Skip to content

Latest commit

 

History

History
124 lines (103 loc) · 7.2 KB

File metadata and controls

124 lines (103 loc) · 7.2 KB

Benchmarks on 1xRTX5060Ti

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"

Model size: 0.5B

FP8

./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%

BF16

./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%

Model size: 1.5B

FP8

./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%

BF16

./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%

Model size: 3B

FP8

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%

BF16

./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%

Model size: 7B

FP8:

./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%

BF16:

./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%