All commands expect the following environment variable for common arguments:
export ARGS="--train-file=data/tiny-stories-qwen/train-*.bin --eval-file=data/tiny-stories-qwen/eval.bin \
--ckpt-interval=10000 --from-scratch --seq-length=1024 --model-dtype=bf16 --opt-m-dtype=bf16 \
--opt-v-dtype=bf16 --gpus=1 --use-cuda-graphs"./train ${ARGS} --model=Qwen2.5-0.5B --matmul-dtype=e4m3 --batch-size=12 --grad-accumulation=42 --lmhead-chunks=6
# [T] step 0 [ 0.1%] | time: 30916 ms | norm 14.364049 | loss 12.123322 | tps 16693 | sol 71.2%
# [T] step 1 [ 0.2%] | time: 31005 ms | norm 15.903967 | loss 11.503590 | tps 16645 | sol 71.0%
# [T] step 2 [ 0.3%] | time: 31023 ms | norm 12.736087 | loss 10.920015 | tps 16635 | sol 71.0%
# [T] step 3 [ 0.5%] | time: 31086 ms | norm 9.621207 | loss 10.510395 | tps 16602 | sol 70.8%
# [T] step 4 [ 0.6%] | time: 31121 ms | norm 7.813529 | loss 10.238544 | tps 16583 | sol 70.8%
# [T] step 5 [ 0.7%] | time: 31127 ms | norm 6.885838 | loss 10.012827 | tps 16580 | sol 70.7%./train ${ARGS} --model=Qwen2.5-0.5B --matmul-dtype=bf16 --batch-size=10 --grad-accumulation=51 --lmhead-chunks=5
# [T] step 0 [ 0.1%] | time: 40002 ms | norm 14.584712 | loss 12.124982 | tps 13055 | sol 85.3%
# [T] step 1 [ 0.2%] | time: 40021 ms | norm 16.088415 | loss 11.504570 | tps 13049 | sol 85.3%
# [T] step 2 [ 0.3%] | time: 40115 ms | norm 12.634532 | loss 10.919158 | tps 13018 | sol 85.1%
# [T] step 3 [ 0.5%] | time: 40136 ms | norm 9.692027 | loss 10.508515 | tps 13011 | sol 85.0%
# [T] step 4 [ 0.6%] | time: 40139 ms | norm 7.875288 | loss 10.236643 | tps 13010 | sol 85.0%
# [T] step 5 [ 0.7%] | time: 40142 ms | norm 6.852544 | loss 10.006608 | tps 13009 | sol 85.0%./train ${ARGS} --model=Qwen2.5-1.5B --matmul-dtype=e4m3 --batch-size=8 --grad-accumulation=64 \
--lmhead-chunks=4 --recompute-block --offload-residual
# [T] step 0 [ 0.1%] | time: 90735 ms | norm 16.107918 | loss 12.241903 | tps 5778 | sol 68.2%
# [T] step 1 [ 0.2%] | time: 90869 ms | norm 20.678112 | loss 11.064617 | tps 5769 | sol 68.1%
# [T] step 2 [ 0.4%] | time: 90888 ms | norm 12.360747 | loss 9.906040 | tps 5768 | sol 68.1%
# [T] step 3 [ 0.5%] | time: 90892 ms | norm 12.093929 | loss 9.316822 | tps 5768 | sol 68.1%
# [T] step 4 [ 0.6%] | time: 90885 ms | norm 13.479617 | loss 9.044628 | tps 5768 | sol 68.1%
# [T] step 5 [ 0.7%] | time: 90887 ms | norm 11.006710 | loss 8.758279 | tps 5768 | sol 68.1%./train ${ARGS} --model=Qwen2.5-1.5B --matmul-dtype=bf16 --batch-size=8 --grad-accumulation=64 \
--lmhead-chunks=4 --recompute-ffn --recompute-att
# [T] step 0 [ 0.1%] | time: 134 s | norm 16.183846 | loss 12.242573 | tps 3907 | sol 78.6%
# [T] step 1 [ 0.2%] | time: 134 s | norm 20.687468 | loss 11.053905 | tps 3899 | sol 78.4%
# [T] step 2 [ 0.4%] | time: 134 s | norm 12.302109 | loss 9.896978 | tps 3899 | sol 78.4%
# [T] step 3 [ 0.5%] | time: 134 s | norm 11.886316 | loss 9.304335 | tps 3899 | sol 78.4%
# [T] step 4 [ 0.6%] | time: 134 s | norm 12.836724 | loss 9.011662 | tps 3899 | sol 78.4%
# [T] step 5 [ 0.7%] | time: 134 s | norm 13.274802 | loss 8.743011 | tps 3899 | sol 78.4%Recomputing swiglu and norms to enable bs=6 turns out to be slower (and in fact, bs=6 is generally slower than bf=4).
./train ${ARGS} --model=Qwen2.5-3B --matmul-dtype=e4m3 --batch-size=16 --grad-accumulation=32 \
--lmhead-chunks=8 --attn-bwd-chunks=4 --recompute-ffn --recompute-att --recompute-norm \
--offload-opt-m --offload-opt-v --offload-master --offload-residual
# [T] step 0 [ 0.1%] | time: 167 s | norm 17.138525 | loss 12.362571 | tps 3124 | sol 70.1%
# [T] step 1 [ 0.2%] | time: 167 s | norm 30.094492 | loss 10.852694 | tps 3128 | sol 70.2%
# [T] step 2 [ 0.4%] | time: 167 s | norm 153.881561 | loss 10.714193 | tps 3129 | sol 70.3%
# [T] step 3 [ 0.5%] | time: 167 s | norm 70.368210 | loss 9.615002 | tps 3129 | sol 70.2%
# [T] step 4 [ 0.6%] | time: 167 s | norm 37.567360 | loss 9.484173 | tps 3129 | sol 70.2%
# [T] step 5 [ 0.7%] | time: 167 s | norm 23.869703 | loss 9.445107 | tps 3129 | sol 70.2%./train ${ARGS} --model=Qwen2.5-3B --matmul-dtype=bf16 --batch-size=16 --grad-accumulation=32 \
--lmhead-chunks=8 --attn-bwd-chunks=4 --recompute-ffn --recompute-qkv --recompute-norm \
--offload-opt-m --offload-opt-v --offload-residual --shard-weights --offload-master
# [T] step 0 [ 0.1%] | time: 260 s | norm 17.327423 | loss 12.364223 | tps 2012 | sol 80.6%
# [T] step 1 [ 0.2%] | time: 260 s | norm 30.029100 | loss 10.842840 | tps 2013 | sol 80.6%
# [T] step 2 [ 0.4%] | time: 260 s | norm 151.518051 | loss 10.669870 | tps 2012 | sol 80.6%
# [T] step 3 [ 0.5%] | time: 260 s | norm 72.001083 | loss 9.619955 | tps 2013 | sol 80.6%
# [T] step 4 [ 0.6%] | time: 260 s | norm 38.074726 | loss 9.484632 | tps 2013 | sol 80.6%
# [T] step 5 [ 0.7%] | time: 260 s | norm 23.123236 | loss 9.422847 | tps 2013 | sol 80.6%./train ${ARGS} --model=Qwen2.5-7B --matmul-dtype=e4m3 --batch-size=32 --grad-accumulation=16 \
--lmhead-chunks=32 --attn-bwd-chunks=32 --recompute-block --offload-opt-m --offload-opt-v \
--offload-master --offload-residual --shard-gradients --offload-gradients --shard-weights \
--persistent-quants --offload-quants
# [T] step 0 [ 0.1%] | time: 360 s | norm 12.318468 | loss 12.652546 | tps 1454 | sol 72.0%
# [T] step 1 [ 0.2%] | time: 359 s | norm 23.801075 | loss 10.818451 | tps 1457 | sol 72.1%
# [T] step 2 [ 0.4%] | time: 359 s | norm 197.750610 | loss 15.778147 | tps 1456 | sol 72.1%
# [T] step 3 [ 0.5%] | time: 359 s | norm 148.652817 | loss 15.137997 | tps 1456 | sol 72.1%
# [T] step 4 [ 0.6%] | time: 359 s | norm 130.939774 | loss 15.104545 | tps 1458 | sol 72.1%
# [T] step 5 [ 0.7%] | time: 359 s | norm 103.617371 | loss 13.061750 | tps 1456 | sol 72.1%./train ${ARGS} --model=Qwen2.5-7B --matmul-dtype=bf16 --batch-size=32 --grad-accumulation=16 \
--lmhead-chunks=32 --attn-bwd-chunks=32 --recompute-block --offload-opt-m --offload-opt-v \
--offload-master --offload-residual --shard-gradients --offload-gradients --shard-weights
# [T] step 0 [ 0.1%] | time: 587 s | norm 12.375262 | loss 12.651280 | tps 892 | sol 81.0%
# [T] step 1 [ 0.2%] | time: 589 s | norm 23.723534 | loss 10.803972 | tps 889 | sol 80.8%
# [T] step 2 [ 0.4%] | time: 589 s | norm 198.276062 | loss 15.801797 | tps 889 | sol 80.8%
# [T] step 3 [ 0.5%] | time: 589 s | norm 149.826492 | loss 15.168220 | tps 889 | sol 80.8%
# [T] step 4 [ 0.6%] | time: 589 s | norm 131.514847 | loss 15.141993 | tps 889 | sol 80.8%
# [T] step 5 [ 0.7%] | time: 589 s | norm 106.500847 | loss 13.073975 | tps 889 | sol 80.8%