As traduções são geradas por tradução automática. Em caso de conflito entre o conteúdo da tradução e da versão original em inglês, a versão em inglês prevalecerá.

Configurar pontos de verificação gerenciados em camadas

Esta seção contém o processo de configuração do ponto de verificação hierárquico gerenciado para a Amazon. SageMaker HyperPod Você aprenderá a habilitar o recurso no cluster e implementar pontos de verificação no código de treinamento.

Pré-requisitos

Antes de configurar o ponto de verificação hierárquico gerenciado, verifique se você tem:

Etapa 1: habilitar o ponto de verificação gerenciado em camadas para seu cluster

Ative o checkpoint gerenciado em camadas por meio do HyperPod APIs ao criar ou atualizar seu cluster. O serviço instala automaticamente o sistema de gerenciamento de memória quando você especifica o parâmetro TieredStorageConfig.

Para novos clusters, você pode usar create-cluster AWS 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"
    }'

O parâmetro InstanceMemoryAllocationPercentage especifica a percentage (int) da memória do cluster a ser alocada ao ponto de verificação. O intervalo é 20-100.

Etapa 2: instalar a biblioteca do Python na imagem de treinamento

Instale a biblioteca SageMaker de pontos de verificação da Amazon e suas dependências em sua imagem de treinamento adicionando-a ao seu Dockerfile:

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

Etapa 3: salve os pontos de verificação em seu ciclo de treinamento

Em seu ciclo de treinamento, você pode salvar pontos de verificação de forma assíncrona usando o DCP. PyTorch Veja a seguir um exemplo de como fazer isso.

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

Etapa 4: pontos de verificação de carga para recuperação

Veja a seguir um exemplo de carregamento de um ponto de verificação.

# 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

Valide suas operações gerenciadas de ponto de verificação hierárquico

Você pode validar suas operações gerenciadas de ponto de verificação em camadas com registros.

Registro em log personalizado (opcional)

Você pode integrar logs de ponto de verificação a outros logs transferindo um agente de log personalizado para a biblioteca. Por exemplo, é possível adicionar um agente de log personalizado ao código de treinamento para que todos os logs da biblioteca também sejam coletados no agente de log de treinamento.

Registro em log aprimorado do serviço (opcional)

Para melhorar a depuração e a visibilidade do serviço, é possível montar o caminho de log do ponto de verificação /var/log/sagemaker_checkpointing dentro do pod até um caminho /var/logs/sagemaker_checkpointing no host. Isso garante que somente os logs específicos da biblioteca sejam coletados separadamente. Isso oferece visibilidade aprimorada à equipe do serviço para depuração e suporte.