Pith. sign in

REVIEW 2 major objections 3 minor 3 cited by

FORGE claims the materialized weight gradient is a staging artifact, not a learning requirement, and eliminates it by fusing the optimizer into the backward pass at tile granularity—provably exact for element-wise optimizers and exact under

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · deepseek-v4-flash

2026-08-02 10:27 UTC pith:GWT7BTE6

load-bearing objection A genuinely new per-tile optimizer fusion with clean single-GPU math and careful benchmarks, but the sequence-parallel exactness claim is overbroad and the artifact isn't public yet. the 2 major comments →

arxiv 2606.22932 v2 pith:GWT7BTE6 submitted 2026-06-22 cs.LG

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

Dikshant Kukreja , Kritarth Prasad , Avinash Anand , Zhengkui Wang , Erik Cambria , Timothy Liu , Aik Beng Ng , Simon See
show 1 more author
Bapi Chatterjee
This is my paper
classification cs.LG
keywords memory-efficient traininggradient eliminationfused optimizerbackward-pass fusionelement-wise optimizerslow-precision trainingtensor parallelismLLM fine-tuning
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

FORGE argues that the materialized gradient tensor that sets the memory ceiling of LLM training is not something learning requires: it is an artifact of running backprop and the optimizer as two separate phases. The method fuses the optimizer step into the backward pass and applies it one 128×128 tile at a time inside GPU registers, so each gradient tile is consumed the instant it is produced and never becomes a tensor. In exact arithmetic the fused step is identical to the standard optimizer for every element-wise-separable rule—AdamW, SGD, RMSprop, Lion, and others—with the only floating-point difference being the GEMM's token-reduction order. Empirically, on a standard 8-billion-parameter transformer at small batch size, FORGE more than halves peak memory and runs about 1.5× faster, and in the bf16/8-bit regimes it preserves full-precision gradient fidelity by never truncating the gradient before the update. The claimed scope has a sharp boundary: weights with more than one gradient consumer in a step, and optimizers that couple coordinates, are excluded.

Core claim

The central claim is that the weight-gradient tensor—whose 2P term sets the memory ceiling in standard training—can be eliminated rather than compressed. FORGE computes the weight gradient of a linear layer as a sequence of 128×128 tiles inside the backward matrix multiply, applies the optimizer arithmetic to each tile in fp32 registers, and writes back only the updated weight and moment tiles. Because an element-wise-separable optimizer updates each coordinate independently, the per-tile schedule lands on exactly the same weights as materializing the full gradient and stepping once (Proposition 1). The same coordinate-wise argument makes the fused update exact under tensor- and sequence-par

What carries the argument

The central object is the per-tile fused gradient-and-optimizer epilogue: the backward weight-gradient matrix multiply is partitioned into 128×128 tiles, each tile accumulates its exact gradient in fp32 registers over token chunks, then the optimizer's element-wise update runs in the same registers, and the tile is discarded before any byte of it reaches global memory. The load-bearing identity is element-wise separability—an optimizer whose update at each coordinate depends only on that coordinate of weight, gradient, and state commutes with any tiling of the weight index set. This turns the O(P) gradient pool into O(TR·TC) per active tile, removes the separate optimizer kernel launch, and

Load-bearing premise

The claim rests on the requirement that a weight's gradient be complete and readable before the fused in-place update fires, which rules out every multi-consumer weight (tied embeddings or LM heads, mixture-of-experts experts with cross-rank reduction) and any data/context-parallel setup where partial gradients must be all-reduced before stepping.

What would settle it

Numerical falsifier: run FORGE with fp32 state on one linear layer and compare its updated weights against a standard optimizer applied to the exact same fp32 gradient with the same token-axis summation order; the two should be bit-identical, and any difference beyond documented GEMM reduction-order non-associativity would refute the exactness proposition. Memory falsifier: train a model with a tied embedding/LM head while fusing all linear layers and measure peak memory during backward; if the parameter-scaled gradient term persists, the single-consumer precondition has been violated.

Watch this falsifier — get emailed when new claim-graph text bears on it.

