Gradient Checkpointing

O gradient checkpointing é uma técnica de treino que poupa memória, armazenando apenas ativações intermediárias selecionadas da passagem direta e recomputando as outras durante a backpropagation. Também chamado de activation checkpointing, ele troca computação adicional por um uso de pico de memória menor. Apesar do nome, a técnica faz o checkpoint de ativações em vez de gradientes de parâmetros ou arquivos de modelo, tornando-a especialmente valiosa quando tensores de ativação impedem que uma rede neural caiba na GPU memory disponível.

Como Funciona o Gradient Checkpointing#

Durante uma passagem direta padrão, uma rede neural calcula tensores intermediários chamados ativações. O sistema de diferenciação automática retém as ativações necessárias para calcular gradientes mais tarde, conforme descrito em PyTorch autograd mechanics. Redes profundas, lotes grandes, imagens de alta resolução e sequências de entrada longas podem fazer com que esses tensores salvos consumam muita memória.

O gradient checkpointing divide a rede em segmentos:

1. A passagem direta armazena entradas ou ativações de limite para segmentos selecionados.
2. Outras ativações intermediárias dentro desses segmentos são descartadas.
3. Durante a passagem reversa, cada segmento com checkpoint executa a passagem direta novamente para reconstruir os valores ausentes.
4. As ativações reconstruídas são usadas imediatamente para calcular gradientes.

O PyTorch activation checkpointing API expõe esse comportamento por meio de torch.utils.checkpoint. Conceitos equivalentes aparecem como JAX gradient checkpointing and rematerialization e TensorFlow tape checkpointing.

A colocação de checkpoints determina a compensação (tradeoff). Fazer o checkpoint de mais regiões geralmente economiza mais memória, mas repete mais operações. A colocação seletiva em torno de blocos pesados em ativação pode proporcionar um equilíbrio melhor do que recomputar toda a rede.

Compensações e Técnicas Relacionadas#

O gradient checkpointing geralmente deixa os parâmetros do modelo, gradientes e estados do otimizador inalterados. Seu principal alvo é a memória de ativação, e o treino torna-se mais lento porque algumas computações diretas são executadas duas vezes. O resultado exato depende da arquitetura, dos limites de checkpoint, da forma do lote e do hardware.

Ele difere de várias técnicas relacionadas:

• Gradient accumulation processa múltiplos microbatches antes de atualizar os pesos do modelo, criando um lote efetivo maior sem carregar cada amostra simultaneamente. Em vez disso, o gradient checkpointing reduz as ativações retidas para cada microbatch.
• Mixed precision usa tipos de dados de menor precisão para operações selecionadas. O PyTorch automatic mixed precision workflow pode reduzir a memória e acelerar operações compatíveis, enquanto o checkpoint adiciona deliberadamente computação.
• A redução de Batch size diminui a memória processando menos amostras juntas. O checkpoint pode permitir que um lote maior ou uma resolução de entrada continue viável.
• Os checkpoints de treino salvam pesos e o estado do otimizador para recuperação ou inferência posterior. O PyTorch model checkpoint guide descreve esse mecanismo de persistência, que não está relacionado à recomputação de ativação.

O código com checkpoint deve ser funcionalmente consistente entre suas execuções diretas originais e recomputadas. Estado mutável, transferências de dispositivos ou aleatoriedade descontrolada dentro de uma região com checkpoint podem causar erros ou gradientes incorretos. Em transformers estilo decodificador, um key-value cache também pode precisar ser desativado durante o treino com checkpoint, porque o estado de inferência em cache pode entrar em conflito com a reconstrução do grafo direto.

Exemplo em PyTorch#

O exemplo a seguir faz o checkpoint de um bloco que consome muita memória durante uma etapa de treino:

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

torch.manual_seed(0)

block = nn.Sequential(
    nn.Linear(1024, 4096),
    nn.ReLU(),
    nn.Linear(4096, 1024),
)
optimizer = torch.optim.AdamW(block.parameters())
inputs = torch.randn(8, 1024, requires_grad=True)
targets = torch.zeros_like(inputs)

optimizer.zero_grad(set_to_none=True)
outputs = checkpoint(block, inputs, use_reentrant=False)
loss = nn.functional.mse_loss(outputs, targets)
loss.backward()
optimizer.step()

Apenas as entradas do bloco e as informações de limite necessárias são retidas; suas ativações internas são reconstruídas durante loss.backward(). Projetos reais devem comparar execuções com e sem checkpoint usando uma métrica como PyTorch peak allocated GPU memory.

Aplicações no Mundo Real#

• Visão computacional de alta resolução: Segmentação médica, deteção aérea e inspeção industrial podem ser treinadas em imagens grandes cujos mapas de características consomem mais memória do que os pesos do modelo. Fazer o checkpoint de estágios selecionados da espinha dorsal (backbone) pode preservar a resolução da imagem ou permitir amostras adicionais por lote em vez de reduzir drasticamente as entradas.

• Treino de transformers de sequência longa: O armazenamento de ativação cresce rapidamente à medida que o comprimento da sequência e a contagem de camadas aumentam. Recomputar blocos de transformers pode fazer com que contextos mais longos ou microbatches maiores caibam no mesmo acelerador. O NVIDIA activation recomputation guide ilustra a recomputação total e seletiva para camadas de transformers.

Orientações Práticas#

Usa o gradient checkpointing quando a análise de perfil (profiling) mostrar que as ativações, em vez dos parâmetros ou do estado do otimizador, dominam a memória. Começa com blocos grandes e repetidos e faz o benchmark do pico de memória, do tempo de iteração e do comportamento de validação antes de expandir a cobertura.

Para as Ultralytics YOLO training configurations documentadas, os controlos de memória de primeira linha incluem precisão mista automática, tamanho de imagem e tamanho de lote físico. O Ultralytics AutoBatch reference explica a seleção automática de lote com base na memória de GPU disponível. Quando o hardware local permanece insuficiente, o Ultralytics Platform cloud training fornece GPUs de nuvem configuráveis para execuções de treino geridas. O gradient checkpointing pode complementar esses controlos quando uma arquitetura PyTorch personalizada requer um gerenciamento mais refinado de memória de ativação.