Comece com a biblioteca de paralelismo de SageMaker modelos v2 - Amazon SageMaker
Etapa 1: Adapte seu script de PyTorch treinamento do FSDPEtapa 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á.

Comece com a biblioteca de paralelismo de SageMaker modelos v2

Nesta página, você aprenderá a usar as APIs v2 da biblioteca de paralelismo de SageMaker modelos e começará a executar um trabalho de treinamento PyTorch Fully Sharded Data Parallel (FSDP) na plataforma de treinamento ou em um cluster. SageMaker SageMaker HyperPod

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

  1. Para SageMaker treinamento, use um dos SageMaker Framework Containers pré-criados para PyTorch v2.0.1 e versões posteriores, que são pré-empacotados com o SMP v2.

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

  3. Estenda os SageMaker Framework Containers pré-criados para PyTorch v2.0.1 e versões posteriores para instalar quaisquer requisitos funcionais adicionais para seu caso de uso. Para saber como estender um contêiner pré-construído, consulteEstenda uma imagem de contêiner predefinida.

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

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

Etapa 1: Adapte seu script de PyTorch treinamento do FSDP

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

SageMaker Training

Em seu 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. EmEtapa 2: iniciar um trabalho de treinamento, você configurará um objeto da classe SageMaker PyTorch estimadora com um dicionário de configuração SMP por meio do distribution argumento da classe estimadora.

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

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

SageMaker HyperPod

Em seu script de treinamento, adicione as duas linhas de código a seguir. EmEtapa 2: iniciar um trabalho de treinamento, você configurará um smp_config.json arquivo para definir as configurações SMP no formato JSON e o carregará em um armazenamento ou sistema de arquivos mapeado com seu cluster. 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 passar diretamente um dicionário de configuração do Parâmetros de configuração do recurso principal do SMP v2 para o torch.sagemaker.init() módulo.

Etapa 2: iniciar um trabalho de treinamento

Saiba como configurar as opções de distribuição do SMP para iniciar um trabalho de treinamento do PyTorch FSDP com os principais recursos do SMP.

SageMaker Training

Ao configurar um objeto iniciador de trabalhos de treinamento da classe estimador de PyTorch estrutura no SDK do SageMaker Python, configure por meio do argumento a seguir. Parâmetros de configuração do recurso principal do SMP v2 distribution

nota

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

nota

No SMP v2, você deve configurar smdistributed with torch_distributed para o distribution argumento do SageMaker PyTorch estimador. Withtorch_distributed, SageMaker runstorchrun, 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 SMP em vez da mais recente, você precisa especificar a imagem do SMP Docker diretamente usando o image_uri argumento em vez do framework_version par e. py_version A seguir está 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 URIs de imagem do SMP Docker, consulte. Estruturas compatíveis

SageMaker HyperPod

Antes de começar, verifique se os pré-requisitos a seguir foram atendidos.

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

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

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

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

  1. Prepare um smp_config.json arquivo que contenha um dicionário deParâmetros de configuração do recurso principal do SMP v2. Certifique-se de carregar esse arquivo JSON para onde você armazena seu script de treinamento ou o caminho que você especificou para o torch.sagemaker.init() módulo na Etapa 1. Se você já passou o dicionário de configuração para o torch.sagemaker.init() módulo 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. Carregue o smp_config.json arquivo em um diretório no seu sistema de arquivos. O caminho do diretório deve corresponder ao caminho especificado na Etapa 1. Se você já passou o dicionário de configuração para o torch.sagemaker.init() módulo no script de treinamento, pode pular esta 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 de 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), clone o GitHub repositório do Awsome Distributed Training 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 o seguinte.

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

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

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

      tail -f ./logs/fsdp_smp_ABCDEF.out