上下文并行性

上下文并行性是一种模型并行性，它沿序列维度对模型激活进行分区。其他序列并行性技术只对 LayerNorm 和 RMSNorm 进行分区，而上下文并行性则不同，它沿着序列维度对网络输入和所有中间激活进行分区。

SMP v2 与 Transformer Engine 集成，可实现上下文并行性，并可与 PyTorch FSDP 和 SMP 张量并行 结合使用。您可以同时启用所有三个并行性系统进行模型训练。上下文并行性有利于训练具有大激活规模和长序列长度的模型。它允许每台设备只计算序列维度上的部分分数和输出，从而加快了注意力分数和注意力输出的计算。虽然张量并行通过沿隐藏维度进行分区也能加快计算速度，但上下文并行性的优势更大，因为计算需求随序列维度的增加而呈四倍增长。

Hugging Face 转换器模型兼容 SMP 上下文并行性

SMP v2 目前可为以下 Hugging Face 转换器模型提供上下文并行性支持。

配置上下文并行性

为 context_parallel_degree 参数设置一个整数值，平均分配集群中 GPU 的数量。例如，如果您有一个 8 GPU 实例，请在 context_parallel_degree 中使用 2、4 或 8。我们建议从较小的 context_parallel_degree 值开始，逐渐增加，直到模型适合 GPU 内存所需的输入序列长度。

以下代码片段显示了如何在训练脚本中添加 SMP 初始化模块 torch.sagemaker.init()，并按照 使用 SageMaker 模型并行性库 v2 中介绍的两步流程，为训练作业启动器设置 JSON 格式的 SMP 配置字典。您无需对 PyTorch 模型或 PyTorch FSDP 配置进行任何更改。有关 context_parallel_degree 参数的更多信息，请参阅 SMP v2 核心功能配置参数。

在您的训练脚本中

作为步骤 1 的一部分，使用 torch.sagemaker.init() 初始化脚本以激活 SMP v2，然后使用 torch.sagemaker.transform API 封装模型。

从 SMP v2.6.0 开始，您可以使用参数 cp_comm_type 来确定使用哪种上下文并行性实现。SMP 库目前支持两种实现方式：p2p 和 all_gather。在注意力实现过程中，p2p 实现使用点对点发送接收调用进行键值积累，并异步运行，允许与计算重叠。而 all_gather 实现则使用 AllGather 集体操作并同步运行。

import torch.sagemaker as tsm
tsm.init()

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model, cp_comm_type="p2p")

SMP 配置

作为步骤 2 的一部分，将以下参数添加到 SageMaker PyTorch 估算器的 SMP 配置字典中。

{   
    ..., # other SMP config parameters
    "context_parallel_degree": 2
}