If this is right

  • For single-consumer linear layers, the 2P gradient term disappears from peak memory, leaving the gradient pool O(TR·TC) per active tile, independent of parameter count, and cutting the per-parameter optimizer-side cost from 18 to 6 bytes with bf16 state.
  • The fused update is exact for any element-wise-separable optimizer in fp32 state, and the coordinate-wise argument proves exactness under tensor- and sequence-parallel sharding; data- and context-parallel training must fall back to a standard optimizer because partial gradients must be summed before the nonlinear update.
  • Because the gradient is consumed in fp32 registers, FORGE avoids the bf16 truncation bias that low-precision optimizers carry, and with stochastic rounding of the weight store it matches an fp32-master baseline where round-to-nearest bf16 in-place steps lose about 6%.
  • The reclaimed memory composes with activation checkpointing and FP8/8-bit quantized states, so downstream compression translates directly into peak-memory savings rather than being capped by a co-resident gradient transient.
  • At the largest tested scale, a 14-billion-parameter model that does not fit on a single accelerator under the fused-AdamW baseline does fit under FORGE, and in a 4-GPU tensor-parallel setup FORGE fits four times the micro-batch, turning the memory saving into throughput.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • We infer that the paper's boundary implies a practical hybrid: any training loop with tied embeddings or mixture-of-experts experts would fuse only single-consumer linears and keep a standard step for the rest, so the claimed memory floor applies to the bulk of transformer linear layers but not to the whole model.
  • A direct ablation that materializes the tile before the update—forcing the exact gradient round-trip that FORGE skips—would isolate how much of the ~1.5× speedup comes from removing the HBM gradient write/read versus from eliminating the optimizer kernel launch; the paper's mechanism suggests both are large but does not separate them.
  • If exactness extends as claimed, then approximate memory-saving optimizers that replace AdamW with factored or global-statistic rules are trading away exact large-batch AdamW dynamics on precisely the layers where an exact element-wise fusion costs nothing; whether the approximations ever pay for themselves in quality remains an open question the paper does not resolve.
  • A testable scaling prediction follows: FORGE's memory saving as a fraction of peak should be largest when the optimizer step dominates, i.e., small batch, short sequence, many parameters, and the paper's tables across model size and sequence length are consistent with that; a prospective study could use the same kernel to predict fit at 70B+ scale.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

2 major / 3 minor

Summary. FORGE proposes to fuse element-wise optimizer steps (AdamW and others) into the per-tile backward weight-gradient GEMM of linear layers, so that each gradient tile is produced, consumed by the optimizer update, and discarded inside GPU registers before any gradient tensor is written to HBM. The paper gives an exact-arithmetic proof that per-tile fusion matches the standard full-gradient optimizer step for any element-wise-separable rule (Proposition 1), characterizes floating-point deviation as backward-GEMM reduction non-associativity, analyzes bf16/INT8 fidelity (bounded INT8 moment error, unbiased stochastic rounding for the bf16 weight store), and reports substantial memory and speed improvements on Llama-3.1-8B, Qwen3 family, and GPT-2 124M, plus a Megatron-LM tensor-parallel integration that fits an 8B model at 4x micro-batch on the same hardware.

Significance. If correct, FORGE addresses a real and relatively underexplored memory bottleneck: the materialized gradient tensor of size O(P), which is untouched by optimizer-state compression and activation recomputation. The core exactness argument is a clean coordinate-wise identity and is verified empirically with careful methodology (N=20 median timings, autotune disclosure, per-cell OOM logs, multi-seed convergence runs). The claimed 20--55% per-rank memory savings in tensor-parallel training and the 8B capability result are meaningful if they survive scrutiny. The main weakness is that the sequence-parallel exactness claim, prominently featured in the abstract and Section 4, is not supported by the proof as written; the single-GPU and pure tensor-parallel claims are not affected.

