Offload di attivazione

In genere, il passaggio in avanti calcola le attivazioni su ogni livello e le conserva nella memoria della GPU fino al termine del passaggio all’indietro per il livello corrispondente. Lo scaricamento di questi tensori sulla memoria della CPU dopo il passaggio in avanti e il loro recupero sulla GPU quando sono necessari può far risparmiare un notevole utilizzo della memoria della GPU. PyTorch supporta l'offload delle attivazioni, ma l'implementazione rimane inattiva mentre le attivazioni vengono GPUs recuperate dalla CPU durante il passaggio all'indietro. Questo causa un notevole peggioramento delle prestazioni quando si utilizza l’offload di attivazione.

SMP v2 è in grado di migliorare questo offload di attivazione. Recupera le attivazioni in anticipo, prima che siano necessarie, affinché la GPU possa iniziare a passarle all’indietro. La funzione di pre-fetching consente di eseguire i progressi dell'allenamento in modo più efficiente senza inattività. GPUs Questo consente di usufruire dei vantaggi derivanti da un minore utilizzo della memoria senza peggioramento delle prestazioni.

È possibile mantenere i PyTorch moduli nativi per l'offload delle attivazioni nello script di formazione. Di seguito è riportato un esempio di struttura di applicazione della funzionalità di offload dell’attivazione SMP nello script. Tieni presente che l’offload di attivazione è applicabile solo se utilizzato insieme a Checkpoint di attivazione. Per saperne di più sugli strumenti nativi di PyTorch checkpoint per l'offload delle attivazioni, consulta:

È possibile applicare la funzionalità di offload dell'attivazione SMP al checkpoint di attivazione. PyTorch A tale scopo, è sufficiente aggiungere i parametri sm_activation_offloading e activation_loading_horizon al dizionario di configurazione di SMP durante Fase 2: avvio di un job di addestramento.

I seguenti frammenti di codice mostrano come aggiungere il modulo di inizializzazione SMP torch.sagemaker.init() allo script di addestramento e impostare il dizionario di configurazione di SMP in formato JSON per il programma di avvio del job di addestramento seguendo il processo in due fasi introdotto in Utilizzo della libreria SageMaker Model Parallelism v2. Non è necessario apportare modifiche al PyTorch modello o alla configurazione FSDP. PyTorch Per ulteriori informazioni sui parametri sm_activation_offloading e activation_loading_horizon, vedi Parametri di configurazione delle funzionalità principali di SMP v2.

Configurazione di SMP

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

Nello script di addestramento

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)