张量并行

训练一个千亿参数的大模型，需要数千GPU协同工作数周。如何高效利用这些GPU，是决定训练成本和速度的关键。分布式训练策略就是这个问题的核心答案。本文将系统解析各种并行策略的原理与工程实践。 一、为什么单GPU不够 1.1 显存瓶颈 一个70B参数的模型： 模型参数（FP16）：140GB 梯度（FP16）：140GB 优化器状态（Adam, FP32）：560GB 激活值：取决于batch和序列长度 总计约840GB+ ——远超单张A100的80GB显存。 1.2 计算瓶颈 训练70B模型一个epoch（1万亿token）： 单A100算力：312 TFLOPS (FP16) 训练所需FLOPS：~6×10²³ 单卡理论时间：~22000天 需要2048张A100并行才能在11天内完成。 二、数据并行（DP） 2.1 基本原理 最直观的并行——每张GPU持有完整模型副本，处理不同数据： GPU 0: 完整模型 + Batch[0:32] GPU 1: 完整模型 + Batch[32:64] GPU 2: 完整模型 + Batch[64:96] GPU 3: 完整模型 + Batch[96:128] 前向传播 → 各GPU独立计算loss 反向传播 → AllReduce同步梯度 更新 → 各GPU同步更新参数 2.2 问题：大模型装不下 数据并行要求每张GPU能装下完整模型+优化器状态+梯度。对于70B模型，需要840GB——单卡80GB根本装不下。 2.3 ZeRO优化：显存突破 DeepSpeed的ZeRO（Zero Redundancy Optimizer）通过切分训练状态来突破显存限制： ZeRO-1: 切分优化器状态 每GPU只存 1/N 的优化器状态 70B模型: 560GB/N → N=8时70GB/卡 ZeRO-2: 切分优化器状态 + 梯度 每GPU只存 1/N 的优化器状态和梯度 70B模型: (560+140)GB/N → N=8时87.5GB/卡 ZeRO-3: 切分优化器状态 + 梯度 + 参数 每GPU只存 1/N 的所有状态 70B模型: (560+140+140)GB/N → N=8时105GB/卡 参数在前向/反向时按需All-Gather # ZeRO-3的参数按需获取 class ZeRO3Layer: def forward(self, x): # 按需All-Gather本层参数 full_weight = self.all_gather(self.shard_weight) output = F.linear(x, full_weight) # 立即释放完整参数 del full_weight return output def backward(self, grad_output): # 按需All-Gather本层参数 full_weight = self.all_gather(self.shard_weight) grad_input = grad_output @ full_weight.T # Reduce-Scatter梯度 self.shard_grad = self.reduce_scatter(grad_output.T @ self.input) return grad_input 三、张量并行（TP） 3.1 原理 将单个层的参数矩阵切分到多张GPU上： ...

为什么需要张量并行？ 当模型参数量超过单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求和。 ...