Les traductions sont fournies par des outils de traduction automatique. En cas de conflit entre le contenu d'une traduction et celui de la version originale en anglais, la version anglaise prévaudra.

Utilisation de l'entraînement élastique sur Amazon SageMaker HyperPod

Elastic Training est une nouvelle SageMaker HyperPod fonctionnalité Amazon qui adapte automatiquement les tâches de formation en fonction de la disponibilité des ressources de calcul et de la priorité de la charge de travail. Les tâches d'entraînement Elastic peuvent démarrer avec un minimum de ressources de calcul requises pour l'entraînement des modèles et évoluer dynamiquement vers le haut ou vers le bas grâce à des points de contrôle automatiques et à une reprise sur différentes configurations de nœuds (taille mondiale). La mise à l'échelle est réalisée en ajustant automatiquement le nombre de répliques parallèles aux données. Pendant les périodes d'utilisation élevée du cluster, les tâches d'entraînement élastiques peuvent être configurées pour être automatiquement réduites en réponse aux demandes de ressources provenant de tâches plus prioritaires, libérant ainsi du calcul pour les charges de travail critiques. Lorsque les ressources sont libérées pendant les périodes creuses, les tâches de formation élastiques sont automatiquement redimensionnées pour accélérer la formation, puis diminuent lorsque les charges de travail les plus prioritaires nécessitent à nouveau des ressources.

Elastic Training repose sur l'opérateur d' HyperPod entraînement et intègre les composants suivants :

Frameworks supportés

Conditions préalables

SageMaker HyperPod Pôle EKS

Vous devez disposer d'un SageMaker HyperPod cluster en cours d'exécution avec l'orchestration Amazon EKS. Pour plus d'informations sur la création d'un cluster HyperPod EKS, consultez :

SageMaker HyperPod Opérateur de formation

Elastic Training est pris en charge dans les versions 1.2 et supérieures de la formation des opérateurs.

Pour installer l'opérateur de formation en tant que module complémentaire EKS, voir : https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-eks-operator-install.html

(Recommandé) Installer et configurer Task Governance et Kueue

Nous vous recommandons d'installer et de configurer Kueue via HyperPod Task Governance pour spécifier les priorités en matière de charge de travail avec Elastic Training. Kueue fournit une meilleure gestion de la charge de travail grâce à la mise en file d'attente, à la hiérarchisation, à la planification des gangs, au suivi des ressources et à une préemption élégante, qui sont essentiels pour fonctionner dans des environnements de formation multi-locataires.

Nous vous recommandons de configurer les composants Kueue suivants :

Pour des configurations plus avancées, vous pouvez également intégrer :

Veuillez vous référer à :

(Recommandé) Configurer les espaces de noms des utilisateurs et les quotas de ressources

Lors du déploiement de cette fonctionnalité sur Amazon EKS, nous vous recommandons d'appliquer un ensemble de configurations fondamentales au niveau du cluster afin de garantir l'isolation, l'équité des ressources et la cohérence opérationnelle entre les équipes.

Configuration de l'espace de noms et de l'accès

Organisez vos charges de travail en utilisant des espaces de noms distincts pour chaque équipe ou projet. Cela vous permet d'appliquer une isolation et une gouvernance précises. Nous vous recommandons également de configurer le mappage AWS IAM vers Kubernetes RBAC pour associer des utilisateurs ou des rôles IAM individuels à leurs espaces de noms correspondants.

Les principales pratiques sont les suivantes :

Contraintes en matière de ressources et de calcul

Pour éviter les conflits de ressources et garantir une planification équitable entre les équipes, appliquez des quotas et des limites au niveau de l'espace de noms :

Ces contraintes contribuent à maintenir la stabilité du cluster et à permettre une planification prévisible des charges de travail de formation distribuées.

Auto-scaling

