使用 SageMaker 模型并行度库 v2 - Amazon SageMaker
第 1 步：调整 PyTorch FSDP训练脚本步骤 2：启动训练作业

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

使用 SageMaker 模型并行度库 v2

在本页上，您将学习如何使用 SageMaker 模型并行度库 v2，APIs并开始在训练平台或集群上运行 PyTorch 完全分片的数据并行 (FSDP) SageMaker 训练作业。 SageMaker HyperPod

使用 SMP v2 运行 PyTorch 训练作业有多种场景。

  1. 要进行 SageMaker 训练，请使用 PyTorch v2.0.1 及更高版本中预先构建的 SageMaker 框架容器之一，该容器已预先打包在 v2 中。SMP

  2. 使用 SMP v2 二进制文件设置 Conda 环境，以便在集群上运行分布式训练工作负载。 SageMaker HyperPod

  3. 扩展适用于 PyTorch v2.0.1 及更高版本的预构建 SageMaker 框架容器，为您的用例安装任何其他功能要求。要了解如何扩展预先构建的容器，请参阅扩展预构建容器。

  4. 您也可以自带 Docker 容器，使用培训工具包手动设置所有 SageMakerSageMaker 训练环境并安装 SMP v2 二进制文件。由于依赖关系的复杂性，这是最不推荐的选项。要了解如何运行自己的 Docker 容器，请参阅调整自己的训练容器。

本入门指南涵盖了前两个场景。

第 1 步：调整 PyTorch FSDP训练脚本

要激活和配置 SMP v2 库，请先在脚本顶部导入和添加torch.sagemaker.init()模块。该模块采用你要准备SMPv2 核心功能配置参数的SMP配置字典步骤 2：启动训练作业。此外，要使用 SMP v2 提供的各种核心功能，您可能需要再做一些更改以适应您的训练脚本。有关调整训练脚本以使用 SMP v2 核心功能的更多详细说明，请参阅。 SageMaker 模型并行度库 v2 的核心功能

SageMaker Training

在训练脚本中，添加以下两行代码，这是开始使用 SMP v2 进行训练的最低要求。在中步骤 2：启动训练作业，您将通过 SageMaker PyTorch估算器类的distribution参数设置带有SMP配置字典的估算器类的对象。

import torch.sagemaker as tsm
tsm.init()
注意

您也可以直接将的配置字典传递SMPv2 核心功能配置参数给torch.sagemaker.init()模块。但是，传递给 PyTorch 估算器的参数具有优先级，并会覆盖为模块指定的参数。步骤 2：启动训练作业 torch.sagemaker.init()

SageMaker HyperPod

在训练脚本中，添加以下两行代码。在中步骤 2：启动训练作业，您将以JSON格式设置SMP配置smp_config.json文件，然后将其上传到与您的 SageMaker HyperPod 集群映射的存储或文件系统。我们建议您将配置文件保存在上传训练脚本的同一目录下。

import torch.sagemaker as tsm
tsm.init("/dir_to_training_files/smp_config.json")
注意

您也可以直接将的配置字典传递SMPv2 核心功能配置参数到torch.sagemaker.init()模块中。

步骤 2：启动训练作业

学习如何配置SMP分发选项，以启动具有SMP核心功能的 PyTorch FSDP训练作业。

SageMaker Training

在 Pyth SageMaker on 中设置PyTorch 框架估算器类的训练作业启动器对象时SDK，请按如下SMPv2 核心功能配置参数方式通过distribution参数进行配置。

注意

从 v2.200 SDK 开始，SMPv2 的distribution配置已集成到 SageMaker Python 中。请务必使用 SageMaker Python 版本 SDK 2.200 或更高版本。

注意

在 SMP v2 中，你应该torch_distributed为 SageMaker PyTorch估计器的distribution参数配置smdistributed为。使用torch_distributed， SageMaker 运行torchrun，这是 Distributed 的默认多节点作业启动器PyTorch 。

