当您使用 SageMaker 模型并行库运行分布式训练任务时，请遵循以下指南。

为特定模型设置正确配置

在扩展模型时，建议您按以下顺序逐一考量。每一项都讨论了使用该库技术的优势以及可能出现的权衡。

提示：如果一个模型仅使用库的部分功能就能很好地适配，增加更多的模型并行或内存节省特性通常不会带来性能提升。

使用大型 GPU 实例

在模型并行领域，最好使用配备大显存的强大实例，以应对模型并行操作（如在多个 GPU 间分区模型）带来的开销。我们推荐使用 ml.p4d 或 ml.p3dn 实例来训练大型深度学习模型。这些实例还配备了弹性结构适配器（EFA），可提供更高的内网带宽，从而支持大规模模型并行训练。

优化器状态分片 (Optimizer State Sharding)

优化器状态分片的效果取决于数据并行的 rank 数量。通常，更高的数据并行度（与计算节点大小成正比）可以提高内存使用效率。

当您计划缩减集群规模时，请务必检查优化器状态分片的配置。例如，一个在 16 个 GPU（如两个 P4d 或 P4de 实例）的计算集群上能够通过优化器状态分片成功运行的大型模型，不一定能在一个 8 GPU 的节点（如单个 P4d 或 P4de 实例）上运行。这是因为 8 个 GPU 的总显存低于 16 个 GPU，并且在 8 个 GPU 上进行分片时，每个 GPU 所需的显存也高于 16 GPU 的情况。因此，增加的显存需求可能无法在较小的集群中得到满足。

激活检查点 (Activation Checkpointing)

通过对一组模块使用激活检查点，可以提高内存效率。您分组的模块越多，内存使用效率就越高。

当对序列模块的层进行检查点设置时，smp.set_activation_checkpointing 函数的 strategy 参数会将这些层组合在一起进行检查点。例如，将两个或更多层组合在一起进行检查点比一次只检查一个层更节省内存，这是一种以额外的计算时间换取内存空间的方法。

张量并行 (Tensor Parallelism)

张量并行的度数应为 2 的幂（2, 4, 8, …, 2^n），且最大度数必须等于每个节点上的 GPU 数量。例如，如果您使用一个有 8 个 GPU 的节点，张量并行的度数可以是 2、4 或 8。我们不推荐使用任意数字（如 3、5、6、7）作为张量并行的度数。

当您使用多个节点时，错误配置张量并行度可能会导致张量并行跨节点运行。这会因跨节点传输激活值而增加显著的通信开销，并可能变得计算成本高昂。

跨节点流水线并行 (Pipeline Parallelism)

流水线并行既可以在单个节点内运行，也可以跨多个节点运行。当您将流水线并行与张量并行结合使用时，我们建议跨多个节点运行流水线并行，并将张量并行限制在单个节点内部。

流水线并行有三个关键参数可供调节：microbatches、active_microbatches 和 prescaled_batch。

  • prescaled_batch：当与张量并行结合使用时，建议激活此选项。这样可以增加每个模型并行组的批量大小，以实现高效的流水线。激活后，训练脚本中设置的批量大小将变为未激活时每个 rank 批量大小的 tp_size 倍。
  • microbatches：增加微批次数量有助于实现高效的流水线和更好的性能。请注意，有效的微批次大小是 batch_size / microbatches。如果在保持批量大小不变的情况下增加微批次数量，则每个微批次处理的样本会变少。
  • active_microbatches：这是流水线处理期间同时进行的最大微批次数量。每个处理中的活动微批次都会占用其激活值和梯度的 GPU 显存。因此，增加此值会占用更多 GPU 显存。如果 GPU 和显存利用率都较低，可以增加此值以获得更好的流水线并行效果。

将激活卸载到 CPU (Activation Offloading)

请确保此功能与激活检查点和流水线并行结合使用。为了确保卸载和预加载在后台进行，请为 microbatches 参数指定一个大于 1 的值。

当卸载激活时，您或许能够增加 active_microbatches 的数量，有时甚至可以使其与 microbatches 的总数相匹配。这取决于哪些模块被设置了检查点以及模型是如何分区的。

参考配置

SageMaker 模型并行训练团队基于 GPT-2 模型（序列长度 512，词汇量 50,000）的实验提供了以下参考点。

模型参数量 实例配置 流水线并行度 张量并行度 优化器状态分片 激活检查点 prescaled_batch 批量大小配置
100 亿 16 x ml.p4d.24xlarge 1 4 True 每个 Transformer 层 True batch_size=40
300 亿 16 x ml.p4d.24xlarge 1 8 True 每个 Transformer 层 True batch_size=32
600 亿 32 x ml.p4d.24xlarge 2 8 True 每个 Transformer 层 True batch_size=56, microbatches=4, active_microbatches=2

您可以根据上述配置来推断您自己模型配置的 GPU 显存使用情况。例如，如果您为一个 100 亿参数的模型增加了序列长度，或者将模型大小增加到 200 亿，您可能需要首先降低批量大小。如果模型仍然无法适配，再尝试增加张量并行的度数。

修改您的训练脚本

在您的训练脚本中使用 SageMaker 模型并行库的功能之前，请先查阅相关的配置技巧和常见陷阱。

为了更快地启动训练任务，建议使用 SageMaker 本地模式。这可以帮助您在 SageMaker 笔记本实例上快速地本地运行训练任务。根据您笔记本实例的规模，您可能需要通过更改模型配置（如隐藏层宽度、Transformer 层数和注意力头数）来调整模型的大小。在将完整模型部署到大型集群进行训练之前，先在笔记本实例上验证缩减后的模型是否能正常运行。

监控和记录训练任务

使用 SageMaker 控制台和 Amazon CloudWatch

要监控系统级指标，如 CPU 内存利用率、GPU 显存利用率和 GPU 利用率，请使用 SageMaker 控制台提供的可视化工具。

  1. 在左侧导航窗格中，选择 训练。
  2. 选择 训练任务。
  3. 在主窗格中，选择您想查看详情的训练任务名称。
  4. 浏览主窗格，找到 监控 部分以查看自动生成的可视化图表。

要查看训练任务日志，请在 监控 部分选择 查看日志。您可以在 CloudWatch 中访问分布式训练任务的日志。如果您启动了多节点分布式训练，您应该会看到多个带有 algo-n-1234567890 格式标签的日志流。其中，algo-1 日志流跟踪来自主节点（第 0 个节点）的训练日志。

权限设置

要运行带有模型并行的 SageMaker 训练任务或示例笔记本，请确保您的 IAM 角色具有正确的权限，例如：

  • AmazonFSxFullAccess：用于使用 FSx for Lustre。
  • AmazonS3FullAccess：用于使用 Amazon S3 作为数据通道。
  • AmazonEC2ContainerRegistryFullAccess：用于使用 Docker、构建您自己的容器并将其推送到 Amazon ECR。
  • AmazonSageMakerFullAccess：用于完全访问 SageMaker 功能套件。

👉 如果你需要 ChatGPT 代充 / Claude / Claude Code / 镜像 / 中转 API：