Configuración de los puntos de comprobación por niveles administrados

En esta sección se explica el proceso de configuración de los puntos de comprobación por niveles administrados 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 comprobación por niveles administrados, asegúrese de que cuenta con lo siguiente:

Paso 1: activación de los puntos de comprobación por niveles administrados para su clúster

Activa los puntos de comprobación por niveles administrados a través de la API de HyperPod al crear o actualizar su clúster. El servicio instala automáticamente el sistema de administración de memoria cuando especifica el parámetro TieredStorageConfig. Para un clúster nuevo, crea:

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

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

Instale la biblioteca de puntos de comprobación de Amazon SageMaker en su imagen de entrenamiento añadiéndola a su Dockerfile:

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

Paso 3: creación de una configuración de protección

Cree un objeto CheckpointConfig para especificar el comportamiento del punto de comprobación. Esto incluye:

En el siguiente ejemplo se muestra una configuración de control de comprobación:

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
)

Paso 4: definición de un escritor de sistema de archivos de SageMaker

Defina su escritor de sistema de archivos de puntos de comprobación. Opcionalmente, puede especificar un número de paso durante la inicialización.

Escritor básico (parámetro step especificado al guardar la llamada):

smWriter = sagemaker_checkpointing.SageMakerTieredStorageWriter(checkpoint_config)

Escritor con parámetro step (parámetros step especificado en la inicialización):

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

Paso 5: grabación de los puntos de comprobación en el bucle de entrenamiento

En su bucle de entrenamiento, guarde los puntos de comprobación utilizando PyTorch DCP con FileSystemWriter.

Uso de PyTorch DCP con FileSystemWriter

Llame al método dist_cp.save() con la entrada FileSystemWriter:

Opción 1: usar checkpoint_id con el formato step (cuando step no está especificado en el escritor)

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

Opción 2: usar el escritor con el parámetro step (checkpoint_id pasa a ser opcional)

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

Paso 6: carga de los puntos de comprobación para la recuperación

Cuando necesite cargar un punto de comprobación, utilice PyTorch DCP con FileSystemReader.

Uso de PyTorch DCP con FileSystemReader

Llame al método de carga DCP con la entrada FileSystemReader:

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

Supervisión y validación

Puede supervisar y validar las operaciones de puntos de comprobación por niveles administrados mediante métricas y 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.