分布式训练

训练一个千亿参数的大模型，需要数千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上： ...

序列并行：解决超长序列的显存瓶颈 张量并行（TP）将模型参数切分到多个GPU，流水线并行（PP）将模型层切分到多个GPU，但两者都没有解决序列长度带来的显存问题。 当序列长度达到128K甚至1M时，即使使用TP和PP，每个GPU上存储的激活值仍然是 O(seq_len × hidden_dim)。序列并行（Sequence Parallelism, SP）正是为了解决这个问题——将序列维度也切分到多个GPU上。 序列并行的核心思想 基本切分 将序列维度切分到N个GPU上，每个GPU只处理 seq_len/N 个token： 原始序列: [t0, t1, t2, ..., t_{L-1}] GPU 0: [t0, t1, ..., t_{L/N-1}] GPU 1: [t_{L/N}, ..., t_{2L/N-1}] ... GPU N-1: [t_{L-L/N}, ..., t_{L-1}] 通信需求 序列切分后，LayerNorm和Dropout等逐token操作可以独立进行（无需通信），但注意力计算需要完整的序列信息——每个token的Query需要与所有token的Key/Value计算注意力。 这是序列并行的核心挑战：如何在不重新聚合完整序列的前提下完成注意力计算？ DeepSpeed-Ulysses 原理 DeepSpeed-Ulysses使用All-to-All通信来实现序列并行注意力： class UlyssesAttention(nn.Module): def __init__(self, d_model, n_heads, n_gpus): super().__init__() self.n_gpus = n_gpus self.local_heads = n_heads // n_gpus # 每个GPU只存储部分注意力头的参数 self.W_q = nn.Linear(d_model, self.local_heads * (d_model // n_heads)) self.W_k = nn.Linear(d_model, self.local_heads * (d_model // n_heads)) self.W_v = nn.Linear(d_model, self.local_heads * (d_model // n_heads)) self.W_o = nn.Linear(self.local_heads * (d_model // n_heads), d_model) def forward(self, x): # x的形状: [batch, local_seq_len, d_model] # 每个GPU持有序列的一部分 # 1. 本地计算QKV Q_local = self.W_q(x) # [batch, local_seq, local_heads * head_dim] K_local = self.W_k(x) V_local = self.W_v(x) # 2. All-to-All通信：重新排列序列和头维度 # 从 [batch, local_seq, local_heads, head_dim] # 到 [batch, global_seq, local_heads, head_dim] Q_global = all_to_all(Q_local, split_dim=1, concat_dim=2) K_global = all_to_all(K_local, split_dim=1, concat_dim=2) V_global = all_to_all(V_local, split_dim=1, concat_dim=2) # 3. 本地注意力计算（现在每个GPU有完整的序列，但只有部分头） attn_out = flash_attention(Q_global, K_global, V_global) # 4. All-to-All通信：将结果按序列切分回去 attn_local = all_to_all(attn_out, split_dim=2, concat_dim=1) # 5. 输出投影 output = self.W_o(attn_local) return output 通信分析 Ulysses需要2次All-to-All通信（前向）+ 2次All-to-All（反向）。 ...

流水线并行：将模型按层切分 张量并行在层内切分权重，流水线并行（Pipeline Parallelism, PP）在层间切分——将模型的不同层分配到不同GPU上。GPU 0持有第1-8层，GPU 1持有第9-16层，依次类推。 这种切分方式让每个GPU只需存储部分层的参数和激活值，极大降低了单GPU显存需求。但它引入了一个新问题：流水线气泡——不同GPU之间需要等待数据传递。 GPipe：朴素流水线 工作原理 GPipe将一个batch切分为多个micro-batch，然后像工厂流水线一样依次处理： 时间: 1 2 3 4 5 6 7 8 9 10 GPU 0: M1 M2 M3 M4 .. .. .. .. .. .. GPU 1: .. M1 M2 M3 M4 .. .. .. .. .. GPU 2: .. .. M1 M2 M3 M4 .. .. .. .. GPU 3: .. .. .. M1 M2 M3 M4 .. .. .. ↑ 前向完成 GPU 3: .. .. .. .. .. .. .. M4' M3' M2' M1' ← 反向 GPU 2: .. .. .. .. .. .. M4' M3' M2' M1' ... M1-M4是4个micro-batch的前向传播，M1’-M4’是反向传播。 ...

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