SageMaker HyperPod sur EKS prend en charge la mise à l'échelle automatique des clusters via Karpenter. Lorsque Karpenter ou un fournisseur de ressources similaire est utilisé conjointement avec Elastic Training, le cluster ainsi que la tâche Elastic Training peuvent évoluer automatiquement une fois qu'une tâche Elastic Training est soumise. En effet, Elastic Training Operator adopte une approche gourmande et demande toujours plus que les ressources de calcul disponibles jusqu'à ce que la limite maximale fixée par la tâche soit atteinte. Cela est dû au fait que l'opérateur Elastic Training demande en permanence des ressources supplémentaires dans le cadre de l'exécution des tâches élastiques, ce qui peut déclencher le provisionnement des nœuds. Les fournisseurs de ressources continues tels que Karpenter répondront aux demandes en agrandissant le cluster de calcul.

Pour que ces mises à l'échelle restent prévisibles et contrôlées, nous vous recommandons de configurer le niveau de l'espace de noms dans les espaces de noms ResourceQuotas dans lesquels les jobs Elastic Training sont créés. ResourceQuotas aident à limiter le maximum de ressources que les jobs peuvent demander, en empêchant une croissance illimitée des clusters tout en permettant un comportement élastique dans des limites définies.

Par exemple, une instance ResourceQuota pour 8 ml.p5.48xlarge aura la forme suivante :

apiVersion: v1
kind: ResourceQuota
metadata:
  name: <quota-name>
  namespace: <namespace-name>
spec:
  hard:
    nvidia.com/gpu: "64"
    vpc.amazonaws.com/efa: "256"
    requests.cpu: "1536"
    requests.memory: "5120Gi"
    limits.cpu: "1536"
    limits.memory: "5120Gi"
      

Construire un conteneur de formation

