Utilizzo della libreria SageMaker Model Parallelism v2 - Amazon SageMaker AI
Fase 1: adattamento dello script di addestramento PyTorch FSDPFase 2: avvio di un job di addestramento

Le traduzioni sono generate tramite traduzione automatica. In caso di conflitto tra il contenuto di una traduzione e la versione originale in Inglese, quest'ultima prevarrà.

Utilizzo della libreria SageMaker Model Parallelism v2

In questa pagina, imparerai come utilizzare le API della libreria SageMaker Model Parallelism v2 e iniziare a eseguire un job di addestramento PyTorch FSDP (Fully Sharded Data Parallel) nella piattaforma SageMaker Training o su un cluster SageMaker HyperPod.

Esistono vari scenari di esecuzione di un job di addestramento su PyTorch con SMP v2.

  1. Per l’addestramento su SageMaker, utilizza uno dei container framework SageMaker predefiniti per PyTorch v2.0.1 e versioni successive, preconfezionati con SMP v2.

  2. Utilizza il file binario SMP v2 per configurare un ambiente Conda per l’esecuzione di un carico di lavoro di addestramento distribuito su un cluster SageMaker HyperPod.

  3. Estendi i container framework SageMaker predefiniti per PyTorch v2.0.1 e versioni successive per installare eventuali requisiti funzionali aggiuntivi per il tuo caso d’uso. Per informazioni su come estendere un container predefinito, consulta Estensione di un container predefinito.

  4. Puoi anche introdurre il tuo container Docker e configurare manualmente tutto l’ambiente SageMaker Training utilizzando il kit di strumenti di SageMaker Training e installare il file binario SMP v2. Questa è l’opzione meno consigliata a causa della complessità delle dipendenze. Per informazioni su come gestire il tuo container Docker, consulta Adattamento del container di addestramento.

Questa guida introduttiva copre i primi due scenari.

Fase 1: adattamento dello script di addestramento PyTorch FSDP

Per attivare e configurare la libreria SMP v2, inizia importando e aggiungendo il modulo torch.sagemaker.init() nella parte superiore dello script. Questo modulo include il dizionario di configurazione di SMP Parametri di configurazione delle funzionalità principali di SMP v2 da utilizzare per la preparazione in Fase 2: avvio di un job di addestramento. Inoltre, per utilizzare le varie funzionalità principali offerte da SMP v2, potrebbe essere necessario apportare qualche altra modifica per adattare lo script di addestramento. Istruzioni più dettagliate sull’adattamento dello script di addestramento per l’utilizzo delle funzionalità principali di SMP v2 sono disponibili all’indirizzo Caratteristiche principali della libreria di parallelismo dei SageMaker modelli v2.

SageMaker Training

Nello script di addestramento, aggiungi le seguenti due righe di codice, che rappresentano il requisito minimo per iniziare l’addestramento con SMP v2. In Fase 2: avvio di un job di addestramento, imposterai un oggetto della classe dello strumento di stima PyTorch SageMaker con un dizionario di configurazione di SMP tramite l’argomento distribution della classe stessa.

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

È anche possibile passare direttamente un dizionario di configurazione di Parametri di configurazione delle funzionalità principali di SMP v2 al modulo torch.sagemaker.init(). Tuttavia, i parametri passati allo strumento di stima PyTorch in Fase 2: avvio di un job di addestramento hanno la priorità e sostituiscono quelli specificati nel modulo torch.sagemaker.init().

SageMaker HyperPod

Nello script di addestramento, aggiungi le seguenti due righe di codice. In Fase 2: avvio di un job di addestramento, dovrai configurare un file smp_config.json per l’impostazione delle configurazioni SMP in formato JSON e caricarlo in uno spazio di archiviazione o in un file system mappato con il tuo cluster SageMaker HyperPod. Ti consigliamo di conservare il file di configurazione nella stessa directory in cui carichi lo script di addestramento.

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

È anche possibile passare direttamente un dizionario di configurazione di Parametri di configurazione delle funzionalità principali di SMP v2 al modulo torch.sagemaker.init().

Fase 2: avvio di un job di addestramento

Scopri come configurare le opzioni di implementazione di SMP per l’avvio di un job di addestramento PyTorch FSDP con le funzionalità principali di SMP.

SageMaker Training

Quando configuri un oggetto launcher del job di addestramento della classe dello strumento di stima del framework PyTorch in SageMaker Python SDK, configura Parametri di configurazione delle funzionalità principali di SMP v2 tramite l’argomento distribution come indicato di seguito.

Nota

La configurazione distribution per SMP v2 è integrata in SageMaker Python SDK a partire dalla v2.200. Assicurati di utilizzare SageMaker Python SDK v2.200 o versioni successive.

Nota

In SMP v2, è necessario configurare smdistributed con torch_distributed per l’argomento distribution dello strumento di stima PyTorch di SageMaker. Con torch_distributed, SageMaker AI esegue torchrun, ovvero il launcher di un processo multi-nodo di 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

Per utilizzare una delle versioni precedenti di PyTorch o SMP anziché la più recente, è necessario specificare l’immagine Docker di SMP utilizzando direttamente l’argomento image_uri anziché la coppia framework_version e py_version. Di seguito è riportato un esempio di 

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

Per trovare gli URI delle immagini Docker di SMP, vedi Framework supportati.

SageMaker HyperPod

Prima di iniziare, verifica che siano soddisfatti i requisiti preliminari indicati di seguito.

  • Una directory condivisa Amazon FSx montata (/fsx) sul tuo cluster HyperPod.

  • Conda installato nella directory condivisa FSx. Per informazioni su come installare Conda, utilizza le istruzioni riportate in Installazione su Linux nella Guida per l’utente di Conda.

  • cuda11.8 o cuda12.1 installato sui nodi principali e di calcolo del cluster HyperPod.

Se tutti i prerequisiti sono soddisfatti, attieniti alle istruzioni sull’avvio di un carico di lavoro con SMP v2 su un cluster HyperPod riportate di seguito.

  1. Prepara un file smp_config.json che contenga un dizionario di Parametri di configurazione delle funzionalità principali di SMP v2. Assicurati di caricare questo file JSON nella posizione in cui archivi lo script di addestramento o nel percorso specificato per il modulo torch.sagemaker.init() nella Fase 1. Se hai già passato il dizionario di configurazione al modulo torch.sagemaker.init() nello script di addestramento durante la Fase 1, puoi saltare questa fase. 

    // 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. Carica il file smp_config.json in una directory del tuo file system. Il percorso della directory deve corrispondere al percorso specificato nella Fase 1. Se hai già passato il dizionario di configurazione al modulo torch.sagemaker.init() nello script di addestramento, puoi saltare questa fase.

  3. Sui nodi di calcolo del cluster, avvia una sessione di terminale con il seguente comando.

    sudo su -l ubuntu

  4. Crea un ambiente Conda sui nodi di calcolo. Il codice seguente è uno script di esempio per creare un ambiente Conda e installare SMP, SMDDP, CUDA e altre dipendenze.

    # 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. Esegui un job di addestramento di prova.

    1. Nel file system condiviso (/fsx), clona il repository GitHub di Awesome Distributed Training e accedi alla cartella 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. Invia un job utilizzando sbatch come segue.

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

      Se l’invio del job ha esito positivo, il messaggio di output di questo comando sbatch dovrebbe essere simile a Submitted batch job ABCDEF.

    3. Controlla il file di log nella directory corrente in logs/.

      tail -f ./logs/fsdp_smp_ABCDEF.out