重要

您正在查看 NeMo 2.0 文档。此版本引入了对 API 的重大更改和一个新的库 NeMo Run。我们目前正在将 NeMo 1.0 的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

上下文并行#

上下文并行 (CP) 是一种通过沿序列维度划分输入张量，在多个 GPU 上并行处理神经网络激活的方法。与序列并行 (SP) 划分特定层的激活不同，CP 划分所有层的激活。

CP 对于训练长上下文模型至关重要，因为它允许模型通过跨多个 GPU 分布序列激活来处理更长的序列。此方法减少了处理长序列的内存占用和计算成本。

启用上下文并行#

要在 NeMo 框架中激活 CP，请在模型配置中设置 context_parallel_size 参数。此参数指定模型序列激活分布在多少个 GPU 上。

将 context_parallel_size 设置为大于 1 的值，以在 Megatron 策略 内部启用序列范围的模型并行。

strategy = nl.MegatronStrategy(
  tensor_model_parallel_size=2,
  pipeline_model_parallel_size=2,
  virtual_pipeline_model_parallel_size=None,
  context_parallel_size=1, # Example to enable Context Parallelism
  sequence_parallel=True,
  expert_model_parallel_size=1,
)

有关模型训练设置的更多详细信息可以在这里找到和修改：NeMo 2.0 预训练。

实施上下文并行#

NeMo 框架利用 Megatron Core 和 Transformer Engine 的功能来高效地实施 CP。在正向传播期间，每个 GPU 处理序列的一个片段，仅存储必要的键和值 (KV) 对。在反向传播中，这些 KV 对使用高级通信方案（如 all-gather 和 reduce-scatter）在 GPU 之间重新组装，这些方案在环形拓扑中转换为点对点通信。此方法显著减少了内存占用，同时保持了计算效率。

访问我们的源代码以获取更多关于实施的见解

在本页