Megatron-LM

为什么需要张量并行？ 当模型参数量超过单GPU显存容量时，模型并行成为必需。张量并行（Tensor Parallelism, TP）是最常用的模型并行方式——将单个层的权重矩阵切分到多个GPU上，每个GPU计算一部分结果，然后通过通信合并。 NVIDIA的Megatron-LM提出了张量并行的标准方案，成为2026年大模型训练的基础设施。 张量并行的核心原理 线性层的切分 一个标准线性层 Y = XW，其中 W ∈ ℝ^{d×h}。将W切分为两个子矩阵有两种方式： 列切分（Column Parallel） 将W按列切分：W = [W₁, W₂]，每个GPU持有 W_i ∈ ℝ^{d×h/N} GPU 1: Y₁ = XW₁ ∈ ℝ^{n×h/N} GPU 2: Y₂ = XW₂ ∈ ℝ^{n×h/N} Y = [Y₁, Y₂] ∈ ℝ^{n×h} 列切分后，每个GPU独立计算Y的一部分列，最终通过All-Gather合并。但更聪明的方式是——不需要显式合并，因为后续操作可以利用分块结构。 行切分（Row Parallel） 将W按行切分：W = [W₁; W₂]，每个GPU持有 W_i ∈ ℝ^{d/N×h} GPU 1: Y₁ = X₁W₁ (X₁是X的前d/N列) GPU 2: Y₂ = X₂W₂ (X₂是X的后d/N列) Y = Y₁ + Y₂ (All-Reduce) 行切分后需要All-Reduce求和。 ...