使用 SageMaker 模型平行化程式庫 v2

在此頁面上，您將了解如何使用 SageMaker 模型平行化程式庫 v2 API，並開始在 SageMaker 訓練平台或 SageMaker HyperPod 叢集中執行 PyTorch 全碎片資料平行化 (FSDP) 訓練任務。

使用 SMP v2 執行 PyTorch 訓練任務有各種情況。

針對 SageMaker 訓練，請使用其中一個預先建置的 PyTorch v2.0.1 和更新版本的 SageMaker Framework Containers，其已預先封裝 SMP v2。
使用 SMP v2 二進位檔案來設定 Conda 環境，以在 SageMaker HyperPod 叢集上執行分散式訓練工作負載。
擴展適用於 PyTorch v2.0.1 和更新版本的預先建置 SageMaker Framework Containers，為您的使用案例安裝任何其他功能需求。若要了解如何擴展預先建置的容器，請參閱延伸預先建置的容器。
您也可以自備 Docker 容器，並使用 SageMaker Training 工具組手動設定所有 SageMaker Training 環境，並安裝 SMP v2 二進位檔案。由於相依性的複雜性，這是最不建議的選項。若要了解如何執行您自己的 Docker 容器，請參閱調整您自己的訓練容器。

本入門指南涵蓋前兩個案例。

步驟 1：調整 PyTorch FSDP 訓練指令碼

若要啟用和設定 SMP v2 程式庫，請從指令碼頂端匯入和新增 torch.sagemaker.init() 模組開始。本單元採用您將在步驟 2：啟動訓練任務中準備的 SMP 組態字典 SMP v2 核心功能組態參數。此外，若要使用 SMP v2 提供的各種核心功能，您可能需要進行一些變更來調整訓練指令碼。有關調整訓練指令碼以使用 SMP v2 核心功能的詳細指示，請參閱SageMaker 模型平行化程式庫第 2 版的核心功能。

步驟 2：啟動訓練任務

了解如何設定 SMP 分佈選項，以使用 SMP 核心功能啟動 PyTorch FSDP 訓練任務。

SageMaker Training

當您在 SageMaker Python SDK 中設定 PyTorch 架構估算器類別的訓練任務啟動器物件時，請透過 distribution 引數設定 SMP v2 核心功能組態參數，如下所示。

注意

從 v2.200 開始，SMP v2 的 distribution 組態已整合在 SageMaker Python SDK 中。請務必使用 SageMaker Python SDK 2.200 版或更新版本。

注意

在 SMP v2 中，您應該使用 torch_distributed 設定 smdistributed 來作為 SageMaker PyTorch 估算器的 distribution 引數。使用torch_distributed 時，SageMaker AI 會執行 torchrun，這是 PyTorch Distributed 的預設多節點任務啟動器。


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 映像檔 URI，請參閱支援的架構。

SageMaker HyperPod

開始之前，請務必先達成以下必要條件。

Amazon FSx 共用目錄掛載 (/fsx) 到您的 HyperPod 叢集。
Conda 安裝在 FSx 共用目錄中。若要了解如何安裝 Conda，請使用《Conda 使用者指南》中的在 Linux 上安裝中的指示。
cuda11.8 或 cuda12.1 安裝在 HyperPod 叢集的前端和運算節點上。

如果所有先決條件都符合，請繼續在 HyperPod 叢集上使用 SMP v2 啟動工作負載的下列指示。

準備包含 SMP v2 核心功能組態參數字典的 smp_config.json 檔案。請務必將此 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
}
```
將 smp_config.json 檔案上傳至檔案系統中的目錄。目錄路徑必須與您在步驟 1 中指定的路徑相符。如果您已將組態字典傳遞至訓練指令碼中的 torch.sagemaker.init() 模組，您可以略過此步驟。
在叢集的運算節點上，使用以下命令啟動終端機工作階段。
```
sudo su -l ubuntu
```

在運算節點上建立 Conda 環境。下列程式碼是建立 Conda 環境和安裝 SMP、SMDDP、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. 在共用檔案系統 (/fsx) 中，複製 Awsome 分散式訓練 GitHub 儲存庫，然後移至 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
```

您的瀏覽器已停用或無法使用 Javascript。

您必須啟用 Javascript，才能使用 AWS 文件。請參閱您的瀏覽器說明頁以取得說明。

文件慣用形式

支援的架構與 AWS 區域

SMP v2 的核心功能