【学习】MiniMind 分布式并行训练扩展:TP/SP/VP/PP/CP
torchrun --standalone --nproc_per_node=2 scripts/demo_tensor_parallel.py
--tp_size 2
--hidden_size 768
--num_hidden_layers 8
--num_attention_heads 8
--num_key_value_heads 4
--vocab_size 6400
--seq_len 340
--batch_size 32
--dtype float32
--check_backward
--optimizer_steps 100
--log_interval 10
--learning_rate 5e-4
--seed 42
--atol 1e-3
forward max logits diff: 6.198883e-06
forward loss diff: 0.000000e+00
backward max grad diff: 1.006993e-08
layer 0:
self_attn.q_proj.weight mean=8.708e-10 max=5.530e-09 rel_l2=2.877e-06
self_attn.k_proj.weight mean=1.237e-09 max=7.014e-09 rel_l2=2.889e-06
self_attn.v_proj.weight mean=1.690e-09 max=1.007e-08 rel_l2=2.366e-06
self_attn.o_proj.weight mean=1.176e-09 max=7.276e-09 rel_l2=2.304e-06
self_attn.q_norm.weight mean=1.133e-09 max=4.540e-09 rel_l2=2.104e-06
self_attn.k_norm.weight mean=1.248e-09 max=4.075e-09 rel_l2=2.375e-06
mlp.gate_proj.weight
…
内容来源: jingyaogong/minimind