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

• URL: http://arxiv.org/abs/2606.22932v1
• Date: Mon, 22 Jun 2026 07:08:31 GMT
• Title: FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training
• Abstract summary: FORGE folds step into the backward pass and applies it one tile at a time, entirely in registers, so each tile is consumed the instant it is produced and never becomes a tensor.<n>For every element-wise update, for every element-wise rule, for every element-wise rule, for every exactness survives and sequence-paralleling.
• Score: 43.28268986528884
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Reverse-mode differentiation computes every weight gradient, writes it to memory, and only then lets the optimizer read it back. This two-phase schedule sets the memory ceiling of modern training: at the seam between the phases, every layer's gradient is live at once. We argue that this materialized gradient is an artifact of how differentiation is staged, not a quantity that learning requires -- and we eliminate it. FORGE folds the optimizer step into the backward pass and applies it one tile at a time, entirely in registers, so each gradient tile is consumed the instant it is produced and never becomes a tensor. The fusion changes only when the update happens, not what it computes: in full precision the fused step is provably exact -- the identical optimizer update, for every element-wise rule -- and that exactness survives tensor- and sequence-parallel sharding; in the bf16 and 8-bit regimes used in practice it is faithful rather than bit-identical, its deviation bounded and, for the weight store, rendered unbiased by stochastic rounding. Because each gradient tile is born and consumed in the same registers, it is never converted down to bf16 to be stored and read back; FORGE thus preserves the full-precision fidelity that both bf16 and 8-bit optimizers lose to that conversion. Nor is the method tied to one architecture or one optimizer: linear layers are ubiquitous, and FORGE reclaims the gradient memory of any of them under any element-wise rule. Empirically FORGE more than halves the memory of an optimizer step and, at the small batch sizes typical of fine-tuning and continued pretraining, runs about 1.5x faster; integrated into tensor-parallel Megatron-LM it fits 8B training at four times the micro-batch a standard optimizer allows on the same GPUs.

Related papers

• HEAT: Faster Fully Homomorphic Inference via Approximations-Weights Co-Adaptation [56.14416807766864]
  Homomorphic encryption allows a server to run a language model directly on encrypted user prompts.<n>Existing approaches fix the iteration counts uniformly across the model rather than tailoring them to each site's error tolerance.<n>We introduce Homomorphic Encryption-Aware Training (HEAT), a fine-tuning method that makes the per-nonline iteration counts learnable.
  arXiv  Detail & Related papers  (2026-09-01T18:02:02Z)
• Zero Sum SVD: Balancing Loss Sensitivity for Low Rank LLM Compression [11.908793753919745]
  We propose textbfZero Sum SVD (textbfZS-SVD), a post-training method that performs singular component selection in whitened coordinates.<n>textbfZS-SVD prunes components across the whole model with a textbfzero sum rule that keeps the cumulative predicted loss change near zero.<n>Experiments show consistent gains across diverse benchmarks and compression ratios.
  arXiv  Detail & Related papers  (2026-02-02T21:51:01Z)
• Backward-Friendly Optimization: Training Large Language Models with Approximate Gradients under Memory Constraints [14.20716202034732]
  Full fine-tuning of Large Language Models (LLMs) is notoriously memory-intensive.<n>We introduce GradLite, a backward-friendly solution that relaxes the requirement of exact gradients.<n>We show that GradLite maintains unbiased estimates with bounded variance, ensuring convergence rates comparable to Adam.
  arXiv  Detail & Related papers  (2025-10-26T00:50:12Z)
• Scaling Recurrent Neural Networks to a Billion Parameters with Zero-Order Optimization [0.0]
  RNNs scale constant in FLOPs and GPU memory with increasing context length.<n>Transformers scale linearly in FLOPs and, at best, linearly in memory during generation.<n>Training large RNNs on long contexts remains impractical because standard optimization methods depend on Backpropagation Through Time (BPTT)
  arXiv  Detail & Related papers  (2025-05-23T13:04:06Z)
• Fine-tuning Quantized Neural Networks with Zeroth-order Optimization [21.0540879091664]
  We propose Quantized Zeroth-order Optimization (QZO), a simple yet effective approach that perturbs the continuous quantization scale for gradient estimation.<n>QZO can reduce the total memory cost by more than 18$times$ for 4-bit LLMs, and enables fine-tuning Llama-2-13B within a single 24GB GPU.
  arXiv  Detail & Related papers  (2025-05-19T17:55:15Z)
• Cut Your Losses in Large-Vocabulary Language Models [102.6981011879656]
  We propose Cut Cross-Entropy (CCE), a method that computes the cross-entropy loss without materializing the logits for all tokens into global memory.<n>CCE reduces the memory footprint of the loss from 24 GB to 1 MB, and the total training-time memory consumption of the head from 28 GB to 1 GB.
  arXiv  Detail & Related papers  (2024-11-13T20:30:15Z)
• Breaking the Memory Barrier: Near Infinite Batch Size Scaling for Contrastive Loss [59.835032408496545]
  We propose a tile-based strategy that partitions the contrastive loss calculation into arbitrary small blocks. We also introduce a multi-level tiling strategy to leverage the hierarchical structure of distributed systems. Compared to SOTA memory-efficient solutions, it achieves a two-order-of-magnitude reduction in memory while maintaining comparable speed.
  arXiv  Detail & Related papers  (2024-10-22T17:59:30Z)
• Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of Language Model [89.8764435351222]
  We propose a new family of unbiased estimators called WTA-CRS, for matrix production with reduced variance. Our work provides both theoretical and experimental evidence that, in the context of tuning transformers, our proposed estimators exhibit lower variance compared to existing ones.
  arXiv  Detail & Related papers  (2023-05-24T15:52:08Z)
• 8-bit Optimizers via Block-wise Quantization [57.25800395197516]
  Statefuls maintain statistics over time, e.g., the exponentially smoothed sum (SGD with momentum) or squared sum (Adam) of past values. This state can be used to accelerate optimization compared to plain gradient descent but uses memory that might otherwise be allocated to model parameters. In this paper, we develop first gradients that use 8-bit statistics while maintaining the performance levels of using 32-bit gradient states.
  arXiv  Detail & Related papers  (2021-10-06T15:43:20Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.