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. L’offload di questi tensori nella memoria della CPU dopo il passaggio in avanti e il loro recupero nella GPU quando sono necessari può ridurre notevolmente l’utilizzo della memoria della GPU. PyTorch supporta le attivazioni di offload, ma l’implementazione fa sì che le GPU restino inattive mentre le attivazioni vengono 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 funzionalità di recupero anticipato consente di eseguire i progressi dell’addestramento in modo più efficiente senza GPU inattive. Questo consente di usufruire dei vantaggi derivanti da un minore utilizzo della memoria senza peggioramento delle prestazioni.

Puoi conservare i moduli PyTorch nativi per scaricare le attivazioni nel tuo script di addestramento. 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 ulteriori informazioni sugli strumenti di checkpoint nativi di PyTorch per l’offload di attivazione, consulta:

Puoi applicare la funzionalità di offload di attivazione SMP sul checkpoint di attivazione di 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 alcuna modifica al modello PyTorch o alla configurazione PyTorch FSDP. 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)