步骤 2:使用 SageMaker Python SDK 启动训练作业 - Amazon SageMaker
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅 中国的 Amazon Web Services 服务入门 (PDF)

步骤 2:使用 SageMaker Python SDK 启动训练作业

SageMaker Python SDK 支持采用 ML 框架的托管训练模型,例如 TensorFlow 和 PyTorch。要使用这些框架之一启动训练作业,您需要定义 SageMaker TensorFlow 估算器SageMaker PyTorch 估算器或 SageMaker 通用估算器来使用修改后的训练脚本和模型并行性配置。

使用 SageMaker TensorFlow 和 PyTorch 估算器

TensorFlow 和 PyTorch 估算器类包含 distribution 参数,您可以用来指定配置参数以用于分布式训练框架。SageMaker 模型并行库内部使用 MPI 来实现混合数据和模型的并行性,因此必须对库使用 MPI 选项。

以下 TensorFlow 或 PyTorch 估算器模板显示了如何配置 distribution 参数,以便将 SageMaker 模型并行库与 MPI 结合使用。

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 模型并行性指定一组参数。

  • 对于 "mpi" 键,传递一个包含以下内容的字典:

    • "enabled" – 必需。设置为 True 以使用 MPI 启动分布式训练作业。

    • "processes_per_host" – 必需。指定 MPI 应在每台主机上启动的进程数。在 SageMaker 中,一台主机是一个 Amazon EC2 ML 实例。SageMaker Python SDK 在模型和数据并行性之间,维护进程与 GPU 的一对一映射。这意味着 SageMaker 将每个进程安排在一个单独的 GPU 上,并且没有 GPU 包含多个进程。如果您使用 PyTorch,则必须通过 torch.cuda.set_device(smp.local_rank()) 将每个进程限制在自己的设备上。要了解更多信息,请参阅PyTorch 的自动拆分

      重要

      process_per_host 不可大于每个实例的 GPU 数量,通常等于每个实例的 GPU 数量。

    • "custom_mpi_options"(可选)– 使用此键传递您可能需要的任何自定义 MPI 选项。如果您没有将任何 MPI 自定义选项传递给键,则默认情况下,MPI 选项将设置为以下标志。

      --mca btl_vader_single_copy_mechanism none
      注意

      您无需为键明确指定此默认标志。如果明确指定了该标志,则分布式模型并行训练作业可能会失败并出现以下错误:

      The following MCA parameter has been listed multiple times on the command line: MCA param: btl_vader_single_copy_mechanism MCA parameters can only be listed once on a command line to ensure there is no ambiguity as to its value. Please correct the situation and try again.
      提示

      如果您使用启用 EFA 的实例类型启动训练作业,例如 ml.p4d.24xlargeml.p3dn.24xlarge,请使用以下标志以获得最佳性能:

      -x FI_EFA_USE_DEVICE_RDMA=1 -x FI_PROVIDER=efa -x RDMAV_FORK_SAFE=1

要使用估算器和 SageMaker 模型并行配置的训练脚本启动训练作业,请运行 estimator.fit() 函数。

使用以下资源详细了解在 SageMaker Python SDK 中使用模型并行性功能:

扩展包含 SageMaker 分布式模型并行库的预构建 Docker 容器

要扩展预构建容器并使用 SageMaker 的模型并行性库,您必须使用适用于 PyTorch 或 TensorFlow 的可用 Amazon Deep Learning Containers (DLC) 映像之一。带 CUDA (cuxyz) 的 TensorFlow(2.3.0 及更高版本)和 PyTorch(1.6.0 及更高版本)DLC 映像中均包括 SageMaker 模型并行性库。有关 DLC 映像的完整列表,请参阅 Amazon Deep Learning Containers GitHub 存储库中的深度学习容器映像

提示

建议您使用包含最新 TensorFlow 或 PyTorch 版本的映像来访问 SageMaker 模型并行性库的最新版本。

例如,您的 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 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 Elastic Container Registry (Amazon ECR),并使用映像 URI (image_uri) 来定义 SageMaker 估算器用于训练。有关更多信息,请参阅 扩展预构建容器

完成 Docker 容器托管并检索了容器的映像 URI 后,请按如下所示创建 SageMaker PyTorch 估算器对象。本例假设您已经定义 smp_optionsmpi_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 训练环境和模型并行库时至少必须包含的代码块。

注意

此自定义 Docker 选项将 SageMaker 模型并行库作为二进制文件,仅适用于 PyTorch。

使用 SageMaker 训练工具包和模型并行库创建 Dockerfile
  1. NVIDIA CUDA 基础映像之一开始使用。

    FROM <cuda-cudnn-base-image>
    提示

    官方的 Amazon Deep Learning Container (DLC) 映像基于 NVIDIA CUDA 基础映像而构建。我们建议您查看 Amazon Deep Learning Container for PyTorch 的官方 Dockerfile,查找您需要安装的库版本以及如何进行配置。官方 Dockerfile 提供了完整的内容,已通过基准测试,并由 SageMaker 和 Deep Learning Container 服务团队管理。在提供的链接中,选择您使用的 PyTorch 版本,选择 CUDA (cuxyz) 文件夹,然后选择 .gpu.sagemaker.gpu 或结尾的 Dockerfile。

  2. 要设置分布式训练环境,您需要为通信和网络设备安装软件,例如 Elastic Fabric Adapter (EFA)NVIDIA Collective Communications Library (NCCL)Open MPI。根据所选的 PyTorch 和 CUDA 版本,您必须安装兼容版本的库。

    重要

    由于 SageMaker 模型并行库在后续步骤中需要使用 SageMaker 数据并行库,因此强烈建议您按照 使用 SageMaker 分布式数据并行库创建您自己的 Docker 容器 中的说明正确设置用于 SageMaker 训练环境,以便用于分布式训练。

    有关使用 NCCL 和 Open MPI 设置 EFA 的更多信息,请参阅开始使用 EFA 和 MPI 以及开始使用 EFA 和 NCCL

  3. 添加以下参数以指定适用于 PyTorch 的 SageMaker 分布式训练包的 URL。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
  4. 安装 SageMaker 模型并行库所需的依赖项。

    1. 安装 METIS 库。

      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
    2. 安装 RAPIDS Memory Manager 库。这需要 CMake 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}
  5. 安装 SageMaker 模型并行库。

    RUN pip install --no-cache-dir -U ${SMD_MODEL_PARALLEL_URL}
  6. 安装 SageMaker 数据并行库。

    RUN SMDATAPARALLEL_PT=1 pip install --no-cache-dir ${SMDATAPARALLEL_BINARY}
  7. 安装 sagemaker-training 工具包。该工具包中包含常用的功能,可满足创建与 SageMaker 训练平台和 SageMaker Python SDK 兼容的容器的需求。

    RUN pip install sagemaker-training
  8. 完成 Dockerfile 的创建后,请参阅调整自己的训练容器,了解如何构建 Docker 容器并将其托管在 Amazon ECR 中。

提示

有关创建自定义 Dockerfile 以用于 SageMaker 中的训练的更多一般信息,请参阅使用自己的训练算法