from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    framework_version=2.2.0,
    py_version="310"
    # image_uri="<smp-docker-image-uri>" # For using prior versions, specify the SMP image URI directly.
    entry_point="your-training-script.py", # Pass the training script you adapted with SMP from Step 1.
    ... # Configure other required and optional parameters
    distribution={
        "torch_distributed": { "enabled": True },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "hybrid_shard_degree": Integer,
                    "sm_activation_offloading": Boolean,
                    "activation_loading_horizon": Integer,
                    "fsdp_cache_flush_warnings": Boolean,
                    "allow_empty_shards": Boolean,
                    "tensor_parallel_degree": Integer,
                    "expert_parallel_degree": Integer,
                    "random_seed": Integer
                }
            }
        }
    }
)
重要

要使用先前版本 PyTorch 或SMP代替最新版本，则需要使用image_uri参数而不是framework_version和py_version对直接指定 SMP Docker 镜像。以下是一个示例 

estimator = PyTorch(
    ...,
    image_uri="658645717510.dkr.ecr.us-west-2.amazonaws.com/smdistributed-modelparallel:2.2.0-gpu-py310-cu121"
)

要查找 SMP Docker 镜像URIs，请参阅支持的框架。

SageMaker HyperPod

在开始之前，请确保是否满足以下先决条件。

  • 安装到您的 HyperPod 集群的 Amazon FSx 共享目录 (/fsx)。

  • Conda 安装在FSx共享目录中。要了解如何安装 Conda，请按照《Conda 用户指南》中在 Linux 上安装中的说明进行操作。

  • cuda11.8或者cuda12.1安装在 HyperPod 集群的头部和计算节点上。

如果所有先决条件都得到满足，请按照以下说明在 HyperPod 集群上使用 SMP v2 启动工作负载。

  1. 准备一个包含字典的smp_config.json文件SMPv2 核心功能配置参数。确保将此JSON文件上传到存储训练脚本的位置，或者上传到您在步骤 1 中为torch.sagemaker.init()模块指定的路径。如果您已经在步骤 1 中将配置字典传递给训练脚本中的torch.sagemaker.init()模块，则可以跳过此步骤。

    // smp_config.json
{
    "hybrid_shard_degree": Integer,
    "sm_activation_offloading": Boolean,
    "activation_loading_horizon": Integer,
    "fsdp_cache_flush_warnings": Boolean,
    "allow_empty_shards": Boolean,
    "tensor_parallel_degree": Integer,
    "expert_parallel_degree": Integer,
    "random_seed": Integer
}

  2. 将smp_config.json文件上传到文件系统中的某个目录。目录路径必须与您在步骤 1 中指定的路径相匹配。如果您已经将配置字典传递给训练脚本中的torch.sagemaker.init()模块，则可以跳过此步骤。

  3. 在集群的计算节点上，使用以下命令启动终端会话。

    sudo su -l ubuntu

  4. 在计算节点上创建 Conda 环境。以下代码是创建 Conda 环境以及安装SMP、SMDDPCUDA、和其他依赖项的示例脚本。

    # Run on compute nodes
SMP_CUDA_VER=<11.8 or 12.1>

source /fsx/<path_to_miniconda>/miniconda3/bin/activate

export ENV_PATH=/fsx/<path to miniconda>/miniconda3/envs/<ENV_NAME>
conda create -p ${ENV_PATH} python=3.10

conda activate ${ENV_PATH}

# Verify aws-cli is installed: Expect something like "aws-cli/2.15.0*"
aws ‐‐version
# Install aws-cli if not already installed
# https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html#cliv2-linux-install

# Install the SMP library
conda install pytorch="2.0.1=sm_py3.10_cuda${SMP_CUDA_VER}*" packaging ‐‐override-channels \
  -c https://sagemaker-distributed-model-parallel.s3.us-west-2.amazonaws.com/smp-2.0.0-pt-2.0.1/2023-12-11/smp-v2/ \
  -c pytorch -c numba/label/dev \
  -c nvidia -c conda-forge

# Install dependencies of the script as below
python -m pip install packaging transformers==4.31.0 accelerate ninja tensorboard h5py datasets \
    && python -m pip install expecttest hypothesis \
    && python -m pip install "flash-attn>=2.0.4" ‐‐no-build-isolation

