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 Vorwärtsdurchlauf in den CPU-Speicher auslagern und sie bei Bedarf wieder auf die GPU laden, kann die GPU-Speicherauslastung erheblich reduziert werden. PyTorch unterstützt das Auslagern von Aktivierungen, aber die Implementierung führt dazu, dass GPUs 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. Die Vorabruffunktion trägt dazu bei, dass Trainingsfortschritte effizienter ausgeführt werden können, ohne dass GPUs im Leerlauf verwendet werden müssen. Dies führt dazu, dass die Vorteile einer geringeren Speichernutzung ohne Leistungseinbußen genutzt werden können.

Sie können die nativen 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 nativen PyTorch-Prüfpunkt-Tools für das Auslagern von Aktivierungen finden Sie unter:

Sie können die Funktion zum Auslagern der SMP-Aktivierung auf das PyTorch-Aktivierungs-Checkpointing anwenden. 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 der PyTorch-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)