Richten Sie verwaltetes Tier-Checkpointing ein

Dieser Abschnitt enthält den Einrichtungsprozess für verwaltetes Tier-Checkpointing für Amazon SageMaker HyperPod. Sie erfahren, wie Sie die Funktion in Ihrem Cluster aktivieren und Checkpointing in Ihrem Trainingscode implementieren.

Voraussetzungen

Bevor Sie Managed Tier Checkpointing einrichten, stellen Sie sicher, dass Sie über Folgendes verfügen:

Schritt 1: Managed-Tier-Checkpointing für Ihren Cluster aktivieren

Aktivieren Sie beim Erstellen oder Aktualisieren Ihres Clusters Managed-Tier-Checkpointing über die HyperPod-API. Der Service installiert das Speicherverwaltungssystem automatisch, wenn Sie den TieredStorageConfig-Parameter angeben. Für neue Cluster-Erstellungen:

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

Der InstanceMemoryAllocationPercentage-Parameter gibt die percentage (int) des Cluster-Speichers an, der für Checkpointing zugewiesen werden soll. Der Bereich liegt zwischen 20 und 100.

Schritt 2: Python-Bibliothek in Ihrem Trainings-Image installieren

Installieren Sie die Checkpointing-Bibliothek von Amazon SageMaker in Ihrem Trainings-Image, indem Sie sie zu Ihrer Dockerfile-Datei hinzufügen:

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

Schritt 3: Eine Checkpoint-Konfiguration erstellen

Erstellen Sie ein CheckpointConfig-Objekt, um das Verhalten des Checkpoints festzulegen. Dies umfasst:

Im folgenden Beispiel wird die Checkpoint-Konfiguration gezeigt.

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
)

Schritt 4: Einen SageMaker-Dateisystem-Writer definieren

Definieren Sie Ihren Checkpointing-Dateisystem-Writer. Optional können Sie bei der Initialisierung auch eine Schrittnummer angeben.

Basic Writer (Schritt wurde im Speicheraufruf angegeben):

smWriter = sagemaker_checkpointing.SageMakerTieredStorageWriter(checkpoint_config)

Writer mit Schrittparameter (Schritt wurde bei der Initialisierung angegeben):

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

Schritt 5: Checkpoints in Ihrer Trainingsschleife speichern

Speichern Sie in Ihrer Trainingsschleife Checkpoints mithilfe von PyTorch-DCP mit FileSystemWriter.

Verwenden von PyTorch-DCP mit FileSystemWriter

Rufen Sie die dist_cp.save()-Methode mit FileSystemWriter als Eingabe auf:

Option 1: checkpoint_id mit Schrittformat verwenden (wenn der Schritt nicht im Writer angegeben ist)

# 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: Writer mit Schrittparameter verwenden (checkpoint_id wird optional)

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

Schritt 6: Checkpoints für Wiederherstellung laden

Wenn Sie einen Checkpoint laden müssen, verwenden Sie PyTorch-DCP mit FileSystemReader.

Verwenden von PyTorch-DCP mit FileSystemReader

Rufen Sie die DCP-Lademethode mit FileSystemReader als Eingabe auf:

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

Überwachung und Validierung

Sie können Ihre Operationen im Rahmen von Managed-Tier-Checkpointing anhand von Metriken und Protokollen überwachen und validieren.

Benutzerdefinierte Protokollierung (optional)

Sie können Checkpointing-Protokolle in andere Protokolle integrieren, indem Sie einen benutzerdefinierten Logger an die Bibliothek übergeben. Beispielsweise können Sie Ihrem Trainingscode einen benutzerdefinierten Logger hinzufügen, sodass alle Protokolle aus der Bibliothek auch im Trainings-Logger gesammelt werden.

Verbesserte Serviceprotokollierung (optional)

Um das Debugging und die Transparenz der Services zu verbessern, können Sie den Checkpointing-Protokollpfad /var/log/sagemaker_checkpointing von Ihrem Pod aus in einen /var/logs/sagemaker_checkpointing-Pfad auf Ihrem Host mounten. Dadurch wird sichergestellt, dass nur bibliotheksspezifische Protokolle separat gesammelt werden. So erhält das Serviceteam verbesserte Transparenz hinsichtlich Debugging und Support.