Configurazione dei checkpoint con livelli gestiti

Questa sezione descrive il processo di configurazione per i checkpoint con livelli gestiti per Amazon SageMaker HyperPod. Scoprirai come abilitare la funzionalità nel cluster e implementare i checkpoint nel codice di addestramento.

Prerequisiti

Prima di configurare i checkpoint con livelli gestiti, assicurati di avere:

Fase 1. Abilita i checkpoint con livelli gestiti per il cluster

Abilita i checkpoint con livelli gestiti tramite l’API HyperPod durante la creazione o l’aggiornamento del cluster. Il servizio installa automaticamente il sistema di gestione della memoria quando specifichi il parametro TieredStorageConfig. Per i nuovi cluster, crea:

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

Il parametro InstanceMemoryAllocationPercentage specifica il valore percentage (int) della memoria del cluster da allocare per i checkpoint. L’intervallo è compreso tra 20 e 100.

Fase 2. Installa la libreria Python nell’immagine di addestramento

Installa la libreria di checkpoint di Amazon SageMaker nell’immagine di addestramento aggiungendola al Dockerfile:

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

Fase 3. Crea una configurazione del checkpoint

Crea un oggetto CheckpointConfig per specificare il comportamento del checkpoint. Questo include:

L’esempio seguente mostra la configurazione di un checkpoint:

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
)

Fase 4. Definisci uno scrittore di file system SageMaker

Definisci il tuo scrittore di file system per checkpoint. Facoltativamente, puoi specificare un numero di fase durante l’inizializzazione.

Scrittore di base (fase specificata nella chiamata di salvataggio):

smWriter = sagemaker_checkpointing.SageMakerTieredStorageWriter(checkpoint_config)

Scrittore con parametro step (fase specificata durante l’inizializzazione):

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

Fase 5. Salva i checkpoint nel ciclo di addestramento

Nel ciclo di addestramento, salva i checkpoint utilizzando PyTorch DCP con FileSystemWriter.

Utilizzo di PyTorch DCP con FileSystemWriter

Chiama il metodo dist_cp.save() con FileSystemWriter come input:

Opzione 1. Utilizzo di checkpoint_id con formato step (quando la fase non è specificata nello scrittore)

# 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
)

Opzione 2. Utilizzo di uno scrittore con parametro step (checkpoint_id diventa facoltativo)

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

Fase 6. Carica i checkpoint per il ripristino

Quando devi caricare un checkpoint, utilizza PyTorch DCP con FileSystemReader.

Utilizzo di PyTorch DCP con FileSystemReader

Chiama il metodo di caricamento DCP con FileSystemReader come input:

# 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
)

Monitoraggio e convalida

Puoi monitorare e convalidare le operazioni di checkpoint con livelli gestiti tramite metriche e log.

Registrazione di log personalizzata (facoltativo)

Puoi integrare i log dei checkpoint con altri log passando un logger personalizzato alla libreria. Ad esempio, puoi aggiungere un logger personalizzato al codice di addestramento in modo che tutti i log della libreria vengano raccolti anche nel logger di addestramento.

Registrazione avanzata di log dei servizi (facoltativo)

Per migliorare il debug e la visibilità del servizio, puoi montare il percorso del log dei checkpoint /var/log/sagemaker_checkpointing dall’interno del pod a un percorso /var/logs/sagemaker_checkpointing sull’host. Questa operazione garantisce che solo i log specifici della libreria vengano raccolti separatamente. Il team di assistenza può quindi avere una maggiore visibilità per il debug e il supporto.