Utilisation de la bibliothèque SageMaker de parallélisme des modèles v2 - Amazon SageMaker AI
Étape 1 : Adapter votre script d’entraînement PyTorch FSDPÉtape 2 : lancer une tâche d’entraînement

Utilisation de la bibliothèque SageMaker de parallélisme des modèles v2

Sur cette page, vous découvrirez comment utiliser les API v2 de la bibliothèque SageMaker de parallélisme des modèles et comment commencer à exécuter une tâche d’entraînement PyTorch de parallélisme entièrement partitionné des données (FSDP) sur la plateforme d’entraînement SageMaker ou sur un cluster SageMaker HyperPod.

Il existe différents scénarios pour exécuter une tâche d’entraînement PyTorch avec SMP v2.

  1. Pour l’entraînement SageMaker, utilisez l’un des conteneurs SageMaker Framework prédéfinis pour PyTorch v2.0.1 et versions ultérieures, qui sont pré-empaquetés avec SMP v2.

  2. Utilisez le fichier binaire SMP v2 pour configurer un environnement Conda afin d’exécuter une charge de travail d’entraînement distribué sur un cluster SageMaker HyperPod.

  3. Étendez les conteneurs SageMaker Framework prédéfinis pour PyTorch v2.0.1 et versions ultérieures afin d’installer toute configuration fonctionnelle supplémentaire requise pour votre cas d’utilisation. Pour découvrir comment étendre un conteneur prédéfini, consultez Extension d’un conteneur préconçu.

  4. Vous pouvez aussi apporter votre propre conteneur Docker, configurer manuellement tous les environnements d’entraînement SageMaker à l’aide de la boîte à outils d’entraînement SageMaker et installer le fichier binaire SMP v2. Il s’agit de l’option la moins recommandée en raison de la complexité des dépendances. Pour découvrir comment exécuter votre propre conteneur Docker, consultez Adaptation de votre propre conteneur d’entraînement.

Ce guide de démarrage couvre les deux premiers scénarios.

Étape 1 : Adapter votre script d’entraînement PyTorch FSDP

Pour activer et configurer la bibliothèque SMP v2, commencez par importer et ajouter le module torch.sagemaker.init() au début du script. Ce module intègre le dictionnaire de configuration SMP de Paramètres de configuration des caractéristiques principales de SMP v2 que vous allez préparer dans Étape 2 : lancer une tâche d’entraînement. De plus, pour utiliser les différentes caractéristiques de base proposées par SMP v2, vous devrez peut-être apporter quelques modifications supplémentaires pour adapter votre script d’entraînement. Des instructions plus détaillées sur l’adaptation de votre script d’entraînement à l’utilisation des caractéristiques de base de SMP v2 sont fournies sur la page Caractéristiques de base de la bibliothèque SageMaker de parallélisme des modèles v2.

SageMaker Training

Dans votre script d’entraînement, ajoutez les deux lignes de code suivantes, qui constituent le minimum requis pour commencer un entraînement avec SMP v2. À l’Étape 2 : lancer une tâche d’entraînement, vous allez configurer un objet de la classe d’estimateur PyTorch SageMaker avec un dictionnaire de configuration SMP via l’argument distribution de la classe d’estimateur.

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

Vous pouvez aussi transmettre directement un dictionnaire de configuration des Paramètres de configuration des caractéristiques principales de SMP v2 dans le module torch.sagemaker.init(). Cependant, les paramètres transmis à l’estimateur PyTorch à l’Étape 2 : lancer une tâche d’entraînement sont prioritaires et remplacent ceux spécifiés pour le module torch.sagemaker.init().

SageMaker HyperPod

Ajoutez les deux lignes de code suivantes à votre script d’entraînement. À l’Étape 2 : lancer une tâche d’entraînement, vous allez configurer un fichier smp_config.json pour définir les configurations SMP au format JSON, et le charger sur un système de stockage ou de fichiers mappé à votre cluster SageMaker HyperPod. Nous vous recommandons de conserver le fichier de configuration dans même le répertoire que celui où vous avez chargé votre script d’entraînement.

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

