Utilize a biblioteca de paralelismo de modelos do SageMaker v2 - SageMaker IA da Amazon
Etapa 1: adaptar o script de treinamento do FSDP do PyTorchEtapa 2: iniciar um trabalho de treinamento

As traduções são geradas por tradução automática. Em caso de conflito entre o conteúdo da tradução e da versão original em inglês, a versão em inglês prevalecerá.

Utilize a biblioteca de paralelismo de modelos do SageMaker v2

Nesta página, você aprenderá a usar as APIs da biblioteca de paralelismo de modelos versão 2 do SageMaker e começará a executar um trabalho de treinamento do paralelismo de dados totalmente fragmentados (FSDP) do PyTorch na plataforma de treinamento do SageMaker ou em um cluster do SageMaker HyperPod.

Existem vários cenários para executar um trabalho de treinamento do PyTorch com o SMP v2.

  1. Para o treinamento do SageMaker, use um dos contêineres predefinidos do SageMaker Framework para PyTorch versão 2.0.1 e versões posteriores, que vêm pré-empacotados com o SMP v2.

  2. Use o arquivo binário do SMP v2 para configurar um ambiente Conda para executar uma workload de treinamento distribuída em um cluster do SageMaker HyperPod.

  3. Amplie os contêineres pré-definidos do SageMaker Framework para PyTorch versão 2.0.1 e versões posteriores para instalar quaisquer requisitos funcionais adicionais para o caso de uso. Para saber como ampliar um contêiner pré-definido, consulte Estenda uma imagem de contêiner predefinida.

  4. Você também pode trazer seu próprio contêiner do Docker e configurar manualmente todo o ambiente do SageMaker Training usando o kit de ferramentas do SageMaker Training e instalar o arquivo binário do SMP v2. Essa é a opção menos recomendada devido à complexidade das dependências. Para saber como executar seu próprio contêiner do Docker, consulte Como adaptar o próprio contêiner de treinamento.

Este guia de introdução aborda os dois primeiros cenários.

Etapa 1: adaptar o script de treinamento do FSDP do PyTorch

Para ativar e configurar a biblioteca de SMP v2, comece por importar e adicionar o módulo torch.sagemaker.init() na parte superior do script. Este módulo inclui o dicionário de configuração do SMP em Parâmetros de configuração do atributo principal do SMP v2, que você preparará em Etapa 2: iniciar um trabalho de treinamento. Além disso, para usar os vários atributos principais oferecidos pelo SMP v2, talvez seja necessário fazer mais algumas alterações para adaptar o script de treinamento. São fornecidas instruções mais detalhadas sobre como adaptar o script de treinamento para usar os principais atributos do SMP v2 em Principais características da biblioteca de paralelismo de SageMaker modelos v2.

SageMaker Training

No script de treinamento, adicione as duas linhas de código a seguir, que é o requisito mínimo para começar a treinar com o SMP v2. Em Etapa 2: iniciar um trabalho de treinamento, você configurará um objeto da função estimadora do PyTorch SageMaker com um dicionário de configuração do SMP por meio do argumento distribution da função estimadora.

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

Você também pode transmitir diretamente um dicionário de configuração do módulo Parâmetros de configuração do atributo principal do SMP v2 para o torch.sagemaker.init(). No entanto, os parâmetros transmitidos para o estimador do PyTorch em Etapa 2: iniciar um trabalho de treinamento são prioridade e substituem os definidos para o módulo torch.sagemaker.init().

SageMaker HyperPod

No script de treinamento, adicione as duas linhas de código a seguir. Em Etapa 2: iniciar um trabalho de treinamento, você configurará um arquivo smp_config.json para definir as configurações do SMP no formato JSON e fará o upload em um armazenamento ou sistema de arquivos mapeado com o cluster do SageMaker HyperPod. Recomendamos que você mantenha o arquivo de configuração no mesmo diretório em que fez o upload do script de treinamento.

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

Você também pode transmitir diretamente um dicionário de configuração do módulo Parâmetros de configuração do atributo principal do SMP v2 para o torch.sagemaker.init().

Etapa 2: iniciar um trabalho de treinamento

Aprenda a configurar as opções de distribuição do SMP para iniciar um trabalho de treinamento do FSDP do PyTorch com os principais atributos do SMP.

SageMaker Training

Ao configurar um objeto inicializador de tarefas de treinamento da função do estimador da estrutura do PyTorch no SageMaker Python SDK, faça a configuração do argumento de Parâmetros de configuração do atributo principal do SMP v2 até distribution conforme instruções a seguir.

nota

A configuração distribution do SMP v2 está integrada ao SageMaker Python SDK a partir da versão 2.200. Certifique-se de usar o SageMaker Python SDK versão 2.200 ou posterior.

nota

No SMP v2, você deve configurar smdistributed com torch_distributed para o argumento distribution do estimador do SageMaker PyTorch. Com torch_distributed, o SageMaker AI executa torchrun, que é o inicializador de tarefas padrão de vários nós do 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 uma das versões anteriores do PyTorch ou do SMP em vez da mais recente, você precisa definir a imagem do Docker do SMP diretamente usando o argumento image_uri em vez do par framework_version e py_version. Veja a seguir um exemplo de: 

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

Para encontrar os URIs de imagem do Docker do SMP, consulte Frameworks compatíveis.

SageMaker HyperPod

Antes de começar a usá-lo, certifique-se de que os seguintes pré-requisitos foram atendidos:

  • Um diretório compartilhado do Amazon FSx montado (/fsx) em seu cluster HyperPod.

  • O Conda instalado no diretório compartilhado do FSx. Para saber como instalar o Conda, use as instruções em Instalação no Linux, no Guia do usuário do Conda.

  • cuda11.8 ou cuda12.1 instalado nos nós principais e de computação do seu cluster HyperPod.

Se todos os pré-requisitos forem atendidos, siga as instruções a seguir sobre como iniciar uma workload com o SMP v2 em um cluster HyperPod.

  1. Prepare um arquivo smp_config.json que contenha um dicionário de Parâmetros de configuração do atributo principal do SMP v2. Certifique-se de fazer upload desse arquivo JSON onde você armazenar o script de treinamento ou o caminho que você especificou para o módulo torch.sagemaker.init(), na Etapa 1. Se você já transmitiu o dicionário de configuração para o módulo torch.sagemaker.init() no script de treinamento na Etapa 1, você pode pular essa etapa. 

    // 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. Faça upload do arquivo smp_config.json em um diretório no seu sistema de arquivos. O caminho do diretório deve corresponder ao caminho definido na Etapa 1. Se você já transmitiu o dicionário de configuração para o módulo torch.sagemaker.init() no script de treinamento, você pode pular essa etapa.

  3. Nos nós de computação do seu cluster, inicie uma sessão de terminal com o comando a seguir.

    sudo su -l ubuntu

  4. Crie um ambiente Conda nos nós de computação. O código a seguir é um exemplo de script de criação de um ambiente Conda e instalação do SMP, SMDDP, CUDA e outras dependências.

    # 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. Execute um trabalho de treinamento de teste.

    1. No sistema de arquivos compartilhado (/fsx), faça um clone do repositório Awsome Distributed Training GitHub e vá até a pasta 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. Envie um trabalho usando sbatch conforme instruções a seguir.

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

      Se o envio do trabalho for bem-sucedido, a mensagem de saída desse comando sbatch deverá ser semelhante a Submitted batch job ABCDEF.

    3. Verifique o arquivo de log no diretório atual em logs/.

      tail -f ./logs/fsdp_smp_ABCDEF.out