Configuration des points de contrôle par niveaux gérés

Cette section contient le processus de configuration des points de contrôle par niveaux gérés 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.

Prérequis

Avant de configurer les points de contrôle par niveaux gérés, assurez-vous de disposer des éléments suivants :

Étape 1 : Activation des points de contrôle par niveaux gérés pour votre cluster

Activez les points de contrôle par niveaux gérés via l’API HyperPod 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 le nouveau cluster, créez :

aws sagemaker update-cluster \
    --cluster-name my-training-cluster \
    --tiered-storage-config {
            "Mode": "Enable"
            "InstanceMemoryAllocationPercentage": percentage
            }

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 contrôle Amazon SageMaker dans votre image d’entraînement en l’ajoutant à votre fichier Docker :

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

Étape 3 : Création d’une configuration de points de contrôle

Créez un objet CheckpointConfig pour spécifier le comportement des points de contrôle. Cela consiste notamment à :

L’exemple suivant montre une configuration des points de contrôle :

from amzn_sagemaker_checkpointing.config.sagemaker_checkpoint_config import SageMakerCheckpointConfig
from amzn_sagemaker_checkpointing.checkpointing.filesystem import SageMakerTieredStorageWriter, SageMakerTieredStorageReader

checkpoint_config = sm_ckpt.CheckpointConfig(
    world_size = 100,
    in_memory_namespace: my-ml-workload, # Logical grouping for checkpoints
    s3_base_path: "s3://bucket-name/checkpointing-path-prefix/",
    s3_every_n_steps: 100,  # Every 100 steps, save to S3
)

Étape 4 : Définition d’un processus d’écriture dans le système de fichiers SageMaker

Définissez le processus d’écriture dans votre système de fichiers de points de contrôle. Vous pouvez éventuellement spécifier un numéro d’étape lors de l’initialisation.

Processus d’écriture de base (étape spécifiée lors de l’appel d’enregistrement) :

smWriter = sagemaker_checkpointing.SageMakerTieredStorageWriter(checkpoint_config)

Processus d’écriture avec paramètre d’étape (étape spécifiée lors de l’initialisation) :

smWriter = sagemaker_checkpointing.SageMakerTieredStorageWriter(
    checkpoint_config, 
    step=step_number
)

Étape 5 : Enregistrement des points de contrôle dans votre boucle d’entraînement

Dans votre boucle d’entraînement, enregistrez les points de contrôle à l’aide de PyTorch DCP avec FileSystemWriter.

Utilisation de PyTorch DCP avec FileSystemWriter

Appelez la méthode dist_cp.save() avec FileSystemWriter comme entrée :

Option 1 : Utilisation de checkpoint_id avec le format d’étape (lorsque l’étape n’est pas spécifiée dans le processus d’écriture)

# Construct checkpoint directory with step number
checkpoint_dir = f"step_number"

dist_cp.save_state_dict(
    state_dict=state_dict,  # state_dict is a dictionary containing model parameters, optimizer state, etc.
    checkpoint_id=checkpoint_dir,  # Should contain step number
    storage_writer=smWriter
)

Option 2 : Utilisation du processus d’écriture avec le paramètre d’étape (checkpoint_id devient facultatif)

dist_cp.save_state_dict(
    state_dict=state_dict,
    storage_writer=smWriter  # Step already specified in writer initialization
)

Étape 6 : Chargement des points de contrôle pour la récupération

Lorsque vous devez charger un point de contrôle, utilisez PyTorch DCP avec FileSystemReader.

Utilisation de PyTorch DCP avec FileSystemReader

Appelez la méthode de chargement DCP avec FileSystemReader comme entrée :

# Define FileSystemReader
smReader = sagemaker_checkpointing.SageMakerTieredStorageReader(
    config=checkpoint_config
)

# Load checkpoint
dist_cp.load_state_dict(
    state_dict=state_dict,
    checkpoint_id=checkpoint_dir,
    storage_reader=smReader
)

Surveillance et validation

Vous pouvez surveiller et valider vos opérations de points de contrôle par niveaux gérés par le biais de métriques et 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.