Configure puntos de control gestionados por niveles - Amazon SageMaker AI
Requisitos previosPaso 1: Habilite los puntos de control gestionados por niveles para su clústerPaso 2: instalación de la biblioteca de Python en la imagen de entrenamientoPaso 3: Guarda los puntos de control en tu ciclo de entrenamientoPaso 4: Cargue los puntos de control para la recuperaciónValide sus operaciones gestionadas de puntos de control escalonados

Las traducciones son generadas a través de traducción automática. En caso de conflicto entre la traducción y la version original de inglés, prevalecerá la version en inglés.

Configure puntos de control gestionados por niveles

Esta sección contiene el proceso de configuración de los puntos de control por niveles gestionados para Amazon. SageMaker HyperPod Aprenderá a habilitar la capacidad en su clúster e implementar los puntos de comprobación en su código de entrenamiento.

Requisitos previos

Antes de configurar los puntos de control por niveles gestionados, asegúrate de:

  • Un HyperPod clúster de Amazon EKS con suficiente memoria de CPU disponible para la asignación de puntos de control

  • PyTorch cargas de trabajo de entrenamiento y trabajos de DCP (ambos son compatibles)

  • Permisos de IAM adecuados para la administración de clústeres, que incluyen:

    • Amazon CloudWatch y Amazon S3 escriben permisos para que el módulo de entrenamiento lea y escriba puntos de control y envíe métricas

    • Estos permisos se pueden configurar mediante la configuración de EKS OIDC

Paso 1: Habilite los puntos de control gestionados por niveles para su clúster

importante

Debe optar por utilizar los puntos de control por niveles gestionados.

Active los puntos de control gestionados por niveles HyperPod APIs al crear o actualizar el clúster. El servicio instala automáticamente el sistema de administración de memoria cuando especifica el parámetro TieredStorageConfig.

Para clústeres nuevos, puede utilizar. 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"
    }'

El parámetro InstanceMemoryAllocationPercentage especifica el percentage (ent) de la memoria del clúster que se va a asignar a los puntos de comprobación. El rango va de 20 a 100.

Paso 2: instalación de la biblioteca de Python en la imagen de entrenamiento

Instala la biblioteca SageMaker de puntos de control de Amazon y sus dependencias en tu imagen de entrenamiento añadiéndola a tu Dockerfile:

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

Paso 3: Guarda los puntos de control en tu ciclo de entrenamiento

En tu ciclo de entrenamiento, puedes guardar los puntos de control de forma asíncrona mediante DCP. PyTorch El siguiente es un ejemplo de cómo hacerlo.

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

Paso 4: Cargue los puntos de control para la recuperación

El siguiente es un ejemplo de cómo cargar un punto de control.

# 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 sus operaciones gestionadas de puntos de control escalonados

Puede validar sus operaciones de puntos de control por niveles gestionadas con registros.

Registro personalizado (opcional)

Puede integrar los registros de puntos de comprobación con otros registros pasando un registrador personalizado a la biblioteca. Por ejemplo, puede añadir un registrador personalizado a su código de entrenamiento para que todos los registros de la biblioteca también se recopilen en el registrador de entrenamiento.

Registro de servicios mejorado (opcional)

Para mejorar la depuración y la visibilidad del servicio, puede montar la ruta del registro de puntos de comprobación /var/log/sagemaker_checkpointing desde su pod hasta una ruta /var/logs/sagemaker_checkpointing del host. Esto garantiza que solo se recopilen por separado los registros específicos de la biblioteca. Esto proporciona al equipo de servicio mejor visibilidad para la depuración y la asistencia.