本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。 # （已存档）SageMaker 模型并行性库 v1.x **重要** 自 2023 年 12 月 19 日起，SageMaker 模型并行性（SMP）库 v2 已发布。为了支持 SMP 库 v2，未来版本将不再支持 SMP v1 功能。以下部分和主题已存档，专门用于使用 SMP 库 v1。有关使用 SMP 库 v2 的信息，请参阅 [SageMaker 模型并行性库 v2](model-parallel-v2.md)。 使用 Amazon SageMaker AI 的模型并行库来训练由于 GPU 内存限制而难以训练的大型深度学习（DL）模型。该库会自动将在多个 GPU 和实例之间高效地拆分模型。使用该库，您可以高效地训练具有数十亿或万亿参数的大型 DL 模型，从而更快地获得目标预测准确性。 使用该库后，您只需最少的代码更改，即可将自己的 TensorFlow 和 PyTorch 模型自动分区到多个 GPU 和多个节点上。您可以通过 SageMaker Python SDK 访问库的 API。 使用以下部分了解有关模型并行性和 SageMaker 模型并行性库的更多信息。此库的 API 文档位于*《SageMaker Python SDK v2.199.0 文档》*的[分布式训练 API](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html) 中。 **Topics** + [模型并行性简介](model-parallel-intro.md) + [支持的框架和 AWS 区域](distributed-model-parallel-support.md) + [SageMaker 模型并行度库的核心功能](model-parallel-core-features.md) + [使用模型并行度运行 SageMaker 分布式训练 Job](model-parallel-use-api.md) + [对具有模型并行性的模型执行检查点操作和微调](distributed-model-parallel-checkpointing-and-finetuning.md) + [亚马逊 SageMaker AI 模型并行库 v1 示例](distributed-model-parallel-examples.md) + [SageMaker 分布式模型并行性最佳实践](model-parallel-best-practices.md) + [SageMaker 分布式模型并行库配置提示和陷阱](model-parallel-customize-tips-pitfalls.md) + [模型并行故障排除](distributed-troubleshooting-model-parallel.md) # 模型并行性简介 模型并行性是一种分布式训练方法，在这种方法中，深度学习模型在单个实例或多个实例中的设备之间分区。本简介页面提供了有关模型并行性的高级概述，描述了它如何帮助克服在训练通常非常大的深度学习模型时出现的问题，并举例说明了模型 SageMaker 并行库为帮助管理模型并行策略和内存消耗而提供的内容。 ## 什么是模型并行性？ 对于计算机视觉和自然语言处理等复杂的任务，增加深度学习模型（层和参数）的大小可以提高准确性。但是，在单个 GPU 的内存中所能容纳的最大模型大小有限制。在训练 DL 模型时，GPU 内存限制可能会在以下方面成为瓶颈： + 它们限制了您可以训练的模型的大小，因为模型的内存占用量与参数数量成比例扩展。 + 它们限制训练期间的每个 GPU 的批次大小，从而降低了 GPU 利用率和训练效率。 为了克服与在单个 GPU 上训练模型相关的限制， SageMaker 提供了模型并行库，以帮助在多个计算节点上高效地分发和训练 DL 模型。此外，借助该库，您可以使用支持 EFA 的设备实现最大限度优化的分布式训练，这些设备具备低延迟、高吞吐量和 OS 旁路等特性，可以增强节点间通信的性能。 ## 在使用模型并行性之前估算内存需求 在使用 SageMaker 模型并行库之前，请考虑以下内容，以了解训练大型 DL 模型的内存需求。 对于使用 AMP (FP16) 和 Adam 优化器的训练作业，每个参数所需的 GPU 内存约为 20 字节，我们可以将其分解如下： + 一个 FP16 参数 \$1 2 字节 + FP16 渐变 \$1 2 字节 + 基于 Adam FP32 优化器的优化器状态 \$1 8 字节 + 大约 4 字节的参数 FP32 副本（`optimizer apply`（OA）操作需要） + 大约 4 字节的渐变 FP32 副本（OA 操作需要） 即使对于具有 100 亿个参数这样的相对较小 DL 模型，它也会需要至少 200 GB 的内存，这比单个 GPU 上典型的可用 GPU 内存（例如，NVIDIA A100 具有 40 GB/80 GB 内存，V100 具有16/32 GB 内存）大得多。请注意，除了模型和优化器状态的内存要求外，还有其他因素也会占用内存，例如在向前传递中生成的激活。所需的内存可能远远超过 200 GB。 对于分布式训练，我们建议您使用分别具有 NVIDIA V100 和 A100 Tensor Core 的 Amazon EC2 P3 和 P4 实例。 GPUs 有关 CPU 核心、RAM、附属存储卷和网络带宽等规格的详细信息，请参阅 [Amazon EC2 实例类型](https://aws.amazon.com/ec2/instance-types/)页面的*加速计算*部分。 即使使用加速型计算实例，很明显，对于具有大约 100 亿个参数的模型（例如 Megatron-LM 和 T5），甚至具有数千亿个参数的更大模型（例如 GPT-3），也无法在每个 GPU 设备中容纳模型副本。 ## 库如何使用模型并行性和内存节省技术 库中包含各种类型的模型并行功能和节省内存的功能，例如优化器状态分片、激活检查点和激活分载。所有这些技术可以结合使用，从而高效地训练由数千亿个参数组成的大型模型。 **Topics** + [分片数据并行度（可用于） PyTorch](#model-parallel-intro-sdp) + [管道并行度（适用于 PyTorch 和） TensorFlow](#model-parallel-intro-pp) + [张量并行度（可用于） PyTorch](#model-parallel-intro-tp) + [优化器状态分片（可用于） PyTorch](#model-parallel-intro-oss) + [激活卸载和检查点操作（可用于） PyTorch](#model-parallel-intro-activation-offloading-checkpointing) + [为模型选择合适的技术](#model-parallel-intro-choosing-techniques) ### 分片数据并行度（可用于） PyTorch 分@@ *片数据并行性*是一种节省内存的分布式训练技术，它在数据并行组中拆分模型的状态（模型参数、梯度和优化器状态）。 GPUs SageMaker [AI 通过实现 MIC 来实现分片数据并行性，MIC 是 **mi** 对 **c** 通信规模进行模拟化的库，并在博客文章中讨论了巨**型**模型训练的近线性缩放。 AWS](https://www.amazon.science/blog/near-linear-scaling-of-gigantic-model-training-on-aws) 您可以将分片数据并行性作为独立策略应用到模型。此外，如果您使用的是配备NVIDIA A100 Tensor Core的最高性能的GPU实例 GPUs`ml.p4d.24xlarge`，则可以利用SMDDP Collectives提供的`AllGather`操作提高的训练速度。 要深入研究分片数据并行性并学习如何对其进行设置，或者将分片数据并行性与其他技术（例如张量并行性和训练）结合使用，请参阅。 FP16 [分片数据并行性](model-parallel-extended-features-pytorch-sharded-data-parallelism.md) ### 管道并行度（适用于 PyTorch 和） TensorFlow *管道并行性*将一组层或操作在到一组设备之间分区，保留每个操作保持不变。为模型分区数 () 指定值时，(`pipeline_parallel_degree`) 的 GPUs总数必须可被模型分区的数量整除。`processes_per_host`要正确设置此项，您必须为 `pipeline_parallel_degree` 和 `processes_per_host` 参数指定正确的值。简单的数学运算如下： ``` (pipeline_parallel_degree) x (data_parallel_degree) = processes_per_host ``` 根据您提供的两个输入参数，库负责计算模型副本（也称为`data_parallel_degree`）的数量。 例如，如果您设置`"pipeline_parallel_degree": 2`并`"processes_per_host": 8`使用具有八个 GPU 工作线程的 ML 实例（例如）`ml.p3.16xlarge`，则库会自动设置跨 GPUs 和四向数据并行度的分布式模型。下图说明了模型如何分布在八个模型中， GPUs 实现四向数据并行性和双向流水线并行性。每个模型副本（我们将其定义为*流水线 parallel 组*并将其标记为`PP_GROUP`）都分为两 GPUs个。模型的每个分区都分配给四个 GPUs，其中四个分区副本位于一个*数据并行组*中并标记为`DP_GROUP`。如果没有张量并行性，管道并行组本质上就是模型并行组。 ![\[模型如何分布在八个模型中， GPUs 实现四向数据并行性和双向流水线并行性。\]](http://docs.aws.amazon.com/zh_cn/sagemaker/latest/dg/images/distributed/model-parallel/smdmp-pipeline-parallel-only.png) 要深入了解管道并行性，请参阅 [SageMaker 模型并行度库的核心功能](model-parallel-core-features.md)。 要开始使用管道并行度运行模型，请参见使用模型并行[库运行 SageMaker 分布式训练 Job](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-use-api.html)。 SageMaker ### 张量并行度（可用于） PyTorch *张量并行性*可跨设备拆分各个层（即 `nn.Modules`），以并行运行。下图显示了一个最简单的示例，演示库如何拆分具有四个层的模型，以实现两路张量并行性 (`"tensor_parallel_degree": 2`)。每个模型副本的层被一分为二，分布为两 GPUs层。在此示例中，模型并行配置还包括`"pipeline_parallel_degree": 1`和`"ddp": True`（在后台使用 PyTorch DistributedDataParallel 包），因此数据并行度变为八。库管理张量分布式模型副本之间的通信。 ![\[一个最简单的示例，演示库如何拆分具有四个层的模型，以实现两路张量并行 ("tensor_parallel_degree": 2)。\]](http://docs.aws.amazon.com/zh_cn/sagemaker/latest/dg/images/distributed/model-parallel/smdmp-tensor-parallel-only.png) 此功能的用处在于，您可以选择特定的层或层的子集来应用张量并行性。要深入了解张量并行性和其他节省内存的功能 PyTorch，以及如何设置流水线和张量并行度的组合，请参阅。[张量并行性](model-parallel-extended-features-pytorch-tensor-parallelism.md) ### 优化器状态分片（可用于） PyTorch 要了解库如何执行*优化器状态分片*，可以考虑包含四个层的简单示例模型。优化状态分片的关键思想是，你不需要在所有状态中复制优化器状态。 GPUs而是跨数据并行秩对优化器状态的单个副本进行分片，不存在跨设备的冗余。例如，GPU 0 存有第一层的优化器状态，下一个 GPU 1 存有 L2 的优化器状态，依此类推。以下动图显示了使用优化器状态分片技术的向后传播。在向后传播结束时，`optimizer apply` (OA) 操作会占用计算和网络时间来更新优化器状态，`all-gather` (AG) 操作会占用计算和网络时间来更新模型参数用于下一次迭代。最重要的是，`reduce` 操作可能与 GPU 0 上的计算重叠，从而带来更高的内存效率和更快的向后传播。在当前实施中，AG 和 OA 操作不与 `compute` 重叠。在 AG 操作期间，它可能会导致计算延长，因此可能需要权衡取舍。 ![\[利用优化器状态分片技术进行后向传播。\]](http://docs.aws.amazon.com/zh_cn/sagemaker/latest/dg/images/distributed/model-parallel/smdmp-optimizer-state-sharding.gif) 有关如何使用此功能的更多信息，请参阅[优化器状态分片](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-optimizer-state-sharding.html)。 ### 激活卸载和检查点操作（可用于） PyTorch 为了节省 GPU 内存，库支持激活检查点，以避免在向前传递期间，将用户指定模块的内部激活存储在 GPU 内存中。库会在向后传递期间重新计算这些激活。此外，激活分载功能将存储的激活分载到 CPU 内存，并在向后传递期间提取回 GPU，以进一步减少激活内存占用。有关如何使用这些功能的更多信息，请参阅[激活检查点](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-activation-checkpointing.html)和[激活分载](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-activation-offloading.html)。 ### 为模型选择合适的技术 有关选择正确技术和配置的更多信息，请参阅[SageMaker 分布式模型并行最佳实践](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-best-practices.html)和[配置提示和陷阱](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-customize-tips-pitfalls.html)。 # 支持的框架和 AWS 区域 在使用 SageMaker 模型并行度库之前，请检查支持的框架和实例类型，并确定您的 AWS 账户和中是否有足够的配额。 AWS 区域 **注意** 要查看该库的最新更新和发行说明，请参阅 *SageMaker Python SDK 文档*中的[SageMaker 模型并行发行说明](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_release_notes/smd_model_parallel_change_log.html)。 ## 支持的框架 SageMaker 模型并行度库支持以下深度学习框架，可在 Deep Learning Contain AWS ers (DLC) 中使用，也可以作为二进制文件下载。 PyTorch SageMaker AI 支持的版本和 SageMaker 模型并行度库 | PyTorch 版本 | SageMaker 模型并行度库版本 | `smdistributed-modelparallel` 集成 DLC 映像 URI | 二进制文件的 URL\$1\$1 | | --- | --- | --- | --- | | v2.0.0 | smdistributed-modelparallel==v1.15.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:2.0.0-gpu-py310-cu118-ubuntu20.04-sagemaker` | https://sagemaker-distributed-model-parallel.s3.us-west-2.amazonaws.com/pytorch-2.0.0/build-artifacts/2023-04-14-20-14/smdistributed\$1modelparallel-1.15.0-cp310-cp310-linux\$1x86\$164.whl | | v1.13.1 | smdistributed-modelparallel==v1.15.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.13.1-gpu-py39-cu117-ubuntu20.04-sagemaker` | https://sagemaker-distributed-model-parallel.s3.us-west-2.amazonaws.com/pytorch-1.13.1/build-artifacts/2023-04-17-15-49/smdistributed\$1modelparallel-1.15.0-cp39-cp39-linux\$1x86\$164.whl | | v1.12.1 | smdistributed-modelparallel==v1.13.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.12.1-gpu-py38-cu113-ubuntu20.04-sagemaker` | https://sagemaker-distributed-model-parallel.s3.us-west-2.amazonaws.com/pytorch-1.12.1/build-artifacts/2022-12-08-21-34/smdistributed\$1modelparallel-1.13.0-cp38-cp38-linux\$1x86\$164.whl | | v1.12.0 | smdistributed-modelparallel==v1.11.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.12.0-gpu-py38-cu113-ubuntu20.04-sagemaker` | https://sagemaker-distributed-model-parallel.s3.us-west-2.amazonaws.com/pytorch-1.12.0/build-artifacts/2022-08-12-16-16-58/smdistributed\$1modelparallel-1.11.0-cp38-cp38-linux\$1x86\$164.whl | | v1.11.0 | smdistributed-modelparallel==v1.10.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.11.0-gpu-py38-cu113-ubuntu20.04-sagemaker` | https://sagemaker-distributed-model-parallel.s3.us-west-2.amazonaws.com/pytorch-1.11.0/build-artifacts/2022-07-11-19-23/smdistributed\$1modelparallel-1.10.0-cp38-cp38-linux\$1x86\$164.whl | | v1.10.2 | smdistributed-modelparallel==v1.7.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.10.2-gpu-py38-cu113-ubuntu20.04-sagemaker` | - | | v1.10.0 | smdistributed-modelparallel==v1.5.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.10.0-gpu-py38-cu113-ubuntu20.04-sagemaker` | - | | v1.9.1 | smdistributed-modelparallel==v1.4.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.9.1-gpu-py38-cu111-ubuntu20.04` | - | | v1.8.1\$1 | smdistributed-modelparallel==v1.6.0 | `763104351884.dkr.ecr..amazonaws.com/pytorch-training:1.8.1-gpu-py36-cu111-ubuntu18.04` | - | **注意** SageMaker 模型并行度库 v1.6.0 及更高版本为提供了扩展功能。 PyTorch有关更多信息，请参阅 [SageMaker 模型并行度库的核心功能](model-parallel-core-features.md)。 \$1\$1 二进制文件用于在自定义容器中安装 SageMaker 模型并行度库。 URLs 有关更多信息，请参阅 [使用 SageMaker 分布式模型并行库创建自己的 Docker 容器](model-parallel-sm-sdk.md#model-parallel-bring-your-own-container)。 TensorFlow SageMaker AI 支持的版本和 SageMaker 模型并行度库 | TensorFlow 版本 | SageMaker 模型并行度库版本 | `smdistributed-modelparallel` 集成 DLC 映像 URI | | --- | --- | --- | | v2.6.0 | smdistributed-modelparallel==v1.4.0 | 763104351884.dkr.ecr..amazonaws.com/tensorflow-training:2.6.0-gpu-py38-cu112-ubuntu20.04 | | v2.5.1 | smdistributed-modelparallel==v1.4.0 | 763104351884.dkr.ecr..amazonaws.com/tensorflow-training:2.5.1-gpu-py37-cu112-ubuntu18.04 | **人工智能和 SageMaker分布式数据并行库支持 SageMaker 的 Hugging Face 变形金刚版本** 适用于 Hu AWS gging Face 的 Deep Learning Contain PyTorch er TensorFlow s 使用 SageMaker 训练容器作为基础图像。要查找 Hugging Face Transformers 库版本以及 PyTorch 配对版本 TensorFlow 和版本，请查看[最新的 Hugging Face 容器和之前的 Hu](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#huggingface-training-containers) ggin [g Fac](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#prior-hugging-face-container-versions) e 容器版本。 ## AWS 区域 SageMaker 数据并行库可在所有使用 Dee [AWS p Learning Contain](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#sagemaker-framework-containers-sm-support-only) er SageMaker s AWS 区域 的地方使用。有关更多信息，请参阅[可用的深度学习容器映像](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#available-deep-learning-containers-images)。 ## 支持的实例类型 SageMaker 模型并行度库需要以下 ML 实例类型之一。 | 实例类型 | | --- | | ml.g4dn.12xlarge | | ml.p3.16xlarge | | ml.p3dn.24xlarge | | ml.p4d.24xlarge | | ml.p4de.24xlarge | 有关实例类型的规格，请参阅 [Amazon EC2 实例类型](https://aws.amazon.com/ec2/instance-types/)页面中的**加速计算**部分。有关实例定价的信息，请参阅 [Amazon A SageMaker I 定价](https://aws.amazon.com/sagemaker/pricing/)。 如果您遇到类似以下内容的错误消息，请按照[请求增加 SageMaker AI 资源的服务配额中的说明进行](https://docs.aws.amazon.com/sagemaker/latest/dg/regions-quotas.html#service-limit-increase-request-procedure)操作。 ``` ResourceLimitExceeded: An error occurred (ResourceLimitExceeded) when calling the CreateTrainingJob operation: The account-level service limit 'ml.p3dn.24xlarge for training job usage' is 0 Instances, with current utilization of 0 Instances and a request delta of 1 Instances. Please contact AWS support to request an increase for this limit. ``` # SageMaker 模型并行度库的核心功能 Amazon SageMaker AI 的模型并行库提供分发策略和节省内存的技术，例如分片数据并行性、张量并行性、按层划分模型以进行管道调度以及检查点。模型并行性策略和技术有助于将大型模型分布在多个设备上，同时优化训练速度和内存使用。该库还提供 Python 帮助程序函数、上下文管理器和封装器函数，用于调整训练脚本以实现模型的自动或手动分区。 在训练作业中实现模型并行性时，将保持 “使用[模型并行度运行分布式 SageMaker 训练作业” 部分中显示的相同的两步工作](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-use-api.html)流程。对于调整训练脚本，您无需或者仅需在训练脚本中添加几行额外的代码。要启动调整后的训练脚本的训练作业，您需要设置分布配置参数以激活节省内存的功能或者传递用于并行度的值。 要开始使用示例，请参阅以下 Jupyter 笔记本，其中演示了如何使用 SageMaker 模型并行度库。 + [PyTorch 示例笔记本](https://github.com/aws/amazon-sagemaker-examples/tree/main/training/distributed_training/pytorch/model_parallel) + [TensorFlow 示例笔记本](https://github.com/aws/amazon-sagemaker-examples/tree/main/training/distributed_training/tensorflow/model_parallel/mnist) 要深入了解库的核心功能，请参阅以下主题。 **注意** SageMaker 分布式训练库可通过 Hugging Face 的 AWS PyTorch深度学习容器获得， TensorFlow 也可以在训练平台 SageMaker 内使用。要使用分布式训练库的功能，我们建议您使用 SageMaker Python SDK。如果你 SageMaker APIs 通过适用于 Python 的 SDK (Boto3) 或，也可以使用 JSON 请求语法手动配置。 AWS Command Line Interface在整篇文档中，说明和示例都侧重于如何将分布式训练库与 SageMaker Python SDK 配合使用。 **重要** SageMaker 模型并行度库支持的所有核心功能，并支持流水线并行性。 PyTorch TensorFlow **Topics** + [分片数据并行性](model-parallel-extended-features-pytorch-sharded-data-parallelism.md) + [模型管道传输](model-parallel-core-features-pipieline-parallelism.md) + [张量并行性](model-parallel-extended-features-pytorch-tensor-parallelism.md) + [优化器状态分片](model-parallel-extended-features-pytorch-optimizer-state-sharding.md) + [激活检查点](model-parallel-extended-features-pytorch-activation-checkpointing.md) + [激活分载](model-parallel-extended-features-pytorch-activation-offloading.md) + [FP16 使用模型并行度进行训练](model-parallel-extended-features-pytorch-fp16.md) + [Support FlashAttention](model-parallel-attention-head-size-for-flash-attention.md) # 分片数据并行性 分@@ *片数据并行性*是一种节省内存的分布式训练技术，它在数据并行组中拆分模型的状态（模型参数、梯度和优化器状态）。 GPUs **注意** 分片数据并行度可在 SageMaker 模型并行度库 v1.11.0 及更高版本 PyTorch 中使用。 将训练作业扩展到大型 GPU 集群时，您可以通过将模型的训练状态分为多个集群来减少模型的每 GPU 内存占用。 GPUs这会带来两个好处：您可以容纳在标准数据并行性下会耗尽内存的更大模型；或者您可以使用腾出的 GPU 内存来增加批次大小。 标准的数据并行度技术在数据并行组 GPUs中复制训练状态，并根据操作执行梯度聚合。`AllReduce`分片数据并行性修改了标准的数据并行分布式训练过程，以考虑优化器状态的分片性质。用于对模型和优化器状态进行分片的一组秩称为*分片组*。*分片数据并行技术对模型的可训练参数以及分片组中相应的梯度和优化器状态进行分片。 GPUs * SageMaker 人工智能通过实现 MIC 实现分片数据并行性， AWS 博客文章巨型模型训练的[近线性扩展](https://www.amazon.science/blog/near-linear-scaling-of-gigantic-model-training-on-aws)对此进行了讨论。 AWS在此实施中，您可以将分片度设置为可配置参数，该参数必须小于数据并行度。在每次向前和向后传递过程中，MIC 都会在 GPUs 整个操作过程中临时重新组合模型参数。`AllGather`在每层向前或向后传递后，MiCS 会再次对参数进行分片以节省 GPU 内存。在向后传递过程中，MIC 会降低梯度，同时在整个操作中 GPUs 将其分片。`ReduceScatter`最后，MiCS 使用优化器状态的局部分片，将局部缩减梯度和分片梯度应用于其对应的局部参数分片。为了降低通信开销， SageMaker 模型并行度库在向前或向后通道中预取即将到来的层，并将网络通信与计算重叠。 模型的训练状态在分片组之间复制。这意味着，在将梯度应用于参数之前，除了在分片组内进行的 `ReduceScatter` 操作之外，还必须跨分片组进行 `AllReduce` 操作。 实际上，分片数据并行性在通信开销和 GPU 内存效率之间进行了权衡。使用分片数据并行性会增加通信成本，但是在每个 GPU 上占用的内存量（不包括激活导致的内存使用量）会除以分片数据并行度，因此可以将更大的模型放入 GPU 集群。 **选择分片数据并行度** 当您为分片数据并行度选择一个值时，该值必须能够整除数据并行度。例如，对于 8 路数据并行性作业，请选择 2、4 或 8 作为分片数据并行度。在选择分片数据并行度时，我们建议您从一个较小的数字开始，然后逐渐增加它，直到模型与所需的批次大小均适合内存。 **选择批次大小** 设置分片数据并行性后，请确保您找到了可以在 GPU 集群上成功运行的最佳训练配置。要训练大型语言模型 (LLM)，请从批次大小 1 开始，然后逐渐增加批次大小，直到达到接收 out-of-memory (OOM) 错误的程度。如果您即使使用最小的批次也会遇到 OOM 错误，请应用更高的分片数据并行度，或者将分片数据并行性和张量并行性结合使用。 **Topics** + [如何将分片数据并行性应用于训练作业](#model-parallel-extended-features-pytorch-sharded-data-parallelism-how-to-use) + [参考配置](#model-parallel-extended-features-pytorch-sharded-data-parallelism-how-to-use-config-sample) + [使用 SMDDP 集合体的分片数据并行性](#model-parallel-extended-features-pytorch-sharded-data-parallelism-smddp-collectives) + [混合精度训练与分片数据并行性](#model-parallel-extended-features-pytorch-sharded-data-parallelism-16bits-training) + [使用张量并行性的分片数据并行性](#model-parallel-extended-features-pytorch-sharded-data-parallelism-with-tensor-parallelism) + [使用分片数据并行性的提示和注意事项](#model-parallel-extended-features-pytorch-sharded-data-parallelism-considerations) ## 如何将分片数据并行性应用于训练作业 要开始使用分片数据并行性，请对训练脚本进行必要的修改，然后使用参数设置 SageMaker PyTorch 估计器。 sharded-data-parallelism-specific另外还可以考虑以参考值和示例笔记本为起点。 ### 调整您的 PyTorch 训练脚本 按照[步骤 1：修改 PyTorch 训练脚本](model-parallel-customize-training-script-pt.md)中的说明使用和模块的`smdistributed.modelparallel.torch`包装器封装模型`torch.nn.parallel`和`torch.distributed`优化器对象。 **（可选）用于注册外部模型参数的额外修改** 如果您的模型使用 `torch.nn.Module` 构建并使用了未在模型类中定义的参数，您应手动将这些参数注册到模块，以便 SMP 收集完整参数。要将参数注册到模块，请使用 `smp.register_parameter(module, parameter)`。 ``` class Module(torch.nn.Module): def __init__(self, *args): super().__init__(self, *args) self.layer1 = Layer1() self.layer2 = Layer2() smp.register_parameter(self, self.layer1.weight) def forward(self, input): x = self.layer1(input) # self.layer1.weight is required by self.layer2.forward y = self.layer2(x, self.layer1.weight) return y ``` ### 设置 SageMaker PyTorch 估算器 在中配置 SageMaker PyTorch 估算器时[第 2 步：使用 SageMaker Python 软件开发工具包启动训练 Job](model-parallel-sm-sdk.md)，请添加分片数据并行度参数。 要开启分片数据并行性，请将`sharded_data_parallel_degree`参数添加到 Estimator 中。 SageMaker PyTorch此参数指定分片训练状态的数量。 GPUs `sharded_data_parallel_degree` 的值必须是介于 1 和数据并行度之间的整数，并且必须能够整除数据并行度。请注意，该库会自动检测数字 GPUs ，因此数据并行度。以下附加参数可用于配置分片数据并行度。 + `"sdp_reduce_bucket_size"`*（int，默认值：5e8）*— 以默认 dtype 的元素数量指定 [PyTorch DDP 渐变桶](https://pytorch.org/docs/stable/notes/ddp.html#internal-design)的大小。 + `"sdp_param_persistence_threshold"`*（整数，默认值：1e6）* – 以每个 GPU 上可以持续存在的元素数量来指定参数张量的大小。分片数据并行性将每个参数张量分成 GPUs 一个数据并行组。如果参数张量中的元素数量小于此阈值，则不会拆分参数张量；这有助于减少通信开销，因为参数张量是跨数据并行复制的。 GPUs + `"sdp_max_live_parameters"`*（整数，默认值：1e9）* – 指定在向前和向后传递期间，可以同时处于重新组合训练状态的最大参数数量。当活动参数的数量达到给定阈值时，提取参数的 `AllGather` 操作将暂停。请注意，增加此参数会增加内存占用。 + `"sdp_hierarchical_allgather"`*（布尔值，默认值：True）* – 如果设置为 `True`，则 `AllGather` 操作按层次运行：首先在每个节点内运行，然后跨节点运行。对于多节点分布式训练作业，分层 `AllGather` 操作会自动激活。 + `"sdp_gradient_clipping"`*（浮点数，默认值：1.0）* – 指定一个阈值，用于在通过模型参数向后传播梯度之前，先裁剪梯度的 L2 范数。当分片数据并行性被激活时，梯度裁剪也被激活。默认阈值为 `1.0`。如果您遇到梯度爆炸问题，请调整此参数。 以下代码演示了如何配置分片数据并行性的示例。 ``` import sagemaker from sagemaker.pytorch import PyTorch smp_options = { "enabled": True, "parameters": { # "pipeline_parallel_degree": 1, # Optional, default is 1 # "tensor_parallel_degree": 1, # Optional, default is 1 "ddp": True, # parameters for sharded data parallelism "sharded_data_parallel_degree": 2, # Add this to activate sharded data parallelism "sdp_reduce_bucket_size": int(5e8), # Optional "sdp_param_persistence_threshold": int(1e6), # Optional "sdp_max_live_parameters": int(1e9), # Optional "sdp_hierarchical_allgather": True, # Optional "sdp_gradient_clipping": 1.0 # Optional } } mpi_options = { "enabled" : True, # Required "processes_per_host" : 8 # Required } smp_estimator = PyTorch( entry_point="your_training_script.py", # Specify your train script role=sagemaker.get_execution_role(), instance_count=1, instance_type='ml.p3.16xlarge', framework_version='1.13.1', py_version='py3', distribution={ "smdistributed": {"modelparallel": smp_options}, "mpi": mpi_options }, base_job_name="sharded-data-parallel-job" ) smp_estimator.fit('s3://my_bucket/my_training_data/') ``` ## 参考配置 SageMaker 分布式训练团队提供了以下参考配置，您可以将其用作起点。您可以根据上述配置进行推断，以实验并估算模型配置的 GPU 内存使用量。 使用 SMDDP 集合体的分片数据并行性 | 模型/参数数量 | 实例数 | 实例类型 | 序列长度 | 全局批次大小 | 小批次大小 | 分片数据并行度 | | --- | --- | --- | --- | --- | --- | --- | | GPT-NEOX-20B | 2 | ml.p4d.24xlarge | 2048 | 64 | 4 | 16 | | GPT-NEOX-20B | 8 | ml.p4d.24xlarge | 2048 | 768 | 12 | 32 | 例如，如果您增加了 200 亿参数模型的序列长度或将模型的大小增加到 650 亿个参数，则您需要先尝试减小批次大小。如果在最小的批次大小（批次大小为 1）模型仍然不适合，请尝试提高模型的并行度。 使用张量并行性和 NCCL 集合体的分片数据并行性 | 模型/参数数量 | 实例数 | 实例类型 | 序列长度 | 全局批次大小 | 小批次大小 | 分片数据并行度 | 张量并行度 | 激活分载 | | --- | --- | --- | --- | --- | --- | --- | --- | --- | | GPT-NEOX-65B | 64 | ml.p4d.24xlarge | 2048 | 512 | 8 | 16 | 8 | Y | | GPT-NEOX-65B | 64 | ml.p4d.24xlarge | 4096 | 512 | 2 | 64 | 2 | Y | 当你想将大型语言模型 (LLM) 拟合到大规模集群中，同时使用序列长度更长的文本数据，从而使用较小的批处理大小，从而处理 GPU 内存使用量来针对更长的文本序列进行训练时，将分片数据并行性和张量并行性结合起来非常有用。 LLMs 要了解更多信息，请参阅[使用张量并行性的分片数据并行性](#model-parallel-extended-features-pytorch-sharded-data-parallelism-with-tensor-parallelism)。 有关案例研究、基准测试和更多配置示例，请参阅博客文章 [Amazon A SageMaker I 模型并行库中的新性能改进](https://aws.amazon.com/blogs/machine-learning/new-performance-improvements-in-amazon-sagemaker-model-parallel-library/)。 ## 使用 SMDDP 集合体的分片数据并行性 SageMaker 数据并行度库提供针对基础架构进行了优化的集体通信基元（SMDDP 集合）。 AWS 它通过使用 Ela [stic Fabric Adapter (EFA) Adapter (EFA)](https://aws.amazon.com/hpc/efa/) 来实现优化，从而产生高吞吐量且对延迟不敏感的集合，将与通信相关的处理分流给 CPU，并腾出 GPU 周期用于计算。 all-to-all-type在大型集群上，与 NCCL 相比，SMDDP 集合体可以将分布式训练性能提高多达40%。有关案例研究和基准测试结果，请参阅博客 [Amazon A SageMaker I 模型并行度库中的新性能改进](https://aws.amazon.com/blogs/machine-learning/new-performance-improvements-in-amazon-sagemaker-model-parallel-library/)。 **注意** SageMaker 模型并行度库 v1.13.0 及更高版本以及数据并行库 v1.6.0 及更高版本中提供了与 SMDDP Collectives 的分片数据并行性。 SageMaker 有关在 SMDDP 集合体中使用分片数据并行性的信息，另请参阅[Supported configurations](#sharded-data-parallelism-smddp-collectives-supported-config)。 在分片数据并行性（这是在大规模分布式训练中常用的技术）中，`AllGather` 集合体用于针对向前和向后传递计算重构分片层参数，与 GPU 计算并行处理。对于大型模型，高效执行 `AllGather` 操作对于避免 GPU 瓶颈问题以及避免降低训练速度至关重要。激活分片数据并行性后，SMDDP 集合体会放入这些对性能至关重要的 `AllGather` 集合体中，从而提高训练吞吐量。 **使用 SMDDP 集合体进行训练** 当您的训练作业已激活分片数据并行性并满足[Supported configurations](#sharded-data-parallelism-smddp-collectives-supported-config)时，SMDDP 集合体会自动激活。在内部，SMDDP Collectives 会优化`AllGather`集体，使其在 AWS 基础设施上保持高性能，而所有其他集体的性能则回归到 NCCL。此外，在不支持的配置下，所有集合体（包括 `AllGather`）都会自动使用 NCCL 后端。 从 SageMaker 模型并行度库版本 1.13.0 起，该`"ddp_dist_backend"`参数已添加到选项中。`modelparallel`此配置参数的默认值为 `"auto"`，这将尽可能使用 SMDDP 集合体，否则回退到 NCCL。要强制库始终使用 NCCL，请将 `"nccl"` 指定为 `"ddp_dist_backend"` 配置参数。 以下代码示例显示了如何使用分片数据并行度设置 PyTorch 估算器，`"ddp_dist_backend"`参数`"auto"`默认设置为，因此可以选择添加。 ``` import sagemaker from sagemaker.pytorch import PyTorch smp_options = { "enabled":True, "parameters": { "partitions": 1, "ddp": True, "sharded_data_parallel_degree": 64 "bf16": True, "ddp_dist_backend": "auto" # Specify "nccl" to force to use NCCL. } } mpi_options = { "enabled" : True, # Required "processes_per_host" : 8 # Required } smd_mp_estimator = PyTorch( entry_point="your_training_script.py", # Specify your train script source_dir="location_to_your_script", role=sagemaker.get_execution_role(), instance_count=8, instance_type='ml.p4d.24xlarge', framework_version='1.13.1', py_version='py3', distribution={ "smdistributed": {"modelparallel": smp_options}, "mpi": mpi_options }, base_job_name="sharded-data-parallel-demo", ) smd_mp_estimator.fit('s3://my_bucket/my_training_data/') ``` **支持的配置** 满足以下所有配置要求后，训练作业中将激活 `AllGather` 操作和 SMDDP 集合体。 + 分片数据并行度大于 1 + `Instance_count` 大于 1 + `Instance_type` 等于 `ml.p4d.24xlarge` + SageMaker 适用于 PyTorch v1.12.1 或更高版本的训练容器 + SageMaker 数据并行度库 v1.6.0 或更高版本 + SageMaker 模型并行度库 v1.13.0 或更高版本 **性能和内存调整** SMDDP 集合体使用额外的 GPU 内存。根据不同的模型训练使用场景，有两个环境变量可用于配置 GPU 内存使用情况。 + `SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` – 在 SMDDP `AllGather` 操作期间，`AllGather` 输入缓冲区被复制到临时缓冲区中，用于节点间通信。`SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` 变量控制此临时缓冲区的大小（以字节为单位）。如果临时缓冲区的大小小于 `AllGather` 输入缓冲区的大小，则 `AllGather` 集合体将回退到使用 NCCL。 + 默认值：16 \$1 1024 \$1 1024 (16 MB) + 可接受值：8192 的任意倍数 + `SMDDP_AG_SORT_BUFFER_SIZE_BYTES` – `SMDDP_AG_SORT_BUFFER_SIZE_BYTES` 变量用于调整临时缓冲区的大小（以字节为单位），以保存从节点间通信中收集的数据。如果临时缓冲区的大小小于 `1/8 * sharded_data_parallel_degree * AllGather input size`，则 `AllGather` 集合体将回退到使用 NCCL。 + 默认值：128 \$1 1024 \$1 1024 (128 MB) + 可接受值：8192 的任意倍数 **缓冲区大小变量调整指南** 环境变量的默认值应该适用于大多数使用场景。我们建议只有在训练遇到 out-of-memory (OOM) 错误时才调整这些变量。 以下列表讨论了一些调整技巧，可以减少 SMDDP 集合体的 GPU 内存占用，同时继续获得带来的性能提升。 + 调整 `SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` + 对于较小的模型，`AllGather` 输入缓冲区的大小会更小。因此，对于参数较少的模型，所需的 `SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` 大小可能会更小。 + `AllGather`输入缓冲区的大小会随着`sharded_data_parallel_degree`增加而减小，因为模型会被分片得更多 GPUs。因此，对于 `sharded_data_parallel_degree` 具有较大值的训练作业，所需的 `SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` 大小可能会更小。 + 调整 `SMDDP_AG_SORT_BUFFER_SIZE_BYTES` + 对于参数较少的模型，从节点间通信中收集的数据量较少。因此，对于参数数量较少的此类模型，所需的 `SMDDP_AG_SORT_BUFFER_SIZE_BYTES` 大小可能会更小。 有些集合体可能会回退为使用 NCCL；因此，您可能无法获得优化 SMDDP 集合体所带来的性能提升。如果有额外的 GPU 内存可供使用，您可以考虑增加 `SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES` 和 `SMDDP_AG_SORT_BUFFER_SIZE_BYTES` 的值以从性能提升中受益。 以下代码显示了如何通过将环境变量附加到 PyTorch 估计器的分布参数`mpi_options`中来配置环境变量。 ``` import sagemaker from sagemaker.pytorch import PyTorch smp_options = { .... # All modelparallel configuration options go here } mpi_options = { "enabled" : True, # Required "processes_per_host" : 8 # Required } # Use the following two lines to tune values of the environment variables for buffer mpioptions += " -x SMDDP_AG_SCRATCH_BUFFER_SIZE_BYTES=8192" mpioptions += " -x SMDDP_AG_SORT_BUFFER_SIZE_BYTES=8192" smd_mp_estimator = PyTorch( entry_point="your_training_script.py", # Specify your train script source_dir="location_to_your_script", role=sagemaker.get_execution_role(), instance_count=8, instance_type='ml.p4d.24xlarge', framework_version='1.13.1', py_version='py3', distribution={ "smdistributed": {"modelparallel": smp_options}, "mpi": mpi_options }, base_job_name="sharded-data-parallel-demo-with-tuning", ) smd_mp_estimator.fit('s3://my_bucket/my_training_data/') ``` ## 混合精度训练与分片数据并行性 要通过半精度浮点数和分片数据并行度进一步节省 GPU 内存，您可以通过向分布式训练配置中添加一个附加参数来激活 16 位[浮点格式 (FP16BF16) 或 Brain 浮点格式](https://en.wikichip.org/wiki/brain_floating-point_format) ()。 **注意** SageMaker 模型并行度库 v1.11.0 及更高版本中提供了具有分片数据并行性的混合精度训练。 **用于使用分片数据并行度进行 FP16 训练** 要使用分片数据并行度运行 FP16 训练，请`"fp16": True"`添加到配置字典中。`smp_options`在训练脚本中，您可以通过 `smp.DistributedOptimizer` 模块在静态和动态损失缩放选项之间进行选择。有关更多信息，请参阅 [FP16 使用模型并行度进行训练](model-parallel-extended-features-pytorch-fp16.md)。 ``` smp_options = { "enabled": True, "parameters": { "ddp": True, "sharded_data_parallel_degree": 2, "fp16": True } } ``` **用于使用分片数据并行度进行 BF16 训练** SageMaker AI 的分片数据并行功能支持数据类型训练。 BF16 BF16 数据类型使用 8 位来表示浮点数的指数，而 FP16 数据类型使用 5 位。保留指数的 8 位可以保持 32 位单精度浮点 () FP32 数的指数的相同表示形式。这使得 FP32 和之间的转换变得 BF16更简单，并且不太容易导致 FP16 训练中经常出现的溢出和下溢问题，尤其是在训练较大的模型时。虽然这两种数据类型总共使用 16 位，但 BF16 格式中指数表示范围的增加是以降低精度为代价的。对于训练大型模型，这种精度缩减通常会视为面向范围和训练稳定性作出的可接受取舍。 **注意** 当前，只有激活分片数据并行性后， BF16 训练才有效。 要使用分片数据并行度运行 BF16 训练，请`"bf16": True`添加到配置字典中。`smp_options` ``` smp_options = { "enabled": True, "parameters": { "ddp": True, "sharded_data_parallel_degree": 2, "bf16": True } } ``` ## 使用张量并行性的分片数据并行性 如果您使用分片数据并行性并且还需要减小全局批次大小，请考虑将[张量并行性](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-tensor-parallelism.html)与分片数据并行性结合使用。在非常大的计算集群（通常包含 128 个节点或更多）上训练采用分片数据并行性的大型模型时，即使每个 GPU 的批次大小很小，也会得到非常大的全局批次大小。这可能会导致收敛问题或计算性能低下问题。当单个批次已经很大并且无法进一步缩小时，有时候仅使用分片数据并行性无法减少每个 GPU 上的批次大小。在这种情况下，将分片数据并行性与张量并行性结合使用，有助于减少全局批次大小。 最佳分片数据并行度和张量并行度的选择，取决于模型的规模、实例类型以及模型能够合理收敛的全局批次大小。我们建议您从低张量 parallel 度开始，以使全局批量大小适合计算集群，从而解决 CUDA out-of-memory 错误并获得最佳性能。请参阅以下两个示例案例，了解张量并行性和分片数据并行性的组合如何帮助您通过对模型并行性进行分组 GPUs 来调整全局批次大小，从而减少模型副本的数量和更小的全局批次大小。 **注意** 此功能可从 SageMaker 模型并行度库 v1.15 中获得，并支持 v1.13.1。 PyTorch **注意** 此功能可通过库的张量并行度功能，面向支持的模型提供。要查找支持的模型列表，请参阅[对 Hugging Face 转换器模型的支持](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-hugging-face.html)。另请注意，在修改训练脚本时，您需要将 `tensor_parallelism=True` 传递给 `smp.model_creation` 参数。要了解更多信息，请参阅 *SageMaker AI 示例 GitHub 存储库[https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/train_gpt_simple.py#L793](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/train_gpt_simple.py#L793)*中的训练脚本。 ### 示例 1 假设我们要 GPUs 在 1536 个集群 GPUs （192 个节点，每个节点 8 个）上训练模型，将分片数据并行度设置为 32 (`sharded_data_parallel_degree=32`)，将每个 GPU 的批处理大小设置为 1，其中每个批次的序列长度为 4096 个令牌。在本例中有 1536 个模型副本，全局批次大小变为 1536，每个全局批次包含大约 600 万个令牌。 ``` (1536 GPUs) * (1 batch per GPU) = (1536 global batches) (1536 batches) * (4096 tokens per batch) = (6,291,456 tokens) ``` 向其添加张量并行性，使其可以减少全局批次大小。在一种配置示例中，可以将张量并行度设置为 8，将每个 GPU 的批次大小设置为 4。这形成了 192 个张量并行组或 192 个模型副本，其中每个模型副本分布在 8 个副本上。 GPUs批次大小为 4 是每次迭代中每个张量并行组的训练数据量；也就是说，每个模型副本每次迭代消耗 4 个批次。在这种情况下，全局批次大小变为 768，每个全局批次包含大约 300 万个令牌。这样，与之前仅使用分片数据并行性的情况相比，全局批次大小减少了一半。 ``` (1536 GPUs) / (8 tensor parallel degree) = (192 tensor parallelism groups) (192 tensor parallelism groups) * (4 batches per tensor parallelism group) = (768 global batches) (768 batches) * (4096 tokens per batch) = (3,145,728 tokens) ``` ### 示例 2 当分片数据并行性和张量并行性均已激活时，库会首先应用张量并行性，并在该维度上对模型分片。对于每个张量并行秩，数据并行性根据 `sharded_data_parallel_degree` 来应用。 例如，假设我们要设置 32 GPUs ，张量 parallel 度为 4（形成一组 4 GPUs），分片数据并行度为 4，最终复制度为 2。该分配基于如下所示张量并行度创建八个 GPU 组：`(0,1,2,3)`、`(4,5,6,7)`、`(8,9,10,11)`、`(12,13,14,15)`、`(16,17,18,19)`、`(20,21,22,23)`、`(24,25,26,27)`、`(28,29,30,31)`。也就是说，四 GPUs 形成一个张量 parallel 群。在这种情况下，张量并行组第 0 个等级 GPUs 的简化数据并行组将是。`(0,4,8,12,16,20,24,28)`简化的数据并行组根据分片数据并行度 4 进行分片，从而产生两个用于数据并行性的复制组。 GPUs`(0,4,8,12)`形成一个分片组，该分片组共同保存第 0 张量 parallel 等级的所有参数的完整副本， GPUs`(16,20,24,28)`然后形成另一个这样的分片组。其他张量并行秩也有类似的分片和复制组。 ![\[图 1：张量并行组。\]](http://docs.aws.amazon.com/zh_cn/sagemaker/latest/dg/images/distributed/model-parallel/sdp_tp_group_tp.jpg) 图 1：（节点、分片数据并行度、张量并行度）=（4, 4, 4）的张量并行组，其中每个矩形代表一个指数从 0 到 31 的 GPU。从 TPG 到 T 0 PG 的 GPUs 形式张量并行度分组。7复制组为（\$1TPG0, TPG4\$1、\$1TPG1, TPG5\$1、\$1TPG2, TPG6\$1 和 \$1TPG3, TPG7\$1）；每个复制组对的颜色相同，但填充方式不同。 ![\[图 2：分片数据并行组。\]](http://docs.aws.amazon.com/zh_cn/sagemaker/latest/dg/images/distributed/model-parallel/sdp_tp_group_sdp.jpg) 图 2：（节点、分片数据并行度、张量并行度）=（4, 4, 4）的分片数据并行组，其中每个矩形代表一个指数从 0 到 31 的 GPU。该 GPUs 表单将从 SDPG 到 SDPG 的数据并行分组划分0。7复制组为（\$1SDPG 0, SDPG 4\$1、\$1SDPG 1, SDPG 5\$1、\$1SDPG 2, SDPG 6\$1 和 \$1SDPG 3, SDPG 7\$1）；每个复制组对的颜色相同，但填充方式不同。 ### 如何激活使用张量并行性的分片数据并行性 要将分片数据并行性与张量并行性结合使用，您需要在创建估计器类的对象时`tensor_parallel_degree`在配置中`distribution`同时设置`sharded_data_parallel_degree`和。 SageMaker PyTorch 您还需要激活 `prescaled_batch`。这意味着，不是每个 GPU 读取自己的批次数据，而是每个张量并行组集体读取具有所选批次大小的一个组合批次。实际上，它不是将数据集分成等于数量 GPUs （或数据并行大小`smp.dp_size()`）的部分，而是将其分成等于 GPUs 除以数的部分`tensor_parallel_degree`（也称为缩小后的数据并行大小，`smp.rdp_size()`）。有关预缩放批处理的更多详细信息，请参阅 *SageMaker Python* SDK 文档中的 Pres [caled Bat](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#prescaled-batch) ch。另请参阅 *SageMaker AI 示例 GitHub 存储库中的 GPT-2 训练脚本[https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/train_gpt_simple.py#L164](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/train_gpt_simple.py#L164)示例*。 以下代码片段显示了基于中上述场景创建 PyTorch 估算器对象的示例。[示例 2](#model-parallel-extended-features-pytorch-sharded-data-parallelism-with-tensor-parallelism-ex2) ``` mpi_options = "-verbose --mca orte_base_help_aggregate 0 " smp_parameters = { "ddp": True, "fp16": True, "prescaled_batch": True, "sharded_data_parallel_degree": 4, "tensor_parallel_degree": 4 } pytorch_estimator = PyTorch( entry_point="your_training_script.py", role=role, instance_type="ml.p4d.24xlarge", volume_size=200, instance_count=4, sagemaker_session=sagemaker_session, py_version="py3", framework_version="1.13.1", distribution={ "smdistributed": { "modelparallel": { "enabled": True, "parameters": smp_parameters, } }, "mpi": { "enabled": True, "processes_per_host": 8, "custom_mpi_options": mpi_options, }, }, source_dir="source_directory_of_your_code", output_path=s3_output_location ) ``` ## 使用分片数据并行性的提示和注意事项 在使用 SageMaker 模型并行度库的分片数据并行度时，请考虑以下几点。 + 分片数据并行性与训练兼容。 FP16 要进行 FP16训练，请参阅一[FP16 使用模型并行度进行训练](model-parallel-extended-features-pytorch-fp16.md)节。 + 分片数据并行性与张量并行性兼容。将分片数据并行性与张量并行性结合使用时，您可能需要考虑以下几点。 + 将分片数据并行性与张量并行性结合使用时，嵌入层也会自动分布在张量并行组中。换而言之，`distribute_embedding` 参数会自动设置为 `True`。有关张量并行性的更多信息，请参阅[张量并行性](model-parallel-extended-features-pytorch-tensor-parallelism.md)。 + 请注意，具有张量并行性的分片数据并行性目前使用 NCCL 集合体作为分布式训练策略的后端。 要了解更多信息，请参阅[使用张量并行性的分片数据并行性](#model-parallel-extended-features-pytorch-sharded-data-parallelism-with-tensor-parallelism)一节。 + 分片数据并行性目前不兼容[管道并行性](model-parallel-intro.md#model-parallel-intro-pp)或[优化器状态分片](model-parallel-extended-features-pytorch-optimizer-state-sharding.md)。要激活分片数据并行性，请关闭优化器状态分片，并将管道并行度设置为 1。 + [激活检查点](model-parallel-extended-features-pytorch-activation-checkpointing.md)和[激活分载](model-parallel-extended-features-pytorch-activation-offloading.md)功能与分片数据并行性兼容。 + 要将分片数据并行性与梯度累积一起使用，请在使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.DistributedModel](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.DistributedModel) 模块包装模型时，将 `backward_passes_per_step` 参数设置为累积步骤数。这可确保跨模型复制组（分片组）的梯度 `AllReduce` 操作发生在梯度累积的边界上。 + 您可以使用库的检查点检查使用分片数据并行度训练的模型，以及。 APIs `smp.save_checkpoint` `smp.resume_from_checkpoint`有关更多信息，请参阅 [对分布式 PyTorch 模型执行检查点操作（适用于 SageMaker 模型并行度库 v1.10.0 及更高版本）](distributed-model-parallel-checkpointing-and-finetuning.md#model-parallel-extended-features-pytorch-checkpoint)。 + 在分片数据并行性下，[https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.delay_param_initialization](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.delay_param_initialization) 配置参数的行为会发生变化。同时启用这两个功能时，在分片方式下，将在模型创建时立即初始化参数，而不是延迟参数初始化，这样每个秩都会初始化并存储自己的参数分片。 + 当分片数据并行性被激活时，该库会在 `optimizer.step()` 调用运行时在内部执行梯度裁剪。您无需使用实用工具 APIs 进行渐变剪裁，例如[https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html](https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html)。要调整梯度裁剪的阈值，可以在构造 SageMaker PyTorch 估计器时通过分布参数配置的参数对其进行设置，如部分所示。`sdp_gradient_clipping` [如何将分片数据并行性应用于训练作业](#model-parallel-extended-features-pytorch-sharded-data-parallelism-how-to-use) # 模型管道传输 模型并行度库 SageMaker的核心功能之一是*管道并行性*，它决定了模型训练期间进行计算的顺序和跨设备处理数据的顺序。Pipelining 是一种通过在不同的数据样本上同时进行计算，从而在模型并行性中实现真正的并行化，并克服顺序 GPUs 计算导致的性能损失的技术。使用管道并行性时，训练作业以管道方式按照微批次执行，以最大限度地提高 GPU 使用率。 **注意** 管道并行度（也称为模型分区）可用于和。 PyTorch TensorFlow有关支持的框架版本，请参阅[支持的框架和 AWS 区域](distributed-model-parallel-support.md)。 ## 管道执行计划 Pipelining 基于将迷你批次拆分为微批，这些微批量被输入到训练管道中， one-by-one并遵循库运行时定义的执行时间表。*微批次*是给定训练小批次的更小的子集。管道计划决定了每个时段由哪个设备执行哪个微批次。 例如，根据流水线计划和模型分区，GPU `i` 可能会在微批量上执行（向前或向后）计算，`b`而 GPU 则在微批处理上`i+1`执行计算`b+1`，从而使两者同时 GPUs 处于活动状态。在单次向前或向后传递期间，根据分区决策，单个微批次的执行流程可能会多次访问同一个设备。例如，位于模型开头的操作，与位于模型末尾的操作可能会放在同一设备上，而两者之间的操作位于不同的设备上，这意味着此设备被访问了两次。 该库提供两种不同的流水线计划，*简单*调度和*交错排*程，可以使用 Pyth SageMaker on SDK 中的`pipeline`参数进行配置。在大多数情况下，交错流水线可以通过 GPUs 更有效地利用流水线来实现更好的性能。 ### 交错管道 在交错管道中，尽可能优先考虑向后执行微批次。这样可以更快地释放内存用于激活，从而更有效地使用内存。它还允许扩大微粒的数量，从而缩短微粒的闲置时间。 GPUs在稳态下，每个设备在向前传递和向后传递之间交替运行。这意味着一个微批次的向后传递可能会在另一个微批次的向前传递完成之前运行。 ![\[超过 2 的交错管道的执行计划示例。 GPUs\]](http://docs.aws.amazon.com/zh_cn/sagemaker/latest/dg/images/distributed/model-parallel/interleaved-pipeline-execution.png) 上图说明了超过 2 的交错管道的执行计划示例。 GPUs在图中，F0 表示微批次 0 的向前传递，B1 表示微批次 1 的向后传递。**更新**表示优化器对参数的更新。 GPU0 始终尽可能优先考虑向后传递（例如，在 F2 之前执行 B0），这样可以清除用于更早激活的内存。 ### 简单管道 相比之下，简单管道则为每个微批次运行向前传递，然后再开始向后传递。这意味着它仅对自身中的向前传递和向后传递阶段进行管道传输。下图说明了其工作原理的示例，超过 2 GPUs. ![\[在管道上的示例，每个微批次都先运行前向传递，然后再开始后向传递。\]](http://docs.aws.amazon.com/zh_cn/sagemaker/latest/dg/images/distributed/model-parallel/simple-pipeline-execution.png) ### 特定框架中的管道执行 使用以下章节来了解特定于框架的管道调度决策 SageMaker的模型并行度库对和的影响。 TensorFlow PyTorch #### 使用管道执行 TensorFlow 下图是使用自动模型拆分由模型并行度库划分的 TensorFlow 图形的示例。拆分图形时，生成的每个子图形都会被复制 B 次（变量除外），其中 B 是微批次的数量。在此图中，每个子图形被复制 2 次 (B=2)。在子图形的每个输入中插入一个 `SMPInput` 运算，在每个输出中插入一个 `SMPOutput` 运算。这些操作与库后端通信，以在彼此之间传输张量。 ![\[使用自动模型拆分由模型并行度库划分的 TensorFlow 图形示例。\]](http://docs.aws.amazon.com/zh_cn/sagemaker/latest/dg/images/distributed/model-parallel/interleaved-pipeline-tf.png) 下图是 2 个子图形的示例，以 B=2 进行拆分，并添加了梯度运算。`SMPInput` 运算的梯度是 `SMPOutput` 运算，反之亦然。这使梯度能够在反向传播期间向后流动。 ![\[2 个子图形的示例，以 B=2 进行拆分，并添加了梯度运算。\]](http://docs.aws.amazon.com/zh_cn/sagemaker/latest/dg/images/distributed/model-parallel/interleaved-pipeline-tf.gif) 此 GIF 演示了一个示例交错管道执行计划，其中 B=2 个微批次和 2 个子图形。每个设备按顺序执行其中一个子图形副本，以提高 GPU 利用率。随着 B 增大，空闲时段的比例趋近于零。每当需要对特定的子图形副本执行（向前或向后）计算时，管道层都会向相应的蓝色 `SMPInput` 运算发出信号以开始执行。 计算出单个小批次中所有微批次的梯度后，该库就会合并各个库的梯度，然后将其应用于参数。 #### 使用管道执行 PyTorch 从概念上讲，流水线也遵循类似的想法。 PyTorch但是，由于 PyTorch 不涉及静态图，因此模型并行度库的 PyTorch 功能使用了更动态的流水线模式。 例如 TensorFlow，每个批次都被分成多个微粒，这些微批次在每个设备上一次执行一个。但是，执行计划由在每个设备上启动的执行服务器来处理。只要当前设备需要放置在另一台设备上的子模块的输出，就会向远程设备的执行服务器发送执行请求，同时将输入张量发送到该子模块。然后，服务器使用给定的输入执行此模块，并将响应返回给当前设备。 由于当前设备在远程子模块执行期间处于空闲状态，当前微批处理的本地执行暂停，库运行时系统会将执行切换到当前设备可以主动处理的另一个微批次。微批次的优先级由所选的管道计划决定。对于交错管道计划，只要有可能，就会优先考虑处于计算后退阶段的微批次。 # 张量并行性 *张量并行性*是模型并行性的一种，它在设备之间拆分特定的模型权重、梯度和优化器状态。管道并行性保持单个权重不变但对权重*集*进行拆分，张量并行性则与之相反，会拆分单个权重。这通常涉及对模型的特定运算、模块或层进行分布式计算。 在单个参数占用大部分 GPU 内存的情况下（例如词汇表很大的大型嵌入表，或者具有大量类的大型 softmax 层），则需要张量并行性。在这种情况下，将这种大张量或运算视为原子单元的效率会很低，并且会阻碍内存负载的平衡。 在仅仅依靠管道并不足以满足要求的超大型模型中，张量并行性也很有用。例如，在 GPT-3 规模的模型中需要对数十个实例进行分区，仅使用微批次管道传输效率低下，因为管道深度会变得过深，开销会变得过大。 **注意** 张量并行度可在 SageMaker 模型并行度库 v1.6.0 及更高版本 PyTorch 中使用。 **Topics** + [张量并行性的工作原理](model-parallel-extended-features-pytorch-tensor-parallelism-how-it-works.md) + [使用张量并行度运行 SageMaker 分布式模型并行训练 Job](model-parallel-extended-features-pytorch-tensor-parallelism-examples.md) + [支持 Hugging Face 转换器模型](model-parallel-extended-features-pytorch-hugging-face.md) + [将管道并行性与张量并行性结合使用时的秩评定机制](model-parallel-extended-features-pytorch-ranking-mechanism.md) # 张量并行性的工作原理 张量并行性在级别 `nn.Modules` 上实现；它在张量并行秩之间对模型中的特定模块分区。这是对管道并行性中使用的*模块集*的现有分区的补充。 通过张量并行性对模块进行分区时，其向前和向后传播是分布式的。库处理设备间的必要通信，以实施这些模块的分布式执行。这些模块在多个数据并行秩之间进行分区。与传统的工作负载分布相反，在使用库的张量并行性时，每个数据并行秩并**没有**完整的模型副本。相反，在所有未分布而保持完整的模块之外，每个数据并行秩可能只有分布式模块的一个分区。 **示例：**考虑跨数据并行秩的张量并行性，其中数据并行度为 4，张量并行度为 2。假设在对模块集进行分区后，您有一个存有以下模块树的数据并行组。 ``` A ├── B | ├── E | ├── F ├── C └── D ├── G └── H ``` 假设模块 B、G 和 H 支持张量并行性。此模型的张量并行分区的一种可能结果可能是： ``` dp_rank 0 (tensor parallel rank 0): A, B:0, C, D, G:0, H dp_rank 1 (tensor parallel rank 1): A, B:1, C, D, G:1, H dp_rank 2 (tensor parallel rank 0): A, B:0, C, D, G:0, H dp_rank 3 (tensor parallel rank 1): A, B:1, C, D, G:1, H ``` 每行代表存储在该 `dp_rank` 中的一组模块，标记 `X:y` 表示模块 `X` 的第 `y` 个部分。注意以下几点： 1. 分区跨数据并行秩的子集进行，我们调用 `TP_GROUP`，而不是整个 `DP_GROUP`，因此精确的模型分区会跨 `dp_rank` 0 和 `dp_rank` 2 复制，类似地跨 `dp_rank` 1 和 `dp_rank` 3 复制。 1. 模块 `E` 和 `F` 不再是模型的一部分，因为它们的父模块 `B` 已分区，并且通常属于 `E` 和 `F` 的任何执行在（已分区）`B` 模块内进行。 1. 尽管 `H` 支持张量并行性，但在此示例中，它没有进行分区，这强调了是否对模块进行分区取决于用户输入。模块支持张量并行性并不一定意味着它需要分区。 ## 库如何调整张量并行度以适应模块 PyTorch `nn.Linear` 对数据并行秩执行张量并行性时，*对于所分区的模块*，参数、梯度和优化器状态的子集在张量并行设备上分区。对于其余模块，张量并行设备以常规的数据并行方式操作。为了执行分区模块，设备首先在同一个张量并行度组中，跨对等设备收集*所有数据样本*的必要部分。然后，设备在所有这些数据样本上运行模块的本地部分，然后进行另一轮同步，该同步既合并了每个数据样本的输出部分，又将组合的数据样本返回到数据样本最初来源的样本。 GPUs 下图显示了在已分区 `nn.Linear` 模块上执行此过程的示例。 ![\[两张图显示了两个张量并行概念。\]](http://docs.aws.amazon.com/zh_cn/sagemaker/latest/dg/images/distributed/model-parallel/tensor-parallel-concept.png) 第一张图显示了一个带有 `nn.Linear` 模块的小模型，带有在两个张量并行性秩上的数据并行性。`nn.Linear` 模块复制到两个并行秩中。 第二张图显示了在拆分 `nn.Linear` 模块时应用于较大模型的张量并行性。每个 `tp_rank` 存有一半的线性模块，以及剩余的全部操作。当线性模块运行时，每个 `tp_rank` 都会收集所有数据样本中相关的一半，并传递到自己的一半 `nn.Linear` 模块中。结果必须是减少分散的（使用求和作为归约运算），这样每个秩都有自己数据样本的最终线性输出。模型的其余部分以典型的数据并行方式运行。 # 使用张量并行度运行 SageMaker 分布式模型并行训练 Job 在本部分中，您将学习： + 如何配置 SageMaker PyTorch 估计器和 SageMaker 模型并行度选项以使用张量并行度。 + 如何使用扩展的 `smdistributed.modelparallel` 模块来调整训练脚本以用于张量并行性。 要了解有关这些`smdistributed.modelparallel`模块的更多信息，请参阅 *SageMaker Python SDK 文档 APIs*中的[并行SageMaker 模型](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html)。 **Topics** + [仅使用张量并行性](#model-parallel-extended-features-pytorch-tensor-parallelism-alone) + [张量并行性与管道并行性相结合](#model-parallel-extended-features-pytorch-tensor-and-pipeline-parallelism) ## 仅使用张量并行性 以下分布式训练选项示例单独激活张量并行性，不使用管道并行性。配置`mpi_options`和`smp_options`字典以为 SageMaker `PyTorch`估计器指定分布式训练选项。 **注意** 扩展的内存节省功能可通过 Deep Learning Containers for 获得 PyTorch，该容器实现了 SageMaker 模型并行度库 v1.6.0 或更高版本。 **配置 SageMaker PyTorch 估算器** ``` mpi_options = { "enabled" : True, "processes_per_host" : 8,               # 8 processes "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none " }                 smp_options = { "enabled":True, "parameters": { "pipeline_parallel_degree": 1,    # alias for "partitions" "placement_strategy": "cluster", "tensor_parallel_degree": 4,      # tp over 4 devices "ddp": True } }                smp_estimator = PyTorch(    entry_point='your_training_script.py', # Specify    role=role,    instance_type='ml.p3.16xlarge',    sagemaker_session=sagemaker_session,    framework_version='1.13.1', py_version='py36',    instance_count=1,    distribution={        "smdistributed": {"modelparallel": smp_options},        "mpi": mpi_options    },    base_job_name="SMD-MP-demo", ) smp_estimator.fit('s3://my_bucket/my_training_data/') ``` **提示** 要查找的完整参数列表`distribution`，请参阅 Pyth SageMaker on SDK 文档中的[模型并行度配置参数](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html)。 **调整您的 PyTorch 训练脚本** 以下示例训练脚本展示了如何根据训练脚本调整 SageMaker 模型并行度库。在此示例中，假设脚本命名为 `your_training_script.py`。 ``` import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchnet.dataset import SplitDataset from torchvision import datasets import smdistributed.modelparallel.torch as smp class Net(nn.Module):     def __init__(self):         super(Net, self).__init__()         self.conv1 = nn.Conv2d(1, 32, 3, 1)         self.conv2 = nn.Conv2d(32, 64, 3, 1)         self.fc1 = nn.Linear(9216, 128)         self.fc2 = nn.Linear(128, 10)     def forward(self, x):         x = self.conv1(x)         x = F.relu(x)         x = self.conv2(x)         x = F.relu(x)         x = F.max_pool2d(x, 2)         x = torch.flatten(x, 1)         x = self.fc1(x)         x = F.relu(x)         x = self.fc2(x)         return F.log_softmax(x, 1) def train(model, device, train_loader, optimizer):     model.train()     for batch_idx, (data, target) in enumerate(train_loader):         # smdistributed: Move input tensors to the GPU ID used by         # the current process, based on the set_device call.         data, target = data.to(device), target.to(device)         optimizer.zero_grad()         output = model(data)         loss = F.nll_loss(output, target, reduction="mean")         loss.backward()         optimizer.step() # smdistributed: Initialize the backend smp.init() # smdistributed: Set the device to the GPU ID used by the current process. # Input tensors should be transferred to this device. torch.cuda.set_device(smp.local_rank()) device = torch.device("cuda") # smdistributed: Download only on a single process per instance. # When this is not present, the file is corrupted by multiple processes trying # to download and extract at the same time if smp.local_rank() == 0:     dataset = datasets.MNIST("../data", train=True, download=False) smp.barrier() # smdistributed: Shard the dataset based on data parallel ranks if smp.dp_size() > 1:     partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())}     dataset = SplitDataset(dataset, partitions=partitions_dict)     dataset.select(f"{smp.dp_rank()}") train_loader = torch.utils.data.DataLoader(dataset, batch_size=64) # smdistributed: Enable tensor parallelism for all supported modules in the model # i.e., nn.Linear in this case. Alternatively, we can use # smp.set_tensor_parallelism(model.fc1, True) # to enable it only for model.fc1 with smp.tensor_parallelism():     model = Net() # smdistributed: Use the DistributedModel wrapper to distribute the # modules for which tensor parallelism is enabled model = smp.DistributedModel(model) optimizer = optim.AdaDelta(model.parameters(), lr=4.0) optimizer = smp.DistributedOptimizer(optimizer) train(model, device, train_loader, optimizer) ``` ## 张量并行性与管道并行性相结合 以下是一个分布式训练选项的示例，该选项支持张量并行性与流水线并行性相结合。 设置`mpi_options`和`smp_options`参数，以便在配置估计器时使用张量并行度指定模型并行度选项。 SageMaker `PyTorch` **注意** 扩展的内存节省功能可通过 Deep Learning Containers for 获得 PyTorch，该容器实现了 SageMaker 模型并行度库 v1.6.0 或更高版本。 **配置 SageMaker PyTorch 估算器** ``` mpi_options = { "enabled" : True, "processes_per_host" : 8,               # 8 processes "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none " }                 smp_options = { "enabled":True, "parameters": { "microbatches": 4, "pipeline_parallel_degree": 2,    # alias for "partitions" "placement_strategy": "cluster", "tensor_parallel_degree": 2,      # tp over 2 devices "ddp": True } }                smp_estimator = PyTorch(    entry_point='your_training_script.py', # Specify    role=role,    instance_type='ml.p3.16xlarge',    sagemaker_session=sagemaker_session,    framework_version='1.13.1', py_version='py36',    instance_count=1,    distribution={        "smdistributed": {"modelparallel": smp_options},        "mpi": mpi_options    },    base_job_name="SMD-MP-demo", ) smp_estimator.fit('s3://my_bucket/my_training_data/')   ``` **调整您的 PyTorch 训练脚本** 以下示例训练脚本展示了如何根据训练脚本调整 SageMaker 模型并行度库。请注意，训练脚本现在包括 `smp.step` 修饰器： ``` import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchnet.dataset import SplitDataset from torchvision import datasets import smdistributed.modelparallel.torch as smp class Net(nn.Module):     def __init__(self):         super(Net, self).__init__()         self.conv1 = nn.Conv2d(1, 32, 3, 1)         self.conv2 = nn.Conv2d(32, 64, 3, 1)         self.fc1 = nn.Linear(9216, 128)         self.fc2 = nn.Linear(128, 10)     def forward(self, x):         x = self.conv1(x)         x = F.relu(x)         x = self.conv2(x)         x = F.relu(x)         x = F.max_pool2d(x, 2)         x = torch.flatten(x, 1)         x = self.fc1(x)         x = F.relu(x)         x = self.fc2(x)         return F.log_softmax(x, 1) # smdistributed: Define smp.step. Return any tensors needed outside. @smp.step def train_step(model, data, target):     output = model(data)     loss = F.nll_loss(output, target, reduction="mean")     model.backward(loss)     return output, loss def train(model, device, train_loader, optimizer):     model.train()     for batch_idx, (data, target) in enumerate(train_loader):         # smdistributed: Move input tensors to the GPU ID used by         # the current process, based on the set_device call.         data, target = data.to(device), target.to(device)         optimizer.zero_grad()         # Return value, loss_mb is a StepOutput object         _, loss_mb = train_step(model, data, target)         # smdistributed: Average the loss across microbatches.         loss = loss_mb.reduce_mean()         optimizer.step() # smdistributed: Initialize the backend smp.init() # smdistributed: Set the device to the GPU ID used by the current process. # Input tensors should be transferred to this device. torch.cuda.set_device(smp.local_rank()) device = torch.device("cuda") # smdistributed: Download only on a single process per instance. # When this is not present, the file is corrupted by multiple processes trying # to download and extract at the same time if smp.local_rank() == 0:     dataset = datasets.MNIST("../data", train=True, download=False) smp.barrier() # smdistributed: Shard the dataset based on data parallel ranks if smp.dp_size() > 1:     partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())}     dataset = SplitDataset(dataset, partitions=partitions_dict)     dataset.select(f"{smp.dp_rank()}") # smdistributed: Set drop_last=True to ensure that batch size is always divisible # by the number of microbatches train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True) model = Net() # smdistributed: enable tensor parallelism only for model.fc1 smp.set_tensor_parallelism(model.fc1, True) # smdistributed: Use the DistributedModel container to provide the model # to be partitioned across different ranks. For the rest of the script, # the returned DistributedModel object should be used in place of # the model provided for DistributedModel class instantiation. model = smp.DistributedModel(model) optimizer = optim.AdaDelta(model.parameters(), lr=4.0) optimizer = smp.DistributedOptimizer(optimizer) train(model, device, train_loader, optimizer) ``` # 支持 Hugging Face 转换器模型 SageMaker 模型并行度库的张量并行度为以下 Hugging Face Transformer 模型提供 out-of-the-box支持： + GPT-2、BERT 和 RoBERTa （在 SageMaker 模型并行度库 1.7.0 及更高版本中可用） + GPT-J（在 SageMaker 模型并行度库 1.8.0 及更高版本中可用） + GPT-neo（在 SageMaker 模型并行度库 v1.10.0 及更高版本中可用） **注意** 对于任何其他转换器模型，您需要使用 [smdistributed.modelparallel.torch.tp\$1register\$1with\$1module()](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch_tensor_parallel.html#smdistributed.modelparallel.torch.tp_register_with_module) API 来应用张量并行性。 **注意** 要使用张量并行度来训练 Hugging Face Transformer 模型，请务必使用 SageMaker具有模型并行度库 v1.7.0 及更高版本的 Hugging Face Deep Lear PyTorch ning Containers。有关更多信息，请参阅[SageMaker 模型并行度库发行说明](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_release_notes/smd_model_parallel_change_log.html)。 ## 现成支持的模型 对于开箱即用库支持的 Hugging Face 变压器模型，您无需手动实现挂钩即可将 Transformer `smdistributed` 转换为 APIs 变压器层。[你可以使用上下文管理器 smdistributed.modelparallel.torch.tensor\$1parallelism () 激活张量并行性，然后用 [smdistributed.modelparallel.torch 包装模型。](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch_tensor_parallel.html#smdistributed.modelparallel.torch.tensor_parallelism) DistributedModel()](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.DistributedModel)。您无需使用 `smp.tp_register` API 手动注册钩子用于张量并行性。 Hugging Face Transformers 与 `smdistributed.modelparallel` 之间的 `state_dict` 转换函数可以如下所示访问。 + `smdistributed.modelparallel.torch.nn.huggingface.gpt2.translate_state_dict_to_hf_gpt2(state_dict, max_seq_len=None)` + `smdistributed.modelparallel.torch.nn.huggingface.gpt2.translate_hf_state_dict_to_smdistributed_gpt2(state_dict)` + `smdistributed.modelparallel.torch.nn.huggingface.bert.translate_state_dict_to_hf_bert(state_dict, max_seq_len=None)` + `smdistributed.modelparallel.torch.nn.huggingface.bert.translate_hf_state_dict_to_smdistributed_bert(state_dict)` + `smdistributed.modelparallel.torch.nn.huggingface.roberta.translate_state_dict_to_hf_roberta(state_dict, max_seq_len=None)` + `smdistributed.modelparallel.torch.nn.huggingface.roberta.translate_hf_state_dict_to_smdistributed_roberta(state_dict)` + `smdistributed.modelparallel.torch.nn.huggingface.gptj.translate_state_dict_to_hf_gptj(state_dict, max_seq_len=None)`（在 SageMaker 模型并行度库 v1.8.0 及更高版本中可用） + `smdistributed.modelparallel.torch.nn.huggingface.gptj.translate_hf_gptj_state_dict_to_smdistributed_gptj`（在 SageMaker 模型并行度库 v1.8.0 及更高版本中可用） + `smdistributed.modelparallel.torch.nn.huggingface.gptneo.translate_state_dict_to_hf_gptneo(state_dict, max_seq_len=None)`（在 SageMaker 模型并行度库 v1.10.0 及更高版本中可用） + `smdistributed.modelparallel.torch.nn.huggingface.gptneo.translate_hf_state_dict_to_smdistributed_gptneo(state_dict)`（在 SageMaker 模型并行度库 v1.10.0 及更高版本中可用） **GPT-2 转换函数的使用示例** 首先包装模型，如以下代码所示。 ``` from transformers import AutoModelForCausalLM with smp.tensor_parallelism(): model = AutoModelForCausalLM.from_config(hf_gpt2_config) model = smp.DistributedModel(model) ``` 对于 `DistributedModel` 对象中的 `state_dict`，您可以使用 `translate_state_dict_to_hf_gpt2` 函数将权重加载到原始 Hugging Face GPT-2 模型中，如以下代码所示。 ``` from smdistributed.modelparallel.torch.nn.huggingface.gpt2 \ import translate_state_dict_to_hf_gpt2 max_seq_len = 1024 # [... code block for training ...] if smp.rdp_rank() == 0: state_dict = dist_model.state_dict() hf_state_dict = translate_state_dict_to_hf_gpt2(state_dict, max_seq_len) # can now call model.load_state_dict(hf_state_dict) to the original HF model ``` **Ro BERTa 翻译函数的用法示例** 同样，给定支持的 HuggingFace 模型`state_dict`，您可以使用`translate_hf_state_dict_to_smdistributed`函数将其转换为可读的格式`smp.DistributedModel`。这在迁移学习使用场景中非常有用，此时预训练的模型加载到 `smp.DistributedModel` 中用于模型并行微调： ``` from smdistributed.modelparallel.torch.nn.huggingface.roberta \ import translate_state_dict_to_smdistributed model = AutoModelForMaskedLM.from_config(roberta_config) model = smp.DistributedModel(model) pretrained_model = AutoModelForMaskedLM.from_pretrained("roberta-large") translated_state_dict = translate_state_dict_to_smdistributed(pretrained_model.state_dict()) # load the translated pretrained weights into the smp.DistributedModel model.load_state_dict(translated_state_dict) # start fine-tuning... ``` # 将管道并行性与张量并行性结合使用时的秩评定机制 此部分解释了模型并行性的秩评定机制如何与张量并行性结合使用。这是从 [SageMaker 模型并行度库的核心功能](model-parallel-core-features.md) 的[秩评定基础知识](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#ranking-basics)中扩展而来的。在张量并行度中，该库引入了三种类型的排名和处理组 APIs：用于`smp.tp_rank()`张量并行等级，用于`smp.pp_rank()`流水线并行等级，以及用于`smp.rdp_rank()`简化数据的并行等级。对应的通信进程组是张量并行组 (`TP_GROUP`)、管道并行组 (`PP_GROUP`) 和缩减数据并行组 (`RDP_GROUP`)。这些组的定义如下： + *张量并行组* (`TP_GROUP`) 是数据并行组中一个可均匀分割的子集，在其中完成模块的张量并行分布。当管道并行度为 1 时，`TP_GROUP` 与*模型并行组* (`MP_GROUP`) 相同。 + *管道并行组* (`PP_GROUP`) 是完成管道并行性的进程组。当张量并行度为 1 时，`PP_GROUP` 与 `MP_GROUP` 相同。 + *缩减数据并行组* (`RDP_GROUP`) 是一组进程，同时容纳相同的管道并行性分区和相同的张量并行分区，并在它们自身中执行数据并行性。之所以将其称为缩减数据并行组，是因为它是整个数据并行性组 `DP_GROUP` 的子集。对于分布在 `TP_GROUP` 中的模型参数，梯度 `allreduce` 运算仅对缩减数据并行组执行，而对于未分布的参数，梯度 `allreduce` 在整个 `DP_GROUP` 上进行。 + 模型并行组 (`MP_GROUP`) 是指一组共同存储整个模型的进程。它由当前进程的 `TP_GROUP` 中，所有秩的 `PP_GROUP` 并集组成。当张量并行度为 1 时，`MP_GROUP` 等于 `PP_GROUP`。它也与先前 `smdistributed` 版本中的 `MP_GROUP` 现有定义一致。请注意，当前 `TP_GROUP` 是当前 `DP_GROUP` 和当前 `MP_GROUP` 的子集。 要详细了解 SageMaker 模型并行度库 APIs 中的通信过程，请参阅 Pyth *SageMaker on* SDK 文档中的[通用 API](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_common_api.html#) 和[PyTorch特定 APIs](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html)的。 ![\[张量并行性的排名机制、参数分布和关联 AllReduce 运算。\]](http://docs.aws.amazon.com/zh_cn/sagemaker/latest/dg/images/distributed/model-parallel/tensor-parallel-ranking-mechanism.png) 例如，考虑具有8的单个节点的处理组 GPUs，其中张量并行度为2，流水线并行度为2，数据并行度为4。上图的顶部居中部分显示了一个包含 4 层的模型的示例。图的左下角和右下部分说明了 GPUs 使用流水线并行性和张量并行度分布在 4 层的 4 层模型，其中中间的两层使用张量并行性。下方的两个图是简单的副本，用于说明不同的组边界线。在 GPUs 0-3 和 4-7 之间复制分区模型以实现数据并行性。左下图显示了`MP_GROUP`、`PP_GROUP` 和 `TP_GROUP` 的定义。右下角的图显示`WORLD`了`RDP_GROUP``DP_GROUP`、和在同一组上 GPUs。为了实现数据并行性，具有相同颜色的层和层切片的梯度通过 `allreduce` 分在一起。例如，第一层（浅蓝色）获取 `DP_GROUP` 上的 `allreduce` 运算，而第二层中的深橙色切片只能获取其进程的 `RDP_GROUP` 中的 `allreduce` 运算。加粗深红色箭头表示张量及其完整 `TP_GROUP` 的批次。 ``` GPU0: pp_rank 0, tp_rank 0, rdp_rank 0, dp_rank 0, mp_rank 0 GPU1: pp_rank 1, tp_rank 0, rdp_rank 0, dp_rank 0, mp_rank 1 GPU2: pp_rank 0, tp_rank 1, rdp_rank 0, dp_rank 1, mp_rank 2 GPU3: pp_rank 1, tp_rank 1, rdp_rank 0, dp_rank 1, mp_rank 3 GPU4: pp_rank 0, tp_rank 0, rdp_rank 1, dp_rank 2, mp_rank 0 GPU5: pp_rank 1, tp_rank 0, rdp_rank 1, dp_rank 2, mp_rank 1 GPU6: pp_rank 0, tp_rank 1, rdp_rank 1, dp_rank 3, mp_rank 2 GPU7: pp_rank 1, tp_rank 1, rdp_rank 1, dp_rank 3, mp_rank 3 ``` 在此示例中，管道并行性跨 GPU 对 (0,1)、(2,3)、(4,5) 和 (6,7) 进行。此外，数据并行度 (`allreduce`) 在 GPUs 0、2、4、6 之间进行，并在 GPUs 1、3、5、7 上独立进行。张量并行性发生在 `DP_GROUP` 的子集上，跨 GPU 对 (0,2)、(1,3)、(4,6) 和 (5,7)。 对于这种管道和张量并行性的混合，`data_parallel_degree` 的算术仍然是 `data_parallel_degree = number_of_GPUs / pipeline_parallel_degree`。该库根据关系 `reduced_data_parallel_degree * tensor_parallel_degree = data_parallel_degree`，进一步计算缩减数据并行度。 # 优化器状态分片 *优化器状态分片*是一种非常有用的内存节省技术，它可以跨数据并行设备组对优化器状态（描述优化器状态的一组权重）进行分片。每当使用有状态优化器（例如 Adam）或优化器（存储参数 FP16 和参数 FP32 副本）时，都可以使用 FP16 优化器状态分片。 **注意** 优化器状态分片可在 SageMaker 模型并行度库 v1.6.0 及更高版本 PyTorch 中使用。 ## 如何使用优化器状态分片 您可以通过在 `modelparallel` 配置设置 `"shard_optimizer_state": True` 来启用*优化器状态分片*。 启用此功能后，库会根据数据并行度对模型参数集进行分区。与第 `i` 个分区对应的梯度，仅在第 `i` 个数据并行秩处缩减。在对 `smp.step` 修饰器函数的第一次调用结束时，被 `smp.DistributedOptimizer` 包装的优化器重新定义了其参数，使其仅限于与当前数据并行秩的分区相对应的参数。重新定义的参数称为*虚拟参数*，它们与原始参数共享底层存储。在第一次调用 `optimizer.step` 期间，优化器状态是根据这些重新定义的参数创建的，这些状态由于原始分区而被分片。优化器更新后，该 AllGather操作（作为`optimizer.step`调用的一部分）在数据 parallel 等级中运行，以实现一致的参数状态。 **提示** 当数据并行度大于 1 且模型的参数超过 10 亿时，优化器状态分片可能很有用。 数据并行性由 `(processes_per_host * instance_count / pipeline_parallel_degree)` 计算得出，`smp.dp_size()` 函数在后台处理大小调整。 **配置 SageMaker PyTorch 估算器** ``` mpi_options = { "enabled" : True, "processes_per_host" : 8, # 8 processes "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none " } smp_options = { "enabled":True, "parameters": { "microbatches": 4, "pipeline_parallel_degree": 2, # alias for "partitions" "placement_strategy": "cluster", "tensor_parallel_degree": 2, # tp over 2 devices "ddp": True, "shard_optimizer_state": True } } ``` **调整您的 PyTorch 训练脚本** 请参阅 *Tensor 并行度与流水线并行度相结合部分*中的[调整 PyTorch 训练脚本](model-parallel-extended-features-pytorch-tensor-parallelism-examples.md#model-parallel-extended-features-pytorch-tensor-and-pipeline-parallelism-script)。无需对脚本进行其他修改。 # 激活检查点 *激活检查点*（或*梯度检查点*）技术通过清除某些层的激活并在向后传递期间重新计算它们，来减少内存使用量。实际上，这是用额外的计算时间来换取内存使用量的减少。如果对模块执行了检查点操作，则在向前传递结束时，该模块的输入和输出将保留在内存中。在向前传递期间，任何本应是模块内部计算一部分的中间张量都会被释放。在有检查点的模块的向后传递过程中，会重新计算这些张量。此时，有检查点的模块之外的层已经完成其向后传递，因此检查点操作的峰值内存使用量可能会更低。 **注意** 此功能可在 SageMaker 模型并行度库 v1.6.0 及更高版本 PyTorch 中使用。 ## 如何使用激活检查点 使用 `smdistributed.modelparallel`，您可以按模块使用激活检查点。对于除 `torch.nn.Sequential` 之外的所有 `torch.nn` 模块，只有当从管道并行性的角度来看，模块树位于一个分区内时，才能对模块树执行检查点操作。对于 `torch.nn.Sequential` 模块，顺序模块内的每个模块树必须完全位于一个分区内，激活检查点才能起作用。使用手动分区时，请注意这些限制。 使用[自动模型分区](https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-core-features.html#model-parallel-automated-model-splitting)时，您可在训练作业日志中找到以 `Partition assignments:` 开头的分区分配日志。如果一个模块在多个秩（例如，一个后代属于一个秩，另一个后代处于不同的秩）上分区，则库会忽略对模块执行检查点的尝试，并发出一条警告消息，说明该模块没有检查点。 **注意** SageMaker 模型并行度库支持重叠和非重叠操作以及检查点`allreduce`操作。 **注意** PyTorch的本机检查点 API 与不兼容。`smdistributed.modelparallel` **示例 1：**以下示例代码演示了当脚本中有模型定义时，如何使用激活检查点操作。 ``` import torch.nn as nn import torch.nn.functional as F from smdistributed.modelparallel.torch.patches.checkpoint import checkpoint class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 32, 3, 1) self.conv2 = nn.Conv2d(32, 64, 3, 1) self.fc1 = nn.Linear(9216, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = self.conv1(x) x = self.conv2(x) x = F.max_pool2d(x, 2) x = torch.flatten(x, 1) # This call of fc1 will be checkpointed x = checkpoint(self.fc1, x) x = self.fc2(x) return F.log_softmax(x, 1) ``` **示例 2：**以下示例代码演示了当脚本中有顺序模型时，如何使用激活检查点操作。 ``` import torch.nn as nn from smdistributed.modelparallel.torch.patches.checkpoint import checkpoint_sequential class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.seq = nn.Sequential( nn.Conv2d(1,20,5), nn.ReLU(), nn.Conv2d(20,64,5), nn.ReLU() ) def forward(self, x): # This call of self.seq will be checkpointed x = checkpoint_sequential(self.seq, x) return F.log_softmax(x, 1) ``` **示例 3：**以下示例代码显示了在从库（例如和 Hugging Face Transformers PyTorch ）导入预建模型时如何使用激活检查点。无论您是否对顺序模型执行检查点操作，请完成以下过程： 1. 使用 `smp.DistributedModel()` 包装模型。 1. 为顺序层定义一个对象。 1. 使用 `smp.set_activation_checkpointig()` 包装顺序层对象。 ``` import smdistributed.modelparallel.torch as smp from transformers import AutoModelForCausalLM smp.init() model = AutoModelForCausalLM(*args, **kwargs) model = smp.DistributedModel(model) # Call set_activation_checkpointing API transformer_layers = model.module.module.module.transformer.seq_layers smp.set_activation_checkpointing( transformer_layers, pack_args_as_tuple=True, strategy='each') ``` # 激活分载 当激活检查点和管道并行性均已启用并且微批次数量大于 1 时，*激活分载*是可以进一步减少内存使用量的附加功能。对于当前未在 CPU 中运行的微批次，激活分载会异步移动与其对应的有检查点的激活。就在 GPU 需要激活以便微批次向后传递之前，此功能会将分载的激活从 CPU 预取回来。 **注意** 此功能可在 SageMaker 模型并行度库 v1.6.0 及更高版本 PyTorch 中使用。 ## 如何使用激活分载 当**微批次的数量大于 1 并启用了激活检查点时**，请使用激活分载以减少内存使用量（请参阅[激活检查点](model-parallel-extended-features-pytorch-activation-checkpointing.md)）。如果不使用激活检查点，则激活分享不会生效。在仅与一个微批次一起使用时，这不会节省内存。 要使用激活分载，请在 `modelparallel` 配置中设置 `"offload_activations": True`。 激活分载将 `nn.Sequential` 模块中的有检查点的激活异步移至 CPU。通过 PCIe 链路传输的数据与 GPU 计算重叠。在计算了特定检查点层的向前传递后，会立即进行分载。在特定微批次的向后传递需要激活之前的片刻，激活将加载回 GPU。CPU-GPU 传输同样与计算重叠。 要调整激活加载回 GPU 的时间提前量，您可以使用配置参数 `"activation_loading_horizon"`（默认设置为 4，必须为大于 0 的 `int`）。较大的激活加载范围会导致激活更早地加载回 GPU。如果范围太大，激活分载所带来的内存节省影响可能会减弱。如果范围太小，激活可能无法及时加载回，从而减少重叠量，导致性能下降。 **提示** 激活分载对于参数超过一千亿的大型模型非常有用。 **配置 SageMaker PyTorch 估算器** ``` mpi_options = { "enabled" : True, "processes_per_host" : 8, # 8 processes "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none " } smp_options = { "enabled":True, "parameters": { "microbatches": 4, "pipeline_parallel_degree": 2, # alias for "partitions" "placement_strategy": "cluster", "tensor_parallel_degree": 2, # tp over 2 devices "ddp": True, "offload_activations": True, "activation_loading_horizon": 4 # optional. default is 4. } } ``` # FP16 使用模型并行度进行训练 要进行 FP16 训练，请对训练脚本和估算器进行以下修改。 **注意** 此功能可在 SageMaker 模型并行度库 v1.10.0 及更高版本 PyTorch 中使用。 **调整您的 PyTorch 训练脚本** 1. 使用 [smdistributed.modelparallel.torch.model\$1creation()](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.model_creation) 上下文管理器包装您的模型。 ``` # fp16_training_script.py import torch import smdistributed.modelparallel.torch as smp with smp.model_creation( dtype=torch.float16 if args.fp16 else torch.get_default_dtype() ): model = ... ``` **提示** 如果您使用的是张量并行性，请将 `tensor_parallelism=smp.tp_size() > 1` 添加到 `smp.model_creation` 上下文管理器中。添加此行还有助于自动检测是否已激活张量并行性。 ``` with smp.model_creation( ... , tensor_parallelism=smp.tp_size() > 1 ): model = ... ``` 1. 当您使用 `smdistributed.modelparallel.torch.DistributedOptimizer` 包装优化器时，请设置 `static_loss_scaling` 或 `dynamic_loss_scaling` 参数。默认情况下，`static_loss_scaling` 设置为 `1.0`，`dynamic_loss_scaling` 设置为 `False`。如果您设置 `dynamic_loss_scale=True`，则可以通过 `dynamic_loss_args` 参数将动态损失缩放选项作为字典输入。在大多数情况下，我们建议您使用带有默认选项的动态损失缩放。[有关优化器包装器函数的更多信息、选项和示例，请参阅 smdistributed.modelparallel.torch。 DistributedOptimizer](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed-modelparallel-torch-distributedoptimizer)API。 以下代码是使用动态损失缩放来包装`Adadelta`优化器对象以进行 FP16 训练的示例。 ``` optimizer = torch.optim.Adadelta(...) optimizer = smp.DistributedOptimizer( optimizer, static_loss_scale=None, dynamic_loss_scale=True, dynamic_loss_args={ "scale_window": 1000, "min_scale": 1, "delayed_shift": 2 } ) ``` **配置 SageMaker PyTorch 估算器** 在创建 SageMaker PyTorch 估计器对象时，将 FP16 参数 (`"fp16"`) 添加到分布配置中以实现模型并行性。有关模型并行性配置参数的完整列表，请参阅 [`smdistributed` 的参数](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#parameters-for-smdistributed)。 ``` from sagemaker.pytorch import PyTorch smp_options = { "enabled": True, "parameters": { "microbatches": 4, "pipeline_parallel_degree": 2, "tensor_parallel_degree": 2, ..., "fp16": True } } fp16_estimator = PyTorch( entry_point="fp16_training_script.py", # Specify your train script ..., distribution={ "smdistributed": {"modelparallel": smp_options}, "mpi": {...} } ) fp16_estimator.fit(...) ``` FP16 训练开始时，模型和优化器`FP16_Optimizer`分别由`FP16_Module`和包装，它们是 [Apex](https://nvidia.github.io/apex/fp16_utils.html#apex-fp16-utils) 实用程序的修改`smdistributed`版本。 `FP16_Module`将模型转换为 FP16 dtype 并处理向前传入。 FP16 **提示** 您可以在 `optimizer.step` 之前通过调用 `clip_master_grads` 来应用梯度剪裁。 ``` optimizer.clip_master_grads(max_norm) # max_norm(float or int): max norm of the gradients ``` **提示** 使用`torch.optim.lr_scheduler`和 FP16 训练时，需要传递`optimizer.optimizer`给 LR 调度器而不是优化器。请参阅以下示例代码。 ``` from torch.optim.lr_scheduler import StepLR scheduler = StepLR( optimizer.optimizer if smp.state.cfg.fp16 else optimizer, step_size=1, gamma=args.gamma ) ``` # Support FlashAttention Suppor FlashAttention t for 是该库的一项功能，仅适用于*分布式变压器模型，分布式变压器*模型是为模型并行训练而封装的 Trans [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed-modelparallel-torch-distributedmodel](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed-modelparallel-torch-distributedmodel)former 模型。此功能还与 [张量并行性](model-parallel-extended-features-pytorch-tensor-parallelism.md) 兼容。 该[FlashAttention](https://github.com/HazyResearch/flash-attention)库仅在设置`attention_head_size`为 8 的倍数且小于 128 的值时才支持模型。因此，在训练分布式变压器并确保其 FlashAttention 正常工作时，应调整参数以使注意力头大小符合要求。有关更多信息，另请参阅*FlashAttention GitHub存储库*中的[安装和功能](https://github.com/HazyResearch/flash-attention#installation-and-features)。 例如，假设您使用 `hidden_width=864` 和 `num_heads=48` 配置转换器模型。的头部大小计算公式 FlashAttention 为`attention_head_size = hidden_width / num_heads = 864 / 48 = 18`。要启用 FlashAttention，您需要将`num_heads`参数调整为 `54``attention_head_size = hidden_width / num_heads = 864 / 54 = 16`，即 8 的倍数。 # 使用模型并行度运行 SageMaker 分布式训练 Job 学习如何使用带有模型并行度库的 Pyth SageMaker on SDK 运行自己的训练脚本的 SageMaker 模型并行训练作业。 运行 SageMaker 训练作业有三种用例场景。 1. 您可以将预先构建的 AWS 深度学习容器之一用于 TensorFlow 和。 PyTorch如果这使您首次使用模型并行库，则建议使用此选项。要查找有关如何运行 SageMaker 模型并行训练作业的教程，请参阅[使用 Amazon A SageMaker I 的模型并行度库进行PyTorch 训练的](https://github.com/aws/amazon-sagemaker-examples/tree/main/training/distributed_training/pytorch/model_parallel)示例笔记本。 1. 您可以扩展预先构建的容器，以处理预构建的 SageMaker Docker 镜像不支持的算法或模型的任何其他功能要求。要查找如何扩展预构建容器的示例，请参阅[扩展预构建容器](prebuilt-containers-extend.md)。 1. 您可以使用[SageMaker 训练工具包](https://github.com/aws/sagemaker-training-toolkit)调整自己的 Docker 容器以使用 SageMaker AI。有关示例，请参阅[调整您自己的训练容器](https://docs.aws.amazon.com/sagemaker/latest/dg/adapt-training-container.html)。 对于前述列表中的选项 2 和 3，请参阅[扩展包含分布式模型并行库的预构建 SageMaker的 Docker 容器](model-parallel-sm-sdk.md#model-parallel-customize-container)，以了解如何在扩展或自定义的 Docker 容器中安装模型并行库。 在所有情况下，您都可以启动训练作业，配置 SageMaker `TensorFlow``PyTorch`或估算器以激活库。要了解更多信息，请参阅以下主题。 **Topics** + [步骤 1：使用 SageMaker分布式模型并行库修改自己的训练脚本](model-parallel-customize-training-script.md) + [第 2 步：使用 SageMaker Python 软件开发工具包启动训练 Job](model-parallel-sm-sdk.md) # 步骤 1：使用 SageMaker分布式模型并行库修改自己的训练脚本 使用本节学习如何自定义训练脚本以使用 Amazon A SageMaker I 模型并行度库的核心功能。要使用特定于库的 API 函数和参数，我们建议您将本文档与 Pyth *SageMaker on SDK* 文档 APIs中的[模型并SageMaker 行库](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html)一起使用。 这些部分中提供的训练脚本示例经过简化，旨在重点介绍使用库时必须进行的更改。有关 end-to-end演示如何在 SageMaker 模型并行度库中使用 TensorFlow 或 PyTorch 训练脚本的可运行笔记本示例，请参阅。[亚马逊 SageMaker AI 模型并行度库 v2 示例](distributed-model-parallel-v2-examples.md) **Topics** + [使用模型并行度库拆分训练脚本的 SageMaker 模型](#model-parallel-model-splitting-using-smp-lib) + [修改 TensorFlow 训练脚本](model-parallel-customize-training-script-tf.md) + [修改 PyTorch 训练脚本](model-parallel-customize-training-script-pt.md) ## 使用模型并行度库拆分训练脚本的 SageMaker 模型 有两种方法可以修改训练脚本以设置模型拆分：自动拆分或手动拆分。 ### 自动模型拆分 使用 SageMaker模型并行度库时，可以利用*自动模型拆分*，也称为*自动模型*分区。该库使用分区算法来平衡内存，最大限度地减少设备之间的通信并优化性能。您可以配置自动分区算法以优化速度或内存。 或者，您可以使用手动模型拆分。除非您非常熟悉模型架构，并且很好地掌握了如何有效地对模型进行分区的知识，否则我们建议自动拆分模型。 #### 工作原理 自动分区发生在第一个训练步骤中，也就是首次调用 `smp.step` 修饰的函数时。在此调用期间，库首先在 CPU RAM 上构造模型的版本（以避免 GPU 内存限制），然后分析模型图形并做出分区决策。根据这个决策，每个模型分区加载到一个 GPU 上，然后才会执行第一步。由于这些分析和分区步骤，第一个训练步骤可能需要较长的时间。 在这两个框架中，库都通过自己的后端管理设备之间的通信，后端针对 AWS 基础架构进行了优化。 自动分区设计适应了框架的特性，库按照在每个框架中更自然的粒度级别进行分区。例如，在中 TensorFlow，可以将每个特定操作分配给不同的设备，而在模块级别中 PyTorch，分配是在模块级别完成的，其中每个模块由多个操作组成。以下部分介绍了每个框架中的具体设计。 ##### 自动拆分模型 PyTorch 在第一个训练步骤中，模型并行性库在内部运行一个跟踪步骤，用于构造模型图形并确定张量和参数配置。在此跟踪步骤之后，库将构造一个树，该树包含模型中的嵌套 `nn.Module` 对象，以及通过跟踪收集的其他数据，例如所存储的 `nn.Parameters` 数量，以及每个 `nn.Module` 的执行时间。 接下来，库从根目录遍历此树，并运行分区算法，将每个 `nn.Module` 分配给一个设备，这可以平衡计算负载（按模块执行时间来衡量）和内存使用量（按总共存储的 `nn.Parameter` 大小和激活来衡量）。如果多个 `nn.Modules` 共享相同的 `nn.Parameter`，则将这些模块放置在相同设备上，以避免维护同一个参数的多个版本。在做出分区决策后，分配的模块和权重就会加载到它们的设备上。 有关如何将`smp.step`装饰器注册到 PyTorch 训练脚本的说明，请参阅[使用自动拆分 PyTorch](model-parallel-customize-training-script-pt.md#model-parallel-customize-training-script-pt-16)。 ##### 自动拆分模型 TensorFlow 模型并行性库分析可训练变量的大小和图形结构，并在内部使用图形分区算法。该算法为每个操作提供了设备分配，目的是最大限度地减少所需的设备间通信量，但要遵守两个限制： + 平衡每个设备中存储的变量数量 + 平衡每个设备中执行的操作数量 如果您为 `optimize` 指定 `speed`（在 Python SDK 的模型并行性参数中），则库会尝试平衡每个设备中的操作数和 `tf.Variable` 对象数量。否则，它会尝试平衡 `tf.Variables` 的总大小。 做出分区决策后，库会创建每个设备需要执行的子图形的串行化表示形式，并将其导入到每个设备上。在分区时，库会将使用相同 `tf.Variable` 的操作，以及属于同一 Keras 层的操作放在同一个设备上。它还尊重由 TensorFlow施加的主机托管限制。举例而言，这意味着如果有两个 Keras 层共享一个 `tf.Variable`，则属于这些层的所有操作都将放在单个设备上。 有关如何将`smp.step`装饰器注册到 PyTorch 训练脚本的说明，请参阅[使用自动拆分 TensorFlow](model-parallel-customize-training-script-tf.md#model-parallel-customize-training-script-tf-23)。 ##### 各个框架的自动模型拆分的比较 在中 TensorFlow，计算的基本单位是 a`tf.Operation`，它将模型 TensorFlow 表示为 `tf.Operation` s 的有向无环图 (DAG)，因此，模型并行度库对此 DAG 进行了分区，使每个节点都转到一台设备。至关重要的是，`tf.Operation` 对象具有足够丰富的可自定义属性，而且它们是通用的，因为每个模型都确保由此类对象的图形组成。 PyTorch 另一方面, 却没有一个足够丰富和普遍的同等行动概念. 具有这些特征的最接近的 PyTorch 计算单位是`nn.Module`，它的粒度级别要高得多，这就是库在此级别上进行分区的原因。 PyTorch ### 手动模型拆分 如果要手动指定如何跨设备对模型进行分区，请使用 `smp.partition` 上下文管理器。有关如何设置上下文管理器以进行手动分区的说明，请参阅以下页面。 + [使用手动拆分 TensorFlow](model-parallel-customize-training-script-tf.md#model-parallel-customize-training-script-tf-manual) + [使用手动拆分 PyTorch](model-parallel-customize-training-script-pt.md#model-parallel-customize-training-script-pt-16-hvd) 要在修改后使用此选项，在步骤 2 中，您需要在 Pyth SageMaker on SDK 的框架估算器类`default_partition`中设置`auto_partition`并定义一个。`False`任何未通过 `smp.partition` 上下文管理器明确放置在分区上的操作都将在 `default_partition` 上执行。在这种情况下，将绕过自动拆分逻辑，并根据您的规范放置每个操作。根据生成的图形结构，模型并行性库会自动创建管道执行计划。 # 修改 TensorFlow 训练脚本 在本节中，您将学习如何修改 TensorFlow 训练脚本以配置用于自动分区和手动分区的 SageMaker模型并行度库。这些示例还包括与 Horovod 集成的示例，用于混合模型和数据并行性。 **注意** 要了解该库支持哪些 TensorFlow 版本，请参阅[支持的框架和 AWS 区域](distributed-model-parallel-support.md)。 [使用自动拆分 TensorFlow](#model-parallel-customize-training-script-tf-23)中列出了为使用库而必须对训练脚本进行的修改。 要了解如何修改训练脚本以在 Horovod 中使用混合模型和数据并行性，请参阅 [使用 TensorFlow 和 Horovod 自动拆分，实现混合模型和数据并行性](#model-parallel-customize-training-script-tf-2.3)。 如果您要使用手动分区，另请参阅 [使用手动拆分 TensorFlow](#model-parallel-customize-training-script-tf-manual)。 以下主题显示了训练脚本的示例，您可以使用这些脚本来配置自动分区和手动分区模型 SageMaker的模型并行度库。 TensorFlow **注意** 默认情况下启用自动分区。除非另行指定，否则示例脚本使用自动分区。 **Topics** + [使用自动拆分 TensorFlow](#model-parallel-customize-training-script-tf-23) + [使用 TensorFlow 和 Horovod 自动拆分，实现混合模型和数据并行性](#model-parallel-customize-training-script-tf-2.3) + [使用手动拆分 TensorFlow](#model-parallel-customize-training-script-tf-manual) + [不支持的框架功能](#model-parallel-tf-unsupported-features) ## 使用自动拆分 TensorFlow 要使用模型并行度库运行 TensorFlow 模型，需要对训练脚本进行 SageMaker以下更改： 1. 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init) 导入和初始化库。 1. 通过从 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_tensorflow.html](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_tensorflow.html) 继承来定义 Keras 模型，而不是从 Keras 模型类继承。从 `smp.DistributedModel` 对象的调用方法返回模型输出。请注意，从调用方法返回的任何张量都将在模型并行设备之间广播，这会产生通信开销，因此在调用方法之外不需要的任何张量（例如中间激活）都不应返回。 1. 在 `tf.Dataset.batch()` 方法中设置 `drop_remainder=True`。这是为了确保批次大小始终可以被微批次数量整除。 1. 例如，使用`smp.dp_rank()`以下方法在数据管道中播种随机操作，`shuffle(ds, seed=smp.dp_rank())`以确保包含不同模型分区的数据样本的一致性。 GPUs 1. 将向前和向后逻辑放在步进函数中，然后用 `smp.step` 进行修饰。 1. 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#StepOutput](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#StepOutput) 方法（例如 `reduce_mean`）对微批次的输出进行后处理。[https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init) 函数必须具有一个取决于 `smp.DistributedModel` 的输出的返回值。 1. 如果有评估步骤，则同样将向前逻辑放在 `smp.step` 修饰的函数中，然后使用 [`StepOutput` API](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#StepOutput) 对输出进行后处理。 要了解有关模型并行度库 API SageMaker 的更多信息，请参阅 AP [I](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html) 文档。 以下 Python 脚本是进行更改后的训练脚本的示例。 ``` import tensorflow as tf # smdistributed: Import TF2.x API import smdistributed.modelparallel.tensorflow as smp # smdistributed: Initialize smp.init() # Download and load MNIST dataset. (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data( "MNIST-data-%d" % smp.rank() ) x_train, x_test = x_train / 255.0, x_test / 255.0 # Add a channels dimension x_train = x_train[..., tf.newaxis] x_test = x_test[..., tf.newaxis] # smdistributed: If needed, seed the shuffle with smp.dp_rank(), and drop_remainder # in batching to make sure batch size is always divisible by number of microbatches train_ds = ( tf.data.Dataset.from_tensor_slices((x_train, y_train)) .shuffle(10000, seed=smp.dp_rank()) .batch(256, drop_remainder=True) ) # smdistributed: Define smp.DistributedModel the same way as Keras sub-classing API class MyModel(smp.DistributedModel): def __init__(self): super(MyModel, self).__init__() # define layers def call(self, x, training=None): # define forward pass and return the model output model = MyModel() loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) optimizer = tf.keras.optimizers.Adam() train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name="train_accuracy") # smdistributed: Define smp.step. Return any tensors needed outside @smp.step def get_grads(images, labels): predictions = model(images, training=True) loss = loss_object(labels, predictions) grads = optimizer.get_gradients(loss, model.trainable_variables) return grads, loss, predictions @tf.function def train_step(images, labels): gradients, loss, predictions = get_grads(images, labels) # smdistributed: Accumulate the gradients across microbatches gradients = [g.accumulate() for g in gradients] optimizer.apply_gradients(zip(gradients, model.trainable_variables)) # smdistributed: Merge predictions and average losses across microbatches train_accuracy(labels, predictions.merge()) return loss.reduce_mean() for epoch in range(5): # Reset the metrics at the start of the next epoch train_accuracy.reset_states() for images, labels in train_ds: loss = train_step(images, labels) accuracy = train_accuracy.result() ``` 如果您已准备好训练脚本，请继续到[第 2 步：使用 SageMaker Python 软件开发工具包启动训练 Job](model-parallel-sm-sdk.md)。如果要运行混合模型和数据并行训练作业，请继续到下一个部分。 ## 使用 TensorFlow 和 Horovod 自动拆分，实现混合模型和数据并行性 您可以将 SageMaker 模型并行度库与 Horovod 配合使用，实现混合模型和数据并行性。要详细了解库如何拆分模型以用于混合并行性，请参阅[管道并行度（适用于 PyTorch 和） TensorFlow](model-parallel-intro.md#model-parallel-intro-pp)。 在这一步中，我们将重点介绍如何修改训练脚本以适应 SageMaker模型并行度库。 要正确设置训练脚本，以便选取要在 [第 2 步：使用 SageMaker Python 软件开发工具包启动训练 Job](model-parallel-sm-sdk.md) 中设置的混合并行度配置，请使用库的帮助程序函数 `smp.dp_rank()` 和 `smp.mp_rank()`，它们分别自动检测数据并行秩和模型并行秩。 要查找该库支持的所有 MPI 原语，请参阅 [Pyth SageMaker on SDK 文档中的 MPI 基础知识](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#mpi-basics)。 脚本中需要进行以下更改： + 添加 `hvd.allreduce` + 按照 Horovod 的要求，在第一个批次之后广播变量 + 使用. 在数据管道中播种洗牌 and/or 分片操作。`smp.dp_rank()` **注意** 使用 Horovod 时，您不可在训练脚本中直接调用 `hvd.init`。相反，你必须在中的 SageMaker Python SDK `modelparallel` 参数`True`中`"horovod"`将其设置为[第 2 步：使用 SageMaker Python 软件开发工具包启动训练 Job](model-parallel-sm-sdk.md)。这使得库可以根据模型分区的设备分配，在内部初始化 Horovod。直接在训练脚本中调用 `hvd.init()` 可能会导致问题。 **注意** 在训练脚本中直接使用 `hvd.DistributedOptimizer` API 可能会导致训练性能和速度不佳，因为 API 会隐式地将 `AllReduce` 操作放入 `smp.step` 中。我们建议您在 `smp.step` 返回的梯度上调用 `accumulate()` 或 `reduce_mean()` 后，直接调用 `hvd.allreduce`，以将模型并行性库与 Horovod 一起使用，如下例所示。 要了解有关模型并行度库 API SageMaker 的更多信息，请参阅 AP [I](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html) 文档。 ``` import tensorflow as tf import horovod.tensorflow as hvd # smdistributed: Import TF2.x API import smdistributed.modelparallel.tensorflow as smp # smdistributed: Initialize smp.init() # Download and load MNIST dataset. (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data( "MNIST-data-%d" % smp.rank() ) x_train, x_test = x_train / 255.0, x_test / 255.0 # Add a channels dimension x_train = x_train[..., tf.newaxis] x_test = x_test[..., tf.newaxis] # smdistributed: Seed the shuffle with smp.dp_rank(), and drop_remainder # in batching to make sure batch size is always divisible by number of microbatches train_ds = ( tf.data.Dataset.from_tensor_slices((x_train, y_train)) .shuffle(10000, seed=smp.dp_rank()) .batch(256, drop_remainder=True) ) # smdistributed: Define smp.DistributedModel the same way as Keras sub-classing API class MyModel(smp.DistributedModel): def __init__(self): super(MyModel, self).__init__() # define layers def call(self, x, training=None): # define forward pass and return model outputs model = MyModel() loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) optimizer = tf.keras.optimizers.Adam() train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name="train_accuracy") # smdistributed: Define smp.step. Return any tensors needed outside @smp.step def get_grads(images, labels): predictions = model(images, training=True) loss = loss_object(labels, predictions) grads = optimizer.get_gradients(loss, model.trainable_variables) return grads, loss, predictions @tf.function def train_step(images, labels, first_batch): gradients, loss, predictions = get_grads(images, labels) # smdistributed: Accumulate the gradients across microbatches # Horovod: AllReduce the accumulated gradients gradients = [hvd.allreduce(g.accumulate()) for g in gradients] optimizer.apply_gradients(zip(gradients, model.trainable_variables)) # Horovod: Broadcast the variables after first batch if first_batch: hvd.broadcast_variables(model.variables, root_rank=0) hvd.broadcast_variables(optimizer.variables(), root_rank=0) # smdistributed: Merge predictions across microbatches train_accuracy(labels, predictions.merge()) return loss.reduce_mean() for epoch in range(5): # Reset the metrics at the start of the next epoch train_accuracy.reset_states() for batch, (images, labels) in enumerate(train_ds): loss = train_step(images, labels, tf.constant(batch == 0)) ``` ## 使用手动拆分 TensorFlow 使用 `smp.partition` 上下文管理器将操作放在特定的分区中。未放在任何 `smp.partition` 上下文中的任何操作都放在 `default_partition` 中。要了解有关模型并行度库 API SageMaker 的更多信息，请参阅 AP [I](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html) 文档。 ``` import tensorflow as tf # smdistributed: Import TF2.x API. import smdistributed.modelparallel.tensorflow as smp # smdistributed: Initialize smp.init() # Download and load MNIST dataset. (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data( "MNIST-data-%d" % smp.rank() ) x_train, x_test = x_train / 255.0, x_test / 255.0 # Add a channels dimension x_train = x_train[..., tf.newaxis] x_test = x_test[..., tf.newaxis] # smdistributed: If needed, seed the shuffle with smp.dp_rank(), and drop_remainder # in batching to make sure batch size is always divisible by number of microbatches. train_ds = ( tf.data.Dataset.from_tensor_slices((x_train, y_train)) .shuffle(10000, seed=smp.dp_rank()) .batch(256, drop_remainder=True) ) # smdistributed: Define smp.DistributedModel the same way as Keras sub-classing API. class MyModel(smp.DistributedModel): def __init__(self): # define layers def call(self, x): with smp.partition(0): x = self.layer0(x) with smp.partition(1): return self.layer1(x) model = MyModel() loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) optimizer = tf.keras.optimizers.Adam() train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name="train_accuracy") # smdistributed: Define smp.step. Return any tensors needed outside @smp.step def get_grads(images, labels): predictions = model(images, training=True) loss = loss_object(labels, predictions) grads = optimizer.get_gradients(loss, model.trainable_variables) return grads, loss, predictions @tf.function def train_step(images, labels): gradients, loss, predictions = get_grads(images, labels) # smdistributed: Accumulate the gradients across microbatches gradients = [g.accumulate() for g in gradients] optimizer.apply_gradients(zip(gradients, model.trainable_variables)) # smdistributed: Merge predictions and average losses across microbatches train_accuracy(labels, predictions.merge()) return loss.reduce_mean() for epoch in range(5): # Reset the metrics at the start of the next epoch train_accuracy.reset_states() for images, labels in train_ds: loss = train_step(images, labels) accuracy = train_accuracy.result() ``` ## 不支持的框架功能 该库不支持以下 TensorFlow 功能： + 当前不支持 `tf.GradientTape()`。您可以改用 `Optimizer.get_gradients()` 或 `Optimizer.compute_gradients()` 来计算梯度。 + 目前不支持 `tf.train.Checkpoint.restore()` API。对于检查点操作，请改用 `smp.CheckpointManager`，它提供相同的 API 和功能。请注意，`smp.CheckpointManager` 的检查点还原应在第一步之后进行。 # 修改 PyTorch 训练脚本 在本节中，您将学习如何修改 PyTorch 训练脚本以配置用于自动分区和手动分区的 SageMaker 模型并行度库。 **注意** 要了解该库支持哪些 PyTorch 版本，请参阅[支持的框架和 AWS 区域](distributed-model-parallel-support.md)。 **提示** 有关演示如何在 SageMaker 模型并行度库中使用 PyTorch 训练脚本的 end-to-end笔记本示例，请参阅。[亚马逊 SageMaker AI 模型并行库 v1 示例](distributed-model-parallel-examples.md) 请注意，默认情况下启用自动分区。除非另行指定，否则以下脚本使用自动分区。 **Topics** + [使用自动拆分 PyTorch](#model-parallel-customize-training-script-pt-16) + [使用手动拆分 PyTorch](#model-parallel-customize-training-script-pt-16-hvd) + [注意事项](#model-parallel-pt-considerations) + [不支持的框架功能](#model-parallel-pt-unsupported-features) ## 使用自动拆分 PyTorch 要使用模型并行度库运行 PyTorch 训练脚本，需要进行 SageMaker以下训练脚本更改： 1. 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init) 导入和初始化库。 1. 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_pytorch.html#smp.DistributedModel](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_pytorch.html#smp.DistributedModel) 包装模型。请注意，从底层 `nn.Module` 对象的 `forward` 方法返回的任何张量都将在模型并行设备之间广播，这会产生通信开销，因此在调用方法之外不需要的任何张量（例如中间激活）都不应返回。 **注意** 要进行 FP16 训练，你需要使用 smdistrib [uted.modelparallel.torch.model\$1creation () 上下文管理器来封装模型。](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html)有关更多信息，请参阅 [FP16 使用模型并行度进行训练](model-parallel-extended-features-pytorch-fp16.md)。 1. 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_pytorch.html#smp.DistributedOptimizer](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_pytorch.html#smp.DistributedOptimizer) 包装优化器。 **注意** 要进行 FP16 训练，您需要设置静态或动态损失比例。有关更多信息，请参阅 [FP16 使用模型并行度进行训练](model-parallel-extended-features-pytorch-fp16.md)。 1. 使用返回的 `DistributedModel` 对象而不是用户模型。 1. 将向前和向后逻辑放在步进函数中，然后用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#smp.init) 进行修饰。 1. 通过 `torch.cuda.set_device(smp.local_rank())` 将每个进程限制在自己的设备上。 1. 在 `smp.step` 调用之前，使用 `.to()` API 将输入张量移至 GPU（参见下面的示例）。 1. 将 `torch.Tensor.backward` 和 `torch.autograd.backward` 替换为 `DistributedModel.backward`。 1. 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#StepOutput](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#StepOutput) 方法（例如 `reduce_mean`）对微批次的输出进行后处理。 1. 如果有评估步骤，则同样将向前逻辑放在 `smp.step` 修饰的函数中，然后使用 [`StepOutput` API](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_common_api.html#StepOutput) 对输出进行后处理。 1. 在 `DataLoader` 中设置 `drop_last=True`。或者，如果批次大小不能被微批次数量整除，则可以手动跳过训练循环中的批次。 要了解有关模型并行度库 API SageMaker 的更多信息，请参阅 AP [I](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html) 文档。 ``` import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchnet.dataset import SplitDataset from torchvision import datasets import smdistributed.modelparallel.torch as smp class GroupedNet(nn.Module): def __init__(self): super(GroupedNet, self).__init__() # define layers def forward(self, x): # define forward pass and return model outputs # smdistributed: Define smp.step. Return any tensors needed outside. @smp.step def train_step(model, data, target): output = model(data) loss = F.nll_loss(output, target, reduction="mean") model.backward(loss) return output, loss def train(model, device, train_loader, optimizer): model.train() for batch_idx, (data, target) in enumerate(train_loader): # smdistributed: Move input tensors to the GPU ID used by the current process, # based on the set_device call. data, target = data.to(device), target.to(device) optimizer.zero_grad() # Return value, loss_mb is a StepOutput object _, loss_mb = train_step(model, data, target) # smdistributed: Average the loss across microbatches. loss = loss_mb.reduce_mean() optimizer.step() # smdistributed: initialize the backend smp.init() # smdistributed: Set the device to the GPU ID used by the current process. # Input tensors should be transferred to this device. torch.cuda.set_device(smp.local_rank()) device = torch.device("cuda") # smdistributed: Download only on a single process per instance. # When this is not present, the file is corrupted by multiple processes trying # to download and extract at the same time dataset = datasets.MNIST("../data", train=True, download=False) # smdistributed: Shard the dataset based on data-parallel ranks if smp.dp_size() > 1: partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())} dataset = SplitDataset(dataset, partitions=partitions_dict) dataset.select(f"{smp.dp_rank()}") # smdistributed: Set drop_last=True to ensure that batch size is always divisible # by the number of microbatches train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True) model = GroupedNet() optimizer = optim.Adadelta(model.parameters(), lr=4.0) # smdistributed: Use the DistributedModel container to provide the model # to be partitioned across different ranks. For the rest of the script, # the returned DistributedModel object should be used in place of # the model provided for DistributedModel class instantiation. model = smp.DistributedModel(model) optimizer = smp.DistributedOptimizer(optimizer) train(model, device, train_loader, optimizer) ``` ## 使用手动拆分 PyTorch 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_pytorch.html#smp.DistributedOptimizer](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/v1.2.0/smd_model_parallel_pytorch.html#smp.DistributedOptimizer) 上下文管理器将模块放置在特定设备中。未放在任何 `smp.partition` 上下文中的任何模块都放在 `default_partition` 中。如果 `auto_partition` 设置为 `False`，则需要提供 `default_partition`。在特定 `smp.partition` 上下文中创建的模块将放置在对应的分区上。 要了解有关模型并行度库 API SageMaker 的更多信息，请参阅 AP [I](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel.html) 文档。 ``` import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchnet.dataset import SplitDataset from torchvision import datasets import smdistributed.modelparallel.torch as smp class GroupedNet(nn.Module): def __init__(self): super(GroupedNet, self).__init__() with smp.partition(0): # define child modules on device 0 with smp.partition(1): # define child modules on device 1 def forward(self, x): # define forward pass and return model outputs # smdistributed: Define smp.step. Return any tensors needed outside. @smp.step def train_step(model, data, target): output = model(data) loss = F.nll_loss(output, target, reduction="mean") model.backward(loss) return output, loss def train(model, device, train_loader, optimizer): model.train() for batch_idx, (data, target) in enumerate(train_loader): # smdistributed: Move input tensors to the GPU ID used by the current process, # based on the set_device call. data, target = data.to(device), target.to(device) optimizer.zero_grad() # Return value, loss_mb is a StepOutput object _, loss_mb = train_step(model, data, target) # smdistributed: Average the loss across microbatches. loss = loss_mb.reduce_mean() optimizer.step() # smdistributed: initialize the backend smp.init() # smdistributed: Set the device to the GPU ID used by the current process. # Input tensors should be transferred to this device. torch.cuda.set_device(smp.local_rank()) device = torch.device("cuda") # smdistributed: Download only on a single process per instance. # When this is not present, the file is corrupted by multiple processes trying # to download and extract at the same time dataset = datasets.MNIST("../data", train=True, download=False) # smdistributed: Shard the dataset based on data-parallel ranks if smp.dp_size() > 1: partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())} dataset = SplitDataset(dataset, partitions=partitions_dict) dataset.select(f"{smp.dp_rank()}") # smdistributed: Set drop_last=True to ensure that batch size is always divisible # by the number of microbatches train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True) model = GroupedNet() optimizer = optim.Adadelta(model.parameters(), lr=4.0) # smdistributed: Use the DistributedModel container to provide the model # to be partitioned across different ranks. For the rest of the script, # the returned DistributedModel object should be used in place of # the model provided for DistributedModel class instantiation. model = smp.DistributedModel(model) optimizer = smp.DistributedOptimizer(optimizer) train(model, device, train_loader, optimizer) ``` ## 注意事项 使用 SageMaker模型并行度库配置 PyTorch 训练脚本时，应注意以下几点： + 如果您使用的优化技术依赖于全局梯度范数，例如整个模型的梯度范数（如 LAMB 优化器或全局梯度剪裁的某些变体），则需要收集模型分区中的所有范数以确保正确性。您可以使用库的通信基本数据类型来完成此操作。 + 模型中 `nn.Modules` 的所有向前方法的所有 `torch.Tensor` 参数都必须在模块输出的计算中使用。换而言之，该库不支持向模块传递模块输出所不依赖的 `torch.Tensor` 参数。 + 传递给 `smp.DistributedModel.backward()` 调用的参数必须依赖于所有模型输出。换而言之，`smp.DistributedModel.forward` 调用的输出，必须在计算提供给 `smp.DistributedModel.backward` 的张量时全部使用。 + 如果您的代码中有 `torch.cuda.synchronize()` 调用，则可能需要在同步调用之前立即调用 `torch.cuda.set_device(smp.local_rank())`。否则，可能会在设备 0 中创建不必要的 CUDA 上下文，这将不必要地消耗内存。 + 由于库放置在不同的设备上的 `nn.Modules` 中，因此模型中的模块不得依赖于任何在 `smp.step` 内部修改的全局状态。在整个训练过程中保持固定的任何状态，或者在 `smp.step` 之外以对所有进程都可见的方式修改的任何状态，都是允许的。 + 使用库时，您无需将模型移至 GPU（例如，使用 `model.to(device)`）。如果您尝试在模型分区之前（在第一次 `smp.step` 调用之前）将模型移动到 GPU，则会忽略移动调用。对于模型中分配给某个秩的部分，库会自动将这个部分移动到其 GPU。一旦开始使用库进行训练，请不要将模型移至 CPU 并使用它，因为对于未分配给进程所持有的分区的模块，它不会有正确的参数。如果您想在使用模型并行度库训练模型后重新训练模型或在没有库的情况下将其用于推理，则推荐的方法是使用我们的检查点 API 保存完整模型并将其加载回常规模块。 PyTorch + 如果您有模块列表，这样一个模块的输出就会传送到另一个模块中，而使用 `nn.Sequential` 替换该列表可以显著提高性能。 + 权重更新 (`optimizer.step()`) 需要在 `smp.step` 外部进行，因为此时整个向后传递已完成并且梯度准备就绪。当使用具有模型和数据并行性的混合模型时，此时，梯 AllReduce 度也可以保证完成。 + 将该库与数据并行度结合使用时，请确保所有数据并行等级上的批次数相同，这样在等待 AllReduce 未参与该步骤的排名时就不会挂起。 + 如果您使用 ml.p4d 实例类型（例如 ml.p4d.24xlarge）启动训练作业，则必须设置数据加载程序变量 `num_workers=0`。例如，您可以如下所示定义 `DataLoader`： ``` dataloader = torch.utils.data.DataLoader( data, batch_size=batch_size, num_workers=0, pin_memory=True, drop_last=True, shuffle=shuffle, ) ``` + `smp.step` 的输入必须是 `DataLoader` 生成的模型输入。这是因为在 `smp.step` 内部按照批次维度拆分输入张量，并对其进行管道传输。这意味着将 `DataLoader` 本身传递给 `smp.step` 函数以在其中生成模型输入，这种方法行不通。 例如，如果您如下所示定义 `DataLoader`： ``` train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True) ``` 您应该访问 `train_loader` 生成的模型输入，并将这些输入传递给 `smp.step` 修饰的函数。不要将 `train_loader` 直接传递给 `smp.step` 函数。 ``` def train(model, device, train_loader, optimizer): model.train() for batch_idx, (data, target) in enumerate(train_loader): ... _, loss_mb = train_step(model, data, target) ... @smp.step def train_step(model, data, target): ... return output, loss ``` + `smp.step` 的输入张量必须使用 `.to()` API 移动到当前设备，此操作必须在 `torch.cuda.set_device(local_rank())` 调用后进行。 例如，您可以如下所示定义 `train` 函数。在使用这些输入张量调用 `train_step` 之前，此函数使用 `.to()` API 将 `data` 和 `target` 添加到当前设备。 ``` def train(model, device, train_loader, optimizer): model.train() for batch_idx, (data, target) in enumerate(train_loader): # smdistributed: Move input tensors to the GPU ID used by the current process, # based on the set_device call. data, target = data.to(device), target.to(device) optimizer.zero_grad() # Return value, loss_mb is a StepOutput object _, loss_mb = train_step(model, data, target) # smdistributed: Average the loss across microbatches. loss = loss_mb.reduce_mean() optimizer.step() ``` 在上面的 `train` 函数中，此 `smp.set` 修饰的函数的输入张量已移至当前设备。模型*无需*移至当前设备。对于模型中分配给某个秩的部分，库会自动将这个部分移动到其 GPU。 ``` @smp.step def train_step(model, data, target): output = model(data) loss = F.nll_loss(output, target, reduction="mean") model.backward(loss) return output, loss ``` ## 不支持的框架功能 的模型并行度库不支持以下 PyTorch 功能： SageMaker + 如果将数据并行性与本机 [PyTorch DDP](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html) 一起使用，则该库不支持[https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html)包装器模块。该库在内部管理与 PyTorch DDP 的集成，包括参数广播和渐变 AllReduce。使用库时，模块缓冲区在训练开始时仅广播一次。如果您的模型具有的模块缓冲区需要在每个步骤中跨数据并行组同步，则可以通过 `torch.distributed` API，使用可以通过 `smp.get_dp_process_group()` 获取的进程组进行同步。 + 对于混合精度训练，不支持 `apex.amp` 模块。使用具有自动混合精度的库的推荐方法是使用 `torch.cuda.amp`，唯一的不同是使用 `smp.amp.GradScaler` 而不是在 Torch 中实施。 + `smp.DistributedModel` 不支持 `torch.jit.ScriptModules` 或 `ScriptFunctions`。 + 不支持 `apex` : `FusedLayerNorm`、`FusedAdam`、`FusedLAMB` 以及来自 `apex` 的 `FusedNovoGrad`。你可以通过`smp.optimizers`和`smp.nn` APIs 改用它们的库实现。 # 第 2 步：使用 SageMaker Python 软件开发工具包启动训练 Job SageMaker Python SDK 支持使用机器学习框架（例如 TensorFlow 和）对模型进行托管训练 PyTorch。[要使用其中一个框架启动训练作业，您需要定义估计器、 SageMaker [TensorFlow 估计器或](https://sagemaker.readthedocs.io/en/v2.199.0/frameworks/tensorflow/sagemaker.tensorflow.html#tensorflow-estimator) SageMaker 通用 E SageMaker [PyTorch stimat](https://sagemaker.readthedocs.io/en/v2.199.0/frameworks/pytorch/sagemaker.pytorch.html#pytorch-estimator) or 以使用修改后的训练脚本和模型并行配置。](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/estimators.html#sagemaker.estimator.Estimator) **Topics** + [使用 SageMaker TensorFlow 和 PyTorch 估算器](#model-parallel-using-sagemaker-pysdk) + [扩展包含分布式模型并行库的预构建 SageMaker的 Docker 容器](#model-parallel-customize-container) + [使用 SageMaker 分布式模型并行库创建自己的 Docker 容器](#model-parallel-bring-your-own-container) ## 使用 SageMaker TensorFlow 和 PyTorch 估算器 TensorFlow 和 PyTorch estimator 类包含`distribution`参数，您可以使用该参数来指定使用分布式训练框架的配置参数。 SageMaker 模型并行库内部使用 MPI 来处理混合数据和模型并行性，因此必须在该库中使用 MPI 选项。 以下 TensorFlow PyTorch 或估算器模板显示了如何配置`distribution`参数，以便在 MPI 中使用模型 SageMaker 并行库。 ------ #### [ Using the SageMaker TensorFlow estimator ] ``` import sagemaker from sagemaker.tensorflow import TensorFlow smp_options = { "enabled":True, # Required "parameters": { "partitions": 2, # Required "microbatches": 4, "placement_strategy": "spread", "pipeline": "interleaved", "optimize": "speed", "horovod": True, # Use this for hybrid model and data parallelism } } mpi_options = { "enabled" : True, # Required "processes_per_host" : 8, # Required # "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none" } smd_mp_estimator = TensorFlow( entry_point="your_training_script.py", # Specify your train script source_dir="location_to_your_script", role=sagemaker.get_execution_role(), instance_count=1, instance_type='ml.p3.16xlarge', framework_version='2.6.3', py_version='py38', distribution={ "smdistributed": {"modelparallel": smp_options}, "mpi": mpi_options }, base_job_name="SMD-MP-demo", ) smd_mp_estimator.fit('s3://my_bucket/my_training_data/') ``` ------ #### [ Using the SageMaker PyTorch estimator ] ``` import sagemaker from sagemaker.pytorch import PyTorch smp_options = { "enabled":True, "parameters": { # Required "pipeline_parallel_degree": 2, # Required "microbatches": 4, "placement_strategy": "spread", "pipeline": "interleaved", "optimize": "speed", "ddp": True, } } mpi_options = { "enabled" : True, # Required "processes_per_host" : 8, # Required # "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none" } smd_mp_estimator = PyTorch( entry_point="your_training_script.py", # Specify your train script source_dir="location_to_your_script", role=sagemaker.get_execution_role(), instance_count=1, instance_type='ml.p3.16xlarge', framework_version='1.13.1', py_version='py38', distribution={ "smdistributed": {"modelparallel": smp_options}, "mpi": mpi_options }, base_job_name="SMD-MP-demo", ) smd_mp_estimator.fit('s3://my_bucket/my_training_data/') ``` ------ 要启用该库，您需要通过 SageMaker 估计器构造函数的`distribution`参数将配置字典传递给`"smdistributed"`和`"mpi"`键。 **SageMaker 模型并行度的配置参数** + 对于 `"smdistributed"` 键，用 `"modelparallel"` 键和以下内部字典传递字典。 **注意** 不支持在一个训练作业中使用 `"modelparallel"` 和 `"dataparallel"`。 + `"enabled"` – 必需。要启用模型并行性，请设置 `"enabled": True`。 + `"parameters"` – 必需。为 SageMaker 模型并行度指定一组参数。 + 有关常用参数的完整列表，请参阅 *SageMaker Python SDK 文档`smdistributed`*中的[参数](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#smdistributed-parameters)。 有关信息 TensorFlow，请参阅[TensorFlow特定参数](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#tensorflow-specific-parameters)。 有关信息 PyTorch，请参阅[PyTorch特定参数](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#pytorch-specific-parameters)。 + `"pipeline_parallel_degree"`（在 `smdistributed-modelparallel 要扩展预先构建 SageMaker的容器并使用其模型并行度库，您必须将其中一个可用的 AWS 深度学习容器 (DLC) 图像用于或。 PyTorch TensorFlow SageMaker 模型并行度库包含在带有 CUDA () 的 TensorFlow （2.3.0 及更高版本）和 PyTorch （1.6.0 及更高版本）DLC 镜像中。`cuxyz`有关 DLC 图像的完整列表，请参阅 Deep Learnin [g Containers * GitHub 存储库中可用的AWS 深度学习容器*镜像](https://github.com/aws/deep-learning-containers/blob/master/available_images.md)。 **提示** 我们建议您使用包含最新版本的图像 TensorFlow或 PyTorch 访问 SageMaker 模型并行度库的最新 up-to-date版本。 例如，您的 Dockerfile 应该包含类似于下文的 `FROM` 语句： ``` # Use the SageMaker DLC image URI for TensorFlow or PyTorch FROM aws-dlc-account-id.dkr.ecr.aws-region.amazonaws.com/framework-training:{framework-version-tag} # Add your dependencies here RUN ... ENV PATH="/opt/ml/code:${PATH}" # this environment variable is used by the SageMaker AI container to determine our user code directory. ENV SAGEMAKER_SUBMIT_DIRECTORY /opt/ml/code ``` 此外，在定义 PyTorch TensorFlow 或估计器时，必须`entry_point`为训练脚本指定。这应该与 `ENV SAGEMAKER_SUBMIT_DIRECTORY` 在 Dockerfile 中标识的路径相同。 **提示** 你必须将此 Docker 容器推送到亚马逊弹性容器注册表 (Amazon ECR) Registry，然后使用图像 URI (`image_uri`) 来定义训练的估算器。 SageMaker 有关更多信息，请参阅 [扩展预构建容器](prebuilt-containers-extend.md)。 托管 Docker 容器并检索容器的镜像 URI 后，按如下方式创建一个 SageMaker `PyTorch`估算器对象。本例假设您已经定义 `smp_options` 和 `mpi_options`。 ``` smd_mp_estimator = Estimator( entry_point="your_training_script.py", role=sagemaker.get_execution_role(), instance_type='ml.p3.16xlarge', sagemaker_session=sagemaker_session, image_uri='your_aws_account_id.dkr.ecr.region.amazonaws.com/name:tag' instance_count=1, distribution={ "smdistributed": smp_options, "mpi": mpi_options }, base_job_name="SMD-MP-demo", ) smd_mp_estimator.fit('s3://my_bucket/my_training_data/') ``` ## 使用 SageMaker 分布式模型并行库创建自己的 Docker 容器 要构建自己的 Docker 容器进行训练并使用 SageMaker 模型并行库，您必须在 Dockerfile 中包含正确的依赖项和 SageMaker分布式并行库的二进制文件。本节提供了在自己的 Docker 容器中正确准备 SageMaker 训练环境和模型 parallel 库时必须包含的最少代码块。 **注意** 这个带有 SageMaker 模型并行库作为二进制文件的自定义 Docker 选项仅适用于。 PyTorch **使用 SageMaker 训练工具包和模型并行库创建 Dockerfile** 1. 从 [NVIDIA CUDA 基础映像](https://hub.docker.com/r/nvidia/cuda)之一开始使用。 ``` FROM ``` **提示** 官方的 AWS 深度学习容器 (DLC) 镜像是基于 [NVIDIA CUDA 基础](https://hub.docker.com/r/nvidia/cuda)镜像构建的。我们建议您查看[AWS 深度学习容器的官方 Dockerfiles](https://github.com/aws/deep-learning-containers/tree/master/pytorch/training/docker)， PyTorch以了解需要安装哪些版本的库以及如何配置它们。官方 Dockerfile 已完成，经过基准测试并由深度学习容器服务团队 SageMaker 和深度学习容器服务团队管理。在提供的链接中，选择您使用的 PyTorch版本，选择 CUDA (`cuxyz`) 文件夹，然后选择以或结尾的 Dockerfile。`.gpu` `.sagemaker.gpu` 1. 要设置分布式训练环境，您需要为通信和网络设备安装软件，例如 [Elastic Fabric Adapter (EFA)](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa.html)、[NVIDIA Collective Communications Library (NCCL)](https://developer.nvidia.com/nccl) 和 [Open MPI](https://www.open-mpi.org/)。根据您选择的 PyTorch 和 CUDA 版本，必须安装兼容版本的库。 **重要** 由于 SageMaker 模型并行库在后续步骤中需要 SageMaker 数据并行库，因此我们强烈建议您按照中的说明正确设置分布式 SageMaker 训练的训练环境。[使用 SageMaker AI 分布式数据并行库创建自己的 Docker 容器](data-parallel-bring-your-own-container.md) 有关使用 NCCL 和 Open MPI 设置 EFA 的更多信息，请参阅[开始使用 EFA 和 MPI](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa-start.html) 以及[开始使用 EFA 和 NCCL](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa-start-nccl.html)。 1. 添加以下参数来指定 URLs SageMaker 分布式训练包的 PyTorch。 SageMaker 模型并行库要求 SageMaker 数据并行库使用跨节点远程直接内存访问 (RDMA)。 ``` ARG SMD_MODEL_PARALLEL_URL=https://sagemaker-distributed-model-parallel.s3.us-west-2.amazonaws.com/pytorch-1.10.0/build-artifacts/2022-02-21-19-26/smdistributed_modelparallel-1.7.0-cp38-cp38-linux_x86_64.whl ARG SMDATAPARALLEL_BINARY=https://smdataparallel.s3.amazonaws.com/binary/pytorch/1.10.2/cu113/2022-02-18/smdistributed_dataparallel-1.4.0-cp38-cp38-linux_x86_64.whl ``` 1. 安装 SageMaker 模型并行库所需的依赖项。 1. 安装 [METIS](http://glaros.dtc.umn.edu/gkhome/metis/metis/overview) 库。 ``` ARG METIS=metis-5.1.0 RUN rm /etc/apt/sources.list.d/* \ && wget -nv http://glaros.dtc.umn.edu/gkhome/fetch/sw/metis/${METIS}.tar.gz \ && gunzip -f ${METIS}.tar.gz \ && tar -xvf ${METIS}.tar \ && cd ${METIS} \ && apt-get update \ && make config shared=1 \ && make install \ && cd .. \ && rm -rf ${METIS}.tar* \ && rm -rf ${METIS} \ && rm -rf /var/lib/apt/lists/* \ && apt-get clean ``` 1. 安装 [RAPIDS Memory Manager 库](https://github.com/rapidsai/rmm#rmm-rapids-memory-manager)。这需要 [CMake](https://cmake.org/)3.14 或更高版本。 ``` ARG RMM_VERSION=0.15.0 RUN wget -nv https://github.com/rapidsai/rmm/archive/v${RMM_VERSION}.tar.gz \ && tar -xvf v${RMM_VERSION}.tar.gz \ && cd rmm-${RMM_VERSION} \ && INSTALL_PREFIX=/usr/local ./build.sh librmm \ && cd .. \ && rm -rf v${RMM_VERSION}.tar* \ && rm -rf rmm-${RMM_VERSION} ``` 1. 安装 SageMaker 模型并行库。 ``` RUN pip install --no-cache-dir -U ${SMD_MODEL_PARALLEL_URL} ``` 1. 安装 SageMaker 数据并行库。 ``` RUN SMDATAPARALLEL_PT=1 pip install --no-cache-dir ${SMDATAPARALLEL_BINARY} ``` 1. 安装 [sagemaker-training 工具包](https://github.com/aws/sagemaker-training-toolkit)。该工具包包含创建与 SageMaker 训练平台和 SageMaker Python SDK 兼容的容器所必需的常用功能。 ``` RUN pip install sagemaker-training ``` 1. 完成 Dockerfile 的创建后，请参阅[调整自己的训练容器](https://docs.aws.amazon.com/sagemaker/latest/dg/adapt-training-container.html)，了解如何构建 Docker 容器并将其托管在 Amazon ECR 中。 **提示** 有关创建用于 SageMaker 人工智能训练的自定义 Dockerfile 的更多一般信息，请参阅[使用自己的训练算](https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo.html)法。 # 对具有模型并行性的模型执行检查点操作和微调 SageMaker 模型并行度库提供了检查点 APIs 功能，用于保存按各种模型并行策略划分的模型状态和优化器状态，并从要重新开始训练和微调的地方加载用于持续训练的检查点。 APIs 还支持部分或全部保存模型和优化器状态的选项。 **Topics** + [对分布式模型执行检查点操作](#distributed-model-parallel-checkpoint) + [微调分布式模型](#distributed-model-parallel-fine-tuning) ## 对分布式模型执行检查点操作 根据 PyTorch 和之间的框架以及您使用的 SageMaker 模型 TensorFlow 并行度库的版本，选择以下主题之一。 **Topics** + [对分布式 PyTorch 模型执行检查点操作（适用于 SageMaker 模型并行度库 v1.10.0 及更高版本）](#model-parallel-extended-features-pytorch-checkpoint) + [对分布式 PyTorch 模型执行检查点操作（适用于 v1.6.0 和 v1.9.0 之间的 SageMaker 模型并行度库）](#model-parallel-extended-features-pytorch-saving-loading-checkpoints) + [对分布式模型执行检查点操作 TensorFlow](#distributed-model-parallel-checkpoint-tensorflow) ### 对分布式 PyTorch 模型执行检查点操作（适用于 SageMaker 模型并行度库 v1.10.0 及更高版本） SageMaker 模型并行度库提供了用于保存和加载分布式模型状态及其优化器状态的全部或部分检查点的检查点。 APIs **注意** 如果您使用 SageMaker 模型并行度库 v1.10.0 或更高版本，则建议使用 PyTorch 此检查点方法。 **部分检查点** 要保存使用模型并行性训练的模型的检查点，请使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.save_checkpoint](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.save_checkpoint) API 并将部分检查点选项设置为 true (`partial=True`)。这将单独保存每个模型分区。在模型状态和优化器状态之外，您还可以通过 `user_content` 参数保存任何其他自定义数据。具有检查点的模型、优化器和用户内容保存为单独的文件。`save_checkpoint` API 调用按以下结构创建检查点文件夹。 ``` - path - ${tag}_partial (folder for partial checkpoints) - model_rankinfo.pt - optimizer_rankinfo.pt - fp16_states_rankinfo.pt - user_content.pt - $tag (checkpoint file for full checkpoints) - user_content_$tag (user_content file for full checkpoints) - newest (a file that indicates the newest checkpoint) ``` 要从部分检查点恢复训练，请使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.resume_from_checkpoint](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.resume_from_checkpoint) API 和 `partial=True`，并指定保存部分检查点时使用的检查点目录和标签。请注意，模型权重的实际加载在模型分区之后进行，在第一次运行经过 `smdistributed.modelparallel.torch.step` 修饰的训练步骤函数期间。 保存部分检查点时，库还会使用 `.pt` 文件扩展名，将模型分区决策保存为文件。反过来，在从部分检查点恢复时，库会将分区决策文件一起加载。分区决策一旦加载，您就无法更改分区。 以下代码片段显示了如何在 PyTorch训练脚本 APIs 中设置检查点。 ``` import smdistributed.modelparallel.torch as smp model = ... model = smp.DistributedModel(model) optimizer = ... optimizer = smp.DistributedOptimizer(optimizer) user_content = ... # additional custom data checkpoint_path = "/opt/ml/checkpoint/model_parallel" # Save a checkpoint. smp.save_checkpoint( path=checkpoint_path, tag=f"total_steps{total_steps}", partial=True, model=model, optimizer=optimizer, user_content=user_content num_kept_partial_checkpoints=5 ) # Load a checkpoint. # This automatically loads the most recently saved checkpoint. smp_checkpoint = smp.resume_from_checkpoint( path=checkpoint_path, partial=True ) ``` **完整检查点** 要保存最终模型构件用于推理用途，请使用 `smdistributed.modelparallel.torch.save_checkpoint` API 和 `partial=False`，这会组合模型分区以创建单个模型构件。请注意，这不会合并优化器状态。 要使用特定权重初始化训练，对于给定的完整模型检查点，您可以使用 `smdistributed.modelparallel.torch.resume_from_checkpoint` API 和 `partial=False`。请注意，这不会加载优化器状态。 **注意** 通常，对于张量并行性，`state_dict` 必须在原始模型实施和 `DistributedModel` 实施之间进行转换。或者，您可以将 `state_dict` 转换函数作为参数提供给 `smdistributed.modelparallel.torch.resume_from_checkpoint`。但是，对于 [现成支持的模型](model-parallel-extended-features-pytorch-hugging-face.md#model-parallel-extended-features-pytorch-hugging-face-out-of-the-box)，库会自动处理此转换。 以下代码显示了如何使用检查点对使用模型并行度训练的 PyTorch 模型 APIs 进行完全检查点的示例。 ``` import smdistributed.modelparallel.torch as smp model = ... model = smp.DistributedModel(model) optimizer = ... optimizer = smp.DistributedOptimizer(optimizer) user_content = ... # additional custom data checkpoint_path = "/opt/ml/checkpoint/model_parallel" # Save a checkpoint. smp.save_checkpoint( path=checkpoint_path, tag=f"total_steps{total_steps}", partial=False, model=model, optimizer=optimizer, user_content=user_content num_kept_partial_checkpoints=5 ) # Load a checkpoint. # This automatically loads the most recently saved checkpoint. smp_checkpoint = smp.resume_from_checkpoint( path=checkpoint_path, partial=False ) ``` ### 对分布式 PyTorch 模型执行检查点操作（适用于 v1.6.0 和 v1.9.0 之间的 SageMaker 模型并行度库） SageMaker 模型并行度库提供了 Python 函数，用于使用张量并行度为训练作业保存部分或全部检查点。以下过程演示如何使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.save](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.save) 和 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.load](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.load)，在使用张量并行性的情况下保存和加载检查点。 **注意** 如果您使用 PyTorch、和 v1.6.0 和 v1.9.0 之间的 SageMaker 模型并行度库[张量并行性](model-parallel-extended-features-pytorch-tensor-parallelism.md)，则建议使用这种检查点方法。 1. 准备一个模型对象，并使用库的包装器函数 `smp.DistributedModel()` 进行包装。 ``` model = MyModel(...) model = smp.DistributedModel(model) ``` 1. 为模型准备一个优化器。一组模型参数是优化器函数所需的可迭代参数。要准备一组模型参数，必须处理`model.parameters()`为单个模型参数 IDs 分配唯一的参数。 如果模型参数可迭代 IDs 中存在重复的参数，则加载检查点优化器状态将失败。要为优化器创建具有唯一 IDs 性的可迭代模型参数，请参阅以下内容： ``` unique_params = [] unique_params_set = set() for p in model.parameters(): if p not in unique_params_set: unique_params.append(p) unique_params_set.add(p) del unique_params_set optimizer = MyOpt(unique_params, ...) ``` 1. 使用库的包装器函数 `smp.DistributedOptimizer()` 包装优化器。 ``` optimizer = smp.DistributedOptimizer(optimizer) ``` 1. 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.save](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.save) 保存模型和优化器状态。根据您保存检查点的方式，选择以下两个选项之一： + **选项 1：**在每个 `mp_rank` 上为单个 `MP_GROUP` 保存一个部分模型。 ``` model_dict = model.local_state_dict() # save a partial model opt_dict = optimizer.local_state_dict() # save a partial optimizer state # Save the dictionaries at rdp_rank 0 as a checkpoint if smp.rdp_rank() == 0: smp.save( {"model_state_dict": model_dict, "optimizer_state_dict": opt_dict}, f"/checkpoint.pt", partial=True, ) ``` 使用张量并行性，该库按以下命名格式保存检查点文件：`checkpoint.pt_{pp_rank}_{tp_rank}`。 **注意** 对于张量并行性，请确保将 if 语句设置 `if smp.rdp_rank() == 0` 为而不是 `if smp.dp_rank() == 0`。当使用张量并行性对优化器状态进行分片时，所有缩减数据并行秩都必须保存自己的优化器状态分区。为检查点操作使用错误的 *if* 语句可能会导致训练作业停滞。有关在`if smp.dp_rank() == 0`不使用张量并行的情况下使用的更多信息，请参阅 Pyth *SageMaker on* SD [K 文档中的保存和加载一般说明](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#general-instruction-for-saving-and-loading)。 + **选项 2：**保存完整模型。 ``` if smp.rdp_rank() == 0: model_dict = model.state_dict(gather_to_rank0=True) # save the full model if smp.rank() == 0: smp.save( {"model_state_dict": model_dict}, "/checkpoint.pt", partial=False, ) ``` **注意** 对于完整检查点操作，请考虑以下内容： 如果您设置 `gather_to_rank0=True`，则除 `0` 之外的所有秩返回空字典。 对于完整检查点，您只能对模型执行检查点操作。目前不支持对优化器状态执行完整检查点操作。 只需将完整模型保存在 `smp.rank() == 0` 即可。 1. 使用 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.load](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.load) 加载检查点。根据您在上一步中的检查点操作方式，选择以下两个选项之一： + **选项 1：**加载部分检查点。 ``` checkpoint = smp.load("/checkpoint.pt", partial=True) model.load_state_dict(checkpoint["model_state_dict"], same_partition_load=False) optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) ``` 如果您知道分区不会更改，则可以在 `model.load_state_dict()` 中设置 `same_partition_load=True` 以实现更快的加载速度。 + **选项 2：**加载完整检查点。 ``` if smp.rdp_rank() == 0: checkpoint = smp.load("/checkpoint.pt", partial=False) model.load_state_dict(checkpoint["model_state_dict"]) ``` `if smp.rdp_rank() == 0` 条件并非必需，但它可以帮助避免不同 `MP_GROUP` 之间的冗余加载。使用张量并行性时，目前不支持完整检查点优化器状态字典。 ### 对分布式模型执行检查点操作 TensorFlow 要在使用 TensorFlow 模型并行度训练时保存模型，请使用 SageMaker 模型并行度库提供的以下函数。 + [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_tensorflow.html#smp.DistributedModel.save_model](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_tensorflow.html#smp.DistributedModel.save_model) + [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_tensorflow.html#smp.CheckpointManager](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_tensorflow.html#smp.CheckpointManager) ## 微调分布式模型 需要在训练脚本中配置微调。以下代码片段显示了使用 Hugging Face Transformers 的 [AutoModelForCausalLM](https://huggingface.co/docs/transformers/main/en/model_doc/auto#transformers.AutoModelForCausalLM) 类的训练脚本的示例结构，其中包含用于注册`smdistributed.model.parallel.torch`模块和微调设置的修改。 **注意** 在激活 [smp.delayed\$1param\$1initionational 函数的情况下微调分布式变压器（封装的变压器模型`smp.DistributedModel()`）需要使用 for](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#smdistributed.modelparallel.torch.delay_param_initialization) Lustre 文件系统配置微调作业。 FSx 如果要使用延迟参数初始化选项对大型模型进行微调，则应设置 for Lustre 文件系统。 FSx ``` import argparse from transformers import AutoModelForCausalLM import smdistributed.modelparallel import smdistributed.modelparallel.torch as smp def parse_args(): parser = argparse.ArgumentParser() # set an arg group for model model_grp = parser.add_argument_group( title="model", description="arguments to describe model configuration" ) ... # set up numerous args to parse from the configuration dictionary to the script for training # add arg for activating fine-tuning model_grp.add_argument( "--fine_tune", type=int, default=0, help="Fine-tune model from checkpoint or pretrained model", ) def main(): """Main function to train GPT.""" args = parse_args() ... # parse numerous args if args.fine_tune > 0 and args.delayed_param > 0 and smp.rank() == 0: pretrained_model = AutoModelForCausalLM.from_pretrained( args.model_name or args.model_dir ) model_state_dict = pretrained_model.state_dict() path = os.path.join(args.model_dir, "fullmodel.pt") torch.save(model_state_dict, path) # create a Transformer model and wrap by smp.model_creation() # with options to configure model parallelism parameters offered by SageMaker AI with smp.model_creation( tensor_parallelism=smp.tp_size() > 1 or args.use_distributed_transformer > 0, zero_init=args.use_distributed_transformer == 0, dtype=dtype, distribute_embedding=args.sharded_data_parallel_degree > 1 and smp.tp_size() > 1, use_alibi=args.alibi > 0, attention_in_fp32=args.attention_in_fp32 > 0, fp32_residual_addition=args.residual_addition_in_fp32 > 0, query_key_layer_scaling=args.query_key_layer_scaling > 0 and args.bf16 < 1, fused_softmax=args.fused_softmax > 0, fused_dropout=args.fused_dropout > 0, fused_bias_gelu=args.fused_bias_gelu > 0, flash_attention=args.flash_attention > 0, ): if args.fine_tune > 0 and args.delayed_param == 0: model = AutoModelForCausalLM.from_pretrained( args.model_name or args.model_dir ) else: model = AutoModelForCausalLM.from_config(model_config) # wrap the model by smp.DistributedModel() to apply SageMaker model parallelism model = smp.DistributedModel( model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation ) # wrap the optimizer by smp.DistributedOptimizer() to apply SageMaker model parallelism optimizer= ... # define an optimizer optimizer = smp.DistributedOptimizer( optimizer, static_loss_scale=None, dynamic_loss_scale=True, dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2}, ) # for fine-tuning, use smp.resume_from_checkpoint() to load a pre-trained model if args.fine_tune > 0 and args.delayed_param > 0: smp.resume_from_checkpoint(args.model_dir, tag="fullmodel.pt", partial=False) ``` 有关训练脚本和 Jupyter 笔记本的完整示例，请参阅 *SageMaker AI [示例存储库 PyTorch中的 GPT-2](https://github.com/aws/amazon-sagemaker-examples/tree/main/training/distributed_training/pytorch/model_parallel/gpt2) 示例 GitHub*。 # 亚马逊 SageMaker AI 模型并行库 v1 示例 本页提供了博客和 Jupyter 笔记本的列表，这些博客和 Jupyter 笔记本提供了实现 SageMaker 模型并行度 (SMP) 库 v1 以在 AI 上运行分布式训练作业的实际示例。 SageMaker ## 博客和案例研究 以下博客将讨论有关使用 SMP 版本 1 的案例研究。 + [亚马逊 A SageMaker I 模型并行度库的新性能改进，Machine Lear](https://aws.amazon.com/blogs/machine-learning/new-performance-improvements-in-amazon-sagemaker-model-parallel-library/) n *AWS ing 博客（202* 2 年 12 月 16 日） + 在 Amazon [AI 上使用分片数据并行度训练具有近线性扩展能力的巨型模型，M SageMaker achine Learning Bl](https://aws.amazon.com/blogs/machine-learning/train-gigantic-models-with-near-linear-scaling-using-sharded-data-parallelism-on-amazon-sagemaker/) *og（2022 AWS 年 10 月* 31 日） ## 示例笔记本 [SageMaker AI 示例 GitHub 存储库中提供了示例](https://github.com/aws/amazon-sagemaker-examples/tree/master/training/distributed_training/)笔记本。要下载示例，请运行以下命令克隆库并转到 `training/distributed_training/pytorch/model_parallel`。 **注意** 克隆并运行以下 SageMaker AI ML 中的示例笔记本 IDEs。 [SageMaker JupyterLab](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated-jl.html)（在 2023 年 12 月之后创建[的工作室](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated.html)中可用） [SageMaker 代码编辑器](https://docs.aws.amazon.com/sagemaker/latest/dg/code-editor.html)（在 2023 年 12 月之后创建的 [Studio](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated.html) 中可用） [Studio Classic](https://docs.aws.amazon.com/sagemaker/latest/dg/studio.html)（可作为 2023 年 12 月之后创建的 [Studio](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated.html) 中的应用程序使用） [SageMaker 笔记本实例](https://docs.aws.amazon.com/sagemaker/latest/dg/nbi.html) ``` git clone https://github.com/aws/amazon-sagemaker-examples.git cd amazon-sagemaker-examples/training/distributed_training/pytorch/model_parallel ``` **适用于 SMP v1 的示例笔记本电脑 PyTorch** + [使用模型并行度库中的分片数据并行技术使用近线性缩放训练 GPT-2 SageMaker ](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/smp-train-gpt-sharded-data-parallel.ipynb) + [使用模型并行度库中的分片数据并行技术对具有近线性缩放的 GPT-2 进行微调 SageMaker ](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt2/smp-fine-tune-gpt-sharded-data-parallel.ipynb) + [使用模型并行库中的分片数据并行技术以近线性缩放方式训练 GPT-neox-20b SageMaker ](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt-neox/smp-train-gpt-neox-sharded-data-parallel.ipynb) + [使用模型并行库中的分片数据并行性和张量并行度技术训练 GPT-J 6B SageMaker ](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/gpt-j/smp-train-gptj-sharded-data-parallel-tp.ipynb) + [使用模型并行度库中的分片数据并行技术使用近线性缩放训练 FLAN-T5 SageMaker ](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/flan-t5/smp-train-t5-sharded-data-parallel.ipynb) + [在模型并行度库中使用分片数据并行技术以近线性缩放方式训练 Falcon SageMaker ](https://github.com/aws/amazon-sagemaker-examples/blob/main/training/distributed_training/pytorch/model_parallel/falcon/smp-train-falcon-sharded-data-parallel.ipynb) **适用于 SMP v1 的示例笔记本电脑 TensorFlow** + [带有 TensorFlow 2.3.1 的 CNN 和 SageMaker 模型并行度库](https://sagemaker-examples.readthedocs.io/en/latest/training/distributed_training/tensorflow/model_parallel/mnist/tensorflow_smmodelparallel_mnist.html) + [HuggingFace 使用 TensorFlow 分布式模型并行库在 AI 上训练 SageMaker ](https://github.com/huggingface/notebooks/blob/master/sagemaker/04_distributed_training_model_parallelism/sagemaker-notebook.ipynb) # SageMaker 分布式模型并行性最佳实践 使用 SageMaker 模型 parallel 库运行分布式训练作业时，请遵循以下准则。 ## 针对给定模型设置合适的配置 在纵向扩展模型时，我们建议您按顺序浏览以下列表。每个列表项都讨论了使用该库技术的优点以及可能出现的权衡取舍。 **提示** 如果使用库的部分功能可以很好地拟合模型，那么添加更多的模型并行度或者使用节省内存的功能通常并不能提高性能。 **使用大型 GPU 实例类型** + 在模型并行性领域，最好使用具有大型 GPU 内存的强大实例来处理模型并行操作产生的开销，例如将模型分成多个模型。 GPUs我们建议使用 `ml.p4d` 或 `ml.p3dn` 实例来训练大型 DL 模型。这些实例还配备了 Elastic Fabric Adapter (EFA)，可提供更高的网络带宽，并支持使用模型并行性进行大规模训练。 **分片优化器状态** + 分片优化器状态的影响取决于数据并行秩的数量。通常，较高的数据并行度（与计算节点的大小成正比）可以提高内存使用效率。 当您想要缩小集群规模时，请确保检查优化器状态分片配置。例如，具有优化器状态分片的大型 DL 模型适合具有 16 的计算集群 GPUs （例如，两个 P4d 或 p4de 实例），可能并不总是适合具有 8 的节点 GPUs （例如，单个 P4d 或 p4de 实例）。这是因为 8 GPUs 的组合内存低于总内存 16，并且在 16 GPUs GPU 场景中分片时每个 GPU 所需的内存也高于在 16 GPU 场景中进行分片所需的每个 GPU 内存。 GPUs 因此，内存需求的增加使其可能不适合较小的集群。 有关更多信息，请参阅 [优化器状态分片](model-parallel-extended-features-pytorch-optimizer-state-sharding.md)。 **激活检查点** + 通过对一组模块使用激活检查点功能，可以提高内存效率。分组的模块越多，内存使用效率就越高。在对层的顺序模块执行检查点操作时，`smp.set_activation_checkpointing` 函数的 `strategy` 参数会将层分组在一起以进行检查点操作。例如，将两个或多个层分组在一起执行检查点操作，比逐个层执行检查点操作更节省内存，但代价是需要额外的计算时间来减少内存使用量。 有关更多信息，请参阅 [激活检查点](model-parallel-extended-features-pytorch-activation-checkpointing.md)。 **张量并行性** + 张量并行度应为二的乘方（2、4、8、...、2 n），其中最大度数必须等于每个节点的 GPUs 数量。例如，如果您使用带有 8 的节点 GPUs，则张量并行度的可能数字为 2、4 和 8。对于张量并行度，我们不建议使用随意的数字（例如 3、5、6 和 7）。当您使用多个节点时，张量并行度配置错误可能会导致需要跨节点来运行张量并行性；这会显著增加跨节点激活通信的开销，并会导致计算成本高昂。 有关更多信息，请参阅 [张量并行性](model-parallel-extended-features-pytorch-tensor-parallelism.md)。 **跨节点的流水线并行性** + 您可以在单个节点内以及跨多个节点来运行管道并行性。当您将管道并行性与张量并行性结合使用时，我们建议在多个节点上运行管道并行性，并将张量并行性保持在单个节点中。 + 管道并行性带有以下三个旋钮：`microbatches`、`active_microbatches` 和 `prescaled_batch`。 + 当您将张量并行性与管道并行性结合使用时，我们建议激活 `prescaled_batch`，这样可以增加每个模型并行组的批次大小，从而实现高效的管道传输。激活 `prescaled_batch` 后，在训练脚本中设置的批次大小，将变为针对每个没有 `prescaled_batch` 的秩设置的批次大小的 `tp_size` 倍。 + 增加 `microbatches` 的数量有助于实现高效的管道和更好的性能。请注意，有效微批次大小等于批次大小除以微批次数量。如果您在保持批次大小不变的同时增加微批次数量，则每个微批次处理的样本数会减少。 + `active_microbatches` 的数量是在管道处理期间同时处理的最大微批次数量。对于正在处理的每个活跃微批次，其激活和梯度都会占用 GPU 内存。因此，增加 `active_microbatches` 会占用更多的 GPU 内存。 + 如果 GPU 和 GPU 内存都未得到充分利用，则在管道处理中增加 `active_microbatches` 可以实现更好的并行化。 + 有关如何将张量并行性与管道并行性结合使用的更多信息，请参阅[张量并行性与管道并行性相结合](model-parallel-extended-features-pytorch-tensor-parallelism-examples.md#model-parallel-extended-features-pytorch-tensor-and-pipeline-parallelism)。 + 要查找上述参数的描述，请参阅 *SageMaker Python SDK 文档`smdistributed`*中的[参数](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smd_model_parallel_general.html#parameters-for-smdistributed)。 **将激活任务分载到 CPU** + 确保此操作与激活检查点和管道并行性结合使用。为确保在后台执行分载和预加载，请在微批次参数中指定大于 1 的值。 + 对激活进行分载时，您也许能够增加 `active_microbatches`，并且有时需要与微批次总数相匹配。这取决于对哪些模块执行了检查点操作以及模型是如何分区的。 有关更多信息，请参阅 [激活分载](model-parallel-extended-features-pytorch-activation-offloading.md)。 ### 参考配置 SageMaker 模型并行度训练团队根据对 GPT-2 模型的实验，提供了以下参考点，序列长度为 512，词汇量为 50,000。 | 模型参数的数量 | 实例类型 | 管道并行性 | 张量并行性 | 优化器状态分片 | 激活检查点 | 预缩放批次 | 批次大小 | | --- | --- | --- | --- | --- | --- | --- | --- | | 100 亿 | 16 ml.p4d.24xlarge | 1 | 4 | True | 每个转换器层 | True | batch\$1size=40 | | 300 亿 | 16 ml.p4d.24xlarge | 1 | 8 | True | 每个转换器层 | True | batch\$1size=32 | | 600 亿 | 32 ml.p4d.24xlarge | 2 | 8 | True | 每个转换器层 | True | batch\$1size=56, microbatches=4, active\$1microbatches=2 | 您可以根据上述配置进行推断，以估算模型配置的 GPU 内存使用量。例如，如果您增加了 100 亿参数模型的序列长度或将模型的大小增加到 200 亿，则可能需要先减小批次大小。如果模型仍然不合适，请尝试提高张量并行度。 ## 修改训练脚本 + 在训练脚本中使用 SageMaker 模型并行库的功能之前，请先查看[SageMaker 分布式模型并行库配置提示和陷阱](model-parallel-customize-tips-pitfalls.md)。 + 要更快地启动训练作业，请使用 [SageMaker AI 本地模式](https://sagemaker.readthedocs.io/en/v2.199.0/overview.html?highlight=local%20mode#local-mode)。这可以帮助您在 SageMaker 笔记本实例上快速在本地运行训练作业。根据运行 SageMaker 笔记本实例的 ML 实例的规模，您可能需要通过更改模型配置（例如隐藏宽度、变压器层数和注意点）来调整模型的大小。在使用大型集群训练完整模型之前，请验证缩减模型在笔记本实例上是否运行良好。 ## 使用 SageMaker AI 控制台和 Amazon 监控和记录训练 Job CloudWatch 要监控 CPU 内存利用率、GPU 内存利用率和 GPU 利用率等系统级指标，请使用通过 [SageMaker AI 控制台](https://console.aws.amazon.com/sagemaker/)提供的可视化效果。 1. 在左侧导航窗格中，选择**训练**。 1. 选择**训练作业**。 1. 在主窗格中，选择要查看更多详细信息的训练作业名称。 1. 浏览主窗格并查找**监控**部分以查看自动生成的可视化对象。 1. 要查看训练作业日志，请在**监控**部分选择**查看日志**。您可以在中访问训练作业的分布式训练作业日志 CloudWatch。如果您启动了多节点分布式训练，则应该会看到多个带有 **algo-n-1234567890** 格式标签的日志流。**algo-1** 日志流跟踪来自主（第 0 个）节点的训练日志。 有关更多信息，请参阅 [用于 CloudWatch 监控和分析训练作业的 Amazon 指标](training-metrics.md)。 ## Permissions 要使用模型并行度或[SageMaker 分布式训练示例笔记本运行训练](https://sagemaker-examples.readthedocs.io/en/latest/training/distributed_training/index.html)作业，请确保您在 IAM 角色中拥有正确的权限，例如： SageMaker + 要[FSx 用于光泽](https://aws.amazon.com/fsx/)，请添加[https://console.aws.amazon.com/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonFSxFullAccess](https://console.aws.amazon.com/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonFSxFullAccess)。 + 要使用 Amazon S3 作为数据通道，请添加 [https://console.aws.amazon.com/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonS3FullAccess](https://console.aws.amazon.com/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonS3FullAccess)。 + 要使用 Docker、构建自己的容器并将其推送到 Amazon ECR，请添加 [https://console.aws.amazon.com/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonEC2ContainerRegistryFullAccess](https://console.aws.amazon.com/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonEC2ContainerRegistryFullAccess)。 + 要获得使用整套 SageMaker 人工智能功能的完全访问权限，请添加[https://console.aws.amazon.com/iam/home#/policies/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonSageMakerFullAccess](https://console.aws.amazon.com/iam/home#/policies/iam/home#/policies/arn%3Aaws%3Aiam%3A%3Aaws%3Apolicy%2FAmazonSageMakerFullAccess)。 # SageMaker 分布式模型并行库配置提示和陷阱 在使用 Amazon A SageMaker I 的模型并行度库之前，请查看以下提示和陷阱。此列表包括适用于各个框架的提示。有关 TensorFlow PyTorch具体提示，请分别参见[修改 TensorFlow 训练脚本](model-parallel-customize-training-script-tf.md)和[修改 PyTorch 训练脚本](model-parallel-customize-training-script-pt.md)。 ## 批次大小和微批次数量 + 当批次大小增加时，库的效率最高。对于模型可以放在单个设备中、但只能用小批次训练的使用场景，在集成库后，可以而且应该增加批次大小。模型并行性可以为大型模型节省内存，使您能够使用以前无法放入内存的批次大小进行训练。 + 选择太小或太大的微批次数量会降低性能。该库在每个设备中按顺序执行各个微批次，因此微批次大小（批次大小除以微批次数）必须足够大，才能充分利用每个 GPU。同时，管道效率会随着微批次数量的增加而提高，因此保持适当的平衡非常重要。通常，好的做法是首先尝试 2 或 4 个微批次，将批次大小增加到内存限制，然后尝试更大的批次大小和微批次数量。随着微批次数量的增加，如果使用交错管道，更大的批次大小可能会变得可行。 + 您的批次大小必须始终可以被微批次数量整除。请注意，根据数据集的大小，有时每个纪元的最后一个批次的大小可能比其余纪元时小，并且这个较小的批次也需要能够被微批次数量整除。如果不是，则可以在`tf.Dataset.batch()`调用`drop_remainder=True`中设置（in TensorFlow），或者`drop_last=True`在`DataLoader`（in PyTorch）中设置，这样就不会使用最后一个小批量。如果您为数据管道使用不同的 API，则只要最后一批次无法被微批次数量整除，您就需要手动跳过最后一个批次。 ## 管理分区 + 如果您使用手动分区，请注意模型中多个操作和模块所使用的参数，例如转换器架构中的嵌入表。为了保证正确性，共享相同参数的模块必须放置在同一个设备中。使用自动分区时，库会自动强制执行此约束。 ## 数据准备 + 如果模型接受多个输入，请确保使用 `smp.dp_rank()`，在数据管道中将随机操作（例如随机排序）设置为种子。如果要在数据并行设备上确定性地对数据集进行分片，请确保分片按照 `smp.dp_rank()` 编制索引。这是为了确保在构成模型分区的所有排序上看到的数据顺序是一致的。 ## 从 `smp.DistributedModel` 返回张量 + 从 (for TensorFlow) 或 `smp.DistributedModel.call` `smp.DistributedModel.forward` (for PyTorch) 函数返回的任何张量都将从计算该特定张量的等级广播到所有其他等级。因此，不应返回调用和转发方法（例如，中间激活）之外不需要的任何张量，因为这会导致不必要的通信和内存开销，并损害性能。 ## `@smp.step` 修饰器 + 如果 `smp.step` 修饰的函数的张量参数没有批次维度，则调用 `smp.step` 时必须在 `non_split_inputs` 列表中提供参数名称。这可以防止库尝试将张量拆分为微批次。有关更多信息，请参阅 API 文档中的 [https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_common_api.html](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_common_api.html)。 ## 延迟参数初始化 对于参数超过 1000 亿的超大型模型，通过 CPU 内存进行权重初始化可能会导致 out-of-memory错误。为了解决这个问题，库提供了 `smp.delay_param_initialization` 上下文管理器。在第一次执行 `smp.step` 修饰的函数时，这会将参数的物理分配延迟到参数移动到 GPU 中时。这样可以避免在训练初始化期间不必要地使用 CPU 内存。在创建模型对象时使用上下文管理器，如以下代码所示。 ``` with smp.delay_param_initialization(enabled=True): model = MyModel() ``` ## 的张量并行度 PyTorch + 如果您使用种子来获得确定性结果，请根据 `smp.dp_rank()` 设置种子（例如，`torch.manual_seed(42 + smp.dp_rank())`）。如果不这样做，`nn.Parameter` 的不同分区将以相同方式初始化，从而影响收敛性。 + SageMaker的模型并行度库使用 NCCL 来实现分发模块所需的集合。特别是对于较小的模型，如果同时在 GPU 上计划了太多 NCCL 调用，则由于 NCCL 占用的额外空间，内存使用量可能会增加。为了抵消这种情况，`smp` 会限制 NCCL 的调用，使在热议给定时间的正在执行的 NCCL 操作数量小于或等于给定限制。默认限制为 8，但可以使用环境变量 `SMP_NCCL_THROTTLE_LIMIT` 进行调整。如果您在使用张量并行性时观察到的内存使用量超出预期，则可以尝试降低此限制。但是，选择过小的限制可能会导致吞吐量损失。要完全禁用节流，您可以设置 `SMP_NCCL_THROTTLE_LIMIT=-1`。 + 以下恒等式在张量并行度为 1 时成立，但当张量并行度大于 1 时不成立：`smp.mp_size() * smp.dp_size() == smp.size()`。这是因为张量并行组既是模型并行性组的一部分，也是数据并行性组的一部分。如果您的代码已有对 `mp_rank`、`mp_size`、`MP_GROUP` 等的引用，并且您只想使用管道并行组，则可能需要将这些引用替换为 `smp.pp_size()`。以下恒等式始终是正确的： + `smp.mp_size() * smp.rdp_size() == smp.size()` + `smp.pp_size() * smp.dp_size() == smp.size()` + `smp.pp_size() * smp.tp_size() * smp.rdp_size() == smp.size()` + 由于 `smp.DistributedModel` 包装器在启用张量并行性时会修改模型参数，因此应在调用 `smp.DistributedModel` 之后使用分布式参数创建优化器。例如，以下内容不起作用： ``` ## WRONG model = MyModel() optimizer = SomeOptimizer(model.parameters()) model = smp.DistributedModel(model)  # optimizer now has outdated parameters!  ``` 而是应该改为使用参数 `smp.DistributedModel` 创建优化器，如下所示： ``` ## CORRECT model = smp.DistributedModel(MyModel()) optimizer = SomeOptimizer(model.optimizers()) ``` + 当通过张量并行性将模块替换为分布式的对应模块时，分布式模块不会从原始模块继承其权重，而是初始化新的权重。举例而言，这意味着如果需要在特定调用中初始化权重（例如，通过 `load_state_dict` 调用），则需要在 `smp.DistributedModel` 调用之后，在进行了模块分布后进行初始化。 + 直接访问分布式模块的参数时，请注意，权重的配置与原始模块不同。例如，  ``` with smp.tensor_parallelism():     linear = nn.Linear(60, 60) # will pass assert tuple(linear.weight.shape) == (60, 60) distributed_linear = smp.DistributedModel(linear) # will fail. the number of input channels will have been divided by smp.tp_size() assert tuple(distributed_linear.module.weight.shape) == (60, 60) ``` + 对于张量并行度，强烈建议使用 `torch.utils.data.distributed.DistributedSampler`。这样可以确保每个数据并行秩接收相同数量的数据样本，从而防止因不同 `dp_rank` 采取的不同步骤数而可能导致的挂起。 + 如果您使用`DistributedDataParallel`类 PyTorch的 `join` API 来处理不同数据 parallel 等级具有不同批次数的情况，则仍需要确保相同的等级`TP_GROUP`具有相同的批次数；否则分布式执行模块中使用的通信集合可能会挂起。只要使用 `join` API，处于不同 `TP_GROUP` 的秩可以有不同的批次数。 + 如果要对使用张量并行性的模型执行检查点操作，请考虑以下几点： + 在使用张量并行性时，为了避免在保存和加载模型时出现停滞和争用情况，请确保在缩减数据并行秩内，从以下模型和优化器状态中调用相应的函数。 + 如果要转换现有的管道并行脚本并为脚本启用张量并行，请确保修改用于 `if smp.rdp_rank() == 0` 块的保存和加载的任何 `if smp.dp_rank() == 0` 块。否则，这可能会导致您的训练作业停滞。 有关对使用张量并行度的模型执行检查点操作的更多信息，请参阅[对分布式模型执行检查点操作](distributed-model-parallel-checkpointing-and-finetuning.md#distributed-model-parallel-checkpoint)。 # 模型并行故障排除 在遇到错误时，您可以根据以下列表尝试对训练作业进行故障排除。如果问题仍存在，请联系 [AWS Support](https://aws.amazon.com/premiumsupport)。 **Topics** + [在 SageMaker 模型并行度库中使用 SageMaker 调试器的注意事项](#distributed-ts-model-parallel-debugger) + [保存检查点](#distributed-ts-model-parallel-checkpoints) + [使用模型并行进行收敛和 TensorFlow](#distributed-ts-model-parallel-tf-convergence) + [分布式训练作业停顿或崩溃](#distributed-ts-model-parallel-training-issues) + [收到 Training Job 的 NCCL 错误 PyTorch](#distributed-ts-model-parallel-nccl-error) + [接受 `RecursionError` Training PyTorch Job](#distributed-ts-model-parallel-super-forward-not-supported) ## 在 SageMaker 模型并行度库中使用 SageMaker 调试器的注意事项 SageMaker 调试器不适用于 SageMaker 模型并行度库。默认情况下，所有任务 SageMaker TensorFlow 和 PyTorch 训练作业均启用调试器，您可能会看到如下所示的错误： ``` FileNotFoundError: [Errno 2] No such file or directory: '/opt/ml/checkpoints/metadata.json.sagemaker-uploading ``` 要修复此问题，请在创建框架时 `estimator` 通过传递 `debugger_hook_config=False` 来禁用 Debugger，如下例所示。 ``` bucket=sagemaker.Session().default_bucket() base_job_name="sagemaker-checkpoint-test" checkpoint_in_bucket="checkpoints" # The S3 URI to store the checkpoints checkpoint_s3_bucket="s3://{}/{}/{}".format(bucket, base_job_name, checkpoint_in_bucket) estimator = TensorFlow( ... distribution={"smdistributed": {"modelparallel": { "enabled": True }}}, checkpoint_s3_uri=checkpoint_s3_bucket, checkpoint_local_path="/opt/ml/checkpoints", debugger_hook_config=False ) ``` ## 保存检查点 在 SageMaker AI 上保存大型模型的检查点时，您可能会遇到以下错误： ``` InternalServerError: We encountered an internal error. Please try again ``` 这可能是由于训练期间将本地检查点上传到 Amazon S3 时的 A SageMaker I 限制所致。要在 SageMaker AI 中禁用检查点，请使用以下示例显式上传检查点。 如果您遇到上述错误，请不要在 SageMaker `estimator`调`checkpoint_s3_uri`用中使用。在为较大的模型保存检查点时，建议将检查点保存到自定义目录中，然后将这些检查点传递给帮助程序函数（作为 `local_path` 参数)。 ``` import os def aws_s3_sync(source, destination): """aws s3 sync in quiet mode and time profile""" import time, subprocess cmd = ["aws", "s3", "sync", "--quiet", source, destination] print(f"Syncing files from {source} to {destination}") start_time = time.time() p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) p.wait() end_time = time.time() print("Time Taken to Sync: ", (end_time-start_time)) return def sync_local_checkpoints_to_s3(local_path="/opt/ml/checkpoints", s3_uri=os.path.dirname(os.path.dirname(os.getenv('SM_MODULE_DIR', '')))+'/checkpoints'): """ sample function to sync checkpoints from local path to s3 """ import boto3 #check if local path exists if not os.path.exists(local_path): raise RuntimeError("Provided local path {local_path} does not exist. Please check") #check if s3 bucket exists s3 = boto3.resource('s3') if not s3_uri.startswith("s3://"): raise ValueError(f"Provided s3 uri {s3_uri} is not valid.") s3_bucket = s3_uri.replace('s3://','').split('/')[0] print(f"S3 Bucket: {s3_bucket}") try: s3.meta.client.head_bucket(Bucket=s3_bucket) except Exception as e: raise e aws_s3_sync(local_path, s3_uri) return def sync_s3_checkpoints_to_local(local_path="/opt/ml/checkpoints", s3_uri=os.path.dirname(os.path.dirname(os.getenv('SM_MODULE_DIR', '')))+'/checkpoints'): """ sample function to sync checkpoints from s3 to local path """ import boto3 #try to create local path if it does not exist if not os.path.exists(local_path): print(f"Provided local path {local_path} does not exist. Creating...") try: os.makedirs(local_path) except Exception as e: raise RuntimeError(f"Failed to create {local_path}") #check if s3 bucket exists s3 = boto3.resource('s3') if not s3_uri.startswith("s3://"): raise ValueError(f"Provided s3 uri {s3_uri} is not valid.") s3_bucket = s3_uri.replace('s3://','').split('/')[0] print(f"S3 Bucket: {s3_bucket}") try: s3.meta.client.head_bucket(Bucket=s3_bucket) except Exception as e: raise e aws_s3_sync(s3_uri, local_path) return ``` 帮助程序函数的用法： ``` #base_s3_uri - user input s3 uri or save to model directory (default) #curr_host - to save checkpoints of current host #iteration - current step/epoch during which checkpoint is saved # save checkpoints on every node using local_rank if smp.local_rank() == 0: base_s3_uri = os.path.dirname(os.path.dirname(os.getenv('SM_MODULE_DIR', ''))) curr_host = os.environ['SM_CURRENT_HOST'] full_s3_uri = f'{base_s3_uri}/checkpoints/{curr_host}/{iteration}' sync_local_checkpoints_to_s3(local_path=checkpoint_dir, s3_uri=full_s3_uri) ``` ## 使用模型并行进行收敛和 TensorFlow 当您使用 SageMaker AI 多节点训练 TensorFlow 和模型并行度库时，损失可能无法按预期收集，因为每个节点上训练输入文件的顺序可能不同。这可能会导致同一模型并行组中处理不同输入文件有不同的排名，从而导致不一致的情况。为防止出现这种情况，请确保输入文件在转换为 TensorFlow 数据集之前，所有等级的顺序都相同。要做到这一点，一种方法是在输入脚本中对输入文件名称排序。 ## 分布式训练作业停顿或崩溃 如果您的训练作业出现了停顿、崩溃或没有响应问题，请阅读以下故障排除项目以确定造成问题的原因。[如果您需要任何进一步的支持，请通过AWS Support 联系 SageMaker 分布式培训团队。](https://aws.amazon.com/premiumsupport) + 如果您看见**分布式训练作业在 NCCL 初始化步骤停顿**，请考虑以下几种情况： + 如果您使用的是启用 EFA 的实例之一（`ml.p4d` 或 `ml.p3dn` 实例）与自定义 VPC 及其子网，请确保使用的安全组在所有端口上具有出站和入站连接。作为一项单独的规则（对于互联网访问），您通常还需要与任意 IP 的出站连接。要查找有关如何为 EFA 通信添加入站和出站规则的说明，请参阅[SageMaker 初始化期间 AI 分布式训练作业停滞](distributed-troubleshooting-data-parallel.md#distributed-ts-data-parallel-efa-sg)。 + 如果您在**对整个模型执行检查点操作时遇到分布式训练作业停顿的问题**，这可能是因为对于具有 `rdp_rank()==0`（使用张量并行性时）或 `dp_rank()==0`（使用管道并行性时）的所有排名，没有在其上进行对模型或优化程序的 `state_dict()` 调用。这些排名需要进行通信来构造所要保存的检查点。在启用了 `shard_optimizer_state` 时，对部分优化程序执行检查点操作时，也可能会出现类似的停顿问题。 有关对具有模型并行性的模型执行检查点操作的更多信息，请参阅[保存和加载的一般说明](https://sagemaker.readthedocs.io/en/v2.199.0/api/training/smp_versions/latest/smd_model_parallel_pytorch.html#general-instruction-for-saving-and-loading)和[对分布式 PyTorch 模型执行检查点操作（适用于 v1.6.0 和 v1.9.0 之间的 SageMaker 模型并行度库）](distributed-model-parallel-checkpointing-and-finetuning.md#model-parallel-extended-features-pytorch-saving-loading-checkpoints)。 + 如果训练作业崩溃并出现 **CUDA 内存不足错误**，这意味着需要调整分布式训练配置以适应 GPU 集群上的模型。有关更多信息和最佳实践，请参阅[针对给定模型设置合适的配置](model-parallel-best-practices.md#model-parallel-best-practices-configuration)。 + 如果训练作业因**无法纠正的 [ECC 错误](https://docs.nvidia.com/deploy/a100-gpu-mem-error-mgmt/index.html)**而崩溃，则表示集群 GPUs 中的一个已失效。如果您需要技术支持，请向 AWS 团队提供作业 ARN，如果可能，请从某个检查点重新启动训练作业。 + 在极少数情况下，以前可以正常工作但已接近 GPU 内存限制的作业配置，以后可能会在其他集群上由于 **CUDA 内存不足错误**而失败。这可能是因为 ECC 错误而导致某些 GPU 的可用内存低于正常情况。 + 运行未在节点中使用全部 GPUs 内容的多节点作业时，可能会发生@@ **网络超时崩溃**。要解决这个问题，请确保将`processes_per_host`参数设置为每个实例中的数量，从而 GPUs 在节点 GPUs 上使用 all。例如，对于 `ml.p3.16xlarge`、`ml.p3dn.24xlarge` 和 `ml.p4d.24xlarge` 实例，此项为 `processes_per_host=8`。 + 如果您发现您的训练作业在数据下载阶段需要很长时间，请确保您为该 SageMaker `Estimator`课程提供的 Amazon S3 路径`checkpoint_s3_uri`对于当前的训练作业是唯一的。如果在同时运行的多个训练作业中重复使用此路径，则所有这些检查点都将在同一个 Amazon S3 路径中上传和下载，这可能会显著增加检查点加载时间。 + 处理大型数据和模型时使用 FSx Lustre。 + 如果您的数据集很大，并且需要很长时间才能获取，我们建议您将数据集保留在 [Lu](https://aws.amazon.com/fsx/lustre/) stre 中FSx 。 + 当训练模型的参数超过 100 亿时，我们建议使用 fo FSx r Lustre 进行检查点检查。 + 创建文件系统后，请务必等待状态变为**可用**，然后再使用它来启动训练作业。 ## 收到 Training Job 的 NCCL 错误 PyTorch 如果您遇到以下错误，这可能是由于进程耗尽 GPU 内存所造成。 ``` NCCL error in: ../torch/lib/c10d/ProcessGroupNCCL.cpp:825, unhandled system error, NCCL version 2.7.8 ncclSystemError: System call (socket, malloc, munmap, etc) failed. ``` 您可以通过减少批处理大小或 `active_microbatches` 来解决这个问题。如果自动分区不能得到均衡的分区，则您可能需要考虑手动分区。有关更多信息，请参阅 [跨节点的流水线并行性](model-parallel-best-practices.md#model-parallel-best-practices-configuration-pipeline-across-nodes)。 ## 接受 `RecursionError` Training PyTorch Job 该库不支持在模块的前向调用内部调用 `super.forward()`。如果您使用 `super.forward()`，则可能会收到以下错误消息。 ``` RecursionError: maximum recursion depth exceeded ``` 要修复错误，您不应调用 `super.forward()`，而是应调用 `super()._orig_forward()`。