Aktivierung, Entladung

Wichtig

In SMP Version 2.2.0 funktioniert die Aktivierungs-Offloading-Funktionalität der SMP Bibliothek nicht. Verwenden Sie stattdessen das native PyTorch Aktivierungs-Offloading.

Normalerweise berechnet der Vorwärtsdurchlauf Aktivierungen auf jeder Ebene und behält sie im GPU Speicher, bis der Rückwärtsdurchlauf für die entsprechende Ebene abgeschlossen ist. Wenn diese Tensoren nach dem Vorwärtsdurchlauf in den CPU Speicher ausgelagert und dann wieder abgerufen werden, GPU wenn sie benötigt werden, kann erheblicher Speicherverbrauch eingespart werden. GPU PyTorch unterstützt das Auslagern von Aktivierungen, aber die Implementierung führt GPUs dazu, dass sie inaktiv sind, während Aktivierungen während des Rückwärtsdurchlaufs abgerufen werden. CPU Dies führt zu erheblichen Leistungseinbußen, wenn das Aktivierungs-Offloading verwendet wird.

SMPv2 verbessert dieses Aktivierungs-Offloading. Aktivierungen werden vorab abgerufen, bevor die Aktivierungen für den GPU Rückwärtsstart dieser Aktivierungen benötigt werden. Mit der Funktion zum Vorabrufen können Trainingsfortschritte effizienter und ohne Leerlauf ausgeführt werden. GPUs Dies führt zu Vorteilen einer geringeren Speicherauslastung ohne Leistungseinbußen.

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

• checkpoint_wrapper.py im Repository PyTorch GitHub

• Checkpointing zur Aktivierung im PyTorch Blog Scaling Multimodal Foundation Models in TorchMultimodal with PyTorch Distributed.

Sie können die Funktion zum Offloading der SMP Aktivierung beim Checkpoint der PyTorch Aktivierung anwenden. Dazu werden die activation_loading_horizon Parameter sm_activation_offloading und währenddessen dem SMP Konfigurationswörterbuch hinzugefügt. Schritt 2: Starten Sie einen Schulungsjob

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 unter beschriebenen zweistufigen Prozess. Erste Schritte mit der SageMaker Modellparallelismus-Bibliothek v2 Sie müssen keine Änderungen an Ihrem Modell oder Ihrer PyTorch Konfiguration vornehmen. PyTorch FSDP Weitere Hinweise zu den Parametern sm_activation_offloading und activation_loading_horizon finden Sie unter SMPKonfigurationsparameter der v2-Kernfunktion.

SMPKonfiguration

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

Im Trainingsskript

Anmerkung

Achten Sie bei der SMP Aktivierung der Aktivierungs-Offloading-Funktion darauf, dass Sie die PyTorch offload_wrapper Funktion auch verwenden und sie auf das Root-Modul anwenden. Die Funktion zum Offloading der SMP Aktivierung verwendet das Root-Modul, um zu ermitteln, wann ein Forward-Durchlauf durchgeführt wird, um mit dem Prefetching zu beginnen.

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)