Checkpointing de gradiente

A passagem de propagação direta (forward pass) normalmente faz o cache de todas as ativações intermediárias para reutilizá-las na passagem de propagação reversa (backward pass). No entanto, as ativações escalam com o tamanho do lote (batch size) e o comprimento da sequência. O checkpointing de gradientes salva apenas certas ativações e descarta as demais. Isso força a passagem de propagação reversa a recalcular algumas das ativações em tempo real conforme necessário.

Normal training:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save ALL activations)
  Backward:  ←uses cached activations everywhere

Gradient checkpointing:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save only at checkpoints, discard the rest)
  Backward:  ←reaches L2, recomputes L2→L3 from scratch, uses it, discards it

O treinamento é tipicamente mais lento porque a passagem de retropropagação recalcula as ativações descartadas, mas o checkpointing reduz a memória de ativação.

Defina gradient_checkpointing=True para habilitar.

Use com gradient accumulation para reduzir ainda mais o uso de memória.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
)

Checkpointing parcial

O gradient checkpointing completo recalcula todas as camadas que podem ser checkpointadas. Se sua execução tiver alguma folga de memória, faça o checkpoint de menos camadas para trocar parte da economia de memória por velocidade.

Passe every_n_layers para gradient_checkpointing_enable() para escolher o intervalo de checkpointing.

every_n_layers=2

Forward:   input  -> [L1] -> [L2] -> [L3] -> [L4] -> [L5] -> [L6]
                      CP     keep     CP     keep     CP     keep

Backward:  output <- [L6] <- [L5] <- [L4] <- [L3] <- [L2] <- [L1]
                     keep   rerun    keep   rerun    keep   rerun

Com every_n_layers=2, a primeira camada e toda segunda camada subsequente são checkpointadas. As camadas checkpointadas descartam suas ativações durante a passagem direta e as recomputam durante a passagem reversa, enquanto as outras camadas mantêm suas ativações na memória.

model.gradient_checkpointing_enable(every_n_layers=2)

O padrão, every_n_layers=1, faz checkpoint de cada camada. Valores maiores fazem checkpoint da primeira camada e, em seguida, a cada n camadas, mantendo as ativações das outras camadas na memória. Por exemplo, every_n_layers=2 faz checkpoint das camadas 1, 3, 5 e assim por diante. Somente módulos que herdam de GradientCheckpointingLayer são contabilizados. Outros módulos que suportam gradient checkpointing permanecem habilitados.

Para usar o checkpointing parcial de gradiente com o Trainer, defina every_n_layers em gradient_checkpointing_kwargs.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"every_n_layers": 4},
)

Transferência para fora das ativações salvas

O gradient checkpointing mantém uma ativação por camada com checkpoint na GPU. Em comprimentos de sequência longos, isso pode consumir uma quantidade considerável de memória (aproximadamente layers x sequence x hidden x bytes_per_element). Defina offload para manter essas ativações em memória de host fixada, com o custo de uma cópia de dispositivo para host durante a passagem direta e uma cópia de host para dispositivo durante a passagem reversa.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"offload": True},
)

Ambas as cópias são executadas no stream de computação, então isso troca um passo mais lento por memória. Recorra a ele quando uma execução não couber de outra forma, e não para acelerá-la.

Próximas etapas

• Leia o documento uso de memória da GPU para entender o que está causando o uso de memória na GPU durante o treinamento.
• Consulte o guia Training with mixed precision para aprender a usar tipos de dados de menor precisão para reduzir ainda mais a memória e acelerar o treinamento.
• Consulte o guia Kernels para aprender a acelerar o treinamento com kernels fundidos personalizados.