知了AI学习平台Logo知了
🧠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

张量并行无需通信。

✓

正确理解

前向和反向需通信汇总部分结果,通信是主要开销。