O que é
Um passo de treinamento tem duas metades. Na passagem para frente, cada camada transforma sua entrada em uma saída, produzindo um conjunto de resultados intermediários ao longo do caminho: as ativações. A passagem para trás então trabalha através das camadas na ordem inversa para calcular os gradientes, e os gradientes de cada camada precisam das ativações daquela camada. A abordagem simples mantém todas elas na memória até que a passagem para trás as alcance. Com exemplos longos, as ativações podem ocupar mais memória do que os próprios pesos do modelo.
Gradient checkpointing, também chamado de checkpointing de ativação ou recomputação, mantém apenas algumas ativações, os checkpoints, e descarta o restante. Quando a propagação reversa precisa de uma ativação descartada, ela executa novamente a propagação direta a partir do checkpoint mais próximo para recriá-la. Chen et al. (2016) demonstraram que isso reduz a memória de ativação de uma rede de n camadas para a ordem de √n, ao custo de cerca de uma propagação direta extra por lote. Em seus experimentos, uma rede de 1.000 camadas passou de 48 GB para 7 GB com cerca de 30% a mais de tempo de execução.
Como funciona em um modelo de linguagem
Nas bibliotecas de transformers, os checkpoints geralmente ficam nos limites das camadas. A entrada de cada camada é mantida; tudo o que é computado dentro da camada é descartado e reconstruído quando a propagação reversa chega lá. Assim, apenas o conjunto completo de ativações de uma camada existe a qualquer momento, em vez de todas elas.
activations per GPT-style layer, 16-bit, efficient attention (Korthikanti et al., 2022)
≈ 34 × batch × tokens × hidden size bytes
illustrative 2B-class model: hidden size 2,048 · 28 layers · batch 2 × 1,024 tokens
all kept 34 × 2 × 1,024 × 2,048 × 28 ≈ 4.0 GB
checkpointed, inputs 2 × 2 × 1,024 × 2,048 × 28 ≈ 0.23 GB
+ one layer rebuilt at a time ≈ 0.14 GBA quantidade de checkpoints a manter é ajustável. Manter mais significa menos para recomputar, mas exige mais memória; manter menos significa o inverso. O checkpointing em cada limite de camada é o padrão habitual porque é simples e já remove a maior parte da memória de ativação: o que resta é uma pequena entrada por camada, além da única camada sendo reconstruída.
O custo é computacional. A propagação direta efetivamente é executada duas vezes: uma vez normalmente e outra vez, parte por parte, durante a propagação reversa. A propagação reversa já custa cerca do dobro da propagação direta, de modo que o trabalho extra é de, no máximo, cerca de um terço de um passo; na prática, o treinamento fica cerca de 20–30% mais lento. O resultado não muda: os valores recomputados são os mesmos, portanto a curva de perda corresponde a uma execução sem checkpointing, exceto por pequenas diferenças numéricas.
Por que vocabulários grandes o tornam mais importante
No final do modelo, cada posição em cada exemplo recebe uma pontuação para cada token no vocabulário: os logits. O tamanho deles não depende da profundidade do modelo, apenas de quantas posições são treinadas e do tamanho do vocabulário. Modelos mais recentes têm vocabulários de 150.000 a 250.000 tokens, de modo que apenas os logits podem ocupar gigabytes.
logits = batch × tokens × vocabulary size × bytes per number
2 × 1,024 × 248,000 × 4 bytes (32-bit) ≈ 2.0 GB before their gradients- O modelo: cerca de 2 bilhões de parâmetros, da família do Qwen3.5, com um vocabulário de aproximadamente 248.000 tokens.
- Os logits: lote de 2 × 1.024 tokens × 248.000 × 4 bytes ≈ 2 GB, e seus gradientes precisam de aproximadamente a mesma quantidade novamente.
- Todo o resto: os pesos de 16 bits, o estado de treinamento do adaptador e as ativações de cada camada, todos competindo pelo que resta.
- Sem o checkpointing, uma execução como essa pode ficar sem memória nos seus primeiros passos. Com ele, as ativações das camadas são reduzidas a uma fração e a mesma execução cabe na memória.
- Nota: o checkpointing não reduz os logits em si. Ele libera o espaço ao redor deles, o que é o que faz a execução caber na memória.
Utilizando-o bem
- Recorra a ele antes de cortar dados. Ativar o checkpointing custa tempo; reduzir o comprimento máximo corta o final de exemplos longos. Quando a memória é o limite, a lentidão geralmente é a melhor troca.
- Gaste a memória que ele libera. Um lote maior ou um comprimento máximo mais longo pode recuperar parte da velocidade perdida por época.
- Combine-o com QLoRA para os maiores modelos que uma máquina pode suportar: o QLoRA reduz os pesos, o checkpointing reduz as ativações.
- Saiba o que ele não afeta. Pesos, gradientes e o estado do otimizador mantêm o mesmo tamanho, e ele não faz nada pela execução do modelo finalizado.
- Verifique o padrão. Muitas receitas de fine-tuning já o ativam. Se a memória for confortável e a velocidade for importante, desativá-lo torna cada passo mais rápido.
Erros comuns
- Reduzir o comprimento máximo primeiro, cortando o final de exemplos longos, quando apenas o checkpointing já os faria caber.
- Esperar que ele reduza os pesos ou o estado do otimizador. Ele afeta apenas as ativações.
- Interpretar a lentidão de 20–30% como uma falha na máquina. É o preço da memória economizada.
- Presumir que ele remove o custo dos logits. Com um vocabulário enorme, um lote menor ou um comprimento máximo menor ainda pode ser necessário.
Referências e leitura complementar
Fontes externas abrem em uma nova aba.
- 01Training Deep Nets with Sublinear Memory CostChen et al., 2016 · arXiv · arxiv.orgO artigo que tornou a recomputação um padrão: memória √n por uma propagação direta extra.
- 02Reducing Activation Recomputation in Large Transformer ModelsKorthikanti et al., 2022 · arXiv · arxiv.orgDe onde vem a estimativa de ativação por camada e como recomputar apenas as partes custosas.
- 03torch.utils.checkpointPyTorch docs · docs.pytorch.org
- 04GPU memory usageHugging Face Transformers docs · huggingface.coPesos, gradientes, estado do otimizador e ativações: para onde vai a memória de treinamento.