🤖 machine learning

FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training

FORGE é um método de treinamento de LLM eficiente em memória que funde a etapa do otimizador na passagem de retropropagação para eliminar gradientes materializados ao processá-los inteiramente em registradores, reduzindo assim pela metade o uso de memória do otimizador, acelerando o treinamento e preservando a fidelidade de precisão total sem alterar a lógica de atualização subjacente.

Autores originais: Dikshant Kukreja, Kritarth Prasad, Avinash Anand, Zhengkui Wang, Erik Cambria, Timothy Liu, Aik Beng Ng, Simon See, Bapi Chatterjee

Publicado 2026-06-23

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Dikshant Kukreja, Kritarth Prasad, Avinash Anand, Zhengkui Wang, Erik Cambria, Timothy Liu, Aik Beng Ng, Simon See, Bapi Chatterjee

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

O Grande Problema: A "Oficina Bagunçada"

Imagine que você está treinando um modelo de IA gigante (como o cérebro de um robô) em um computador. Para ensinar este robô, o computador precisa realizar uma quantidade massiva de cálculos matemáticos.

Na forma padrão de fazer isso (chamada de "Diferenciação de modo reverso"), o computador segue um processo rigoroso de duas etapas:

Calcular os Erros: Ele analisa os dados, descobre onde o robô errou e escreve uma lista gigante de "notas de correção" (gradientes) para cada parte do cérebro do robô. Ele escreve essas notas em um quadro branco gigante (a memória do computador).
Aplicar os Ajustes: Somente após todo o quadro branco estar cheio é que o computador pega uma borracha e um marcador para realmente atualizar o cérebro do robô com base nessas notas.

O Gargalo: O problema é que o computador tem que manter esse quadro branco gigante de notas em sua memória enquanto também está segurando o cérebro do robô e as notas para o próximo passo. Este "quadro branco" ocupa tanto espaço que muitas vezes preenche a memória do computador antes mesmo de o treinamento começar. É como tentar consertar um carro em uma garagem, mas você precisa manter o projeto inteiro do carro, as ferramentas e o manual de reparo espalhados pelo chão ao mesmo tempo. Se a garagem for pequena demais, você não consegue colocar o carro lá dentro.

A Solução: FORGE (O Método do "Ajuste Instantâneo")

Os autores deste artigo, FORGE (Fused On-Register Gradient Elimination), dizem: "Por que escrever as notas no quadro branco?"

Eles argumentam que a lista gigante de notas é apenas um artefato de como fazemos a matemática, não algo que o processo de aprendizado realmente necessita.

Como o FORGE funciona:
Em vez de escrever as notas e depois lê-las novamente, o FORGE faz a matemática e o ajuste ao mesmo tempo, na área de armazenamento mais rápida e pequena do computador (chamada de "registradores").

A Analogia: Imagine um mestre cuca preparando um prato complexo.
- Jeito Antigo: O chef pica uma cebola, escreve "cebola picada" em um bloco de notas, coloca a cebola em uma tigela e então passa para o próximo ingrediente. Depois que todos os ingredientes estão picados e listados, o chef lê o bloco de notas e mistura tudo. O bloco de notas ocupa espaço na bancada.
- Jeito FORGE: O chef pica uma cebola e a joga imediatamente na panela. Depois, ele pica uma cenoura e a joga imediatamente. Ele nunca escreve nada. Ele nunca precisa de um bloco de notas. A bancada permanece limpa.

Por que isso é um Grande Avanço

O artigo afirma três grandes benefícios deste método de "sem notas":

1. Economiza Espaço Massivo (Memória)
Como o computador nunca escreve a lista gigante de notas na memória principal, ele libera uma enorme quantidade de espaço.

O Resultado: Em um chip de computador padrão (H200), eles conseguiram treinar um modelo de 8 bilhões de parâmetros usando 53% menos memória.
A Analogia: É como ser capaz de colocar uma mesa de jantar para 10 pessoas em um estúdio pequeno porque você parou de precisar armazenar as cadeiras extras no corredor.

2. É Realmente Mais Rápido
Como o computador não precisa parar para escrever notas, esperar e depois lê-las novamente, todo o processo acelera.

O Resultado: As etapas de treinamento rodam cerca de 1,5 vezes mais rápido.
A Analogia: É a diferença entre um entregador que deixa um pacote, corre de volta para o caminhão para pegar o próximo e repete o processo (Jeito Antigo), versus um motorista que tem uma esteira rolante que carrega o pacote diretamente no caminhão enquanto ele o pega (FORGE).

3. Não Perde a Precisão
Geralmente, quando você tenta economizar espaço pulando etapas, você perde precisão (o robô aprende um pouco menos bem). Os autores provam que, como estão fazendo a matemática nos "registradores" de maior qualidade do computador (precisão total) antes de descartar os dados, o aprendizado é matematicamente idêntico ao método antigo e lento.

O Resultado: O robô aprende tão bem quanto, mas de forma muito mais eficiente.