超长上下文的分布式挑战 随着LLM上下文长度从32K扩展到1M甚至更长，单GPU的显存和计算能力已经远远不够。一个1M token的注意力矩阵在FP16下需要2TB显存——即使分布在8个GPU上，每张卡也需要256GB。 更关键的是，标准注意力计算需要在所有Query-Key对之间计算注意力分数，这在分布式环境下意味着大量的跨GPU通信。Ring Attention通过优雅的环形通信模式解决了这个问题。 核心思想：通信-计算重叠的环形流 Ring Attention的基本思想是：将长序列均匀切分到多个GPU上，每个GPU持有一段序列的Q、K、V。然后通过环形通信传递K/V块，同时与本地Q计算部分注意力。 环形拓扑 假设有N个GPU，编号为0到N-1。GPU i 持有序列的第 i 段的Q_i、K_i、V_i。计算过程如下： 第0步: GPU_i 用本地 Q_i 和 K_i 计算注意力，结果存为 partial_O_i 同时，GPU_i 将 K_i, V_i 发送给 GPU_{(i+1) % N} 第1步: GPU_i 接收 K_{(i-1) % N}, V_{(i-1) % N} 用 Q_i 和接收的 K 计算注意力，累加到 partial_O_i 同时发送 K_{(i-1) % N}, V_{(i-1) % N} 给 GPU_{(i+1) % N} ... (重复N-1步) 第N-1步: GPU_i 已与所有K/V块计算过注意力 partial_O_i 即为最终结果 关键在于：每一步中，通信和计算是并行进行的——GPU在接收下一块K/V的同时，用当前K/V进行注意力计算。 ...

引言 联邦学习（Federated Learning, FL）自2016年Google提出以来，已经走过了10年。从最初解决手机输入法隐私问题，到2026年已经成为跨行业、跨组织的AI协作标准。 2026年的联邦学习不仅在算法上更加成熟，在系统架构、隐私保护、产业应用等方面也取得了显著进展。本文将系统梳理这些进展。 一、联邦学习2026全景 1.1 技术成熟度 方面 2016 2026 通信效率 基础 高效压缩+异步 异构处理 差 多种算法应对 隐私保护 基础DP DP+SMPC+HE组合 拜占庭鲁棒 无 多种鲁棒聚合 系统支持 研究原型 生产级框架 1.2 应用规模 Google：20亿+设备参与联邦学习 Apple：iOS系统级联邦学习基础设施 医疗联盟：跨医院联合训练诊断模型 金融联盟：跨银行联合训练反欺诈模型 二、2026年核心技术进展 2.1 高效通信 通信是联邦学习的主要瓶颈。2026年的进展： 梯度压缩 class GradientCompression: def __init__(self, compression_rate=0.01): self.compression_rate = compression_rate # 只传输1%的梯度 def compress(self, gradient): """梯度压缩""" # 方法1: 稀疏化（只传输最大的k个梯度值） flat_grad = gradient.flatten() k = int(len(flat_grad) * self.compression_rate) threshold = np.sort(np.abs(flat_grad))[-k] mask = np.abs(flat_grad) >= threshold compressed = flat_grad * mask indices = np.where(mask)[0] values = compressed[indices] return {"indices": indices, "values": values, "shape": gradient.shape} def decompress(self, compressed): """梯度解压""" gradient = np.zeros(compressed["shape"]) gradient[compressed["indices"]] = compressed["values"] return gradient 量化 def quantize(gradient, bits=8): """梯度量化""" # 将32位浮点数量化为8位整数 min_val, max_val = gradient.min(), gradient.max() scale = (max_val - min_val) / (2**bits - 1) quantized = np.round((gradient - min_val) / scale).astype(np.uint8) return {"quantized": quantized, "min": min_val, "scale": scale} def dequantize(quantized_data): """梯度反量化""" return quantized_data["quantized"] * quantized_data["scale"] + quantized_data["min"] 异步聚合 不是等待所有客户端，而是来一个聚合一个： ...

Ring Attention：突破 GPU 显存墙 当上下文窗口从 4K 扩展到 1M Token 时，单 GPU 的 KV Cache 存储和注意力计算已经无法满足需求。Ring Attention（环形注意力）通过将注意力计算分布到多 GPU，实现了超长上下文的训练和推理。 一、长上下文的挑战 1.1 注意力计算复杂度 标准注意力的计算和显存需求都是 $O(n^2)$： 对于 $n = 1,048,576$（1M Token）： 注意力矩阵：$10^{12}$ 元素 = 8 TB（FP16） 单 GPU A100 80GB 显存：只能放下 0.01 的注意力矩阵 即使使用 Flash Attention 减少 HBM 访问，单 GPU 也无法存储完整的注意力矩阵。 1.2 KV Cache 显存需求 每个 Token 需要存储 Key 和 Value 向量。以 DeepSeek V4（MLA 压缩后）为例： $$\text{KV Cache/Token} = 512 \times 2 \times 2 \text{ bytes} = 2 \text{ KB}$$ ...