# Install the SMDDP wheel
SMDDP_WHL="smdistributed_dataparallel-2.0.2-cp310-cp310-linux_x86_64.whl" \
  && wget -q https://smdataparallel.s3.amazonaws.com/binary/pytorch/2.0.1/cu118/2023-12-07/${SMDDP_WHL} \
  && pip install ‐‐force ${SMDDP_WHL} \
  && rm ${SMDDP_WHL}

# cuDNN installation for Transformer Engine installation for CUDA 11.8
# Please download from below link, you need to agree to terms 
# https://developer.nvidia.com/downloads/compute/cudnn/secure/8.9.5/local_installers/11.x/cudnn-linux-x86_64-8.9.5.30_cuda11-archive.tar.xz

tar xf cudnn-linux-x86_64-8.9.5.30_cuda11-archive.tar.xz \
    && rm -rf /usr/local/cuda-$SMP_CUDA_VER/include/cudnn* /usr/local/cuda-$SMP_CUDA_VER/lib/cudnn* \
    && cp ./cudnn-linux-x86_64-8.9.5.30_cuda11-archive/include/* /usr/local/cuda-$SMP_CUDA_VER/include/ \
    && cp ./cudnn-linux-x86_64-8.9.5.30_cuda11-archive/lib/* /usr/local/cuda-$SMP_CUDA_VER/lib/ \
    && rm -rf cudnn-linux-x86_64-8.9.5.30_cuda11-archive.tar.xz \
    && rm -rf cudnn-linux-x86_64-8.9.5.30_cuda11-archive/

# Please download from below link, you need to agree to terms 
# https://developer.download.nvidia.com/compute/cudnn/secure/8.9.7/local_installers/12.x/cudnn-linux-x86_64-8.9.7.29_cuda12-archive.tar.xz \
# cuDNN installation for TransformerEngine installation for cuda12.1
tar xf cudnn-linux-x86_64-8.9.7.29_cuda12-archive.tar.xz \
    && rm -rf /usr/local/cuda-$SMP_CUDA_VER/include/cudnn* /usr/local/cuda-$SMP_CUDA_VER/lib/cudnn* \
    && cp ./cudnn-linux-x86_64-8.9.7.29_cuda12-archive/include/* /usr/local/cuda-$SMP_CUDA_VER/include/ \
    && cp ./cudnn-linux-x86_64-8.9.7.29_cuda12-archive/lib/* /usr/local/cuda-$SMP_CUDA_VER/lib/ \
    && rm -rf cudnn-linux-x86_64-8.9.7.29_cuda12-archive.tar.xz \
    && rm -rf cudnn-linux-x86_64-8.9.7.29_cuda12-archive/
    
# TransformerEngine installation
export CUDA_HOME=/usr/local/cuda-$SMP_CUDA_VER
export CUDNN_PATH=/usr/local/cuda-$SMP_CUDA_VER/lib
export CUDNN_LIBRARY=/usr/local/cuda-$SMP_CUDA_VER/lib
export CUDNN_INCLUDE_DIR=/usr/local/cuda-$SMP_CUDA_VER/include
export PATH=/usr/local/cuda-$SMP_CUDA_VER/bin:$PATH
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-$SMP_CUDA_VER/lib

python -m pip install ‐‐no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v1.0

  5. 运行测试训练作业。

    1. 在共享文件系统 (/fsx) 中，克隆 Awsome Distributed Train GitHub ing 存储库，然后转到该3.test_cases/11.modelparallel文件夹。

      git clone https://github.com/aws-samples/awsome-distributed-training/
cd awsome-distributed-training/3.test_cases/11.modelparallel

    2. 使用以下方式提交sbatch作业。

      conda activate <ENV_PATH>
sbatch -N 16 conda_launch.sh

      如果作业提交成功，则此sbatch命令的输出消息应类似于Submitted batch job ABCDEF。

    3. 检查下当前目录下的日志文件logs/。

      tail -f ./logs/fsdp_smp_ABCDEF.out