Since this isn't stable right now and will break in python 3.11.
Tested using for f in $(find . -name '*_test.py'); do python $f; done. All the tests pass except:
ops/topology_test.py, but I see it fail the same way in main so I'm pretty confident this is unrelated.
memory_test.py and parallelism_test.py fail with one GPU because they (unsurprisingly) fail to initialize torch dist, but pass with multiple GPUs as expected.
Pretty sure I got it all:
megablocks $ rg 'turbo'
megablocks $ rg 'uantize'
megablocks $ rg 'bits'
megablocks/layers/moe.py
113: # Calculate the number of bits needed to represent the expert indices
megablocks/layers/dmoe.py
24: # Calculate the number of bits needed to represent the column indices
Since this isn't stable right now and will break in python 3.11.
Tested using
for f in $(find . -name '*_test.py'); do python $f; done
. All the tests pass except:ops/topology_test.py
, but I see it fail the same way inmain
so I'm pretty confident this is unrelated.memory_test.py
andparallelism_test.py
fail with one GPU because they (unsurprisingly) fail to initialize torch dist, but pass with multiple GPUs as expected.Pretty sure I got it all: