🤖 machine learning

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

FORGE is a memory-efficient LLM training method that fuses the optimizer step into the backward pass to eliminate materialized gradients by processing them entirely in registers, thereby halving optimizer memory usage, accelerating training, and preserving full-precision fidelity without altering the underlying update logic.

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

Published 2026-06-23

📖 5 min read🧠 Deep dive

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Cluttered Workshop"

Imagine you are training a giant AI model (like a robot brain) on a computer. To teach this robot, the computer has to do a massive amount of math.

In the standard way of doing this (called "Reverse-mode differentiation"), the computer follows a strict two-step process:

Calculate the Mistakes: It looks at the data, figures out where the robot went wrong, and writes down a giant list of "correction notes" (gradients) for every single part of the robot's brain. It writes these notes down on a giant whiteboard (the computer's memory).
Apply the Fixes: Only after the entire whiteboard is full does the computer pick up an eraser and a marker to actually update the robot's brain based on those notes.

The Bottleneck: The problem is that the computer has to hold that giant whiteboard of notes in its memory while it is also holding the robot's brain and the notes for the next step. This "whiteboard" takes up so much space that it often fills up the computer's memory before the training can even start. It's like trying to fix a car in a garage, but you have to keep the entire blueprint of the car, the tools, and the repair manual spread out on the floor at the same time. If the garage is too small, you can't fit the car in.

The Solution: FORGE (The "Instant Fix" Method)

The authors of this paper, FORGE (Fused On-Register Gradient Elimination), say: "Why write the notes down on the whiteboard at all?"

They argue that the giant list of notes is just an artifact of how we do the math, not something the learning process actually needs.

How FORGE works:
Instead of writing the notes down and then reading them back, FORGE does the math and the fixing at the same time, in the computer's fastest, smallest storage area (called "registers").

The Analogy: Imagine a master chef cooking a complex dish.
- Old Way: The chef chops an onion, writes down "chopped onion" on a notepad, puts the onion in a bowl, then moves to the next ingredient. Once all ingredients are chopped and listed, the chef reads the notepad and mixes them all together. The notepad takes up counter space.
- FORGE Way: The chef chops an onion and immediately tosses it into the pot. Then they chop a carrot and toss it in immediately. They never write anything down. They never need a notepad. The counter stays clear.

Why This is a Big Deal

The paper claims three major benefits from this "no-notes" approach:

1. It Saves Massive Space (Memory)
Because the computer never writes the giant list of notes to the main memory, it frees up a huge amount of space.

The Result: On a standard computer chip (H200), they were able to train an 8-billion-parameter model using 53% less memory.
The Analogy: It's like being able to fit a 10-person dining table into a studio apartment because you stopped needing to store the extra chairs in the hallway.

2. It's Actually Faster
Because the computer doesn't have to stop to write notes, wait, and then read them back, the whole process speeds up.

The Result: The training steps run about 1.5 times faster.
The Analogy: It's the difference between a delivery driver who drops a package off, runs back to the truck to get the next one, and repeats (Old Way), versus a driver who has a conveyor belt that loads the package directly onto the truck as they grab it (FORGE).

3. It Doesn't Lose Accuracy
Usually, when you try to save space by skipping steps, you lose precision (the robot learns slightly less well). The authors prove that because they are doing the math in the computer's highest-quality "registers" (full precision) before discarding the data, the learning is mathematically identical to the old, slow method.

The Result: The robot learns just as well, but much more efficiently.

The "Tile" Trick

How do they do this without chaos? They break the giant math problem into small, manageable chunks called "tiles" (like 128x128 squares).

They process one tile: Calculate the error, fix the brain, and throw away the error immediately.
Then they move to the next tile.
Because they do this tile-by-tile, the computer never has to hold the entire list of errors at once.

What This Allows Us to Do

The paper shows that with FORGE:

You can train larger models on the same computer.
You can use larger "batches" (teaching the robot more examples at once) without running out of memory.
In one specific test, they were able to train a model on four GPUs that would have required eight GPUs using the old method.

Summary

FORGE is a new way to train AI that stops the computer from cluttering its memory with temporary "correction notes." By calculating the mistakes and fixing the model instantly in the fastest part of the chip, it cuts memory usage in half and speeds up training, all without making the AI any less smart. It turns a crowded, slow workshop into a streamlined, high-speed assembly line.

Technical Summary: FORGE (Fused On-Register Gradient Elimination)

1. The Problem: The Gradient Floor

Current memory-efficient Large Language Model (LLM) training methods primarily target three of the four memory pools on a GPU: optimizer states (via compression), activations (via recomputation), and model weights. However, they leave the fourth pool—the materialized weight gradient—untouched.

In standard reverse-mode differentiation, the backward pass computes the full gradient tensor ( $\nabla W$ ) for every layer, writes it to High-Bandwidth Memory (HBM), and only then does the optimizer read it back to update weights. This two-phase schedule creates a "gradient floor" where, at the boundary between the backward pass and the optimizer step, the gradients for all layers must coexist in memory. For an 8B parameter model, this transient gradient pool alone consumes approximately 16 GiB, often dictating whether a training configuration fits on a single GPU. Existing methods that fuse backward and optimizer steps (e.g., LOMO, FlashOptim) operate at the granularity of entire parameters or layers, meaning the full gradient tensor for a parameter is still materialized in memory before the update occurs.

2. Methodology: FORGE

