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.

Configurer des points de contrôle hiérarchisés gérés

Cette section décrit le processus de configuration du point de contrôle hiérarchisé géré pour Amazon. SageMaker HyperPod Vous allez apprendre à activer cette fonctionnalité sur votre cluster et à implémenter les points de contrôle dans votre code d’entraînement.

Conditions préalables

Avant de configurer le point de contrôle hiérarchisé géré, assurez-vous d'avoir :

Étape 1 : activer le point de contrôle hiérarchisé géré pour votre cluster

Activez le point de contrôle hiérarchisé géré HyperPod APIs lors de la création ou de la mise à jour de votre cluster. Le service installe automatiquement le système de gestion de la mémoire lorsque vous spécifiez le paramètre TieredStorageConfig.

Pour les nouveaux clusters, vous pouvez utiliser create-clusterAWS CLI.

aws sagemaker create-cluster \
    --cluster-name cluster-name \
    --orchestrator "Eks={ClusterArn=eks-cluster-arn}" \
    --instance-groups '{
        "InstanceGroupName": "instance-group-name",
        "InstanceType": "instance-type",
        "InstanceCount": instance-count,
        "LifeCycleConfig": {
            "SourceS3Uri": "s3-path-to-lifecycle-scripts",
            "OnCreate": "lifecycle-script-name"
        },
        "ExecutionRole": "instance-group-iam-role",
        "ThreadsPerCore": threads-per-core,
        "InstanceStorageConfigs": [
            { "EbsVolumeConfig": {"VolumeSizeInGB": volume-size} }
        ]
    }' \
    --vpc-config '{
        "SecurityGroupIds": ["security-group-ids"],
        "Subnets": ["subnets"]
    }' \
    --tiered-storage-config '{
        "Mode": "Enable"
    }'

Le paramètre InstanceMemoryAllocationPercentage spécifie le percentage (int) de mémoire du cluster à allouer pour les points de contrôle. La plage est comprise entre 20 et 100.

Étape 2 : Installation de la bibliothèque Python dans votre image d’entraînement

Installez la bibliothèque de points de SageMaker contrôle Amazon et ses dépendances dans votre image d'entraînement en l'ajoutant à votre Dockerfile :

# Add this line to your training image Dockerfile
RUN pip install amzn-sagemaker-checkpointing s3torchconnector tenacity torch boto3 s3torchconnector

Étape 3 : Enregistrez les points de contrôle dans votre boucle d'entraînement

Dans votre boucle d'entraînement, vous pouvez enregistrer des points de contrôle de manière asynchrone à l'aide du DCP. PyTorch Voici un exemple sur la façon de procéder.

import torch
import torch.distributed as dist
from torch.distributed.checkpoint import async_save, load
from amzn_sagemaker_checkpointing.checkpointing.filesystem.filesystem import (
    SageMakerTieredStorageWriter,
    SageMakerTieredStorageReader
)

# Initialize distributed training
dist.init_process_group(backend="nccl")

# Configure checkpointing
checkpoint_config = SageMakerCheckpointConfig(
    # Unique ID for your training job 
    # Allowed characters in ID include: alphanumeric, hyphens, and underscores
    namespace=os.environ.get('TRAINING_JOB_NAME', f'job-{int(time.time())}'),

    # Number of distributed processes/available GPUs
    world_size=dist.get_world_size(),

    # S3 storage location, required for SageMakerTieredStorageReader for read fallbacks
    # Required for SageMakerTieredStorageWriter when save_to_s3 is True
    s3_tier_base_path="s3://my-bucket/checkpoints"
)

# Your model and optimizer
model = MyModel()
optimizer = torch.optim.AdamW(model.parameters())

# Training loop
future = None
in_memory_ckpt_freq = 10
s3_ckpt_freq = 50

for training_step in range(1000):
    # ... training code ...
    
    # Save checkpoint
    if (training_step % in_memory_ckpt_freq == 0 or 
        training_step % s3_ckpt_freq == 0):
        # Create state dictionary
        state_dict = {
            "model": model.state_dict(),
            "optimizer": optimizer.state_dict(),
            "step": training_step,
            "epoch": epoch
        }
        
        # Create storage writer for current step
        checkpoint_config.save_to_s3 = training_step % s3_ckpt_freq == 0
        storage_writer = SageMakerTieredStorageWriter(
            checkpoint_config=checkpoint_config,
            step=training_step
        )

        # wait for previous checkpoint to get completed
        if future is not None:
            exc = future.exception()
            if exc:
                print(f"Failure in saving previous checkpoint:{str(exc)}")
                # Handle failures as required
            else:
                result = future.result()
                # Process results from save, if required
        
        # Async save checkpoint using PyTorch DCP
        future = async_save(state_dict=state_dict, storage_writer=storage_writer)
        
        # Continue training while checkpoint saves in background

Étape 4 : Charger les points de contrôle pour la récupération

Voici un exemple de chargement d'un point de contrôle.

# Create state dictionary template
state_dict = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "step": 0,
    "epoch": 0
}

# Load latest checkpoint
storage_reader = SageMakerTieredStorageReader(checkpoint_config=checkpoint_config)
load(state_dict, storage_reader=storage_reader)

# Load specific checkpoint step
storage_reader = SageMakerTieredStorageReader(
    checkpoint_config=checkpoint_config, 
    step=500 # Or don't pass step if you have to load the latest available step.
)
try:
    load(state_dict, storage_reader=storage_reader)
except BaseException as e:
    print(f"Checkpoint load failed: {str(e)}")
    # Add additional exception handling

Validez vos opérations de point de contrôle hiérarchisé gérées

Vous pouvez valider vos opérations de point de contrôle hiérarchisé gérées à l'aide de journaux.

Journalisation personnalisée (facultatif)

Vous pouvez intégrer les journaux de points de contrôle à d’autres journaux en transmettant un enregistreur personnalisé à la bibliothèque. Par exemple, vous pouvez ajouter un enregistreur personnalisé à votre code d’entraînement afin que tous les journaux de la bibliothèque soient également collectés dans l’enregistreur d’entraînement.

Journalisation des services améliorée (facultatif)

Pour améliorer le débogage et la visibilité des services, vous pouvez monter le chemin du journal de points de contrôle /var/log/sagemaker_checkpointing depuis votre pod vers un chemin /var/logs/sagemaker_checkpointing sur votre hôte. Cela garantit que seuls les journaux spécifiques à la bibliothèque sont collectés séparément. Cela fournit à l’équipe de service une meilleure visibilité pour le débogage et le support.