major comments (2)
  1. [§4 / Appendix A, Proposition 2(i)] The sequence-parallel exactness claim is not established. The proof argues that under sequence parallelism the full-sequence input is all-gathered before the weight-gradient GEMM, so the contraction over tokens is complete. But the weight-gradient contraction is G = Σ_t ΔY[t,:]^T X[t,:]; completing it requires both operands to be full-sequence. All-gathering X alone does not supply the matching ΔY rows. In standard Megatron sequence parallelism the upstream gradient ΔY is token-sharded at the wgrad point (or, depending on dataflow, requires an additional all-gather to make it full-sequence), and the paper's integration statement that 'no communication is added' rules out an all-gather of ΔY. Thus the abstract's claim that exactness 'survives ... sequence-parallel sharding' is overbroad. Please either restrict the exactness claim to tensor parallelism, or specify a sequence-parallel dataf
  2. [§2.2, Eq. (6) and Remark 1] The headline memory equation M_FORGE = 2P + kP + 0 + Σ|aℓ| + O(TR·TC) and the abstract's claim that the gradient floor is eliminated apply only to the subset of weights that are single-consumer and fused. Tied embeddings/LM heads, MoE experts with cross-rank reduction, and any weight with more than one backward consumer remain on the standard optimizer and retain a materialized gradient term. The body does acknowledge this in Remark 1 and in the Limitations, but the unconditional form of Eq. (6) and of the abstract's 'eliminates the materialized gradient' phrasing is likely to mislead. Please state the scope explicitly wherever the global memory equation is presented.
minor comments (3)
  1. [§2.2, Algorithm 1] Algorithm 1 is titled 'fused backward and AdamW kernel' but does not compute the input gradient ΔX = ΔY·W, which Remark 1 says must be read before the in-place weight update. As written, the algorithm only produces the weight-gradient tile and the update. Please include the ΔX computation (or state explicitly that it is performed in a separate preceding kernel) so that the pseudocode matches the described backward computation.
  2. [§3.2 / Appendix F] The 0.001 nats average convergence-parity claim is computed on one continued-pretraining corpus (OpenMathInstruct-2) and averaged over a family of models. This is disclosed in the Limitations, but it would be helpful to state at the point of the claim in §3.2 that the number is corpus-specific and is below the expected BF16-vs-FP32 precision floor, rather than a universal algorithmic invariant.
  3. [Appendix J, Tables 13 and 15] Several B200 reference rows are marked as measured under Protocol A (N=5, mean), while the remaining rows use Protocol B (N=20, median). This is disclosed, but because the headline tables mix the two protocols on the same page, the reader may inadvertently compare Protocol A absolute values against Protocol B values. A visible separator or a 'not directly comparable' annotation on the affected rows would improve clarity.

Circularity Check

0 steps flagged

No circularity: FORGE's exactness is a self-contained algebraic partition argument; the sequence-parallel concern is a correctness gap, not circularity.

full rationale

The derivation chain is self-contained. Proposition 1 proves per-tile exactness from the definition of element-wise-separable optimizers and the fact that any partition of coordinates sums to the same total gradient; this is a direct algebraic identity, not a fitted value or a prediction forced by a construction. The memory bound in Proposition 3 follows from the register-resident lifetime of each tile by definition of the algorithm, and the empirical results are comparisons against external baselines (fused AdamW, bnb 8-bit, GaLore, APOLLO, AdaLomo, FlashOptim) rather than recoveries of fitted parameters. The low-precision and INT8 propositions are standalone mathematical bounds with explicitly stated assumptions. No load-bearing self-citation appears: the paper's references are to standard prior work, and no 'uniqueness theorem' or ansatz-carrying citation is invoked to force the choice of tile fusion. The skeptical concern about Proposition 2's sequence-parallel exactness is a potential correctness gap in the proof (all-gathering the input does not by itself complete the token contraction if the upstream gradient is also token-sharded), but this is not circularity: it is an unsupported claim, not a reduction of the conclusion to its own inputs. Therefore the circularity score is 0.

Axiom & Free-Parameter Ledger

2 free parameters · 4 axioms · 0 invented entities

All assumptions are explicitly stated in the paper. The only hand-chosen numbers are performance-oriented kernel parameters and the INT8 block size; none of them is fitted to make the main exactness claim true.

free parameters (2)
  • Tile geometry (TR, TC, TK, num_warps) = TR=TC=128, TK=64, num_warps=8 (autotuned on H200)
    Chosen from a curated 34-point search space for speed; affects the O(TR·TC) register working set and performance, but not the exactness of the update.
  • INT8 quantization block size for moments = qblock=64
    Design choice to make blocks 32x finer than bnb's 2048; directly controls the bounded INT8 state error in Proposition 4.
axioms (4)
  • domain assumption Optimizer update is element-wise-separable (Definition 1): each coordinate's new weight depends only on that coordinate of (W, G, S).
    Central theorem (Proposition 1) applies only to this class; Adafactor/LAMB/Muon are excluded, as the paper states in §2.3 and Appendix N.
  • domain assumption Single-consumer precondition for in-place weight updates (Remark 1).
    Exactness under backward-time update requires the layer's input gradient to be read before overwriting W and no second consumer of the old weight in the same step; tied embeddings/MoE fall back to the standard optimizer.
  • standard math Real-arithmetic identity of token-chunk summation (Proposition 1, Step 1); floating-point non-associativity is the only discrepancy (Corollary 1).
    The proof treats arithmetic as exact; in practice the tile-local reduction order may differ from cuBLAS, which the authors classify as GEMM-level noise.
  • standard math Weight gradient of a linear layer equals ΔY^T X (Eq. 3).
    Standard reverse-mode calculus for a dense layer, used as the starting identity for all proofs.

pith-pipeline@v1.3.0-alltime-deepseek · 38341 in / 15451 out tokens · 141628 ms · 2026-08-02T10:27:58.826720+00:00 · methodology

0 comments
read the original 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.

Figures

Figures reproduced from arXiv: 2606.22932 by Aik Beng Ng, Avinash Anand, Bapi Chatterjee, Dikshant Kukreja, Erik Cambria, Kritarth Prasad, Simon See, Timothy Liu, Zhengkui Wang.

