Offload di attivazione

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

SMPv2 migliora questo offload di attivazione. Pre-recupera le attivazioni in anticipo prima che siano necessarie per avviare il trasferimento GPU all'indietro di tali attivazioni. La funzione di preacquisizione consente di eseguire i progressi dell'allenamento in modo più efficiente senza inattività. GPUs 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'SMPattivazione 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:

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

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

SMPconfigurazione

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

Nel copione di allenamento

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)