大模型长序列训练并行策略：PCP与DCP原理、实现与调优

1. 大模型训练为什么绕不开并行策略这件事单卡跑大模型的时代早就过去了。现在随便一个能拿得出手的开源模型参数量动辄几十亿起步激活值、梯度、优化器状态加起来显存需求轻松突破单张卡的物理上限。你可能会想那就多卡呗。但多卡不是简单地把模型切开往每张卡上一扔就完事切法不同通信开销、显存占用、计算效率的差异可能是数量级的。这就引出了并行策略的三大主流方向数据并行DP、流水线并行PP、张量并行TP。数据并行是把不同的数据批次分给不同的卡每张卡持有完整的模型副本流水线并行是把模型按层切成若干段不同段放在不同卡上像工厂流水线一样依次处理张量并行则是把单个矩阵运算拆开比如一个大的矩阵乘法拆成几块分给不同的卡同时算。这三种策略各有各的适用场景。数据并行实现简单但显存瓶颈明显——每张卡都得存一份完整的模型参数和优化器状态。流水线并行能解决显存问题但存在“气泡”问题也就是流水线填充和排空阶段有些卡在空转。张量并行通信频繁对带宽要求极高一般只在单机内用 NVLink 这种高速互联才划算。那 PCP 和 DCP 是什么它们属于序列并行这个更细分的领域。序列并行解决的是一个很具体的问题当序列长度变得很长时激活值在序列维度上的显存占用会急剧膨胀。传统的做法是把序列维度也切开来分给不同的卡这就是序列并行。PCP 和 DCP 是两种不同的序列切分策略分别对应不同的切分维度和通信模式。我见过不少团队在长序列场景下踩坑模型能跑起来但效率极低GPU 利用率长期在 30% 以下排查半天发现是并行策略选错了。所以这篇文章不只是讲概念我会把 PCP 和 DCP 的底层逻辑、适用边界、实操配置和踩坑经验都摊开来讲。2. PCP 与 DCP 的底层切分逻辑拆解2.1 从注意力计算的显存瓶颈说起要理解 PCP 和 DCP得先搞清楚 Transformer 里显存到底被什么吃掉了。以一个标准的自注意力层为例输入序列长度为 S隐藏维度为 H注意力头数为 A。计算过程中需要存储的中间张量包括 Query、Key、Value 三个矩阵每个的大小是 S × H。注意力分数矩阵的大小是 S × S × A这个才是真正的显存杀手。当 S 从 2K 涨到 32K 时注意力分数矩阵的显存占用会膨胀 256 倍。这就是为什么长序列场景下必须做序列并行——不切开的话单卡根本放不下。序列并行的核心思路是把序列维度 S 切成若干块每张卡只负责其中一块的计算。但注意力机制有个麻烦的地方每个 token 的 Query 需要和所有 token 的 Key 做点积。如果你把序列切开了每张卡只有部分 Key那 Query 怎么和完整的 Key 做注意力2.2 PCP 的切分方式与通信模式PCP 的全称是Parallel Context Processing有些框架里也叫 Context Parallelism。它的切分逻辑是把序列维度切成 N 份每张卡持有 S/N 个 token 的完整隐藏状态。但在计算注意力时每张卡需要拿到所有其他卡上的 Key 和 Value。具体来说PCP 的通信模式是这样的在前向传播时每张卡先用自己的 Query 和本地的 Key 算局部注意力然后通过 All-Gather 操作收集其他卡上的 Key再算全局注意力。反向传播时需要把梯度的贡献通过 Reduce-Scatter 操作分发回各卡。这种方式的通信量是 O(S × H) 级别的因为每张卡都需要拿到完整的 Key 和 Value 矩阵。当序列很长时All-Gather 的通信开销会变得相当可观。但 PCP 的好处是实现相对简单对模型结构的改动小基本上只需要在注意力层插入通信操作就行。2.3 DCP 的切分方式与通信模式DCP 的全称是Distributed Context Parallelism它的切分维度比 PCP 更细。DCP 不仅切分序列维度还把注意力头维度也纳入切分范围。也就是说每张卡负责一部分序列位置和一部分注意力头的计算。DCP 的通信模式更复杂一些。因为每张卡只持有部分注意力头所以在计算注意力时需要做头维度的 All-Gather。同时由于序列也被切分了还需要做序列维度的通信。但 DCP 的优势在于每张卡需要存储的 Key 和 Value 矩阵更小因为只存部分头的部分序列。从通信量来看DCP 的通信模式可以分解为两个阶段第一阶段是头维度的 All-Gather通信量是 O(S × H / N)第二阶段是序列维度的 All-Gather通信量也是 O(S × H / N)。总通信量和 PCP 相当但 DCP 的显存占用更低因为每张卡只存部分头的 KV。2.4 两种策略的对比与选型依据对比维度PCPDCP切分维度仅序列维度序列维度 注意力头维度单卡 KV 显存O(S × H)O(S × H / N)通信模式All-Gather Reduce-Scatter两阶段 All-Gather实现复杂度较低较高适用场景中等序列长度头数较少超长序列头数较多对模型结构的侵入性小中等选型的时候我一般会看两个指标序列长度和注意力头数。如果序列长度在 8K 到 32K 之间头数不超过 32PCP 基本够用实现简单调试也方便。如果序列长度超过 64K或者头数在 64 以上DCP 的显存优势就体现出来了。还有一个容易被忽略的因素是通信带宽。PCP 的 All-Gather 是单次大通信对带宽的瞬时要求高DCP 是两阶段通信单次通信量小但次数多对延迟更敏感。如果你的集群是 NVLink 互联两者差别不大如果是跨机通信DCP 的分阶段模式可能更友好一些。3. 在训练框架里落地 PCP 和 DCP 的实操路径3.1 环境准备与依赖检查在动手改代码之前先把环境理清楚。我假设你用的是 PyTorch 生态因为目前大多数训练框架都是基于 PyTorch 的。需要确认的依赖包括PyTorch 版本不低于 2.0因为 2.0 之后对分布式通信原语的支持更完善NCCL 版本不低于 2.14低版本在 All-Gather 大张量时可能有性能问题如果用的是 Megatron-LM 或 DeepSpeed需要确认它们对序列并行的支持程度检查命令很简单python -c import torch; print(torch.__version__) python -c import torch.distributed as dist; print(dist.is_nccl_available())如果 NCCL 不可用需要先装 NCCL 库并确保 PyTorch 编译时链接了它。这一步经常被忽略尤其是用 conda 装 PyTorch 的时候有时候会装到 CPU 版本。3.2 PCP 的核心代码实现PCP 的实现主要是在注意力层插入通信操作。以下是一个简化版的实现思路import torch import torch.distributed as dist class PCPAttention(torch.nn.Module): def __init__(self, hidden_size, num_heads, pcp_group): super().__init__() self.hidden_size hidden_size self.num_heads num_heads self.head_dim hidden_size // num_heads self.pcp_group pcp_group self.pcp_size dist.get_world_size(pcp_group) def forward(self, hidden_states): # hidden_states: [batch, seq_len // pcp_size, hidden_size] batch, local_seq, _ hidden_states.shape q self.query_proj(hidden_states) k self.key_proj(hidden_states) v self.value_proj(hidden_states) # All-Gather Key 和 Value k_gathered self._all_gather_along_seq(k) v_gathered self._all_gather_along_seq(v) # 计算注意力 attn_output self._compute_attention(q, k_gathered, v_gathered) return self.output_proj(attn_output) def _all_gather_along_seq(self, tensor): # tensor: [batch, local_seq, hidden_size] world_size self.pcp_size gathered [torch.zeros_like(tensor) for _ in range(world_size)] dist.all_gather(gathered, tensor, groupself.pcp_group) return torch.cat(gathered, dim1)这段代码的关键点在于_all_gather_along_seq方法。它把每张卡上的局部 Key 和 Value 收集起来拼成完整的序列。注意这里用的是all_gather而不是gather因为每张卡都需要完整的 Key 和 Value 来计算注意力。反向传播的时候PyTorch 的自动微分机制会自动处理梯度。但需要注意的是all_gather的反向操作是reduce_scatterPyTorch 会自动处理这个转换。如果你手动实现了通信需要确保反向传播的梯度正确。3.3 DCP 的两阶段通信实现DCP 的实现要复杂一些因为涉及两个维度的切分。以下是一个简化版的实现class DCPAttention(torch.nn.Module): def __init__(self, hidden_size, num_heads, dcp_group): super().__init__() self.hidden_size hidden_size self.num_heads num_heads self.head_dim hidden_size // num_heads self.dcp_group dcp_group self.dcp_size dist.get_world_size(dcp_group) self.local_heads num_heads // self.dcp_size def forward(self, hidden_states): batch, local_seq, _ hidden_states.shape # 每张卡只计算部分头 q self.query_proj(hidden_states) # [batch, local_seq, local_heads * head_dim] k self.key_proj(hidden_states) v self.value_proj(hidden_states) # 第一阶段头维度 All-Gather k_head_gathered self._all_gather_along_head(k) v_head_gathered self._all_gather_along_head(v) # 第二阶段序列维度 All-Gather k_full self._all_gather_along_seq(k_head_gathered) v_full self._all_gather_along_seq(v_head_gathered) # 计算注意力 attn_output self._compute_attention(q, k_full, v_full) return self.output_proj(attn_output) def _all_gather_along_head(self, tensor): # tensor: [batch, local_seq, local_heads * head_dim] world_size self.dcp_size gathered [torch.zeros_like(tensor) for _ in range(world_size)] dist.all_gather(gathered, tensor, groupself.dcp_group) return torch.cat(gathered, dim-1) def _all_gather_along_seq(self, tensor): world_size self.dcp_size gathered [torch.zeros_like(tensor) for _ in range(world_size)] dist.all_gather(gathered, tensor, groupself.dcp_group) return torch.cat(gathered, dim1)DCP 的关键在于两个 All-Gather 的顺序。先做头维度的 All-Gather再做序列维度的 All-Gather。这个顺序不能反因为头维度的通信量更小先做可以减少后续序列维度通信的数据量。3.4 配置参数与调优建议在实际配置中有几个参数需要特别注意pcp_size / dcp_size这个参数决定了序列被切成多少份。一般来说设置为 GPU 数量的约数比较合理。比如 8 卡机器可以设成 2、4、8。设成 8 意味着每张卡只处理 1/8 的序列通信开销最大但显存占用最低。通信组划分PCP 和 DCP 都需要独立的通信组。如果你的集群同时用了数据并行和序列并行需要仔细规划通信组的划分。我一般建议把序列并行的通信组限制在单机内因为跨机的 All-Gather 延迟太高。梯度累积步数序列并行会增加单步的计算时间适当增加梯度累积步数可以提高吞吐。但要注意梯度累积步数增加会线性增加显存占用需要和序列并行度配合调整。混合精度训练PCP 和 DCP 都支持混合精度但需要注意 All-Gather 操作在 FP16 下可能有精度损失。我一般建议在通信前把张量转成 FP32通信后再转回 FP16。这会增加通信量但精度更有保障。4. 实测中遇到的性能陷阱与排查过程4.1 通信瓶颈的定位方法有一次在 8 卡 A100 上跑一个 32K 序列的模型用了 PCP4 的配置结果发现 GPU 利用率只有 40% 左右。一开始怀疑是计算瓶颈用 nsys profiling 之后发现All-Gather 操作占了总时间的 60% 以上。排查过程是这样的先用torch.cuda.Event在通信前后打时间戳确认通信确实是瓶颈。然后检查 NCCL 的配置发现 NCCL 的 buffer size 设得太小默认是 2MB对于大张量的 All-Gather 来说不够用。改成 16MB 之后通信时间下降了约 30%。import os os.environ[NCCL_BUFFSIZE] 16777216 # 16MB os.environ[NCCL_ALGO] Ring # 对于 All-GatherRing 算法通常比 Tree 更快还有一个容易忽略的点是通信和计算的重叠。PyTorch 默认是同步通信也就是说通信的时候计算单元在空转。可以通过torch.distributed的异步通信接口来实现重叠handle dist.all_gather(gathered, tensor, groupgroup, async_opTrue) # 在等待通信完成的同时可以做一些不依赖通信结果的计算 handle.wait()但要注意异步通信会增加显存占用因为需要同时保留通信前后的张量。4.2 显存溢出的常见原因DCP 配置下显存溢出最常见的原因是头维度切分不均衡。比如 32 个头DCP8每张卡分 4 个头这没问题。但如果 30 个头DCP8除不尽就需要做 padding 或者不均衡切分容易导致某些卡显存溢出。另一个原因是中间激活值没有及时释放。在 PCP 的实现中All-Gather 之后的完整 Key 和 Value 矩阵会一直保留到反向传播这部分显存占用是 O(S × H) 级别的。如果序列很长这部分显存可能比模型参数还大。解决办法是用梯度检查点Gradient Checkpointing。把注意力层的中间激活值不保存反向传播时重新计算。这会增加约 30% 的计算量但能省下大量显存。from torch.utils.checkpoint import checkpoint def forward(self, hidden_states): return checkpoint(self._attention_forward, hidden_states)4.3 数值精度问题的处理序列并行下注意力计算的数值精度问题比单卡更突出。原因是 All-Gather 之后的 Key 和 Value 矩阵来自不同的卡如果各卡的计算精度不一致拼接后的结果可能有细微差异。我遇到过一次 loss 突然变成 NaN 的情况排查后发现是某张卡的 Key 矩阵中有异常值。解决办法是在 All-Gather 之前做一次数值范围检查def _check_numerical_stability(tensor, name): if torch.isnan(tensor).any(): raise ValueError(fNaN detected in {name}) if torch.isinf(tensor).any(): raise ValueError(fInf detected in {name}) max_val tensor.abs().max().item() if max_val 1e4: print(fWarning: large value {max_val} in {name})另外注意力分数的计算建议用 FP32即使模型其他部分用 FP16。因为注意力分数涉及 exp 运算FP16 的动态范围不够容易溢出。4.4 负载不均衡的识别与调整序列并行下如果序列长度不能被并行度整除就会出现负载不均衡。比如序列长度 1000PCP8每张卡分 125 个 token这没问题。但如果序列长度 1001就需要某些卡多处理一个 token。短序列下这个问题不明显但长序列下一个 token 的差异可能导致某张卡的计算时间明显更长。解决办法是padding 到并行度的整数倍虽然会浪费一些计算但能保证负载均衡。还有一个隐蔽的负载不均衡来源是注意力掩码。如果用了 causal mask序列前面的 token 注意力计算量小后面的 token 计算量大。在序列并行下如果某张卡恰好分到了序列末尾的部分它的计算量会比其他卡大。解决办法是交错切分而不是连续切分。也就是说第 i 张卡分到的 token 是 i, iN, i2N, ... 而不是连续的 i*S/N 到 (i1)*S/N。5. 从单机到集群PCP 和 DCP 的扩展性考量5.1 单机内与跨机通信的差异单机内用 NVLink带宽能到 600GB/s 甚至更高PCP 和 DCP 的通信开销基本可以忽略。但跨机走 InfiniBand 或者以太网带宽可能只有 100GB/s 到 400GB/s延迟也从微秒级涨到毫秒级。这意味着在跨机场景下PCP 的单次大 All-Gather 可能成为瓶颈。我实测过在 4 机 32 卡的集群上PCP8 的配置下All-Gather 的通信时间占了总时间的 45%。换成 DCP8 之后因为通信分成了两阶段单次通信量小了重叠效果更好通信时间降到了 30% 左右。所以我的建议是单机内优先用 PCP跨机场景考虑 DCP。当然这不是绝对的还要看具体的序列长度和模型结构。5.2 与数据并行、流水线并行的组合实际训练中序列并行很少单独使用通常是和数据并行、流水线并行组合。比如 32 卡集群可以配置成 DP2, PP2, PCP8或者 DP4, PP1, DCP8。组合的时候需要注意通信组的嵌套关系。数据并行的通信组是全局的流水线并行的通信组是跨阶段的序列并行的通信组是阶段内的。如果通信组划分不当可能出现死锁或者通信冲突。我一般用这样的配置逻辑# 假设 world_size 32 dp_size 2 pp_size 2 pcp_size 8 # 先划分流水线组 pp_groups [] for i in range(dp_size * pcp_size): pp_groups.append(list(range(i, world_size, dp_size * pcp_size))) # 再划分序列并行组 pcp_groups [] for i in range(dp_size * pp_size): pcp_groups.append(list(range(i * pcp_size, (i 1) * pcp_size)))这种划分方式保证了序列并行的通信组在单机内假设每机 8 卡减少了跨机通信。5.3 扩展性测试的指标与方法评估 PCP 和 DCP 的扩展性我一般看三个指标吞吐量Tokens/s这是最直接的指标。固定模型大小和序列长度增加并行度看吞吐量是否线性增长。如果增长曲线变平说明通信瓶颈出现了。通信占比用 profiling 工具测出通信时间占总时间的比例。这个比例超过 40% 就说明通信是瓶颈了需要考虑优化通信或者调整并行策略。显存利用率每张卡的显存占用是否均衡。如果某张卡显存占用明显高于其他卡说明切分不均衡需要调整切分策略。测试的时候要注意预热。分布式训练的第一次迭代通常很慢因为要建立通信连接、初始化 NCCL 等。我一般会跑 10 个迭代取后 5 个的平均值。6. 几个容易被忽略的工程细节6.1 随机种子与可复现性序列并行下随机种子的设置比单卡复杂。因为每张卡处理的数据不同如果随机种子设置不当可能导致各卡的 dropout mask 不一致影响模型效果。正确的做法是按数据并行维度设置随机种子而不是按全局 rank。也就是说同一个数据并行组内的卡应该用相同的随机种子不同数据并行组用不同的种子。def set_seed(seed, dp_rank): torch.manual_seed(seed dp_rank) torch.cuda.manual_seed(seed dp_rank) np.random.seed(seed dp_rank)6.2 检查点保存与加载序列并行下保存检查点需要注意只保存一份完整的模型参数而不是每张卡各存一份。因为序列并行切分的是激活值模型参数在每张卡上是完整的对于 PCP或者部分完整的对于 DCP。对于 DCP因为头维度被切分了每张卡只有部分头的参数。保存的时候需要做 All-Gather 把参数拼完整或者保存分片参数并在加载时重新切分。我一般用分片保存的方式每张卡保存自己那部分参数加载的时候按相同的切分方式加载。这样避免了 All-Gather 的开销但需要保证加载时的并行度配置和保存时一致。6.3 与混合精度训练的配合混合精度训练下序列并行需要注意**梯度缩放GradScaler**的配合。因为 All-Gather 操作在 FP16 下可能有精度损失如果梯度缩放因子设置不当可能导致梯度下溢。我的经验是在通信前把张量转成 FP32通信后再转回 FP16。虽然增加了通信量但精度更有保障。如果通信带宽是瓶颈可以考虑用 BF16 代替 FP16BF16 的动态范围更大不容易溢出。# 通信前转 FP32 k_fp32 k.float() k_gathered all_gather(k_fp32) # 通信后转回 FP16 k_gathered k_gathered.half()6.4 调试工具与日志策略分布式训练的调试比单卡麻烦得多因为你看不到所有卡的状态。我一般会在关键位置加日志但要注意只让 rank 0 打印否则日志会刷屏。def log_rank0(msg): if dist.get_rank() 0: print(f[Rank 0] {msg})另外推荐用TensorBoard或者WandB记录每张卡的显存占用、通信时间、计算时间等指标。这些指标能帮你快速定位瓶颈。还有一个技巧是用小规模配置复现问题。比如 8 卡的问题可以尝试用 2 卡复现。如果 2 卡能复现调试起来就简单多了。7. 一些个人实践中的体会PCP 和 DCP 不是银弹它们解决的是特定场景下的特定问题。如果你的序列长度不超过 4K说实话没必要上序列并行数据并行加张量并行基本够用。序列并行的通信开销和实现复杂度都不低用错了反而拖慢训练。我自己的经验是序列长度超过 16K 再考虑序列并行超过 64K 再考虑 DCP。而且一定要先做 profiling确认瓶颈确实在显存或者注意力计算上而不是盲目上并行策略。还有一个体会是通信优化比并行策略选择更重要。同样的 PCP4 配置NCCL 参数调优前后性能可能差 30% 以上。所以在上序列并行之前先把 NCCL 的 buffer size、算法选择、通信重叠这些基础优化做好。最后分布式训练的调试一定要有耐心。一个问题可能涉及通信、计算、显存、精度多个方面排查的时候要系统性地排除不要东一榔头西一棒子。我一般会按照“先确认瓶颈在哪再定位具体原因最后验证修复效果”的流程来走这样效率最高。