Offload di attivazione

Importante

In SMP v2.2.0, la funzionalità di offload di attivazione della libreria SMP non funziona. Utilizzate invece l'offload di attivazione nativo. PyTorch

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 nella memoria della CPU dopo il passaggio in avanti e il loro recupero nella GPU quando sono necessari può far risparmiare un notevole utilizzo della memoria della GPU. PyTorch supporta l'offload delle attivazioni, ma l'implementazione fa sì che le GPU restino inattive mentre le attivazioni vengono recuperate dalla CPU durante il passaggio all'indietro. Ciò causa un notevole peggioramento delle prestazioni quando si utilizza l'offload di attivazione.

SMP v2 migliora questo offload di attivazione. Pre-recupera le attivazioni in anticipo prima che siano necessarie affinché la GPU inizi a trasmettere tali attivazioni all'indietro. La funzione di pre-fetching consente di eseguire i progressi della formazione in modo più efficiente senza GPU inattive. Ciò si traduce nell'offrire vantaggi derivanti da un minore utilizzo della memoria senza un 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. Si noti che l'offload di attivazione è applicabile solo se utilizzato insieme a. Checkpoint di attivazione Per saperne di più sugli strumenti nativi dei PyTorch checkpoint per l'offload dell'attivazione, consulta:

• checkpoint_wrapper.py nel repository PyTorch GitHub

• Activation Checkpointing nel PyTorch blog Scaling Multi-modal Foundation Models in with Distributed. TorchMultimodal PyTorch

È possibile applicare la funzionalità di offload dell'attivazione SMP al checkpoint di attivazione. PyTorch Ciò viene fatto aggiungendo i activation_loading_horizon parametri sm_activation_offloading and al dizionario di configurazione SMP durante. Fase 2: Avviare un lavoro di formazione

I seguenti frammenti di codice mostrano come aggiungere il modulo di inizializzazione SMP allo script di formazione e torch.sagemaker.init() impostare il dizionario di configurazione SMP in formato JSON per Training Job Launcher seguendo il processo in due fasi introdotto in. Inizia con la libreria di parallelismo dei SageMaker modelli v2 Non è necessario apportare modifiche al modello o alla configurazione FSDP. PyTorch PyTorch Per ulteriori informazioni sui parametri sm_activation_offloading e activation_loading_horizon, vedi Parametri di configurazione delle funzionalità principali di SMP v2.

Configurazione SMP

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

Nello script di addestramento

Nota

Durante l'attivazione della funzionalità di offload di attivazione SMP, assicuratevi di utilizzare anche la PyTorch offload_wrapper funzione e di applicarla al modulo root. La funzionalità di offload di attivazione SMP utilizza il modulo root per determinare quando viene eseguito il forward pass per avviare il prefetching.

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)