

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

# 使用 SageMaker 模型并行度库 v2
<a name="model-parallel-use-api-v2"></a>

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

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

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

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

1. 扩展适用于 PyTorch v2.0.1 及更高版本的预构建 SageMaker 框架容器，为您的用例安装任何其他功能要求。要了解如何扩展预构建的容器，请参阅 [扩展预构建容器](prebuilt-containers-extend.md)。

1. 您也可以自带 Docker 容器，使用培训[工具包](https://github.com/aws/sagemaker-training-toolkit)手动设置所有 SageMaker训练环境并安装 SMP v2 二进制文件。SageMaker 由于依赖关系的复杂性，这是最不推荐的选项。要了解如何运行自己的 Docker 容器，请参阅[调整自己的训练容器](https://docs.amazonaws.cn/sagemaker/latest/dg/adapt-training-container.html)。

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

**Topics**
+ [第 1 步：调整你的 PyTorch FSDP 训练脚本](#model-parallel-adapt-pytorch-script-v2)
+ [步骤 2：启动训练作业](#model-parallel-launch-a-training-job-v2)

## 第 1 步：调整你的 PyTorch FSDP 训练脚本
<a name="model-parallel-adapt-pytorch-script-v2"></a>

要激活和配置 SMP v2 库，首先要在脚本顶部导入并添加 `torch.sagemaker.init()` 模块。本模块采用您将在 [步骤 2：启动训练作业](#model-parallel-launch-a-training-job-v2) 中准备的 [SMP v2 核心功能配置参数](distributed-model-parallel-v2-reference.md#distributed-model-parallel-v2-reference-init-config) 的 SMP 配置字典。此外，为了使用 SMP v2 提供的各种核心功能，您可能还需要对训练脚本进行一些修改。有关如何调整训练脚本以使用 SMP v2 核心功能的更详细说明，请参阅 [SageMaker 模型并行度库 v2 的核心功能](model-parallel-core-features-v2.md)。

------
#### [ SageMaker Training ]

在训练脚本中添加以下两行代码，这是开始使用 SMP v2 进行训练的最低要求。在中[步骤 2：启动训练作业](#model-parallel-launch-a-training-job-v2)，你将通过 SageMaker `PyTorch`估计器类的`distribution`参数设置一个带有 SMP 配置字典的估计器类的对象。

```
import torch.sagemaker as tsm
tsm.init()
```

**注意**  
您也可以直接将 [SMP v2 核心功能配置参数](distributed-model-parallel-v2-reference.md#distributed-model-parallel-v2-reference-init-config) 的配置字典传递给 `torch.sagemaker.init()` 模块。但是，传递给 PyTorch 估算器的参数具有优先级，并会覆盖为模块指定的参数。[步骤 2：启动训练作业](#model-parallel-launch-a-training-job-v2) `torch.sagemaker.init()`

------
#### [ SageMaker HyperPod ]

将在训练脚本中添加以下两行代码。在中[步骤 2：启动训练作业](#model-parallel-launch-a-training-job-v2)，您将以 JSON 格式设置一个用于设置 SMP 配置的`smp_config.json`文件，然后将其上传到与您的 SageMaker HyperPod 集群映射的存储或文件系统。我们建议将配置文件保存在上传训练脚本的同一目录下。

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

**注意**  
您也可以直接将 [SMP v2 核心功能配置参数](distributed-model-parallel-v2-reference.md#distributed-model-parallel-v2-reference-init-config) 的配置字典传递给 `torch.sagemaker.init()` 模块。

------

## 步骤 2：启动训练作业
<a name="model-parallel-launch-a-training-job-v2"></a>

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

------
#### [ SageMaker Training ]

在 Pyth SageMaker on SDK 中设置[PyTorch 框架估算](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/sagemaker.pytorch.html)器类的训练作业启动器对象时，请按以下[SMP v2 核心功能配置参数](distributed-model-parallel-v2-reference.md#distributed-model-parallel-v2-reference-init-config)方式通过`distribution`参数进行配置。

**注意**  
从 v2.200 开始，SMP v2 的`distribution`配置已集成在 Pyt SageMaker hon SDK 中。请务必使用 SageMaker Python 软件开发工具包 v2.200 或更高版本。

**注意**  
在 SMP v2 中，应将`smdistributed`估计`torch_distributed`器的`distribution`参数配置为。 SageMaker `PyTorch`使用`torch_distributed`， SageMaker AI 运行`torchrun`，这是 Distributed 的默认多节点任务启动器[PyTorch 。](https://pytorch.org/tutorials/beginner/dist_overview.html)

```
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`参数而不是和对直接指定 SMP Docker 镜像。`framework_version` `py_version`以下是一个示例   

```
estimator = PyTorch(
    ...,
    image_uri="658645717510.dkr.ecr.us-west-2.amazonaws.com/smdistributed-modelparallel:2.2.0-gpu-py310-cu121"
)
```
要查找 SMP Docker 镜像 URIs，请参阅。[支持的框架](distributed-model-parallel-support-v2.md#distributed-model-parallel-supported-frameworks-v2)

------
#### [ SageMaker HyperPod ]

开始之前，请确保满足以下先决条件。
+ 安装到您的 HyperPod 集群的 Amazon FSx 共享目录 (`/fsx`)。
+ Conda 安装在 FSx 共享目录中。要了解如何安装 Conda，请按照*《Conda 用户指南》*中的[在 Linux 上安装](https://docs.conda.io/projects/conda/en/latest/user-guide/install/linux.html)的说明进行操作。
+ `cuda11.8`或者`cuda12.1`安装在 HyperPod 集群的头部和计算节点上。

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

1. 准备一个包含 [SMP v2 核心功能配置参数](distributed-model-parallel-v2-reference.md#distributed-model-parallel-v2-reference-init-config) 字典的 `smp_config.json` 文件。确保将此 JSON 文件上传到存储训练脚本的位置，或上传到[步骤 1](#model-parallel-adapt-pytorch-script-v2) 中指定的 `torch.sagemaker.init()` 模块路径。如果您已经在[步骤 1](#model-parallel-adapt-pytorch-script-v2) 中将配置字典传递给训练脚本中的 `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
   }
   ```

1. 将 `smp_config.json` 文件上传到文件系统中的一个目录。目录路径必须与您在[步骤 1](#model-parallel-adapt-pytorch-script-v2) 中指定的路径一致。如果您已经将配置字典传递给训练脚本中的 `torch.sagemaker.init()` 模块，则可以跳过此步骤。

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

   ```
   sudo su -l ubuntu
   ```

1. 在计算节点上创建 Conda 环境。下面的代码是创建 Conda 环境并安装 SMP、[SMDDP](data-parallel.md)、CUDA 和其他依赖关系的示例脚本。

   ```
   # 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
   ```

1. 运行测试训练作业。

   1. 在共享文件系统 (`/fsx`) 中，克隆 [Awsome Distributed Train GitHub ing 存储库](https://github.com/aws-samples/awsome-distributed-training/)，然后转到该`3.test_cases/11.modelparallel`文件夹。

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

   1. 使用如下 `sbatch` 提交作业。

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

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

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

      ```
      tail -f ./logs/fsdp_smp_ABCDEF.out
      ```

------