context_parallel 包 - NVIDIA 文档

上下文并行概述

图 1：使用 TP2CP2 运行的 Transformer 层。注意力机制旁边的通信用于 CP，其他的用于 TP。（AG/RS：前向传播中的 all-gather 和反向传播中的 reduce-scatter，RS/AG：前向传播中的 reduce-scatter 和反向传播中的 all-gather，/AG：前向传播中无操作，反向传播中 all-gather）。

上下文并行（“CP”）是一种在序列长度维度上的并行化方案。与之前的 SP（序列并行）仅分割 Dropout 和 LayerNorm 激活的序列不同，CP 沿着序列维度划分网络输入和所有激活。使用 CP，除了注意力机制之外的所有模块（例如，Linear、LayerNorm 等）都可以像往常一样工作，无需任何更改，因为它们没有 token 间的操作。至于注意力机制，每个 token 的 Q（查询）需要与同一序列中所有 token 的 KV（键和值）一起计算。因此，CP 需要在 GPU 之间进行额外的 all-gather 操作，以收集完整的 KV 序列。相应地，reduce-scatter 应该应用于反向传播中 KV 的激活梯度。为了减少激活内存占用，每个 GPU 在前向传播中仅存储序列块的 KV，并在反向传播中再次收集 KV。KV 通信发生在 GPU 和其在其他 TP 组中的对应部分之间。all-gather 和 reduce-scatter 在底层被转换为环形拓扑中的点对点通信。交换 KV 还可以利用 MQA/GQA 来减少通信量，因为它们对于 KV 只有一个或几个注意力头。

例如，在图 1 中，假设序列长度为 8K，则每个 GPU 处理 4K 个 token。GPU0 和 GPU2 组成一个 CP 组，它们彼此交换 KV。GPU1 和 GPU3 之间也发生同样的事情。CP 类似于 Ring Attention，但通过以下方式提供更好的性能：（1）利用最新的 OSS 和 cuDNN flash attention kernels；（2）消除由低三角形因果掩码导致的不必要计算，并在 GPU 之间实现最佳负载平衡。

上下文并行优势

图 2：在各种 TP+CP 组合下，175B GPT 相对于完全重计算（即 TP8CP1）的加速比。

LLM 在长上下文（即长序列长度）下遇到 OOM（内存不足）问题，这是因为激活的内存占用线性增加。在反向传播中重新计算激活可以避免 OOM，但也会引入显著的开销（完全重计算约为 30%）。扩大 TP（张量模型并行）也可以解决 OOM 问题，但这可能会使计算（例如，Linear）太短，以至于无法重叠通信延迟。需要明确的是，无论是否发生 OOM，使用更大的 TP 扩展到更多 GPU 都可能遇到重叠问题。

CP 可以更好地解决这些问题。使用 CP，每个 GPU 仅计算序列的一部分，从而将计算和通信都减少 CP 倍。因此，无需担心它们之间的重叠。每个 GPU 的激活内存占用也减少 CP 倍，因此不再有 OOM 问题。如图 2 所示，TP 和 CP 的组合可以通过消除重计算开销并在计算和通信之间实现最佳权衡来获得最佳性能。

启用上下文并行

GPT 已添加 CP 支持。所有共享 GPT 代码路径的模型也应该能够从 CP 中受益，例如 Llama。CP 可以与 TP（张量模型并行）、PP（流水线模型并行）和 DP（数据并行）一起工作，其中 GPU 总数等于 TPxCPxPPxDP。CP 也可以与不同的注意力变体一起工作，包括 MHA/MQA/GQA、单向和双向掩码。

通过在命令行中简单地设置 context_parallel_size=<CP_SIZE> 即可启用 CP。默认的 context_parallel_size 为 1，这意味着 CP 已禁用。使用 CP 需要 Megatron-Core (>=0.5.0) 和 Transformer Engine (>=1.1)。