Figure 1
Figure 1. Figure 1: FORGE (8-bit state) on Llama-3.1-8B (H200, BS=1, SEQ=512, BF16-everywhere) cuts peak [PITH_FULL_IMAGE:figures/full_fig_p002_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Measured within-step GPU memory (memory_allocated) over one forward+backward+optimizer cycle (Llama-3.1-8B). Fused AdamW (red) holds its peak through backward and the optimizer step, where the gradient and fp32 moments are co-resident (∼64 GB); FORGE (blue) sheds memory through backward—each gradient tile is consumed in place—and never spikes. With activation checkpointing, FORGE +AC (green) stays flat at … view at source ↗
Figure 3
Figure 3. Figure 3: FORGE mechanism: per-tile data flow inside the fused backward-and-optimizer step. The weight [PITH_FULL_IMAGE:figures/full_fig_p004_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: Convergence parity on continued pretraining over OpenMathInstruct-2 (BS [PITH_FULL_IMAGE:figures/full_fig_p007_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Headline single-GPU comparison on Llama-3.1-8B (H200 141 GB, BS=1, SEQ=512, BF16- [PITH_FULL_IMAGE:figures/full_fig_p008_5.png] view at source ↗
Figure 6
Figure 6. Figure 6: From-scratch pretraining parity on GPT-2 124M (nanoGPT) over FineWeb-Edu (sample-10BT, [PITH_FULL_IMAGE:figures/full_fig_p009_6.png] view at source ↗
Figure 7
Figure 7. Figure 7: Composability with FP8 and activation checkpointing at BS [PITH_FULL_IMAGE:figures/full_fig_p011_7.png] view at source ↗
Figure 8
Figure 8. Figure 8: Multi-model scaling on H200 (BS=1, SEQ=512, BF16-everywhere): peak memory of fused AdamW vs. FORGE across the Qwen3 family and Llama-3.1-8B, with FORGE’s step-time speedup over fused AdamW labeled above each model. FORGE holds peak memory to 0.47–0.73× the fused-AdamW footprint and runs 1.07–1.52× faster; at Qwen3-14B fused AdamW OOMs at this configuration while FORGE runs at 150.1 ms / 63.2 GiB. optimizer… view at source ↗
Figure 9
Figure 9. Figure 9: Distributed sweeps on 8×A100, FORGE vs. fused-CUDA-AdamW, across model size, tensor-parallel degree, micro-batch, and sequence length. Top: per-rank peak-memory saving (20–55%), largest where optimizer state dominates and at small batch/sequence. Bottom: step-time ratio vs. fused AdamW (the dashed line marks parity)—within 0.8–1.2× throughout—parity at no speed cost. Every reclaimed gigabyte is micro-batch… view at source ↗
Figure 10
Figure 10. Figure 10: Llama-style 8B on four A100-40GB (TP= 4, SEQ=1024): per-rank peak memory by micro￾batch. foreach AdamW OOMs at every micro-batch and fused-CUDA-AdamW fits only micro-batch 1 (34.0 GiB / 305.1 ms); FORGE stays under the 40 GiB card all the way to micro-batch 4 (24.6 GiB)—four times the micro-batch on the same hardware—at per-step latencies 241.3/346.7/504.0 ms for micro-batch 1/2/4. At micro-batch 4 FORGE … view at source ↗
Figure 11
Figure 11. Figure 11: Convergence parity across the remaining Qwen3 sizes on OpenMathInstruct-2 continued pretraining [PITH_FULL_IMAGE:figures/full_fig_p023_11.png] view at source ↗
Figure 12
Figure 12. Figure 12: Per-family training-loss curve (left sub-panel), peak GPU memory (center sub-panel), and average [PITH_FULL_IMAGE:figures/full_fig_p035_12.png] view at source ↗
Figure 13
Figure 13. Figure 13: Optimizer sweep (continued from Figure [PITH_FULL_IMAGE:figures/full_fig_p036_13.png] view at source ↗
Figure 14
Figure 14. Figure 14: Optimizer sweep (continued from Figure [PITH_FULL_IMAGE:figures/full_fig_p038_14.png] view at source ↗

Review history (2 revisions) →

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Forward citations

Cited by 3 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

  1. Reference Traces for Auditing Invisible Weight Updates and Guiding Exact-Budget Protection

    cs.LG 2026-07 conditional novelty 6.5

    Weight updates that fall below half a ULP freeze coordinates deterministically, and freeze time is predictable a priori from a high-precision trajectory and mantissa length alone.

  2. MMOE: Modernizing Diffusion Transformers with Efficient Expert Design

    cs.CV 2026-07 conditional novelty 5.0

    Composing MoE++ lightweight experts with attention-residual reuse inside SiT yields faster FID convergence and better quality-cost tradeoff than dense or plain sparse baselines at fixed single-node budget.

  3. Reference Traces for Auditing Invisible Weight Updates and Guiding Exact-Budget Protection

    cs.LG 2026-07 conditional novelty 5.0

    A matched high-precision reference trace forecasts invisible weight-update events in controlled low-precision runs, and a pre-frozen 10% fp32-master mask recovers ~61% of the RNE-vs-fp32 loss gap on three unseen targets.

Reference graph

Works this paper leans on

14 extracted references · 1 linked inside Pith · cited by 2 Pith papers

  1. [1]

    AdamW (Loshchilov & Hutter, 2019): first and second moments, decoupled weight decay

    2019

  2. [2]

    NAdam: AdamW with Nesterov momentum correction onˆm

  3. [3]

    RAdam: AdamW with rectified-variance term

  4. [4]

    Lion: sign-of-momentum update with a single moment buffer

  5. [5]

    RMSprop: second moment only, no first moment

  6. [6]

    AdaGrad: cumulative squared gradient (no decay)

  7. [7]

    SGD: vanilla gradient descent (no moment buffers)

  8. [8]

    SGD with momentum: single moment, no second moment. (ii) Linear-layer fusion only (standard optimizer step).These five couple coordinates through a row, column, block, or global statistic, so the update cannot be formed inside one tile; FORGE runs the tiled linear-layer backward and applies the optimizer as a standard step. They are included to show the l...

  9. [10]

    Adam-mini (Zhang et al., 2025): AdamW with shared per-block second moment

    2025

  10. [11]

    Adafactor (Shazeer & Stern, 2018): factored (row×column) second-moment statistics

    2018

  11. [12]

    AdaLomo (Lv et al., 2024): low-memory AdamW with factored second moment

    2024

  12. [13]

    LAMB: layer-wise adaptive scaling (per-tensor trust ratio)

  13. [14]

    SM3: per-dimension running max for the cumulative moment. 33 Each fully-fused family replaces the body of Phase 2 of Algorithm 1 with∼10 lines of Triton arithmetic; the wgrad mainloop (Phase 1) and the state read/write pattern (Phase 3) are unchanged. Cross-element preconditioners (Muon’s Newton–Schulz orthogonalization (Jordan et al., 2024), Shampoo, SOA...

    2024

  14. [2025]

    update-then-sum

    doi: 10.48550/arxiv.2406.16793. Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. GaLore: Memory-efficient LLM training by gradient low-rank projection. InInternational Conference on Machine Learning, 2024. doi: 10.48550/arxiv.2403.03507. Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Les...

    Pith/arXiv arXiv doi:10.48550/arxiv.2406.16793 2024