使用SageMaker分布式模型并行 API - Amazon SageMaker
AWS 文档中描述的 AWS 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅中国的 AWS 服务入门

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

使用SageMaker分布式模型并行 API

要使用 Amazon SageMaker的分布式模型并行库,您必须为支持的框架之一创建训练脚本,并使用 SageMaker Python 开发工具包启动训练作业。要了解如何将库集成到训练脚本中,请参阅使用 SageMaker的分布式模型并行库修改您的训练脚本。库的 API 文档位于 SageMaker Python 开发工具包中。请参阅 SageMaker的分布式模型并行 API 文档

SageMaker 支持以下训练环境配置。

  1. 您可以使用预构建的 TensorFlow 或 PyTorch容器。建议为模型并行库的新用户使用此选项,在示例 MNIST with PyTorch 1.6 和 Amazon SageMaker 的分布式模型并行库中对此进行了演示。

  2. 您可以扩展SageMaker预构建的容器,以满足预构建的 SageMaker Docker 映像不支持的算法或模型的任何其他功能要求。要查看如何扩展预构建容器的示例,请参阅扩展预构建容器

  3. 您可以使用 SageMaker 训练工具包SageMaker调整自己的 Docker 容器以使用 。有关示例,请参阅调整您自己的训练容器

有关上述列表中的选项 2 和 3,请参阅 扩展或调整包含 的分布式模型并行库SageMaker的 Docker 容器 以了解如何在扩展或自定义的 Docker 容器中安装模型并行库。

在所有情况下,您使用 SageMaker Python 开发工具包 TensorFlow 或 PyTorch 启动作业Estimator,以初始化库并启动训练作业。有关更多信息,请参阅以下部分:使用 SageMaker Python 开发工具包启动训练作业

使用 SageMaker Python 开发工具包启动训练作业

SageMaker Python 开发工具包支持 TensorFlow和 PyTorch模型的托管训练。要使用这些框架之一启动训练作业,您可以定义 TensorFlowEstimatorPyTorchEstimator

TensorFlow 和 PyTorchEstimator对象包含一个 distribution 参数,该参数用于启用和指定参数以初始化 SageMaker的分布式模型并行库。库在内部使用 MPI,因此,要使用模型并行度,必须使用 distribution 参数启用 MPI。

以下是如何使用 库启动新的 PyTorch 训练作业的示例。

