Uso de la biblioteca de paralelismo de modelos de SageMaker v2 - Amazon SageMaker AI
Paso 1: adaptación del script de entrenamiento PyTorch FSDPPaso 2: inicialización de un trabajo de entrenamiento

Uso de la biblioteca de paralelismo de modelos de SageMaker v2

En esta página, aprenderá a utilizar las API de la biblioteca de paralelismo de modelos de SageMaker v2 y empezará a ejecutar un trabajo de entrenamiento de Paralelismo de datos totalmente particionados (FSDP) de PyTorch en la plataforma SageMaker Training o en un clúster de SageMaker HyperPod.

Existen varios escenarios para ejecutar un trabajo de entrenamiento de PyTorch con SMP v2.

  1. Para el entrenamiento de SageMaker, utilice uno de los contenedores SageMaker Framework prediseñados para PyTorch v2.0.1 y versiones posteriores, que vienen preempaquetados con SMP v2.

  2. Utilice el archivo binario SMP v2 para configurar un entorno de Conda para ejecutar una carga de trabajo de entrenamiento distribuida en un clúster de SageMaker HyperPod.

  3. Amplíe los contenedores prediseñados de SageMaker Framework para PyTorch v2.0.1 y versiones posteriores para instalar cualquier requisito funcional adicional para su caso de uso. Para aprender a ampliar un contenedor prediseñado, consulte Ampliar una contenedor precompilado.

  4. También puede utilizar su propio contenedor de Docker y configurar manualmente todo el entorno de SageMaker Training mediante el kit de herramientas de SageMaker Training e instalar el archivo binario de SMP v2. Esta es la opción menos recomendada debido a la complejidad de las dependencias. Para obtener información sobre cómo ejecutar su propio contenedor de Docker, consulte Adaptación del contenedor de entrenamiento propio.

Esta guía de introducción cubre los dos primeros escenarios.

Paso 1: adaptación del script de entrenamiento PyTorch FSDP

Para activar y configurar la biblioteca de SMP v2, comience importando y añadiendo el módulo torch.sagemaker.init() en la parte superior del script. Este módulo utiliza el diccionario de configuración de SMP de Parámetros de configuración de las características esenciales de SMP v2 que preparará en Paso 2: inicialización de un trabajo de entrenamiento. Además, para utilizar las diversas características principales que ofrece SMP v2, es posible que tenga que realizar algunos cambios más para adaptar el script de entrenamiento. Encontrará instrucciones más detalladas sobre cómo adaptar el script de entrenamiento para utilizar las características esenciales de SMP v2 en Características principales de la biblioteca de paralelismo de modelos de SageMaker v2.

SageMaker Training

En el script de entrenamiento, añada las dos líneas de código siguientes, que son el requisito mínimo para empezar a entrenar con SMP v2. En Paso 2: inicialización de un trabajo de entrenamiento configurará un objeto de la clase del estimador de PyTorch de SageMaker con un diccionario de configuración de SMP a través del argumento distribution de la clase de estimador.

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

También puede pasar directamente un diccionario de configuración de la Parámetros de configuración de las características esenciales de SMP v2 al módulo de torch.sagemaker.init(). Sin embargo, los parámetros pasados al estimador PyTorch en Paso 2: inicialización de un trabajo de entrenamiento tienen prioridad y anulan los especificados en el módulo de torch.sagemaker.init().

SageMaker HyperPod

En el script de entrenamiento, agregue las dos líneas de código siguientes. En Paso 2: inicialización de un trabajo de entrenamiento, configurará un archivo smp_config.json para configurar las configuraciones de SMP en formato JSON y lo cargará en un almacenamiento o un sistema de archivos asignado con el clúster de SageMaker HyperPod. Le recomendamos que guarde el archivo de configuración en el mismo directorio en el que carga el script de entrenamiento.

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

También puede pasar directamente un diccionario de configuración de la Parámetros de configuración de las características esenciales de SMP v2 al módulo de torch.sagemaker.init().

Paso 2: inicialización de un trabajo de entrenamiento

Aprenda a configurar las opciones de distribución de SMP para inicializar un trabajo de entrenamiento de PyTorch FSDP con características principales de SMP.

SageMaker Training

Cuando configure un objeto lanzador de tareas de entrenamiento de la clase estimador de marco de PyTorch en el SageMaker Python SDK, configure Parámetros de configuración de las características esenciales de SMP v2 a través del argumento distribution de la siguiente manera.

nota

La configuración de distribution para SMP v2 está integrada en el SageMaker Python SDK a partir de la v2.200. Asegúrese de utilizar el SageMaker Python SDK v2.200 o posteriores.

nota

En SMP v2, debe configurar smdistributed con torch_distributed para el argumento distribution del estimador de SageMaker PyTorch. Con torch_distributed, SageMaker AI ejecuta torchrun, que es el lanzador de trabajos multinodo predeterminado de 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
                }
            }
        }
    }
)
importante

Para usar una de las versiones anteriores de PyTorch o SMP en lugar de la última, debe especificar la imagen de Docker de SMP directamente utilizando el argumento image_uri en lugar del par framework_version y py_version. A continuación se muestra un ejemplo de 

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

Para buscar los URI de imagen de Docker de SMP, consulte Marcos admitidos.

SageMaker HyperPod

Antes de empezar, asegúrese de que se cumplen los siguientes requisitos previos.

  • Un directorio compartido de Amazon FSx montado (/fsx) en el clúster de HyperPod.

  • Conda instalado en el directorio compartido de FSx. Para obtener información sobre cómo instalar Conda, siga las instrucciones de Installing on Linux de la Conda User Guide.

  • cuda11.8 o cuda12.1 instalado en los nodos principal y de cómputo del clúster de HyperPod.

Si se cumplen todos los requisitos previos, siga estas siguientes instrucciones para iniciar una carga de trabajo con SMP v2 en un clúster de HyperPod.

  1. Prepare un archivo smp_config.json que contenga un diccionario de Parámetros de configuración de las características esenciales de SMP v2. Asegúrese de cargar este archivo JSON al lugar donde guarda el script de entrenamiento o la ruta que ha especificado para el módulo torch.sagemaker.init() en el paso 1. Si ya ha pasado el diccionario de configuración al módulo torch.sagemaker.init() del script de entrenamiento del paso 1, puede omitir este paso. 

    // 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. Cargue el archivo smp_config.json en un directorio del sistema de archivos. La ruta del directorio debe coincidir con la ruta que ha especificado en el paso 1. Si ya ha pasado el diccionario de configuración al módulo torch.sagemaker.init() del script de entrenamiento, puede omitir este paso.

  3. En los nodos de computación del clúster, inicie una sesión de terminal con el siguiente comando.

    sudo su -l ubuntu

  4. Cree un entorno de Conda en los nodos de cómputo. El siguiente código es un script de ejemplo para crear un entorno de Conda e instalar SMP, SMDDP, CUDA y otras dependencias.

    # 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. Ejecute un trabajo de entrenamiento de prueba.

    1. En el sistema de archivos compartidos (/fsx), clone el repositorio de GitHub de Awsome Distributed Training y vaya a la carpeta 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. Envíe un trabajo utilizando sbatch de este modo.

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

      Si el envío del trabajo se realiza correctamente, el mensaje de salida de este comando sbatch debería ser similar a Submitted batch job ABCDEF.

    3. Compruebe el archivo de registro en el directorio actual en logs/.

      tail -f ./logs/fsdp_smp_ABCDEF.out