Gradient checkpointing e a matemática da memória

01 · Conceito

Conceito

Suponha que o sharding já tenha reduzido o estado persistente por rank sob uma receita de optimizer explicitamente escolhida. No papel o job cabe. Então uma sequência de 8.192 tokens com batch size um morre no primeiro backward pass. Nada na aritmética do sharding estava necessariamente errado; outra categoria — ativações salvas — dominou o pico.

A backpropagation precisa de valores do forward pass. O autograd portanto retém ativações, cópias normalizadas, saídas de projeção e intermediários de attention até que suas funções de backward rodem, e num transformer profundo esses valores salvos rotineiramente excedem os pesos. O activation checkpointing — normalmente chamado de gradient checkpointing, embora nada nele envolva um arquivo em disco — reduz isso esquecendo e recomputando deliberadamente.

Acerte o tamanho primeiro, porque a intuição que as pessoas carregam está muito errada. A estimativa tentadora conta o residual stream: batch vezes sequência vezes hidden vezes camadas vezes bytes do dtype.

Cinco gibibytes. Confortável. Agora conte o que o autograd de fato guarda numa camada. Tome um modelo uniformemente de full attention com as dimensões deste — a conta real do híbrido exige custear as camadas DeltaNet à parte, o que é um exercício próprio; o que vem a seguir é o método e a ordem de grandeza. A entrada residual e sua cópia normalizada têm 5120 de largura cada. Numa camada de full attention as projeções não são quadradas: a projeção de query vai de 5120 para 12288, porque carrega a porta de saída junto com os 24 heads de dimensão 256 (lição 4.2); keys e values vão de 5120 para 1024 cada (4 KV heads), a saída da attention é 6144, e a projeção de saída retorna a 5120. Isso dá aproximadamente 35.800 valores por token. O bloco feed-forward é ainda maior: uma entrada normalizada de 5120 de largura, depois os tensores gate, up e do produto com gate, de largura 17.408 cada, depois uma saída de down-projection de 5120 — cerca de 62.500 valores. Somando:

Isso é 19 vezes o residual stream sozinho, e os três tensores de 17.408 de largura da FFN são a maior parte da diferença. Escale:

Noventa e seis gibibytes de ativações salvas para uma sequência com batch size um — mais de dez vezes o estado persistente que seu sharding tanto se esforçou para reduzir, e mais do que um cartão de 80 GB comporta. Esse é o número do modelo uniformemente de attention que postulamos, e é com ele que se raciocina. Para o Qwen3.8-27B especificamente é uma aproximação de sinal desconhecido: a lição 4.16 mostrou que só 16 das suas 64 camadas são de full attention, e as outras 48 são Gated DeltaNet, cujas ativações de treino são uma contabilidade à parte — um estado de inferência de tamanho fixo não implica armazenamento barato no backward, porque o mixer recorrente ainda materializa intermediários por token que o autograd precisa guardar. A metade da FFN, que é a maior parte do total, é idêntica nas 64 camadas de qualquer forma. A estimativa só do residual errou por um fator de aproximadamente dezenove, e erra na direção que faz um job ser agendado e depois morto.

Agora aplique o checkpointing, e deixe a arquitetura escolher a fronteira. Divida a rede em segmentos; durante o forward pass guarde apenas o tensor de fronteira de cada segmento em vez de tudo dentro dele; durante o backward, rode de novo o forward daquele segmento a partir da fronteira salva, reconstrua seus intermediários, compute os gradientes e libere-os. As 64 camadas deste modelo formam 16 super-blocos repetidos de três camadas de Gated DeltaNet mais uma camada de full attention, o que é uma segmentação natural e perfeitamente balanceada. Cada fronteira é um hidden state:

Durante o backward, um super-bloco é rematerializado por vez, custando o equivalente aos intermediários de quatro camadas — ainda precificados à taxa do modelo uniformemente de attention, mas essa taxa é apenas um ponto de comparação hipotético para o híbrido, não um teto conhecido nem um pico medido:

Dentro daquele modelo didático de full attention, a memória de pico de ativações salvas vira aproximadamente $1.25+6.0\approx 7.3$ GiB em vez de 96 GiB. A razão de cerca de treze vezes pertence apenas a essa contagem por formatos; ela não se transfere ao híbrido real até medir seu grafo de autograd de DeltaNet, kernels fundidos, allocator e política de recomputação.

O tradeoff é direto: menor pico de memória de ativações por computação extra. Não é simplesmente dobrar o tempo de treino, porque backward, comunicação e trabalho não checkpointado continuam de qualquer forma e a recomputação pode se sobrepor de modo diferente. Meça tempo de passo e throughput realizado, não o aumento teórico de FLOPs.

Um orçamento de memória deve manter as categorias separadas:

O checkpointing ataca $M_{activations}$ e nada mais. Ele não sharda parâmetros nem muda estado persistente do optimizer. Os dtypes e componentes exatos desse estado dependem da receita, e a lição 9.3 assume o orçamento de memória do Qwen. Um job dominado por estado persistente precisa de ZeRO ou FSDP; um dominado por ativações de sequência longa precisa de checkpointing.

O comprimento de sequência é onde este modelo fica interessante. A estimativa acima usou 8.192 tokens; o contexto nativo é 262.144. As ativações nas partes position-wise da rede crescem linearmente com $T$, então a mesma aritmética no contexto completo é trinta e duas vezes maior e absolutamente inviável sem checkpointing e o sequence parallelism da lição 5.8 juntos. Attention comum acrescentaria um termo quadrático em $T$ por cima, que é precisamente por que kernels de attention eficientes em memória, que nunca materializam a matriz completa de probabilidades, são obrigatórios nesses comprimentos — e por que só 16 das 64 camadas enfrentam esse termo, com as outras 48 varrendo posições com um estado recorrente de tamanho fixo.

Duas armadilhas de correção merecem ser nomeadas. Operações estocásticas precisam reproduzir sua aleatoriedade: se o dropout sorteou uma máscara no forward original e outra durante a recomputação, o backward derivaria uma função diferente, então utilitários de checkpoint dos frameworks preservam e restauram o estado do RNG a um custo pequeno. E módulos com estado, efeitos colaterais, caches mutáveis e chamadas externas dependentes de dados podem tornar a recomputação silenciosamente incorreta em vez de ruidosamente quebrada.

A composição acrescenta o resto da complexidade. O pipeline parallelism faz a recomputação competir com outros microbatches pelos mesmos dispositivos. O FSDP pode precisar de um novo all-gather dos pesos durante a recomputação, a menos que os parâmetros tenham permanecido materializados, trocando comunicação por memória. Compiladores fundem e reordenam regiões de formas que mudam o que um segmento sequer é. A unidade efetiva é o escalonamento distribuído inteiro, nunca uma função isolada.

O modelo mental duradouro é uma troca tempo–espaço: guarde estado de fronteira suficiente para reconstruir o grafo do forward, descarte o miolo volumoso e pague compute para recriá-lo pouco antes de seus gradientes serem necessários. A decisão decorre de uma decomposição medida de para onde a memória de fato foi — que, num modelo com um bloco feed-forward de 17.408 de largura, quase nunca é onde uma primeira estimativa a coloca.