Utilice la biblioteca de paralelismo de SageMaker modelos v2 - Amazon SageMaker AI
Paso 1: Adapta tu PyTorch guion de formación sobre el FSDPPaso 2: inicialización de un trabajo de entrenamiento

Las traducciones son generadas a través de traducción automática. En caso de conflicto entre la traducción y la version original de inglés, prevalecerá la version en inglés.

Utilice la biblioteca de paralelismo de SageMaker modelos v2

En esta página, aprenderás a usar la biblioteca de paralelismo de SageMaker modelos, versión 2, APIs y comenzarás a ejecutar un trabajo de formación de PyTorch Fully Sharded Data Parallel (FSDP) en la plataforma de formación o en un clúster. SageMaker SageMaker HyperPod

Existen varios escenarios para ejecutar un PyTorch trabajo de formación con SMP v2.

  1. Para la SageMaker formación, utilice uno de los contenedores SageMaker Framework prediseñados para la PyTorch versión 2.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 formación distribuida en un clúster. SageMaker HyperPod

  3. Amplíe los contenedores SageMaker Framework prediseñados para la PyTorch versión 2.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 traer su propio contenedor de Docker y configurar manualmente todo el entorno de SageMaker formación mediante el kit de herramientas de SageMaker formación e instalar el archivo binario 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: Adapta tu PyTorch guion de formación sobre el 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 SageMaker modelos 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. EnPaso 2: inicialización de un trabajo de entrenamiento, configurará un objeto de la clase SageMaker PyTorch estimadora con un diccionario de configuración SMP a través del distribution argumento de la clase estimadora.

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 que se pasan al PyTorch estimador Paso 2: inicialización de un trabajo de entrenamiento tienen prioridad y anulan los especificados en el módulo. torch.sagemaker.init()

SageMaker HyperPod

En el script de entrenamiento, agregue las dos líneas de código siguientes. EnPaso 2: inicialización de un trabajo de entrenamiento, configurará un smp_config.json archivo para configurar las configuraciones de SMP en formato JSON y lo cargará en un almacenamiento o un sistema de archivos mapeado con su clúster. 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 iniciar un trabajo de formación sobre el PyTorch FSDP con las funciones principales del SMP.

SageMaker Training

Al configurar un objeto lanzador de tareas de entrenamiento de la clase PyTorch Framework Estimator en el SDK de SageMaker Python, configúrelo Parámetros de configuración de las características esenciales de SMP v2 mediante el distribution argumento de la siguiente manera.

nota

La distribution configuración de SMP v2 está integrada en el SDK de SageMaker Python a partir de la versión 2.200. Asegúrese de usar el SDK de SageMaker Python v2.200 o posterior.

nota

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

Para usar una de las versiones anteriores PyTorch o SMP en lugar de la última, debe especificar la imagen de Docker SMP directamente utilizando el image_uri argumento en lugar del par and. framework_version 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 una imagen de Docker SMP, consulte. URIs Marcos admitidos

SageMaker HyperPod

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

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

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

  • cuda11.8o cuda12.1 instalado en los nodos principal y de procesamiento de su HyperPod clúster.

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

  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 compartido (/fsx), clona el GitHub repositorio de Awsome Distributed Training y ve 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