在 PyTorch 训练脚本中使用该SMDDP库 - Amazon SageMaker
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅 中国的 Amazon Web Services 服务入门 (PDF)

本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。

在 PyTorch 训练脚本中使用该SMDDP库

从 SageMaker 分布式数据并行度 (SMDDP) 库 v1.4.0 开始,您可以将该库用作分布式软件包的后端选项。PyTorch 要使用SMDDPAllReduceAllGather集合操作,您只需要在训练脚本的开头导入SMDDP库,并在流程组初始化期间SMDDP将其设置为 PyTorch 分布式模块的后端即可。使用单行后端规范,您可以保持所有原生 PyTorch 分布式模块和整个训练脚本不变。以下代码片段展示了如何使用该SMDDP库作为 PyTorch基于分布式训练包的后端:PyTorch 分布式数据 parallel (DDP)PyTorch 完全分片的数据并行度 (FSDP) 和威震天-DeepSpeed。DeepSpeed

为了 PyTorch DDP或 FSDP

按如下方式初始化进程组。

import torch.distributed as dist import smdistributed.dataparallel.torch.torch_smddp dist.init_process_group(backend="smddp")
注意

(仅适用于 PyTorch DDP作业)smddp后端目前不支持使用创建子进程组。torch.distributed.new_group() API您也不能将smddp后端与其他进程组后端(例如NCCL和)同时使用。Gloo

对于我们的 DeepSpeed 威震天-DeepSpeed

按如下方式初始化进程组。

import deepspeed import smdistributed.dataparallel.torch.torch_smddp deepspeed.init_distributed(dist_backend="smddp")
注意

要在中SMDDPAllGather使用mpirun基于启动器(smdistributedpytorchddpSMDDP使用 SageMaker Python 启动分布式训练作业 SDK,还需要在训练脚本中设置以下环境变量。

export SMDATAPARALLEL_OPTIMIZE_SDP=true

有关编写 PyTorch FSDP训练脚本的一般指导,请参阅 PyTorch文档中的使用完全分片数据并行进行高级模型训练 (FSDP)

有关编写 PyTorch DDP训练脚本的一般指导,请参阅 PyTorch 文档中的分布式并行数据入门

调整完训练脚本后,继续到 SMDDP使用 SageMaker Python 启动分布式训练作业 SDK