FORGE (Fused On-Register Gradient Elimination) eliminates the materialized gradient tensor entirely by fusing the optimizer step directly into the weight-gradient computation at a tile granularity (specifically $128 \times 128$ tiles).

Core Mechanism

Instead of computing the full gradient matrix and storing it, FORGE processes the weight matrix one tile at a time:

Tile-wise Gradient Accumulation: For a specific weight tile, the kernel accumulates the partial products of the upstream gradient ( $\Delta Y$ ) and the input activation ( $X$ ) directly into FP32 GPU registers.
In-Register Optimization: Once the tile's gradient is fully accumulated in registers, the optimizer update (e.g., AdamW) is applied immediately to that tile, also within the registers.
Immediate Discard: The gradient tile is consumed the instant it is produced. It is never written to HBM as a tensor.
Write-Back: Only the updated weight tile and the (optionally quantized) optimizer state tiles (moments) are written back to global memory.

Precision and Fidelity

Exactness: For any element-wise-separable optimizer (where the update for a weight depends only on that weight's coordinate, e.g., AdamW, SGD, Lion), the per-tile update is mathematically identical to a standard full-gradient step. The only deviation is the order of floating-point summation over tokens, which is inherent to any GEMM implementation and does not constitute an approximation of the optimizer logic.
Low-Precision Fidelity: In BF16 or INT8 training regimes, FORGE preserves full-precision fidelity in two ways:
- The gradient is never truncated to BF16 before optimization; it is computed and consumed in FP32 registers.
- To prevent the "dead zone" of BF16 (where small updates are rounded to zero), FORGE employs stochastic rounding for weight write-backs, ensuring unbiased accumulation of sub-ULP updates.
- For INT8 optimizer states, FORGE uses smaller block sizes (64 vs. 2048 in standard 8-bit optimizers) to minimize quantization error.

Distributed Training Compatibility

FORGE is provably exact under Tensor Parallelism (TP) and Sequence Parallelism (SP) because the local gradient for a shard is complete at backward time. However, it is incompatible with Data Parallelism (DP) or Context Parallelism in its current form, as these require an all-reduce of partial gradients before the update can be applied, which would necessitate materializing the gradient tensor.

3. Key Contributions

Structural Elimination, Not Compression: FORGE identifies the materialized gradient as an artifact of training scheduling rather than a mathematical necessity. It eliminates the $2P$ (parameter count) gradient memory term, reducing the gradient pool from $O(P)$ to $O(T_R T_C)$ (tile size), independent of model size.
Provable Exactness: The paper proves that for element-wise optimizers, the tile-fused schedule produces updates identical to standard AdamW on FP32 gradients. This exactness holds under tensor and sequence parallelism.
Full-Precision Fidelity at Low Precision: By keeping gradients in FP32 registers until the update, FORGE avoids the bias introduced by truncating gradients to BF16 before optimization. Combined with stochastic rounding, it matches FP32-master convergence even when using BF16 or INT8 states.
Generality: The method is architecture-agnostic for linear layers and supports eight fully fused element-wise optimizers (AdamW, NAdam, RAdam, Lion, RMSprop, AdaGrad, SGD, SGD-momentum).

4. Experimental Results

Evaluated on NVIDIA H200 and B200 GPUs across models ranging from 0.6B to 14B parameters (Llama-3.1-8B, Qwen3 family):

Memory Reduction: FORGE reduces peak memory by 38–53% compared to fused AdamW. For Llama-3.1-8B, peak memory drops from ~75 GiB to 35.4 GiB (a 53% reduction).
Throughput: FORGE achieves a 1.52× speedup in step time (99.1 ms vs. 150.9 ms) by eliminating the separate optimizer kernel launch and the HBM round-trip for gradients.
Scalability:
- On a single H200, FORGE trains Qwen3-14B (which causes Out-Of-Memory errors with standard fused AdamW) at 63.2 GiB.
- In a 4-GPU A100-40GB setup, FORGE enables a micro-batch size of 4 (fitting 24.6 GiB), whereas standard fused AdamW only fits micro-batch 1. This allows FORGE to achieve 2.42× higher throughput than the standard baseline on the same hardware.
Convergence: FORGE tracks fused AdamW convergence within 0.001 nats (average deviation) on continued pretraining tasks, confirming that the low-precision deviations are within run-to-run noise.
Composability: FORGE composes cleanly with FP8 quantization and activation checkpointing, absorbing the dequantization costs that typically penalize standard optimizers.

5. Significance and Claims

The paper claims that FORGE represents a structural shift rather than an incremental optimization. By removing the gradient floor, it converts optimizer memory overhead into usable headroom for larger batch sizes, longer sequences, or larger models on existing hardware.

Capability Boundary: The authors state that the gradient floor elimination is not just a constant-factor improvement but a capability boundary. For instance, standard training configurations for 14B models on single H200s are impossible due to the gradient floor, whereas FORGE makes them feasible.
Orthogonality: The method is orthogonal to other memory-saving techniques (like state quantization or activation recomputation) and can be combined with them to further reduce memory footprints.
Limitations: The method is currently limited to element-wise separable optimizers and does not support data parallelism without modification. It also requires the input gradient to be read before the weight update, excluding tied embeddings or MoE experts with multiple gradient consumers in a single step.

In summary, FORGE demonstrates that the "gradient floor" is an unnecessary artifact of current training loops. By fusing the optimizer into the backward pass at the tile level, it achieves significant memory savings and speedups while maintaining mathematical exactness for a wide class of optimizers.