Zubnet AIAprenderWiki › Gradient Checkpointing
Treinamento

Gradient Checkpointing

Também conhecido como: Activation Checkpointing, Rematerialização
Uma técnica de economia de memória que troca computação por memória durante o treinamento. Em vez de armazenar todas as ativações intermediárias do forward pass (necessárias para backpropagation), gradient checkpointing armazena ativações apenas em certas camadas de "checkpoint" e recalcula as outras durante o backward pass. Isso reduz o uso de memória em até 5–10x ao custo de ~30% mais computação.

Por que isso importa

Gradient checkpointing é o que torna possível fazer fine-tuning de modelos grandes em memória GPU limitada. Sem ele, um modelo 7B poderia precisar de 80+ GB só para ativações durante o treinamento, excedendo a capacidade de uma única GPU. Com gradient checkpointing, o mesmo modelo pode ser ajustado em uma GPU de consumo de 24GB. É a otimização de memória mais comumente usada para treinamento.

Em profundidade

Durante o forward pass, as ativações de entrada de cada camada são necessárias durante o backward pass para computar gradientes. Normalmente, todas as ativações são armazenadas em memória. Com gradient checkpointing, apenas ativações de certas camadas são armazenadas. Durante o backward pass, quando uma ativação não armazenada é necessária, o forward pass é re-executado a partir do checkpoint mais próximo para recalculá-la. Isso troca ~30% de computação extra (recalculando ativações) por ~5x de economia de memória (não armazenando todas).

Posicionamento de Checkpoints

O posicionamento ideal de checkpoints depende da arquitetura do modelo. A abordagem mais simples: fazer checkpoint a cada N camadas (ex.: a cada 3 blocos Transformer). Mais sofisticado: analisar os tamanhos de ativação por camada e posicionar checkpoints para minimizar memória total limitando a recomputação. Alguns frameworks (torch.utils.checkpoint do PyTorch) tornam isso tão simples quanto envolver uma chamada de camada em uma função de checkpoint.

Combinando com Outras Técnicas

Gradient checkpointing se compõe com outras otimizações de memória: precisão mista (FP16/BF16 reduz pela metade o tamanho das ativações), acumulação de gradiente (lotes menores reduzem pico de memória) e FSDP/DeepSpeed (fragmenta parâmetros entre GPUs). Juntas, essas técnicas podem reduzir o footprint de memória de um modelo em 10–50x comparado a treinamento ingênuo em FP32, permitindo treinar modelos muito maiores que a memória de qualquer GPU individual. Essa pilha de otimizações é padrão para fine-tuning de modelos 7B+.

Conceitos relacionados

← Todos os termos
← GQA Gradiente Descendente →