HyperPod L'opérateur de formation utilise un PyTorch lanceur personnalisé fourni via le package Python d' HyperPod Elastic Agent (https://www.piwheels.org/project/hyperpod-elastic-agent/). Les clients doivent installer l'agent élastique et remplacer la torchrun commande par hyperpodrun pour lancer la formation. Pour plus de détails, veuillez consulter :

https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-eks-operator-install.html#sagemaker -eks-operator-elastic-agent

Un exemple de conteneur de formation :

FROM ...

...

RUN pip install hyperpod-elastic-agent
ENTRYPOINT ["entrypoint.sh"]

# entrypoint.sh ...
hyperpodrun --nnodes=node_count --nproc-per-node=proc_count \
  --rdzv-backend hyperpod \
 # Optional ...
 # Other torchrun args
 # pre-traing arg_group
 --pre-train-script pre.sh --pre-train-args "pre_1 pre_2 pre_3" \
 # post-train arg_group
 --post-train-script post.sh --post-train-args "post_1 post_2 post_3" \
 training.py --script-args
  

Modification du code d'entraînement

SageMaker HyperPod fournit un ensemble de recettes déjà configurées pour être exécutées avec Elastic Policy.

Pour activer l'entraînement élastique pour les scripts d' PyTorch entraînement personnalisés, vous devez apporter des modifications mineures à votre boucle d'entraînement. Ce guide explique les modifications nécessaires pour garantir que votre tâche de formation réponde aux événements de mise à l'échelle élastique qui se produisent lorsque la disponibilité des ressources de calcul change. Pendant tous les événements élastiques (par exemple, des nœuds sont disponibles ou des nœuds sont préemptés), la tâche d'entraînement reçoit un signal d'événement élastique qui est utilisé pour coordonner un arrêt progressif en enregistrant un point de contrôle, et en reprenant l'entraînement en redémarrant à partir de ce point de contrôle enregistré avec une nouvelle configuration mondiale. Pour activer l'entraînement élastique avec des scripts d'entraînement personnalisés, vous devez :

Détecter les événements Elastic Scaling

Dans votre boucle d'entraînement, vérifiez la présence d'événements élastiques à chaque itération :

from hyperpod_elastic_agent.elastic_event_handler import elastic_event_detected

def train_epoch(model, dataloader, optimizer, args):
    for batch_idx, batch_data in enumerate(dataloader):
        # Forward and backward pass
        loss = model(batch_data).loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        # Handle checkpointing and elastic scaling
        should_checkpoint = (batch_idx + 1) % args.checkpoint_freq == 0
        elastic_event = elastic_event_detected()
        
        # Save checkpoint if scaling-up or scaling down job
        if should_checkpoint or elastic_event:
            save_checkpoint(model, optimizer, scheduler, 
                            checkpoint_dir=args.checkpoint_dir, 
                            step=global_step)
              
            if elastic_event:
                print("Elastic scaling event detected. Checkpoint saved.")
                return
    

Mettre en œuvre la sauvegarde des points de contrôle et le chargement des points de contrôle

Remarque : Nous vous recommandons d'utiliser le point de contrôle PyTorch distribué (DCP) pour enregistrer les états du modèle et de l'optimiseur, car le DCP permet de reprendre à partir d'un point de contrôle dont la taille du monde varie. D'autres formats de points de contrôle peuvent ne pas prendre en charge le chargement de points de contrôle sur différentes tailles de monde. Dans ce cas, vous devrez implémenter une logique personnalisée pour gérer les changements dynamiques de taille du monde.

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict

def save_checkpoint(model, optimizer, lr_scheduler, user_content, checkpoint_path):
    """Save checkpoint using DCP for elastic training."""
    state_dict = {
        "model": model,
        "optimizer": optimizer,
        "lr_scheduler": lr_scheduler,
        **user_content
    }
      
    dcp.save(
        state_dict=state_dict,
        storage_writer=dcp.FileSystemWriter(checkpoint_path)
    )

def load_checkpoint(model, optimizer, lr_scheduler, checkpoint_path):
    """Load checkpoint using DCP with automatic resharding."""
    state_dict = {
        "model": model,
        "optimizer": optimizer,
        "lr_scheduler": lr_scheduler
    }
      
    dcp.load(
        state_dict=state_dict,
        storage_reader=dcp.FileSystemReader(checkpoint_path)
    )
      
    return model, optimizer, lr_scheduler
    

(Facultatif) Utiliser des chargeurs de données dynamiques

Si vous ne vous entraînez que pour une seule époque (c'est-à-dire un seul passage dans l'ensemble de données), le modèle doit voir chaque échantillon de données exactement une fois. Si la tâche de formation s'arrête en milieu d'époque et reprend avec une taille de monde différente, les échantillons de données précédemment traités seront répétés si l'état du chargeur de données n'est pas conservé. Un chargeur de données dynamique empêche cela en enregistrant et en rétablissant la position du chargeur de données, garantissant ainsi que les exécutions reprises se poursuivent à partir de l'événement Elastic Scaling sans retraiter aucun échantillon. Nous vous recommandons d'utiliser StatefulDataLoader, qui remplace automatiquement ces ajouts state_dict() et ces load_state_dict() méthodes, afin de permettre un contrôle à mi-époque du processus de chargement des données. torch.utils.data.DataLoader

Soumettre des offres de formation élastique

HyperPod L'opérateur de formation définit un nouveau type de ressource -hyperpodpytorchjob. Elastic Training étend ce type de ressource et ajoute les champs surlignés ci-dessous :

apiVersion: sagemaker.amazonaws.com/v1
kind: HyperPodPyTorchJob
metadata:
  name: elastic-training-job
spec:
  elasticPolicy:
    minReplicas: 1
    maxReplicas: 4
    # Increment amount of pods in fixed-size groups
    # Amount of pods will be equal to minReplicas + N * replicaIncrementStep
    replicaIncrementStep: 1           
    # ... or Provide an exact amount of pods that required for training
    replicaDiscreteValues: [2,4,8]     

    # How long traing operator wait job to save checkpoint and exit during
    # scaling events. Job will be force-stopped after this period of time
    gracefulShutdownTimeoutInSeconds: 600

    # When scaling event is detected:   
    # how long job controller waits before initiate scale-up.
    # Some delay can prevent from frequent scale-ups and scale-downs
    scalingTimeoutInSeconds: 60

    # In case of faults, specify how long elastic training should wait for
    # recovery, before triggering a scale-down
    faultyScaleDownTimeoutInSeconds: 30
  ...
  replicaSpecs:
    - name: pods
      replicas: 4           # Initial replica count
      maxReplicas: 8        # Max for this replica spec (should match elasticPolicy.maxReplicas)
      ...
  

Utilisation de kubectl

Vous pouvez ensuite lancer l'entraînement élastique à l'aide de la commande suivante.

kubectl apply -f elastic-training-job.yaml

Utilisation de SageMaker recettes

Les jobs d'entraînement Elastic peuvent être lancés via SageMaker HyperPod des recettes.

Outre les recettes statiques, les recettes élastiques ajoutent les champs suivants pour définir les comportements élastiques :

Politique élastique

Le elastic_policy champ définit la configuration au niveau de la tâche pour la tâche d'entraînement élastique. Il possède les configurations suivantes :

Voici un exemple de définition de ce champ, que vous pouvez également trouver dans le dépôt Hyperpod Recipe dans la recette : recipes_collection/recipes/fine-tuning/llama/llmft_llama3_1_8b_instruct_seq4k_gpu_sft_lora.yaml

<static recipe>
...
elastic_policy:
  is_elastic: true
  min_nodes: 1
  max_nodes: 16
  use_graceful_shutdown: true
  scaling_timeout: 600
  graceful_shutdown_timeout: 600
      

Configuration de la balance

Le scale_config champ définit les configurations principales à chaque échelle spécifique. Il s'agit d'un dictionnaire clé-valeur, où key est un entier représentant l'échelle cible et la valeur est un sous-ensemble de la recette de base. À <key> grande échelle, nous les utilisons <value> pour mettre à jour les configurations spécifiques de la base/static recette. Voici un exemple de ce champ :

scale_config:   
...
  2:
    trainer:
      num_nodes: 2
    training_config:
      training_args:
        train_batch_size: 128
        micro_train_batch_size: 8
        learning_rate: 0.0004
  3:
    trainer:
      num_nodes: 3
    training_config:
      training_args:
        train_batch_size: 128
        learning_rate: 0.0004
        uneven_batch:
          use_uneven_batch: true
          num_dp_groups_with_small_batch_size: 16
          small_local_batch_size: 5
          large_local_batch_size: 6
 ...
      

La configuration ci-dessus définit la configuration d'entraînement aux échelles 2 et 3. Dans les deux cas, nous utilisons le taux d'apprentissage4e-4, la taille du lot de128. Mais à l'échelle 2, nous utilisons un micro_train_batch_size de 8, tandis qu'à l'échelle 3, nous utilisons une taille de lot inégale car la taille du lot de train ne peut pas être divisée de manière égale sur 3 nœuds.

Taille de lot inégale

Ce champ permet de définir le comportement de distribution des lots lorsque la taille globale du lot ne peut pas être divisée de manière égale par le nombre de rangs. Ce n'est pas spécifique à l'entraînement élastique, mais il permet une granularité d'échelle plus fine.

Surveillez la formation sur MLflow

Les tâches de recette Hyperpod favorisent l'observabilité via MLflow. Les utilisateurs peuvent spécifier les configurations MLflow dans la recette :

training_config:
  mlflow:
    tracking_uri: "<local_file_path or MLflow server URL>"
    run_id: "<MLflow run ID>"
    experiment_name: "<MLflow experiment name, e.g. llama_exps>"
    run_name: "<run name, e.g. llama3.1_8b>"

      

Ces configurations sont mappées à la configuration MLflow correspondante. Voici un exemple de tableau de bord MLflow pour une tâche d'entraînement élastique.

Après avoir défini les recettes élastiques, nous pouvons utiliser les scripts de lancement, par exemple launcher_scripts/llama/run_llmft_llama3_1_8b_instruct_seq4k_gpu_sft_lora.sh pour lancer une tâche d'entraînement élastique. Ceci est similaire au lancement d'une tâche statique à l'aide de la recette Hyperpod.

Use-case exemples et limites

Scale-up lorsque davantage de ressources sont disponibles

Lorsque davantage de ressources sont disponibles sur le cluster (par exemple, d'autres charges de travail sont terminées). Au cours de cet événement, le contrôleur de formation intensifiera automatiquement la tâche de formation. Ce comportement est expliqué ci-dessous.

Pour simuler une situation dans laquelle davantage de ressources seront disponibles, nous pouvons soumettre une tâche prioritaire, puis libérer des ressources en supprimant la tâche hautement prioritaire.

# Submit a high-priority job on your cluster. As a result of this command
# resources will not be available for elastic training
kubectl apply -f high_prioriy_job.yaml

# Submit an elastic job with normal priority
kubectl apply -f hyperpod_job_with_elasticity.yaml

# Wait for training to start....

# Delete high priority job. This command will make additional resources available for
# elastic training
kubectl delete -f high_prioriy_job.yaml

# Observe the scale-up of elastic job
    

Comportement attendu :

Remarque : Lorsque Kueue n'est pas utilisé, l'opérateur de formation ignore les deux premières étapes. Il tente immédiatement de créer les modules supplémentaires nécessaires à la nouvelle taille du monde. Si des ressources suffisantes ne sont pas disponibles dans le cluster, ces pods resteront dans l'état En attente jusqu'à ce que la capacité soit disponible.

Préemption par un travail hautement prioritaire

Les tâches élastiques peuvent être réduites automatiquement lorsqu'une tâche hautement prioritaire nécessite des ressources. Pour simuler ce comportement, vous pouvez soumettre une tâche d'entraînement élastique, qui utilise le maximum de ressources disponibles dès le début de la formation, puis soumettre une tâche hautement prioritaire et observer un comportement de préemption.

# Submit an elastic job with normal priority
kubectl apply -f hyperpod_job_with_elasticity.yaml

# Submit a high-priority job on your cluster. As a result of this command
# some amount of resources will be   
kubectl apply -f high_prioriy_job.yaml

# Observe scale-down behaviour
    

Lorsqu'une tâche hautement prioritaire nécessite des ressources, Kueue peut préempter les charges de travail Elastic Training de moindre priorité (plusieurs objets de charge de travail peuvent être associés à une tâche Elastic Training). Le processus de préemption suit la séquence suivante :

La préemption peut entraîner la suspension de l'ensemble de la tâche de formation, ce qui n'est peut-être pas souhaitable pour tous les flux de travail. Pour éviter la suspension complète des tâches tout en permettant une mise à l'échelle élastique, les clients peuvent configurer deux niveaux de priorité différents au sein d'une même tâche de formation en définissant deux replicaSpec sections :

Cette configuration permet une préemption partielle, où seule la capacité élastique est récupérée, ce qui permet de maintenir la continuité de la formation tout en favorisant un partage équitable des ressources dans les environnements multi-locataires. Exemple :

apiVersion: sagemaker.amazonaws.com/v1
kind: HyperPodPyTorchJob
metadata:
  name: elastic-training-job
spec:
  elasticPolicy:
    minReplicas: 2
    maxReplicas: 8
    replicaIncrementStep: 2
  ...
  replicaSpecs:
    - name: base
      replicas: 2
      template:
        spec:
          priorityClassName: high-priority # set high-priority to avoid evictions
           ...
    - name: elastic
      replicas: 0
      maxReplicas: 6
      template:
        spec:
          priorityClassName: low-priority. # Set low-priority for elastic part
           ...
    

Gestion de l'expulsion des modules, des pannes des modules et de la dégradation du matériel :

L'opérateur de HyperPod formation comprend des mécanismes intégrés permettant de reprendre le processus d'entraînement en cas d'interruption inattendue. Les interruptions peuvent survenir pour diverses raisons, telles que des échecs de code de formation, des expulsions de pods, des pannes de nœuds, une dégradation du matériel et d'autres problèmes d'exécution.

Dans ce cas, l'opérateur tente automatiquement de recréer les modules concernés et de reprendre l'entraînement au dernier point de contrôle. Si la reprise n'est pas immédiatement possible, par exemple en raison d'une capacité inutilisée insuffisante, l'opérateur peut poursuivre ses progrès en réduisant temporairement la taille du monde et en réduisant la charge de formation élastique.

Lorsqu'une tâche d'entraînement élastique tombe en panne ou perd des répliques, le système se comporte comme suit :

Ce mécanisme garantit que la formation peut se poursuivre avec un minimum de temps d'arrêt, même en cas de pression sur les ressources ou de défaillances partielles de l'infrastructure, tout en tirant parti de l'évolutivité élastique.

Utilisez l'entraînement élastique avec d'autres HyperPod fonctionnalités

Elastic Training ne prend actuellement pas en charge les fonctionnalités d'entraînement sans point de contrôle, les points de contrôle hiérarchisés HyperPod gérés ou les instances Spot.