O Truque do "Tile" (Ladrilho)

Como eles fazem isso sem causar o caos? Eles dividem o grande problema matemático em pequenos pedaços gerenciáveis chamados "tiles" (como quadrados de 128x128).

Eles processam um tile: Calcula o erro, ajusta o cérebro e descarta o erro imediatamente.
Depois, passam para o próximo tile.
Como fazem isso tile por tile, o computador nunca precisa segurar a lista inteira de erros ao mesmo tempo.

O Que Isso Nos Permite Fazer

O artigo mostra que com o FORGE:

Você pode treinar modelos maiores no mesmo computador.
Você pode usar "lotes" (batches) maiores (ensinando mais exemplos ao robô de uma só vez) sem ficar sem memória.
Em um teste específico, eles foram capazes de treinar um modelo em quatro GPUs que exigiria oito GPUs usando o método antigo.

Resumo

FORGE é uma nova maneira de treinar IA que impede o computador de entulhar sua memória com "notas de correção" temporárias. Ao calcular os erros e ajustar o modelo instantaneamente na parte mais rápida do chip, ele corta o uso de memória pela metade e acelera o treinamento, tudo sem tornar a IA menos inteligente. Ele transforma uma oficina lotada e lenta em uma linha de montagem ágil e de alta velocidade.

Resumo Técnico: FORGE (Fused On-Register Gradient Elimination)

1. O Problema: O Piso de Gradiente (The Gradient Floor)

Os métodos atuais de treinamento de Large Language Models (LLMs) eficientes em memória visam principalmente três dos quatro pools de memória em uma GPU: estados do otimizador (via compressão), ativações (via recomputação) e pesos do modelo. No entanto, eles deixam o quarto pool — o gradiente de peso materializado — intocado.

No modo de diferenciação reversa padrão, a passagem de retropropagação (backward pass) computa o tensor de gradiente completo ( $\nabla W$ ) para cada camada, escreve-o na Memória de Alta Largura de Banda (HBM) e só então o otimizador o lê de volta para atualizar os pesos. Esse cronograma de duas fases cria um "piso de gradiente" onde, na fronteira entre a passagem de backward e o passo do otimizador, os gradientes de todas as camadas devem coexistir na memória. Para um modelo de 8B parâmetros, esse pool de gradiente transitório sozinho consome aproximadamente 16 GiB, ditando frequentemente se uma configuração de treinamento cabe em uma única GPU. Métodos existentes que fundem os passos de backward e otimizador (ex: LOMO, FlashOptim) operam na granularidade de parâmetros ou camadas inteiras, o que significa que o tensor de gradiente completo para um parâmetro ainda é materializado na memória antes que a atualização ocorra.

2. Metodologia: FORGE

O FORGE (Fused On-Register Gradient Elimination) elimina o tensor de gradiente materializado por completo, fundindo o passo do otimizador diretamente na computação do gradiente de peso em uma granularidade de tile (especificamente tiles de $128 \times 128$ ).

Mecanismo Central

Em vez de computar a matriz de gradiente completa e armazená-la, o FORGE processa a matriz de pesos um tile por vez:

Acumulação de Gradiente por Tile: Para um tile de peso específico, o kernel acumula os produtos parciais do gradiente ascendente ( $\Delta Y$ ) e da ativação de entrada ( $X$ ) diretamente em registradores FP32 da GPU.
Otimização em Registro (In-Register Optimization): Uma vez que o gradiente do tile é totalmente acumulado nos registradores, a atualização do otimizador (ex: AdamW) é aplicada imediatamente a esse tile, também dentro dos registradores.
Descarte Imediato: O tile de gradiente é consumido no instante em que é produzido. Ele nunca é escrito na HBM como um tensor.
Escrita de Retorno (Write-Back): Apenas o tile de peso atualizado e os tiles de estado do otimizador (momentos) — opcionalmente quantizados — são escritos de volta na memória global.

Precisão e Fidelidade

Exatidão: Para qualquer otimizador elemento-a-elemento separável (onde a atualização de um peso depende apenas das coordenadas desse peso, ex: AdamW, SGD, Lion), a atualização por tile é matematicamente idêntica a um passo de gradiente total padrão. A única divergência é a ordem da soma de ponto flutuante sobre os tokens, o que é inerente a qualquer implementação de GEMM e não constitui uma aproximação da lógica do otimizador.
Fidelidade de Baixa Precisão: Em regimes de treinamento BF16 ou INT8, o FORGE preserva a fidelidade de precisão total de duas maneiras:
- O gradiente nunca é truncado para BF16 antes da otimização; ele é computado e consumido em registradores FP32.
- Para evitar a "zona morta" do BF16 (onde atualizações pequenas são arredondadas para zero), o FORGE emprega arredondamento estocástico para as escritas de volta dos pesos, garantindo a acumulação não enviesada de atualizações de sub-ULP.
- Para estados de otimizador INT8, o FORGE utiliza blocos menores (64 vs 2048 em otimizadores de 8 bits padrão) para minimizar o erro de quantização.