sagemaker_session = sagemaker.session.Session(boto_session=session) mpi_options = { "enabled" : True, "processes_per_host" : 8, "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none " } smp_options = { "enabled":True, "parameters": { "microbatches": 4, "placement_strategy": "spread", "pipeline": "interleaved", "optimize": "speed", "partitions": 2, "ddp": True, } } smd_mp_estimator = PyTorch( entry_point="pt_mnist.py", # Pick your train script source_dir='utils', role=role, instance_type='ml.p3.16xlarge', sagemaker_session=sagemaker_session, framework_version='1.6.0', # You must set py_version to py36 py_version='py36', instance_count=1, distribution={ "smdistributed": {"modelparallel": smp_options}, "mpi": mpi_options }, base_job_name="SMD-MP-demo", ) smd_mp_estimator.fit('s3://my_bucket/my_training_data/')

要启用该库"mpi",需要将包含键 "smdistributed" 和 的词典作为 Python 开发工具包中 TensorFlow 和 PyTorch 评估程序构造函数的 distribution 参数传递。对于 "mpi" 键,必须传递一个 dict,其中包含:

  • "enabled" True :使用 MPI 启动训练作业。

  • "processes_per_host":指定 MPI 应在每台主机上启动的进程数。在 中SageMaker,一个主机是一个 Amazon EC2 ml 实例。SageMaker Python 开发工具包跨模型和数据并行度维护进程和 GPUs 之间的一对一映射。这意味着, 在单个独立的 GPU 上SageMaker计划每个进程,并且没有 GPU 包含多个进程。如果您使用的是 PyTorch则必须通过 将每个进程限制为其自己的设备torch.cuda.set_device(smp.local_rank())。要了解更多信息,请参阅“PyTorch”。

    重要

    process_per_host 必须小于每个实例的 GPUs数,通常等于每个实例的 GPUs数。

    例如,如果您使用一个具有4-way模型并行度和2-way数据并行度的实例,则 processes_per_host 应为 2 x 4 = 8。因此,您必须选择一个至少具有 8 个 GPUs的实例,例如 p316xlarge。

    下图说明了如何在 8 个 GPUs 中分配2-way数据并行度和4-way模型并行度:在 4 个 GPUs 中对模型进行分区,并将每个分区添加到 2 GPUs。

  • "custom_mpi_options":使用此键可传递您可能需要的任何自定义 MPI 选项。为避免 Docker 警告导致训练日志被占用,我们建议使用以下标记。

    --mca btl_vader_single_copy_mechanism none

对于 "smdistributed" 键,必须传递一个具有唯一键 的词典"modelparallel"。不支持在同一训练作业"modelparallel"中使用 "dataparallel" 和 。

对于 "modelparallel" 词典,必须传递一个内部词典,它通过设置 来启用模型并行库"enabled": True,以及一个 dict "parameters" 来自定义库。在参数中, "partitions" 键是必需的,它指定请求多少个模型分区。要了解有关您可以在 中使用的值的更多信息parameters,请参阅SageMaker分布式数据并行 API 文档

要使用SageMaker模型并行配置的训练脚本启动训练作业,请使用 Estimator.fit() 函数。您可以使用SageMaker笔记本实例或本地启动SageMaker训练作业。

使用以下资源了解有关将 SageMaker Python 开发工具包与这些框架结合使用的更多信息:

扩展或调整包含 的分布式模型并行库SageMaker的 Docker 容器

要扩展预构建的容器或调整您自己的容器以使用 SageMaker的分布式模型并行库,您必须使用以下 Pytorch 或 TensorFlow GPU 通用框架基本映像之一:

  • 763104351884.dkr.ecr。<region>.amazonaws.com/tensorflow-training:2.4.11-gpu-py37-cu110-ubuntu184

  • 763104351884.dkr.ecr。<region>.amazonaws.com/tensorflow-training:2.3.11-gpu-py37-cu110-ubuntu184

  • 763104351884.dkr.ecr。<region>.amazonaws.com/pytorch-training:1.6.00-gpu-py36-cu110-ubuntu184

  • 763104351884.dkr.ecr。<region>.amazonaws.com/pytorch-training:1.7.11-gpu-py36-cu110-ubuntu184

例如,如果您在 us-east-1 中使用 Pytorch 1.7.1,则 Dockerfile 应包含类似于以下内容的FROM语句:

# SageMaker PyTorch image FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.7.1-gpu-py36-cu110-ubuntu18.04 # Add your dependencies here RUN ... ENV PATH="/opt/ml/code:${PATH}" # this environment variable is used by the SageMaker PyTorch container to determine our user code directory. ENV SAGEMAKER_SUBMIT_DIRECTORY /opt/ml/code # /opt/ml and all subdirectories are utilized by SageMaker, use the /code subdirectory to store your user code. COPY cifar10.py /opt/ml/code/cifar10.py

此外,在定义 PyTorch 或 TensorFlow 时estimator,您必须entry_point为训练脚本指定 。这应该与在 Dockerfile ENV SAGEMAKER_SUBMIT_DIRECTORY 中使用 标识的路径相同。

您必须将此 Docker 容器推送到 Amazon Elastic Container Registry (Amazon ECR),并使用映像 URI (image_uri) 定义评估程序。有关更多信息,请参阅扩展预构建容器

例如,如果您的 Dockerfile 是使用前面的代码块定义的并使用 nametag 要将其推送到 Amazon ECR,您可以定义 PyTorchestimator如下所示。此示例假定您已定义了 smp_optionsmpi_options

smd_mp_estimator = PyTorch( entry_point='/opt/ml/code/pt_mnist.py', # Identify source_dir='utils', role=role, instance_type='ml.p3.16xlarge', sagemaker_session=sagemaker_session, image_uri='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/')