Vous pouvez aussi transmettre directement un dictionnaire de configuration des Paramètres de configuration des caractéristiques principales de SMP v2 dans le module torch.sagemaker.init().

Étape 2 : lancer une tâche d’entraînement

Découvrez comment configurer les options de distribution SMP pour lancer une tâche d’entraînement PyTorch FSDP avec les caractéristiques de base SMP.

SageMaker Training

Lorsque vous configurez un objet lanceur de tâches d’entraînement de la classe d’estimateur du cadre PyTorch dans le kit SageMaker Python SDK, configurez Paramètres de configuration des caractéristiques principales de SMP v2 via un argument distribution comme suit.

Note

La configuration distribution pour SMP v2 est intégrée dans le kit SageMaker Python SDK à partir de la version 2.200. Assurez-vous d’utiliser le kit SageMaker Python SDK v2.200 ou version ultérieure.

Note

Dans SMP v2, vous devez configurer smdistributed avec torch_distributed comme argument distribution de l’estimateur SageMaker PyTorch. Avec torch_distributed, SageMaker AI exécute torchrun, qui est le lanceur de tâches multi-nœuds par défaut 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
                }
            }
        }
    }
)
Important

Pour utiliser l’une des versions précédentes de PyTorch ou de SMP au lieu de la dernière, vous devez spécifier l’image Docker SMP directement en utilisant l’argument image_uri au lieu de la paire framework_version et py_version. Voici un exemple : 

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

Pour trouver les URI des images Docker SMP, consultez Cadres pris en charge.

SageMaker HyperPod

Avant de commencer, vérifiez que vous respectez les prérequis suivants.

  • Répertoire partagé Amazon FSx monté (/fsx) sur votre cluster HyperPod.

  • Conda installé dans le répertoire partagé FSx. Pour découvrir comment installer Conda, suivez les instructions de la page Installing on Linux dans le Guide de l’utilisateur Conda.

  • cuda11.8 ou cuda12.1 installé sur le nœud principal et les nœuds de calcul de votre cluster HyperPod.

Si toutes le prérequis sont respectés, suivez les instructions suivantes pour lancer une charge de travail avec SMP v2 sur un cluster HyperPod.

  1. Préparez un fichier smp_config.json contenant un dictionnaire de Paramètres de configuration des caractéristiques principales de SMP v2. Assurez-vous de charger ce fichier JSON dans l’endroit où vous stockez votre script d’entraînement ou au chemin que vous avez spécifié pour accéder au module torch.sagemaker.init() à l’Étape 1. Si vous avez déjà transmis le dictionnaire de configuration au module torch.sagemaker.init() dans le script d’entraînement à l’Étape 1, vous pouvez ignorer cette étape. 

    // 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. Chargez le fichier smp_config.json dans un répertoire de votre système de fichiers. Le chemin du répertoire doit correspondre à celui que vous avez spécifié à l’Étape 1. Si vous avez déjà transmis le dictionnaire de configuration au module torch.sagemaker.init() dans le script d’entraînement, vous pouvez ignorer cette étape.

  3. Sur les nœuds de calcul de votre cluster, lancez une session de terminal avec la commande suivante.

    sudo su -l ubuntu

  4. Créez un environnement Conda sur les nœuds de calcul. Le code suivant est un exemple de script de création d’un environnement Conda et d’installation de SMP, SMDDP, CUDA et d’autres dépendances.

    # 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. Exécutez une tâche d’entraînement test.

    1. Dans le système de fichiers partagés (/fsx), clonez le référentiel GitHub Awsome Distributed Training et accédez au dossier 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. Soumettez une tâche à l’aide de sbatch en procédant comme suit.

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

      Si la soumission de la tâche est réussie, le message de sortie de cette commande sbatch doit être similaire à Submitted batch job ABCDEF.

    3. Consultez le fichier journal dans le répertoire en cours sous logs/.

      tail -f ./logs/fsdp_smp_ABCDEF.out