Aktivierung, Entladung

In der Regel werden beim Vorwärtsdurchlauf Aktivierungen auf jeder Ebene berechnet und im GPU-Speicher belassen, bis der Rückwärtsdurchlauf für die entsprechende Ebene abgeschlossen ist. Wenn Sie diese Tensoren nach dem Forward-Durchlauf in den CPU-Speicher auslagern und sie bei Bedarf wieder auf die GPU abrufen, kann der GPU-Speicherverbrauch erheblich gesenkt werden. PyTorch unterstützt das Auslagern von Aktivierungen, aber die Implementierung führt dazu, dass GPUs sie inaktiv sind, während Aktivierungen während des Rückwärtsdurchlaufs von der CPU abgerufen werden. Dies führt zu erheblichen Leistungseinbußen, wenn das Auslagern von Aktivierungen verwendet wird.

SMP v2 verbessert die Aktivierungsauslagerung. Es ruft Aktivierungen im Voraus ab, bevor sie benötigt werden, damit die GPU mit dem Rückwärtsdurchlauf dieser Aktivierungen beginnen kann. Mit der Funktion zum Vorabrufen können Trainingsfortschritte effizienter und ohne Leerlauf ausgeführt werden. GPUs Dies führt dazu, dass die Vorteile einer geringeren Speichernutzung ohne Leistungseinbußen genutzt werden können.

Sie können die systemeigenen PyTorch Module zum Auslagern von Aktivierungen in Ihrem Trainingsskript behalten. Im Folgenden finden Sie eine Beispielstruktur für die Anwendung der Funktion zum Auslagern der SMP-Aktivierung in Ihrem Skript. Beachten Sie, dass das Auslagern von Aktivierungen nur in Kombination mit Checkpointing bei der Aktivierung verwendet wird. Weitere Informationen zu den systemeigenen PyTorch Checkpoint-Tools für das Offloading von Aktivierungen finden Sie unter:

Sie können die SMP-Aktivierungsauslagerungsfunktion beim Aktivierungs-Checkpointing anwenden. PyTorch Dazu werden die Parameter activation_loading_horizon und sm_activation_offloading während Schritt 2: Starten eines Trainingsjobs dem SMP-Konfigurationswörterbuch hinzugefügt.

Die folgenden Codefragmente zeigen, wie Sie das SMP-Initialisierungsmodul torch.sagemaker.init() zu Ihrem Trainingsskript hinzufügen und das SMP-Konfigurationswörterbuch im JSON-Format für den Trainingsjob-Launcher einrichten. Dabei folgen Sie dem in Verwendung der SageMaker-Modellparallelitätsbibliothek v2 beschriebenen zweistufigen Prozess. Sie müssen keine Änderungen an Ihrem PyTorch Modell oder PyTorch Ihrer FSDP-Konfiguration vornehmen. Weitere Hinweise zu den Parametern sm_activation_offloading und activation_loading_horizon finden Sie unter Konfigurationsparameter für die Kernfunktionen von SMP v2.

SMP-Konfiguration

{
    "activation_loading_horizon": 2,
    "sm_activation_offloading": True
}

Im Trainingsskript

import torch.sagemaker as tsm
tsm.init()

# Native PyTorch module for activation offloading
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing, 
    offload_wrapper,
)

model = FSDP(...)

# Activation offloading requires activation checkpointing.
apply_activation_checkpointing(
    model,
    check_fn=checkpoint_transformer_layers_policy,
)

model = offload_wrapper(model)