🧠AI 冷知识
你知道吗?
这些有趣的 AI 小知识可能让你大吃一惊
💡
切分矩阵而非切分数据Megatron-LM由NVIDIA开发,核心是张量并行(Tensor Parallelism)——把单个矩阵乘法拆分到多张GPU上。关键洞察是:注意力层的QKV投影矩阵可以按列切分(每个GPU算部分头),FFN层可以先按行再按列切分,只在必要的地方做一次All-Reduce通信。这种切分方式让通信量最小化,通信可以和计算重叠。Megatron曾训练了5300亿参数的Megatron-Turing NLG,是当时最大的密集模型之一。
来源:Shoeybi et al., "Megatron-LM: Training Multi-Billion Parameter Language Models", 2019
一句话总结
💡
Megatron-LM 通过张量并行切分矩阵乘法和注意力,结合流水线并行训练超大 Transformer 模型。
常见误区
这些坑别踩
✗
误区 1
张量并行就是简单拆分权重。
✓
正确理解
需在切分维度上配合行/列并行并在恰当地点 AllReduce,保证正确性。
✗
误区 2
Megatron 只支持 Transformer。
✓
正确理解
其张量并行思想可推广,但工程针对 Transformer 深度优化。
✗
误区 3
张量并行无需通信。
✓
正确理解
前向和反向需通信汇总部分结果,通信是主要开销。