Compatibilidade com Treinamento Distribuído

O FORGE é comprovadamente exato sob Paralelismo de Tensor (TP) e Paralelismo de Sequência (SP) porque o gradiente local para um shard é completo no momento do backward. No entanto, é incompatível com Paralelismo de Dados (DP) ou Paralelismo de Contexto em sua forma atual, pois estes exigem um all-reduce de gradientes parciais antes que a atualização possa ser aplicada, o que exigiria a materialização do tensor de gradiente.

3. Principais Contribuições

Eliminação Estrutural, Não Compressão: O FORGE identifica o gradiente materializado como um artefato do agendamento de treinamento, e não uma necessidade matemática. Ele elimina o termo de memória de gradiente $2P$ (contagem de parâmetros), reduzindo o pool de gradiente de $O(P)$ para $O(T_R T_C)$ (tamanho do tile), independente do tamanho do modelo.
Exatidão Comprovável: O artigo prova que, para otimizadores elemento-a-elemento, o agendamento fundido por tile produz atualizações idênticas ao AdamW padrão em gradientes FP32. Essa exatidão mantém-se sob paralelismo de tensor e de sequência.
Fidelidade de Precisão Total em Baixa Precisão: Ao manter os gradientes em registradores FP32, o FORGE evita o viés introduzido pelo truncamento de gradientes para BF16 antes da otimização. Combinado com o arredondamento estocástico, ele iguala a convergência do mestre FP32 mesmo usando estados BF16 ou INT8.
Generalidade: O método é agnóstico à arquitetura para camadas lineares e suporta oito otimizadores elemento-a-elemento totalmente fundidos (AdamW, NAdam, RAdam, Lion, RMSprop, AdaGrad, SGD, SGD-momentum).

4. Resultados Experimentais

Avaliado em GPUs NVIDIA H200 e B200 em modelos variando de 0.6B a 14B parâmetros (família Llama-3.1-8B, Qwen3):

Redução de Memória: O FORGE reduz o pico de memória em 38–53% comparado ao AdamW fundido. Para o Llama-3.1-8B, o pico de memória cai de ~75 GiB para 35.4 GiB (uma redução de 53%).
Throughput (Vazão): O FORGE alcança um aceleração de 1.52× no tempo de passo (99.1 ms vs 150.9 ms) ao eliminar o lançamento de kernel separado do otimizador e a ida e volta (round-trip) do gradiente na HBM.
Escalabilidade:
- Em uma única H200, o FORGE treina o Qwen3-14B (que causa erros de Out-Of-Memory com o AdamW fundido padrão) com 63.2 GiB.
- Em uma configuração de 4 GPUs A100-40GB, o FORGE permite um micro-batch size de 4 (cabendo em 24.6 GiB), enquanto o AdamW fundido padrão só comporta micro-batch 1. Isso permite que o FORGE alcance um throughput 2.42× maior que o baseline padrão no mesmo hardware.
Convergência: O FORGE acompanha a convergência do AdamW fundido dentro de 0.001 nats (desvio médio) em tarefas de pré-treinamento contínuo, confirmando que as desvios de baixa precisão estão dentro do ruído de execução.
Componibilidade: O FORGE compõe-se perfeitamente com quantização FP8 e checkpointing de ativação, absorvendo os custos de desquantização que tipicamente penalizam otimizadores padrão.

5. Significância e Alegações

O artigo afirma que o FORGE representa uma mudança estrutural em vez de uma otimização incremental. Ao remover o piso de gradiente, ele converte o overhead de memória do otimizador em headroom utilizável para lotes maiores, sequências mais longas ou modelos maiores em hardware existente.

Fronteira de Capacidade: Os autores afirmam que a eliminação do piso de gradiente não é apenas uma melhoria de fator constante, mas uma fronteira de capacidade. Por exemplo, configurações de treinamento padrão para modelos de 14B em H200s únicas são impossíveis devido ao piso de gradiente, enquanto o FORGE as torna viáveis.
Ortogonalidade: O método é ortogonal a outras técnicas de economia de memória (como quantização de estado ou recomputação de ativação) e pode ser combinado com elas para reduzir ainda mais a pegada de memória.
Limitações: O método é atualmente limitado a otimizadores separáveis elemento-a-elemento e não suporta paralelismo de dados sem modificações. Também requer que o gradiente de entrada seja lido antes da atualização de peso, excluindo embeddings vinculados (tied embeddings) ou especialistas de MoE com múltiplos consumidores de gradiente em um único passo.

Em resumo, o FORGE demonstra que o "piso de gradiente" é um artefato desnecessário dos loops de treinamento atuais. Ao fundir o otimizador no processo de backward ao nível de tile, ele alcança economias de memória e acelerações significativas, mantendo a exatidão matemática para uma ampla classe de otimizadores.