ArXiv: 2306.09782

🎯 Pitch

Full parameter fine-tuning of a 65B model is possible on a single machine with eight consumer-grade RTX 3090 GPUs—not just parameter-efficient methods. The key is LOMO, a new optimizer that fuses gradient computation and parameter updates, slashing memory use to 10.8% of standard approaches and running over 11× faster than AdamW, all while matching or beating LoRA's accuracy.


1. Executive Summary

This paper introduces LOw-Memory Optimization (LOMO), a new optimizer that fuses gradient computation with parameter updates in a single backward-pass step, enabling full parameter fine-tuning of large language models on consumer-grade GPUs. The work targets the LLaMA family (7B to 65B) and evaluates downstream performance on the SuperGLUE benchmark suite, comparing LOMO against Zero-shot inference and LoRA parameter-efficient fine-tuning. LOMO reduces memory usage to 10.8% of the standard DeepSpeed approach — from 102.20GB down to 14.58GB for a 7B model — by eliminating the storage of gradient tensors and optimizer states (removing AdamW's momentum and variance buffers entirely in favor of plain SGD). The method achieves an 11× throughput improvement over AdamW on a single GPU (769.92 vs. 67.37 tokens per GPU per second for LLaMA-7B) while roughly matching or exceeding LoRA's accuracy, establishing that full-parameter fine-tuning of a 65B model is feasible on a single machine with 8 RTX 3090 GPUs only when the fine-tuning domain shares the smooth-loss-surface properties of the pre-training distribution.

2. Context and Motivation

The Core Problem: Full-Parameter Fine-Tuning Is the Gold Standard but Inaccessible

The fundamental tension this paper addresses is a de facto segregation in the LLM community: researchers with abundant compute can pursue full-parameter fine-tuning — the acknowledged most powerful adaptation method — while those with limited resources are constrained to parameter-efficient fine-tuning (PEFT) methods like LoRA (Hu et al., 2022) or Prefix-Tuning (Li and Liang, 2021). This is not an incidental inconvenience; it is a structural barrier that determines who can meaningfully participate in LLM research and who is locked out.

The paper states this explicitly in the opening paragraph of Section 1:

"Tuning LLMs often requires expensive GPU resources, such as 8×80GB devices, making it difficult for small labs and companies to participate in this area of research."

The precision of "8×80GB" is specifying that training a 65B-parameter model typically requires multiple high-end datacenter GPUs (e.g., A100s or H100s), while many academic groups and smaller organizations operate with consumer-grade hardware like the RTX 3090 (24GB each). The gap between these hardware tiers is not incremental — it is categorical. A single A100-80GB provides roughly 3.3× the memory of an RTX 3090, and the cluster-scale resources of industry labs multiply this by orders of magnitude.

Why Full-Parameter Fine-Tuning Matters

The motivation for pursuing full-parameter tuning — despite its cost — is empirical rather than philosophical. The paper cites prior work establishing that full-parameter fine-tuning is "a more powerful approach than parameter-efficient fine-tuning" (Ding et al., 2022; Sun et al., 2023). This is not a speculative claim; it reflects a consistent finding across the delta-tuning literature: when you tune more parameters, you generally achieve better downstream performance, provided you have sufficient data to avoid overfitting. PEFT methods achieve impressive efficiency by restricting updates to a small subset of parameters (typically 0.1–1% of the total), but this restriction inherently caps the model's capacity to adapt to new tasks — it can only reconfigure within the low-rank or prefix subspace defined by the PEFT architecture.

The practical implications are significant:

  • Scientific equity: When full-parameter fine-tuning is economically gated, research agendas are shaped by who has access to hardware, not by who has the best ideas. This concentrates LLM research in a small number of well-resourced labs.
  • Reproducibility and iteration speed: Full-parameter fine-tuning enables studying how representations change during adaptation, which is central to understanding transfer learning and catastrophic forgetting. PEFT methods obscure these dynamics because most parameters remain frozen.
  • Task performance ceilings: For challenging downstream tasks where the pre-training distribution is far from the target distribution, the restricted capacity of PEFT methods may impose a hard performance ceiling that no amount of hyperparameter tuning can overcome.

Where Existing Approaches Fall Short

The paper identifies four classes of prior work that attempt to address the memory bottleneck, each with specific limitations:

Parameter-Efficient Fine-Tuning (PEFT): Efficiency at the Cost of Expressiveness

Methods like LoRA and Prefix-Tuning dramatically reduce the number of trainable parameters, which reduces both memory usage and computational cost. However, as the paper notes, these methods "do not offer a practical solution for full parameter fine-tuning" (Section 1). The distinction is crucial: PEFT is not a more efficient way to do full-parameter fine-tuning; it is a fundamentally different training regime that trades model capacity against resource requirements. The paper does not position LOMO as an alternative to PEFT — it acknowledges that "these two methods are not conflicting or mutually exclusive" (Section 4.3.2) — but rather as a way to reclaim the full expressiveness of the base model for researchers who cannot afford it.

Adam and Adaptive Optimizers: The Dominant But Memory-Hungry Default

The near-universal default for training transformer models is the Adam optimizer family (Kingma and Ba, 2015; Loshchilov and Hutter, 2019). Adam maintains two moving-average buffers per parameter — momentum and variance — each stored in full precision (FP32) during mixed-precision training. This means the optimizer states alone consume 2× the parameter count in bytes (actually more, since parameters are stored in half-precision (FP16) while optimizer states are FP32). For a 7B model, the parameters occupy roughly 14GB in FP16, but Adam's optimizer states demand an additional 56GB in FP32 (7B × 4 bytes × 2 buffers). Figure 2 in the paper quantifies this visually: 73.7% of total memory is consumed by optimizer states when using AdamW for LLaMA-7B training — parameters and gradients together account for only about 26% of the memory footprint.

The paper's insight here is not that Adam is memory-intensive (this is well-known), but rather that for fine-tuning specifically, the benefits of adaptive optimization might not justify the memory cost. This is a non-obvious claim that the paper develops through both theoretical analysis and empirical results.

DeepSpeed ZeRO: Partitioning Without Eliminating the Root Cause

ZeRO (Rajbhandari et al., 2020) and its variants (ZeRO-Offload, ZeRO-Infinity) partition optimizer states, gradients, and parameters across GPUs, enabling training of larger models by distributing the memory load. This is the "standard approach" the paper references when claiming a 10.8% memory reduction. However, ZeRO addresses the symptom (insufficient per-GPU memory) rather than the cause (excessive per-parameter memory consumption). It requires multi-GPU setups with high-bandwidth interconnects, which consumer hardware often lacks (the paper's experiments use PCIe-interconnected RTX 3090s, where communication overhead becomes a bottleneck — Section 4.2). Furthermore, ZeRO introduces communication costs that scale with model size and GPU count, partially offsetting its memory benefits.

Gradient Checkpointing and Heterogeneous Memory: Partial Solutions

Activation checkpointing (Chen et al., 2016) reduces activation memory at the cost of recomputation, which the paper acknowledges as orthogonal and complementary to LOMO (the paper combines them to reduce activation memory from 45.61GB to 1.79GB; Table 1). Heterogeneous training systems (Rhu et al., 2016; Pudipeddi et al., 2020; Ren et al., 2021b) offload tensors to CPU or NVMe memory, but these introduce latency from slower memory tiers and require careful orchestration of data movement. The paper notes these techniques can be "effectively combined with LOMO" (Section 2), positioning them as complementary rather than competitive.

The Gap: No Method Eliminates the Need to Store All Gradients

The crucial observation that motivates LOMO is that none of the existing approaches challenge a fundamental assumption of the standard training loop: all gradients must be computed before any parameters are updated. This assumption forces the system to store gradient tensors for every parameter simultaneously — an additional memory burden equal to the parameter count (for FP16 gradients, ~14GB for a 7B model). If gradients could be applied immediately as they are computed, the gradient storage requirement would drop from O(P) to O(1) — storing only the gradient of the current parameter.

This is the key technical gap the paper fills. The existing memory-saving ecosystem (ZeRO, checkpointing, offloading) operates within the standard two-phase paradigm: forward pass → backward pass (accumulate all gradients) → update all parameters. LOMO challenges this paradigm by fusing the backward pass with parameter updates, making gradient storage near-zero.

How This Paper Positions Itself

The paper positions LOMO at the intersection of three design choices that together produce a qualitatively different resource profile:

  1. Replacing Adam with SGD: This eliminates optimizer states entirely (56GB savings for a 7B model), justified by a theoretical argument that the loss surface of LLMs is sufficiently smooth and the fine-tuning regime sufficiently local that adaptive optimization's curvature-aware adjustments are unnecessary. The paper develops this argument through the "Implicit Batch Size" analysis (Section 3.1.2), showing that under a smooth loss surface assumption, sequential SGD updates approximate a larger-batch update.

  2. Fusing gradient computation with parameter updates: This eliminates gradient storage (14GB savings for a 7B model in FP16), implemented via hook functions injected into PyTorch's backward pass. This is the algorithmic contribution that distinguishes LOMO from simply "using SGD instead of Adam."

  3. Integrating with existing techniques: LOMO is explicitly designed to compose with activation checkpointing, mixed-precision training, and ZeRO-style partitioning — the combination achieves the headline 10.8% memory compared to the standard AdamW + DeepSpeed baseline.

The paper's framing is deliberately pragmatic rather than fundamental. It does not claim that SGD is universally superior to Adam, or that fused gradient updates are always preferable. The scope is carefully limited to fine-tuning LLMs on natural language tasks (Section 3.1.1):

"Note that this holds only when we teach the LLMs natural language-based tasks (or code-based if pre-trained with code). A synthetic loss function unrelated to pre-training tasks will indeed face the large curvature problem."

This boundary condition is crucial: LOMO works because fine-tuning starts from a pre-trained model sitting in a smooth region of the loss landscape. Training from scratch, or fine-tuning on tasks far from the pre-training distribution, might require the curvature adaptation that Adam provides.

The paper also distinguishes itself from contemporaneous memory-efficient optimizers. MeZO (Malladi et al., 2023) uses zeroth-order optimization — estimating gradients via finite differences from two forward passes — which avoids backpropagation entirely but introduces gradient estimation error. LOMO, by contrast, uses exact first-order gradients and sacrifices no fidelity to the SGD update. GaLore (Zhao et al., 2024) performs low-rank gradient decomposition, which approximates the full gradient. LOMO, again, uses the exact gradient — the memory savings come from when gradients are stored, not from what gradients are computed.

The Conceptual Through-Line: Memory Proportional to Inference

A revealing way to understand the paper's ambition is through the implicit benchmark it sets: the inference memory floor. During inference, only the model parameters and activations for the current batch need to reside in memory — no gradients, no optimizer states, no intermediate computation graphs. The paper repeatedly emphasizes that LOMO's memory usage "is merely equivalent to the usage of inference" (Section 3.2) and "commensurate with memory usage during inference" (Section 4.1). This is an important framing because it establishes a theoretical lower bound: you cannot use less memory for training than for inference, since training subsumes the forward pass. By pushing memory usage to this floor, LOMO achieves what the paper argues is the best possible outcome without model compression — further reductions would require quantizing parameters themselves, which the paper identifies as future work (Section 5).

3. Technical Approach

3.1 Reader Orientation

LOMO is a gradient computation strategy and optimizer that rewires the standard training loop so that parameters are updated one-by-one during the backward pass, rather than waiting until all gradients have been computed. It solves the problem of excessive GPU memory consumption during full-parameter fine-tuning by eliminating the need to store gradient tensors for every parameter simultaneously — the memory footprint drops from storing all gradients (O(P), where P is the number of parameters) to storing only the single largest gradient tensor (O(1)).

3.2 Big-Picture Architecture (Diagram in Words)

The LOMO system has four integrated components, each addressing a different memory pressure point:

  1. SGD as the sole optimizer — replaces AdamW entirely, eliminating the momentum and variance buffers that consume 73.7% of memory in standard training (storing FP32 copies of two moment estimates per parameter). This is justified theoretically by arguing that the smooth loss surface of LLMs makes adaptive optimization unnecessary during fine-tuning.

  2. Fused backward-pass updates — the core algorithmic innovation. Rather than computing all gradients first (backward pass → gradient storage → optimizer step), LOMO registers hooks on the computational graph that trigger a parameter update immediately after that parameter's gradient is computed, then discard the gradient. This reduces gradient memory from storing all parameter gradients to storing only the current layer's gradient.

  3. Stability mechanisms for mixed-precision training — because LOMO breaks the standard two-phase (compute-all-gradients, then update) paradigm, operations that normally require access to all gradients simultaneously — specifically gradient norm computation and dynamic loss scaling — require re-engineering. The paper proposes value-based clipping as a lightweight alternative to norm-based clipping and a two-pass backward scheme when norm information is genuinely needed.

  4. Compatibility layer with existing techniques — LOMO is explicitly designed to compose with activation checkpointing (reducing activation memory), mixed-precision training (FP16 forward/backward with FP32 master weights), and ZeRO-style parameter partitioning (for multi-GPU setups). The memory savings are multiplicative: each technique addresses a different memory category.

The information flow during a single training step with LOMO is: forward pass (compute activations and store checkpoints if using activation checkpointing) → backward pass begins → for each layer from last to first, compute the layer's gradient → immediately apply the SGD update to that layer's parameters → discard the gradient → proceed to the previous layer. After the backward pass completes, no gradient tensors remain in memory.

3.3 Roadmap for the Deep Dive

  • First, the optimizer choice (SGD over Adam) — because this is the foundational decision that eliminates optimizer states and determines what update rule LOMO must implement. We need to understand why SGD is acceptable before we can understand how LOMO implements it memory-efficiently.
  • Second, the fused gradient computation and update mechanism — this is the core algorithmic contribution. We will trace exactly how hooks are injected, what happens layer-by-layer during the backward pass, and why the memory savings materialize.
  • Third, the implicit batch size argument — a theoretical analysis that explains why sequential SGD updates on a smooth loss surface approximate larger-batch updates, providing analytical support for the empirical claim that SGD is stable for LLM fine-tuning.
  • Fourth, the stability mechanisms — gradient clipping alternatives, precision preservation through dynamic loss scaling and full-precision casting, and the two-pass backward strategy required when gradient norm information is unavoidable.
  • Fifth, the composition with other memory-saving techniques — how activation checkpointing, mixed-precision training, and ZeRO partitioning interact with LOMO, and why the memory categories they address are essentially orthogonal.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a systems paper whose core insight is that the standard two-phase training loop (backward to compute all gradients, then update all parameters) is a convention rather than a necessity — and that breaking this convention eliminates an entire category of memory usage. The theoretical contribution is the argument that SGD is sufficient for LLM fine-tuning under smoothness assumptions. The engineering contribution is the hook-based implementation that realizes fused updates in PyTorch.


Why SGD Replaces Adam: The Smoothness Argument and Optimizer State Elimination

The standard default for training transformer models is the Adam optimizer family. Adam maintains two exponential moving averages per parameter: the first moment estimate mtm_t (momentum, tracking the mean gradient direction) and the second moment estimate vtv_t (tracking the uncentered variance of gradients). During mixed-precision training (FP16 parameters, FP16 gradients, FP32 optimizer states), these buffers are stored in full precision (FP32), consuming:

Moptimizer=2×P×4 bytes=8P bytesM_{\text{optimizer}} = 2 \times P \times 4 \text{ bytes} = 8P \text{ bytes}

where PP is the number of parameters and each FP32 value occupies 4 bytes.

For a 7B-parameter model, this is 8×7×109=568 \times 7 \times 10^9 = 56 GB — for optimizer states alone, separate from parameters (~14 GB in FP16) and gradients (~14 GB in FP16). Figure 2 in the paper quantifies this visually: optimizer states consume 73.7% of total training memory when using AdamW.

The paper's first design decision is to replace Adam with plain stochastic gradient descent (SGD). The SGD update rule is:

θt+1=θtαL(θt)\boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - \alpha \nabla \mathcal{L}(\boldsymbol{\theta}_t)

where θt\boldsymbol{\theta}_t is the parameter vector at step tt, α\alpha is the learning rate, and L(θt)\nabla \mathcal{L}(\boldsymbol{\theta}_t) is the gradient of the loss with respect to the parameters evaluated at θt\boldsymbol{\theta}_t.

What it computes: the simplest possible first-order update — take a step in the negative gradient direction scaled by the learning rate. No momentum, no variance normalization, no per-parameter adaptation. The only state is the learning rate α\alpha, which is a single scalar, not a per-parameter buffer.

Why this form over Adam: the paper develops a three-part argument (Section 3.1.1) that the known weaknesses of SGD relative to adaptive methods are mitigated in the specific regime of LLM fine-tuning:

Argument 1: Large curvature is unlikely. The paper cites prior work (Hao et al., 2019) as empirical evidence that large pre-trained language models have smooth loss surfaces. The intuition is that the pre-training process across massive, diverse corpora produces parameters that sit in broad, flat basins — small perturbations to the parameters do not dramatically change the loss. This means the per-parameter adaptive learning rates that Adam computes (dividing by vt\sqrt{v_t}) are less valuable because the gradient magnitudes are expected to be well-behaved across parameter dimensions. The paper is explicit about the boundary condition: "Note that this holds only when we teach the LLMs natural language-based tasks (or code-based if pre-trained with code). A synthetic loss function unrelated to pre-training tasks will indeed face the large curvature problem."

Argument 2: Local optimum is sufficient for fine-tuning. The goal of fine-tuning is adaptation, not optimization from scratch. The pre-trained parameters already encode substantial knowledge about language. A local optimum near the pre-trained initialization is expected to perform well because it preserves most of the pre-trained representations while adjusting to the target task. The paper argues this explicitly: "a local optimum is often a good enough solution, and the limited training data (compared to pre-training corpus) makes it difficult to push the model to a faraway global optimum."

Argument 3: Saddle points are distant from the initialization. Saddle points — where the gradient is zero but the point is not a local minimum — are a well-known challenge for SGD because the optimizer can stall in flat regions with zero first-order information. The paper argues that since the pre-trained model starts in a "valley" (region of low loss) and fine-tuning data is typically drawn from a distribution related to pre-training, the initialization is far from saddle points that might trap SGD. The argument is: "if we do not change the parameter too far from the pre-trained value," saddle points — which "typically appear on ridges and have a distance from valleys" — are unlikely to be encountered.

Why this is controversial and the paper acknowledges it: the paper does not claim that SGD is universally better than Adam, or even that it is always equal. It presents the argument as a deliberate trade — sacrifice the adaptive optimization machinery for massive memory savings, with the justification that the specific regime (fine-tuning, natural language, large pre-trained models) is forgiving enough that the sacrifice is acceptable. Section 3.1.1 concludes: "there is no guarantee that SGD is a powerful optimizer compared to modern optimizers. Our intention is to create a simple and practical solution for fine-tuning LLMs and identify its flaws to continually improve it."


The Implicit Batch Size Argument: Why Sequential SGD on a Smooth Surface Approximates Larger Batches

Beyond the qualitative smoothness discussion, the paper provides a mathematical argument (Section 3.1.2) to explain why SGD with small batches is stable when fine-tuning LLMs. The argument shows that under a smoothness assumption, two sequential SGD steps on two data points approximate one SGD step on the combined batch — meaning sequential updates effectively simulate larger batch sizes.

Setup: We have a pre-trained model f()f(\cdot) with parameters θ\boldsymbol{\theta}, a training set D={d1,d2,,dn}\mathcal{D} = \{d_1, d_2, \ldots, d_n\}, and a loss function L\mathcal{L}. Consider two training examples did_i and djd_j.

Case 1: One SGD step on the combined batch (two data points):

θ=θα[L(di,f(di,θ))+L(dj,f(dj,θ))]\boldsymbol{\theta}' = \boldsymbol{\theta} - \alpha[\nabla \mathcal{L}(d_i, f(d_i, \boldsymbol{\theta})) + \nabla \mathcal{L}(d_j, f(d_j, \boldsymbol{\theta}))]

where α\alpha is the learning rate. The update is the sum of the two individual gradients computed at the same initial parameters θ\boldsymbol{\theta}.

Case 2: Two sequential SGD steps (one per data point):

Step 1: θ1=θαL(di,f(di,θ))\boldsymbol{\theta_1} = \boldsymbol{\theta} - \alpha \nabla \mathcal{L}(d_i, f(d_i, \boldsymbol{\theta}))

Step 2: θ2=θ1αL(dj,f(dj,θ1))\boldsymbol{\theta_2} = \boldsymbol{\theta_1} - \alpha \nabla \mathcal{L}(d_j, f(d_j, \boldsymbol{\theta_1}))

In the sequential case, the gradient for djd_j is evaluated at θ1\boldsymbol{\theta_1} (the parameters after updating on did_i), not at the original θ\boldsymbol{\theta}. This is the source of the discrepancy — the sequential update uses stale gradient information for the second example.

Bridging the gap via the mean value theorem: The paper applies the differential mean value theorem to relate L(dj,f(dj,θ1))\mathcal{L}(d_j, f(d_j, \boldsymbol{\theta_1})) to L(dj,f(dj,θ))\mathcal{L}(d_j, f(d_j, \boldsymbol{\theta})):

L(dj,f(dj,θ1))=L(dj,f(dj,θ))+L(dj,ξ)(f(dj,θ1)f(dj,θ))\mathcal{L}(d_j, f(d_j, \boldsymbol{\theta_1})) = \mathcal{L}(d_j, f(d_j, \boldsymbol{\theta})) + \nabla \mathcal{L}(d_j, \xi)(f(d_j, \boldsymbol{\theta_1}) - f(d_j, \boldsymbol{\theta}))

where ξ\xi is some point on the line segment between f(dj,θ)f(d_j, \boldsymbol{\theta}) and f(dj,θ1)f(d_j, \boldsymbol{\theta_1}). Substituting this back into the expression for θ2\boldsymbol{\theta_2} and manipulating:

θ2=θα[L(di,f(di,θ))+L(dj,f(dj,θ))]α[L(dj,ξ)(f(dj,θ1)f(dj,θ))]\boldsymbol{\theta_2} = \boldsymbol{\theta} - \alpha[\nabla \mathcal{L}(d_i, f(d_i, \boldsymbol{\theta})) + \nabla \mathcal{L}(d_j, f(d_j, \boldsymbol{\theta}))] - \alpha \nabla[\nabla \mathcal{L}(d_j, \xi)(f(d_j, \boldsymbol{\theta_1}) - f(d_j, \boldsymbol{\theta}))]

What this structure reveals: the first two terms (multiplied by α\alpha) exactly match the combined-batch update in Case 1. The third term is an error term — a second-order correction involving the change in the model's output on djd_j between θ\boldsymbol{\theta} and θ1\boldsymbol{\theta_1}, multiplied by a Hessian-like term [L(dj,ξ)]\nabla[\nabla \mathcal{L}(d_j, \xi) \cdots].

The smoothness assumption makes the error term negligible. If the loss surface is smooth, then two things are true:

  1. The change f(dj,θ1)f(dj,θ)f(d_j, \boldsymbol{\theta_1}) - f(d_j, \boldsymbol{\theta}) — how much the model's prediction on example djd_j changes after a single SGD step on example did_i — is small. This follows from parameter smoothness: a small parameter change produces a small output change.
  2. The second derivative term [L(dj,ξ)]\nabla[\nabla \mathcal{L}(d_j, \xi) \cdots] is bounded (smoothness means second derivatives exist and are moderate).

Under these conditions, the error term is small, and the sequential update θ2\boldsymbol{\theta_2} approximates the batch update θ\boldsymbol{\theta}'. The paper's conclusion:

"It suggests that utilizing SGD optimizer over a smooth loss surface could imply a larger batch size."

Why this form of argument matters: this is not a generic "SGD works" claim. It is a specific mechanism: each sequential SGD step on individual examples approximates a larger effective batch size because the model changes slowly enough that gradient staleness — the fact that later examples in the sequence see different parameters than earlier ones — does not substantially bias the update. This is important because larger batch sizes are associated with more stable training (less noisy gradient estimates). The paper is arguing that sequential SGD on LLMs inherently achieves the stability of larger batches because the loss surface is smooth enough that sequential updates don't diverge significantly from batched updates.

A subtle point about the analysis: the paper uses two data points for clarity, but the argument generalizes: a sequence of kk single-example SGD steps approximates a single step on a batch of kk examples, up to error terms involving second derivatives and the total parameter displacement over the sequence. For LLMs, where each step produces small parameter changes (due to small learning rates and smoothness), this approximation can be good for non-trivial kk.


The Core Algorithm: Fused Gradient Computation and Parameter Update

The paper's central technical innovation is LOMO's fused update (Algorithm 1 in the paper). The standard training loop in deep learning frameworks like PyTorch separates gradient computation from parameter updates:

  1. Forward pass: compute activations and loss.
  2. Backward pass: traverse the computational graph in reverse, computing gradients for all parameters and storing them in per-parameter tensors.
  3. Optimizer step: iterate over all parameters and apply the update rule (e.g., p = p - lr * p.grad for SGD).

This separation means gradients for all PP parameters exist simultaneously in memory, consuming P×sizeof(dtype)P \times \text{sizeof(dtype)} bytes (typically 2 bytes for FP16 or 4 bytes for FP32). For a 7B model with FP16 gradients, this is ~14 GB — equal to the parameter memory itself.

LOMO's key insight: the update rule ppαLpp \leftarrow p - \alpha \cdot \frac{\partial \mathcal{L}}{\partial p} for a single parameter pp requires only that parameter and its gradient. It does not require gradients for any other parameter. Therefore, if we can execute the update immediately after computing Lp\frac{\partial \mathcal{L}}{\partial p} — before the backward pass proceeds to compute gradients for other parameters — we can discard that gradient and reuse its memory.

The paper describes this as:

"The fusion version is p=plrLpp = p - lr * \frac{\partial \mathcal{L}}{\partial p}."

This is a single fused operation: compute gradient and update in place, with no intermediate storage.

Implementation via backward hooks (Algorithm 1): PyTorch's autograd engine computes gradients layer-by-layer in reverse topological order during the backward pass. LOMO registers backward hooks — callback functions that PyTorch invokes immediately after computing the gradient for a specific parameter — on every trainable parameter in the model. The hook function does three things:

  1. Retrieves the gradient tensor for the associated parameter (computed by autograd for that parameter).
  2. Applies the SGD update in place: param.data -= lr * grad.
  3. Sets the gradient to None to release the memory immediately.

The paper's Algorithm 1 traces this explicitly:

for l = L, ..., 1 do                               // backward pass, layer by layer
    θ_l ← [θ_i for θ_i in layer l]                 // collect parameters for layer l
    g_l ← ∂ℓ / ∂θ_l                                 // autograd computes gradient for this layer
    θ_l ← θ_l - α * g_l                             // immediate in-place update
    g_l ← None                                      // discard gradient tensor
end for

What happens physically in memory: as the backward pass proceeds from the last layer (closest to the loss) to the first layer (closest to the input), gradients are computed for each layer's parameters. At any moment, only the gradient tensors for the current layer exist in memory, plus the gradients for activations that autograd needs for the continuing backward propagation. Once a layer's parameters are updated and its gradients discarded, that memory is freed. By the time the backward pass reaches the first layer, gradients for all subsequent layers have already been applied and freed.

The memory ceiling is now the largest single gradient tensor, not the sum of all gradient tensors. For transformer models, the largest parameter matrices are typically the feed-forward network projection layers or the attention output projections. The memory for gradients drops from approximately P×sizeof(dtype)P \times \text{sizeof(dtype)} to maxlθl×sizeof(dtype)\max_l |\theta_l| \times \text{sizeof(dtype)}, where θl|\theta_l| is the number of parameters in layer ll. For a transformer with uniform layer sizes, this is approximately a factor of LL (number of layers) reduction — for LLaMA-7B with 32 layers, the largest single parameter matrix might be ~1/32 of the total.

A subtlety about shared parameters: the paper notes in a footnote (footnote 2) that "we should inject different hook functions accordingly if some of them share the weight." Weight sharing (e.g., tied embeddings) means the same parameter tensor appears in multiple parts of the computational graph. Autograd would compute multiple gradient contributions that need to be summed. LOMO handles this by only updating the parameter after all its gradient contributions have been accumulated, requiring hook logic that tracks pending contributions.

Why this cannot be implemented with standard PyTorch APIs directly: the paper states "we cannot implement the exact immediate update with current APIs. Instead, we store at most one parameter's gradient in memory and update each parameter one by one along with the backward propagation." The issue is that PyTorch's autograd engine may compute gradients for multiple parameters in the same layer simultaneously (e.g., for different linear layers within an attention block), and the hook mechanism fires per-parameter but does not guarantee precise ordering. LOMO's implementation manages this by buffering gradients at the granularity needed to ensure correctness — at worst, storing a few parameter gradients for one layer rather than strictly one-at-a-time — but the memory savings are effectively the same because the buffer is bounded by the largest layer, not the entire model.

Relationship to inference memory: with LOMO and activation checkpointing, the only persistent memory consumers during training are the model parameters (in FP16, ~14 GB for 7B), the master parameter copies in FP32 if using mixed precision (~28 GB), and the activations for the current checkpoint segment (~1-2 GB with checkpointing). This is essentially the inference footprint plus the FP32 master weights. The paper emphasizes this repeatedly: "the memory usage of the forward + backward process should not be less than the forward process alone," establishing this as the theoretical minimum for full-parameter training without model compression.


Gradient Clipping Without All Gradients: Value-Based Clipping and the Two-Pass Alternative

Gradient clipping is a standard technique to stabilize training by preventing individual gradient steps from being excessively large. The standard approach computes the L2 norm (Euclidean norm) of the complete gradient vector:

g2=i=1Pgi2\|\mathbf{g}\|_2 = \sqrt{\sum_{i=1}^{P} g_i^2}

where g\mathbf{g} is the vector of all parameter gradients and PP is the total number of parameters. If this norm exceeds a threshold CC, every gradient value is scaled by C/g2C / \|\mathbf{g}\|_2, preserving the direction while capping the magnitude.

The problem for LOMO: computing the global L2 norm requires access to all gradient values simultaneously — you need to square each gig_i, sum them, and take the square root, and only then can you decide whether and how to scale. But LOMO updates parameters and discards gradients one layer at a time, so when updating layer ll, it has already forgotten the gradient values for layers l+1l+1 through LL (already updated and discarded) and does not yet know the gradients for layers 11 through l1l-1 (not yet computed).

Value-based clipping (lightweight alternative): instead of computing the global norm, clip each gradient value individually to a range [c,c][-c, c]:

giclipped=clamp(gi,c,c)g_i^{\text{clipped}} = \text{clamp}(g_i, -c, c)

where clamp(x,a,b)=max(a,min(x,b))\text{clamp}(x, a, b) = \max(a, \min(x, b)). This operation is local to each gradient element and can be performed in the backward hook immediately before the update, requiring no knowledge of other gradients.

Why this is imperfect: the paper acknowledges the limitation explicitly. Clipping by value can change the direction of the gradient vector, not just its magnitude. The paper gives the example: a two-dimensional vector [1.3,0.8][1.3, 0.8] clipped to 1.01.0 becomes [1.0,0.8][1.0, 0.8], which points in a different direction (different ratio of components). Norm-based clipping preserves direction — [1.3,0.8][1.3, 0.8] scaled to have the same norm as [1.0,0.8][1.0, 0.8] would be scaled proportionally, keeping the 1.3:0.8 ratio.

When value-based clipping is acceptable: the paper reports empirically that "clipping by values performs worse when the learning rate is high because truncations happened more often in that case. However, clipping by values performs well for medium and small learning rates." The paper suggests this as a practical guideline: "we suggest using clipping by values for a learning rate less than 1×1031 \times 10^{-3}." At low learning rates, individual gradient elements are smaller, so the clipping threshold is rarely triggered, and when it is, the distortion to gradient direction is small because only a few extreme values are affected.

The two-pass backward for gradient norm (exact but slower): when gradient norm information is genuinely necessary (e.g., for monitoring, or for tasks where value-based clipping is insufficient), the paper proposes a two-pass scheme:

  1. First backward pass: compute gradients for all parameters but do not update them. Accumulate the squared gradient norm: igi2\sum_i g_i^2. At the end, compute the scaling factor.
  2. Second backward pass: recompute gradients (or retrieve them if stored), apply the scaling factor, and update parameters using the fused LOMO update.

The paper notes: "The memory usage leaves unchanged but sacrifices the speed." The first backward pass can discard gradients after accumulating their squared norm — only the running sum needs to be stored, not the individual gradients — but it still requires traversing the entire computational graph, effectively doubling the backward computation.

The "controversial" grouped approximation: the paper briefly mentions (in a subsection titled "A controversial solution") the possibility of approximating the global gradient norm using only a subset of parameters — for example, computing the norm within each group of adjacent layers and clipping that group independently. This is "controversial" because different parameter groups would receive different scaling factors, which is equivalent to applying per-group dynamic learning rates. The paper cites Sun et al. (2020a) to suggest this might actually be beneficial — "it is not always appropriate to use the same learning rate for all parameters in SGD" — but leaves this as unexplored future work.


Dynamic Loss Scaling and Precision Preservation

Mixed-precision training (FP16 for most operations, FP32 for critical accumulations) is essential for both speed and memory efficiency on modern GPUs. However, FP16 has a limited dynamic range (minimum positive normalized value ~6×1086 \times 10^{-8}), and gradients during training can become very small, causing underflow — values flushed to zero, losing gradient signal entirely.

Loss scaling (Micikevicius et al., 2018) is the standard solution: multiply the loss by a large factor SS before the backward pass, which scales all gradients proportionally. After the backward pass but before the parameter update, divide the gradients by SS. This shifts small gradient values into the representable range of FP16 without changing the effective update (since the scaling is undone before the update).

Dynamic loss scaling adapts SS during training: if no gradient overflow (NaN or Inf values) is detected for a window of steps, SS is doubled (allowing smaller gradients to be represented). If overflow is detected, the current step is skipped entirely and SS is halved.

The problem for LOMO: overflow detection normally happens after all gradients are computed — you inspect all gradient tensors for NaN/Inf, and if any overflow, you discard the step (do not update) and reduce SS. But LOMO updates parameters as gradients are computed, so by the time overflow is detected in a later layer, earlier layers have already been updated — you cannot "skip the step" for them retroactively.

LOMO's solution: the two-pass backward for loss scaling. The paper integrates dynamic loss scaling by performing two backward passes:

  1. First backward pass (scouting): compute gradients without updating parameters. Check for NaN/Inf in any gradient tensor. If overflow is detected, reduce SS, and skip the second pass entirely (no updates). If no overflow, proceed to pass 2 and optionally double SS for future steps.
  2. Second backward pass (update): if no overflow was detected, recompute gradients and apply LOMO's fused updates as normal.

Why this is acceptable (the integration argument): the paper notes that "these two backward passes for dynamic loss scaling can be executed simultaneously with gradient normalization." If gradient norm clipping is also needed and also requires a two-pass scheme (as described above), then both purposes can be served by the same pair of backward passes — the first pass detects overflow and computes the gradient norm, and the second pass applies the clipping factor and performs the fused updates. The cost is still two backward passes, not four.

Precision casting during updates: when applying the SGD update in the backward hook, the gradient and the associated parameter are temporarily cast to FP32 for the actual arithmetic:

paramFP32paramFP32α×gradFP32\text{param}_{\text{FP32}} \leftarrow \text{param}_{\text{FP32}} - \alpha \times \text{grad}_{\text{FP32}}

This prevents the accumulation of rounding errors that would occur if the subtraction were performed in FP16. The master copy of parameters (the FP32 version maintained alongside the FP16 version used in the forward pass) is what actually gets updated. The FP16 parameter is then refreshed from the FP32 master for the next forward pass. This is standard practice in mixed-precision training and does not introduce additional memory overhead beyond the FP32 master weights — which would exist in any mixed-precision training setup, with or without LOMO.


Composition with Existing Memory-Saving Techniques

LOMO is not designed to replace the entire memory-saving ecosystem; it targets specific memory categories (gradients and optimizer states) and is explicitly orthogonal to techniques that target other categories.

Activation checkpointing (Chen et al., 2016): during the forward pass, only a subset of layer activations (the "checkpoints") are retained. During the backward pass, when activations for a non-checkpointed layer are needed, they are recomputed from the nearest preceding checkpoint with a partial forward pass. This trades computation (~33% extra forward pass cost for the standard square-root schedule) for reduced activation memory (from O(L)O(L) to O(L)O(\sqrt{L}) for a network with LL layers, in terms of stored activations).

LOMO combines with activation checkpointing seamlessly because checkpointing only affects when activations are available during the backward pass, not how gradients are handled once computed. Table 1 shows the effect quantitatively: for LLaMA-7B with sequence length 512 and batch size 8, activation memory drops from 45.61 GB (no checkpointing) to 1.79 GB (with checkpointing), while LOMO handles the gradient and optimizer state elimination independently.

ZeRO-3 parameter partitioning (Rajbhandari et al., 2020): ZeRO stage 3 partitions not just optimizer states and gradients (stages 1 and 2) but also the model parameters themselves across GPUs. Each GPU holds only a fraction of each layer's parameters. During the forward pass, parameters are gathered via all-gather communication when needed for their layer's computation, then discarded. During backward, the same gathering happens again for gradient computation.

LOMO composes with ZeRO-3 because LOMO's fused update operates on a per-parameter basis — when a parameter's gradient is computed, LOMO updates it immediately. With ZeRO-3, this means each GPU updates only its partition of parameters, and the memory savings from LOMO (no gradient storage, no optimizer states) multiply with the memory savings from ZeRO-3 (only 1/N1/N of parameters per GPU for an NN-GPU setup).

The paper uses ZeRO-3 in its multi-GPU experiments (Section 4.2) for the AdamW and SGD baselines but notes that LOMO requires fewer GPUs than SGD for the same model size because its per-GPU memory footprint is lower. For example, LLaMA-13B requires 8 GPUs with SGD but only 2 GPUs with LOMO (Table 2).

Heterogeneous memory (ZeRO-Offload): the paper mentions that offloading to CPU or NVMe memory (Ren et al., 2021b; Rajbhandari et al., 2021) is also orthogonal — LOMO reduces the amount of data that would need to be offloaded, making offloading-based approaches more efficient, but does not depend on them.


The Training Procedure End-to-End

Combining all components, a single training step with LOMO (with activation checkpointing and dynamic loss scaling) proceeds as follows:

  1. Forward pass (FP16): the input batch is processed through the model. Activation checkpoints are stored at designated layers. All other activations are discarded after their downstream computations complete. The loss L\mathcal{L} is computed in FP32 to preserve accuracy.

  2. Dynamic loss scaling (part 1): the loss is multiplied by the current scale factor SS. The scaled loss Lscaled=SL\mathcal{L}_{\text{scaled}} = S \cdot \mathcal{L} is the starting point for the backward pass.

  3. First backward pass (scouting): the autograd engine traverses the computational graph in reverse. Gradients are computed but not stored persistently. At each layer, the hook checks for NaN or Inf in the computed gradients. If overflow is detected, the flag is set. No parameter updates occur in this pass. If gradient norm computation is needed, squared gradient norms are accumulated into a running scalar.

  4. Overflow check: if any overflow was detected during the first backward pass, the step is skipped entirely. The loss scale SS is halved, and training proceeds to the next batch (no second backward pass). If no overflow was detected, proceed to step 5.

  5. Second backward pass (update): the autograd engine traverses the graph again. This time, at each layer:

    • Activations needed for computing that layer's gradients are retrieved (either from stored checkpoints or recomputed via the activation checkpointing logic).
    • The gradient for each parameter in the layer is computed.
    • If gradient norm clipping is active, the precomputed global norm from step 3 is used to compute the scaling factor C/g2C / \|\mathbf{g}\|_2, and each gradient is scaled accordingly. If value-based clipping is used instead, each gradient element is clamped individually.
    • The gradient is cast to FP32, and the parameter update is applied: param_fp32 -= lr * (grad_fp32 / S). The division by SS undoes the loss scaling, restoring the correct gradient magnitude.
    • The FP16 parameter is refreshed from the updated FP32 master.
    • The gradient tensor is set to None and its memory is freed.
    • The backward pass continues to the previous layer.
  6. Post-step: the loss scale SS may be doubled if no overflow has occurred for a consecutive window of steps (the paper's implementation uses the standard heuristic from Micikevicius et al., 2018).

After the backward pass completes, no gradient tensors remain in GPU memory. The only persistent memory consumers are the model parameters (both FP16 and FP32 copies) and the activation checkpoints (if using checkpointing). This is the "inference-equivalent" memory floor the paper aims for.


LOMO and Parameter-Efficient Fine-Tuning: Orthogonal Mechanisms

Section 4.3.2 of the paper explicitly tests the combination of LOMO with LoRA (Hu et al., 2022). LoRA works by freezing the pre-trained weights and injecting trainable low-rank decomposition matrices (AA and BB) into specific layers, so that the effective weight becomes W+BAW + BA where only AA and BB (much smaller than WW) are trained. This reduces the number of trainable parameters dramatically (typically to 0.1-1% of the full parameter count).

Why combine them? LOMO reduces memory for full-parameter training; LoRA reduces the number of trainable parameters. They address memory from different angles:

  • LoRA means fewer parameters need gradients computed at all (since most parameters are frozen), reducing both computation and gradient memory.
  • LOMO means the gradients for the trainable parameters (whether many or few) are not stored persistently.

The paper's experiments (Figure 3) show that "LOMO + LoRA" (injecting LoRA modules while fine-tuning the pre-trained model weights using LOMO) consistently outperforms LoRA alone. This means LOMO is tuning the full parameters while LoRA adds additional capacity — the two are additive rather than alternative. The paper notes: "LOMO does not compromise the performance of LoRA; rather, it facilitates better model tuning for downstream tasks."

The practical implication: for researchers with extremely limited resources, the combination provides a spectrum. Use LoRA alone for the most memory-constrained scenarios; use LOMO alone for full-parameter fine-tuning when possible; use LOMO + LoRA for maximum performance when the combined memory footprint is acceptable. The paper does not benchmark this combined memory footprint separately but notes the conceptual orthogonality.


Summary of Design Choices and Their Justifications

  • SGD over Adam: justified by smoothness of LLM loss surfaces and the locality of fine-tuning, enabling elimination of optimizer states (73.7% of memory in standard training). The sacrifice is potential training instability on non-smooth or far-from-pre-training tasks — a recognized boundary condition.
  • Fused gradient computation and update: enabled by backward hooks in autograd, reduces gradient memory from O(P)O(P) to O(1)O(1) (largest single tensor). The sacrifice is the inability to use standard gradient operations that require all gradients simultaneously — addressed by value-based clipping and two-pass schemes.
  • Value-based gradient clipping over norm-based clipping: avoids the need for a full gradient norm computation at the cost of occasionally distorting gradient direction. Empirically acceptable at learning rates below 1×1031 \times 10^{-3}.
  • Two-pass backward for dynamic loss scaling: avoids corrupting parameters with overflow gradients at the cost of doubling backward computation. Can be combined with gradient norm computation in the same pair of passes.
  • FP32 casting during parameter updates: preserves numerical accuracy in mixed-precision training, standard practice that LOMO preserves rather than modifying.
  • Orthogonality to activation checkpointing and ZeRO: LOMO targets gradient and optimizer state memory; these target activation and parameter memory respectively. The savings multiply rather than overlap, which is why the paper can claim the combined system reaches 10.8% of the baseline memory.

4. Key Insights and Innovations

Innovation 1: Full-Parameter Fine-Tuning Memory Can Be Reduced to the Inference Floor — and That Floor Is a Hard Bound

The paper's most conceptually significant contribution is not any single technique but rather a diagnostic reframing of the memory problem itself. Before LOMO, the dominant approach to memory reduction was partitioning: spread the problem across more GPUs (ZeRO), offload to slower memory tiers (ZeRO-Offload), or reduce what you train (PEFT). All of these accept the fundamental memory profile of the training loop as given and work around it. LOMO challenges the profile itself.

The key framing move appears in the paper's repeated emphasis that LOMO's memory usage "is merely equivalent to the usage of inference" (Section 3.2) and that "the memory usage of the forward + backward process should not be less than the forward process alone." This establishes a theoretical lower bound that no training algorithm can beat without model compression: training must include at minimum the forward pass, so it must consume at least the inference memory. By showing that the combination of SGD (eliminating optimizer states) and fused gradient updates (eliminating gradient storage) reaches exactly this bound, the paper demonstrates that the memory problem for full-parameter fine-tuning has been solved to its theoretical limit on the optimizer side. Further improvements must come from quantizing parameters themselves (which the paper flags as future work, Section 5) — compression, not better optimization.

This is a fundamental conceptual contribution rather than an incremental engineering improvement because it converts a fuzzy goal ("reduce memory usage as much as possible") into a precise target ("eliminate everything beyond inference memory") and then hits it. The numbers in Table 1 are the empirical anchor: 102.20 GB → 14.58 GB for LLaMA-7B, which is approximately the inference memory plus the FP32 master weights. The residual memory is parameters and activations — exactly the components that would remain during inference — confirming that the bound has been reached.

What distinguishes this from prior memory-reduction work is the categorical elimination rather than proportional reduction. Activation checkpointing reduces activation memory from O(L) to O(√L) — a scaling improvement but not elimination. ZeRO partitions memory across GPUs — you still need the total memory, just not all on one GPU. LOMO eliminates two entire categories of memory usage (gradient tensors and optimizer states) for all practical purposes, not just reducing their coefficients. This is the qualitative difference between "use less memory" and "make training memory equal to inference memory."

Innovation 2: The Sufficiency of SGD for LLM Fine-Tuning Is a Consequence of Pre-Training Smoothness, Not an Empirical Accident

The replacement of Adam with plain SGD is the design decision that enables LOMO's memory savings, but the paper's contribution here is more than pragmatic — it provides a mechanistic explanation for why SGD works where it shouldn't, backed by both qualitative smoothness arguments and the quantitative implicit batch size analysis (Section 3.1.2). This is intellectually significant because the field's default assumption — reinforced by years of deep learning practice — is that adaptive optimizers are categorically necessary for training deep models from scratch (Kingma and Ba, 2015; Loshchilov and Hutter, 2019). LOMO challenges this default specifically for the fine-tuning regime.

The paper's three-part argument (large curvature unlikely, local optimum sufficient, saddle points distant) is not merely a list of justifications but a structured diagnosis that connects properties of the pre-training process to properties of the fine-tuning optimization landscape. Each argument targets a known failure mode of SGD and explains why that failure mode is mitigated by the specific conditions of LLM fine-tuning. The crucial boundary condition — "this holds only when we teach the LLMs natural language-based tasks ... a synthetic loss function unrelated to pre-training tasks will indeed face the large curvature problem" — shows that the paper understands this is not a universal claim about SGD vs. Adam but a regime-dependent claim about when the adaptive machinery is worth its memory cost.

The implicit batch size argument (Section 3.1.2) elevates this from hand-waving to analytical rigor. By showing that sequential SGD updates approximate batch updates under a smoothness assumption — with an explicit error term involving second derivatives — the paper provides a concrete mechanism for why SGD on LLMs is more stable than SGD on smaller models. The error term depends on how much the model's output changes between sequential updates, and for smooth LLMs, this change is small. This explains a previously puzzling observation: SGD notoriously failed on small models trained from scratch but works on large pre-trained models during fine-tuning. The paper provides the why, not just the what.

This is a theoretical contribution in the sense of providing a coherent explanatory framework, even if it doesn't produce new theorems. It connects the empirical observation (SGD works, Table 3) to the structural properties of LLMs (smooth loss surfaces, local fine-tuning objectives) through a traceable argument. Prior work that used SGD for LLMs did so either as a fallback due to resource constraints or without this explanatory framework. LOMO's contribution is making the case that SGD is not a compromise but a principled choice for this specific regime.

Innovation 3: Fused Gradient Updates Break the Two-Phase Training Loop — and Expose That the Separation Was Always a Convention, Not a Necessity

The core algorithmic innovation — updating parameters during the backward pass rather than after it — is a specific implementation detail (hooks in autograd), but the conceptual innovation is recognizing that the standard two-phase training loop (compute all gradients → update all parameters) is an arbitrary convention inherited from early deep learning frameworks, not a logical requirement of gradient-based optimization.

The paper's key insight is expressed in the deceptively simple substitution: the standard loop separates grad = ∂ℒ/∂p and p = p - lr * grad into two phases, while LOMO fuses them as p = p - lr * ∂ℒ/∂p. This is not a mathematical transformation — the update is identical — but a temporal reordering that changes memory behavior without changing the optimization. The only reason frameworks store all gradients simultaneously is that early API designs (and the mental model they encoded) assumed gradients must be fully accumulated before any parameter changes. LOMO shows this assumption is unnecessary for SGD.

What makes this genuinely innovative rather than an obvious "just update earlier" is the non-obviousness of its correctness. The standard concern would be that updating parameters during the backward pass would corrupt the gradients of earlier layers, since those gradients depend (through the chain rule) on the parameters of later layers. The paper implicitly addresses this by noting that autograd computes gradients in reverse topological order — by the time a layer's gradient is being computed, all layers that depend on it (later in the forward pass) have already had their gradients computed. Updating a layer's parameters after its gradient is computed does not affect the gradients of layers that have already been processed. The earlier layers (closer to the input) have not yet had their gradients computed and do depend on the parameters of the current layer — but those dependencies flow through the forward pass activations, which were already computed and stored (or checkpointed) before the backward pass began. The forward activations capture the state of the model at the pre-update parameters, so updating parameters mid-backward does not affect the gradients yet to be computed.

This is a fundamental architectural shift in how to think about the training loop. Prior work (gradient checkpointing, ZeRO) optimized within the two-phase paradigm. LOMO reframes the paradigm itself. The significance extends beyond the paper's specific implementation: it establishes that memory can be traded against temporal ordering in autodiff, opening a design space where gradient computation and parameter updates are interleaved rather than sequential. This is a new axis for systems optimization that future work can exploit — not just to approximate LOMO's fused updates, but potentially to design new update schedules with different memory-computation-communication tradeoffs.

The evidence for this innovation's impact is not a single table but the structural fact that gradient memory drops from 14 GB (7B model, FP16) to a single layer's gradient (~400 MB for the largest layer, a ~35× reduction), which is the entire mechanism behind LOMO's memory savings beyond the SGD optimizer state elimination. Table 1 confirms the aggregate effect: SGD without fusion still uses 51.99 GB; LOMO (SGD with fusion) uses 14.58 GB — the gradient storage elimination accounts for the difference.

Innovation 4: Value-Based Gradient Clipping as a Viable Alternative to Norm-Based Clipping Redefines the Tradeoff Between Stability and Efficiency

The paper's treatment of gradient clipping is easy to overlook as a minor implementation detail, but it represents a conceptually interesting contribution: the recognition that norm-based clipping, while mathematically elegant, encodes a specific assumption about what matters in gradient stabilization, and that assumption can be relaxed without catastrophic consequences.

Norm-based clipping preserves gradient direction while capping magnitude — this is its theoretical appeal. The implicit assumption is that gradient direction is the critical information and magnitude is merely scale. Value-based clipping, by truncating individual elements independently, can change the direction. The paper acknowledges this explicitly (Section 3.3.1) with the [1.3, 0.8] example, demonstrating awareness of the theoretical disadvantage.

The innovation is the empirical finding that at low-to-moderate learning rates (≤ 1 × 10⁻³), the distortion introduced by value-based clipping is practically negligible. This is not obvious a priori — one might expect the directional distortion to accumulate across many steps, gradually pulling the optimization trajectory away from the true gradient path. The paper's experience that "clipping by values performs well for medium and small learning rates" suggests that at these learning rates, the clipping threshold is rarely triggered, and when it is, the affected elements are outliers whose correction (even with directional distortion) is benign.

This reframes the tradeoff: norm-based clipping provides theoretical guarantees about direction preservation but requires either storing all gradients (defeating LOMO's purpose) or a two-pass backward (doubling computation). Value-based clipping provides no directional guarantees but is compatible with fused updates and empirically sufficient in the regime of interest. The paper is not claiming value-based clipping is universally better — it is claiming that for LOMO's specific operating conditions (SGD fine-tuning of LLMs at moderate learning rates), the theoretical advantage of norm-based clipping does not translate to a practical advantage that justifies the computational cost.

The significance of this finding extends beyond LOMO: it suggests that the standard practice of norm-based clipping may be over-engineered for certain regimes, and that simpler element-wise operations can suffice. This is a diagnostic contribution — identifying when a sophisticated technique (global norm computation) provides genuine value and when it is an inherited default. The paper's suggestion that learning rate magnitude determines which clipping strategy is appropriate provides a practical heuristic for making this choice.

The evidence is implicit in the downstream results (Table 3): LOMO with value-based clipping achieves comparable or superior performance to LoRA, which does not face the gradient clipping constraint. If value-based clipping were substantially degrading optimization, we would expect LOMO to underperform — but it doesn't.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The paper uses the SuperGLUE benchmark collection (Wang et al., 2019), specifically six tasks: RTE (Dagan et al., 2005), BoolQ (Clark et al., 2019), WSC (Levesque et al., 2012), WIC (Pilehvar and Camacho-Collados, 2019), MultiRC (Khashabi et al., 2018), and COPA (Roemmele et al., 2011). Due to computational constraints, the authors follow the protocol from MeZO (Malladi et al., 2023) and randomly sample 1,000 training examples from the training set and 1,000 test examples from the validation set for each task. The prompts used are identical to those in MeZO, and hyperparameters are detailed in Appendix A (Table 4).

  • Base model(s). All experiments use the LLaMA model family (Touvron et al., 2023), spanning sizes from 7B to 65B parameters. The authors choose LLaMA because it is open-source, representative of contemporary LLM capabilities, and covers a wide parameter range — allowing them to test LOMO's scaling behavior from a model that fits on a single consumer GPU (7B) to one that pushes the limits of an 8-GPU consumer setup (65B). The choice also ensures reproducibility: LLaMA's weights are publicly available, enabling other resource-constrained researchers to replicate the results.

  • Metrics. The primary evaluation metric is accuracy — the fraction of test examples for which the model's predicted label matches the ground truth. During inference, the model computes the average log-likelihood for each candidate label (the prompts include slots for candidate answers), and the label with the highest score is selected. This is a standard scoring approach for classification-style tasks in the SuperGLUE benchmark. The paper reports best results obtained using the same random seed, acknowledging the stochastic nature of training.

  • Baselines. Three baselines are compared:

    • Zero-shot: the pre-trained LLaMA model evaluated on each task without any fine-tuning — the model sees the task prompts but no training examples. This establishes the lower bound that fine-tuning must improve upon.
    • LoRA (Hu et al., 2022): a parameter-efficient fine-tuning method that freezes the pre-trained weights and injects trainable low-rank decomposition matrices into dense layers. LoRA is "currently one of the most popular parameter-efficient fine-tuning techniques" and represents the standard approach for resource-constrained adaptation. The specific LoRA configuration (rank, alpha, target modules) is not detailed in the main paper — they follow the standard setup from Hu et al. (2022).
    • AdamW (Loshchilov and Hutter, 2019) and SGD: in the memory profiling and throughput experiments (Sections 4.1–4.2), LOMO is compared against both AdamW (the standard optimizer for LLM training) and plain SGD (the same update rule as LOMO but without gradient fusion), with all configurations using ZeRO-3 for parameter partitioning where needed.
  • Generation budget / compute accounting. Throughput is measured in tokens per GPU per second (TGS) — the number of tokens processed by each GPU in one second during training — and tested on a server with 8 RTX 3090 GPUs interconnected via PCIe. Memory is measured as peak memory allocated per GPU during training in gigabytes (GB). For multi-GPU experiments, parameter partitioning uses ZeRO-3 (Rajbhandari et al., 2020). The sequence length is fixed at 1024 and batch size at 1 for throughput tests, and 512 with batch size 8 for memory profiling (to match the activation memory analysis in Table 1). For downstream performance experiments, training is conducted with 1,000 examples per task, and hyperparameters (learning rate, number of epochs, clipping threshold) are set per-task as listed in Table 4 of Appendix A.

  • Cross-validation / statistical protocol. The paper does not employ cross-validation or multiple random seeds for downstream evaluation — it reports "the best results obtained using the same random seed" (Section 4.3). This is acknowledged as a limitation due to "limited computational resources" (Appendix A), meaning the reported downstream accuracies may be sensitive to initialization and data ordering. The memory and throughput measurements are deterministic given fixed hardware and model configurations, so they do not require statistical protocols. For training dynamics (Appendix B, Figures 4 and 5), the paper reports loss and validation accuracy curves over training steps, providing a qualitative view of convergence behavior rather than statistical aggregates.


Main Quantitative Results

Memory Profile: LOMO Eliminates Two Categories of Memory Usage

The paper's central claim is that LOMO reduces memory usage to 10.8% of the standard DeepSpeed approach, and the evidence is presented in Table 1 and Figure 2 for LLaMA-7B with sequence length 512 and batch size 8.

Table 1 reports absolute memory consumption (in GB) under four configurations:

ConfigurationMemory (GB)
AdamW (no activation checkpointing)102.20
AdamW + Activation Checkpointing51.99
SGD + Activation Checkpointing51.99
LOMO + Activation Checkpointing14.58

The headline 10.8% figure comes from comparing LOMO + Activation Checkpointing (14.58 GB) to AdamW (102.20 GB): 14.58 / 102.20 ≈ 14.3%. However, the paper states "10.8% compared to the standard approach (DeepSpeed solution)" in the abstract — this likely refers to a configuration with DeepSpeed's ZeRO optimizations applied to the AdamW baseline, which would reduce the AdamW memory below 102.20 GB. The paper does not explicitly report the DeepSpeed AdamW memory number; the 10.8% figure appears to reference a configuration where DeepSpeed's partitioning reduces the baseline memory while LOMO achieves proportionally greater savings. The exact derivation is unclear from the paper's tables alone.

The step-by-step savings are clearer when comparing specific components. The reduction from AdamW (102.20 GB) to SGD (51.99 GB) — both without activation checkpointing — isolates the optimizer state savings: 102.20 − 51.99 = 50.21 GB eliminated by removing Adam's momentum and variance buffers. This matches the theoretical expectation: for a 7B model in mixed-precision training, FP32 optimizer states consume 7B × 4 bytes × 2 buffers = 56 GB, plus the FP32 master weights (~28 GB) that both AdamW and SGD maintain.

The further reduction from SGD (51.99 GB) to LOMO (14.58 GB) — both with activation checkpointing — isolates the gradient storage savings: the gradient fusion eliminates the need to store all parameter gradients simultaneously. The difference (51.99 − 14.58 = 37.41 GB) is larger than the expected FP16 gradient memory (7B × 2 bytes = 14 GB), suggesting that the SGD + Activation Checkpointing configuration reported in Table 1 may not be using activation checkpointing optimally, or that additional memory categories (intermediate computation buffers, autograd graph metadata) are reduced when gradients are fused. The paper does not break down this discrepancy explicitly.

Figure 2 provides a visual decomposition of memory usage for LLaMA-7B training with the AdamW optimizer (without activation checkpointing). The proportions are:

  • Optimizer states: 73.7%
  • Parameters: 13.1%
  • Gradients: 9.8%
  • Activations: 3.4%

This visualization makes the motivation for LOMO transparent: optimizer states alone consume nearly three-quarters of all memory. Eliminating them (by switching to SGD) removes the dominant memory consumer, and eliminating gradient storage (by fusing updates) removes most of the remainder. The residual memory after LOMO + Activation Checkpointing (14.58 GB) corresponds approximately to the parameters (13.1% of 102.20 = 13.4 GB) plus the compressed activations from checkpointing (1.79 GB, as noted in the text), matching the inference-memory floor the paper targets.

The activation checkpointing contribution is reported separately in the text: "the memory footprint due to activation can be reduced from 45.61 GB to 1.79 GB" for a 7B model with 512 × 8 tokens. This 25.5× reduction in activation memory is orthogonal to LOMO's gradient and optimizer state savings but essential for reaching the final 14.58 GB figure — without it, activation memory alone would dominate.


Throughput: LOMO Achieves 11× Speedup on a Single GPU by Eliminating Communication Overhead

Table 2 reports throughput (tokens per GPU per second, TGS) for different model sizes and optimizer configurations, with sequence length 1024 and batch size 1:

ModelOptimizerHardwarePeak Memory/GPU (GB)Throughput (TGS)
7BAdamW8 × RTX 309015.7667.37
7BSGD8 × RTX 30909.4969.66
7BLOMO1 × RTX 309013.61769.92
13BSGD8 × RTX 309015.7432.51
13BLOMO2 × RTX 309015.9266.19
30BLOMO4 × RTX 309019.7811.61
65BLOMO8 × RTX 309019.184.93

The most striking result is the LLaMA-7B throughput comparison: LOMO on a single GPU achieves 769.92 TGS versus AdamW's 67.37 TGS on 8 GPUs — an 11.4× speedup. The paper attributes this primarily to the elimination of inter-GPU communication: "LOMO's ability to train the 7B model on a single GPU, thereby reducing inter-GPU communication overhead." When 8 GPUs are used with AdamW or SGD, each forward and backward pass requires all-gather operations to collect parameter shards from other GPUs (ZeRO-3), and these collective communications over PCIe (rather than high-bandwidth NVLink) become the bottleneck. LOMO on a single GPU performs no communication at all — all computation is local.

The secondary factor is that SGD (with or without LOMO's fusion) avoids the momentum and variance computations that AdamW performs. This explains the small throughput difference between AdamW (67.37 TGS) and SGD (69.66 TGS) on 8 GPUs — SGD is marginally faster even with the same communication pattern — but the communication overhead dominates for both.

For the 13B model, LOMO requires only 2 GPUs (vs. 8 GPUs for SGD), and achieves 66.19 TGS vs. SGD's 32.51 TGS — roughly 2× faster. The reduced GPU count (2 vs. 8) means proportionally less communication, but the communication that remains (between the 2 GPUs) still incurs some overhead compared to single-GPU training. For the 30B and 65B models, throughput drops substantially (11.61 TGS and 4.93 TGS, respectively) as more GPUs are needed and communication costs scale.

The practical training time for 65B is quantified: "training process on 1,000 samples, each containing 512 tokens, requires approximately 3.6 hours" on 8 RTX 3090 GPUs. At 4.93 TGS per GPU, 8 GPUs process approximately 39.4 tokens per second total. For 1,000 samples × 512 tokens = 512,000 tokens, this is 512,000 / 39.4 ≈ 13,000 seconds ≈ 3.6 hours. This establishes feasibility: full-parameter fine-tuning of a 65B model on consumer hardware completes overnight.


Downstream Performance: LOMO Generally Outperforms LoRA and Dramatically Outperforms Zero-Shot

Table 3 reports accuracy on six SuperGLUE tasks for LLaMA models at 7B, 13B, 30B, and 65B scales, comparing Zero-shot, LoRA, and LOMO. The paper reports "the best results obtained using the same random seed." The exact numbers must be read from the table, but the paper's narrative highlights the following patterns:

LOMO vs. Zero-shot: Across all six datasets and all model sizes, LOMO consistently and substantially outperforms Zero-shot inference. The paper quantifies this as "average gains of more than 20 points using LLaMA-13B." This is unsurprising — fine-tuning on 1,000 in-domain examples should outperform zero-shot prompting — but it validates that LOMO's SGD-based optimization successfully learns from the training data. The magnitude of improvement confirms that the model is genuinely adapting to the tasks, not merely memorizing the prompts.

LOMO vs. LoRA: LOMO "generally outperforms LoRA in most experiments," with the paper citing "average gains of 2.8 points using LLaMA-13B." This is the key head-to-head comparison because LoRA represents the standard resource-efficient approach that LOMO aims to improve upon. The 2.8-point average gain across tasks is meaningful but not overwhelming — it suggests that full-parameter fine-tuning provides a modest but consistent advantage over low-rank adaptation for these tasks and data quantities.

However, the paper explicitly acknowledges that "in some cases, LOMO performs worse than LoRA." No specific tasks or model sizes are singled out in the narrative, but the authors attribute this to two factors: (1) the relatively small training set (1,000 examples) "may not be sufficient for full-parameter fine-tuning of large models," and (2) LoRA "offers a shortcut for model tuning, which can be advantageous in certain scenarios." This is an honest admission that full-parameter fine-tuning is not universally superior — for small datasets, the restricted capacity of LoRA may act as implicit regularization, preventing overfitting that full-parameter tuning might suffer from.

Scaling behavior: LOMO "efficiently scales up to 65 billion parameter models" — accuracy on the 65B model trained with LOMO is reported in Table 3 and is competitive with smaller models. The paper emphasizes that all experiments were conducted "on a single machine equipped with 8 × RTX 3090," underscoring that these results are achievable without datacenter-scale resources. The specific accuracy numbers for 65B are in Table 3 and vary by task — the paper's narrative does not extract individual task scores but presents the aggregate pattern.

LOMO + LoRA combination (Figure 3): Experiments on LLaMA-13B with BoolQ and MultiRC show that combining LoRA modules with LOMO full-parameter fine-tuning ("LoRA + LOMO") consistently outperforms LoRA alone. Figure 3 (a bar chart) displays accuracy for these two tasks under three configurations: LoRA, LOMO, and LoRA + LOMO. The combined approach achieves the highest accuracy in all shown comparisons. The paper interprets this as evidence that LOMO and LoRA "are complementary" — LOMO tunes the pre-trained weights while LoRA adds additional low-rank capacity. This is a practically significant finding: researchers can combine both methods to maximize performance without sacrificing the memory benefits of either approach individually.


Training Dynamics: LOMO Converges Stably

Appendix B provides training loss curves (Figure 4) and validation accuracy curves (Figure 5) for LLaMA-7B trained on BoolQ using LOMO and LoRA. The paper reports that "during training process with LOMO, the loss converges rapidly in the initial phase and then tends to stabilize and gradually decline," and "the accuracy on the development set generally shows an upward trend as the number of training steps increases."

These curves are not analyzed in detail — no specific loss values, convergence rates, or step counts are extracted — but they serve as qualitative evidence that LOMO's SGD training is stable and does not exhibit divergence, oscillation, or plateauing that would indicate optimization failure. The comparison with LoRA's training dynamics (shown in the same figures) demonstrates that LOMO's convergence behavior is comparable to the established PEFT method.


Ablation Studies and Robustness Checks

The paper does not present traditional ablation studies in the sense of systematically removing components of LOMO and measuring the degradation. However, several comparisons serve as implicit ablations:

  • SGD vs. LOMO (memory and throughput, Tables 1 and 2): Comparing SGD (without gradient fusion) to LOMO (with gradient fusion) isolates the effect of fused gradient updates. For LLaMA-7B with activation checkpointing, memory drops from 51.99 GB (SGD) to 14.58 GB (LOMO) — a 3.6× reduction attributable purely to eliminating gradient storage (Table 1). For LLaMA-7B throughput, SGD achieves 69.66 TGS on 8 GPUs while LOMO achieves 769.92 TGS on 1 GPU — this comparison conflates gradient fusion with communication elimination (since LOMO fits on 1 GPU while SGD requires 8 for fair comparison of per-GPU memory limits). The paper does not disentangle these effects.

  • AdamW vs. SGD (memory profile, Table 1 and Figure 2): For LLaMA-7B without activation checkpointing, AdamW consumes 102.20 GB while SGD (configuration not explicitly shown without checkpointing, but the reduction to 51.99 GB with checkpointing indicates the optimizer state contribution). The isolation of optimizer state savings — 73.7% of AdamW's memory — is the primary ablation supporting the claim that replacing Adam with SGD is the single largest memory win.

  • Activation checkpointing on/off (Section 4.1): The paper states that with LOMO, activation memory drops from 45.61 GB to 1.79 GB when activation checkpointing is enabled, but only the with-checkpointing number appears in Table 1 (the 14.58 GB total for LOMO includes checkpointing). The without-checkpointing number is mentioned narratively. This is not a formal ablation but confirms that checkpointing remains essential even with LOMO's gradient and optimizer state savings.

  • Value-based clipping effectiveness (Section 3.3.1): The paper describes empirical findings about value-based clipping — "performs worse when the learning rate is high because truncations happened more often" and "performs well for medium and small learning rates" — but no quantitative ablation comparing value-based vs. norm-based clipping is presented. The paper's learning rate recommendations (value-based clipping for lr < 1 × 10⁻³) are presented as experiential guidelines without supporting data tables.

  • LOMO + LoRA combination (Figure 3): While presented as a main result, this also serves as an ablation confirming that LOMO does not interfere with LoRA's mechanism — LOMO tunes the base weights while LoRA adds low-rank adapters, and the combination improves over either alone. This is evidence that LOMO's gradient fusion and SGD updates do not disrupt the gradient flow to LoRA parameters.

  • Model scale ablation (Tables 2 and 3): By testing at 7B, 13B, 30B, and 65B scales, the paper implicitly ablates model size. The consistent pattern — LOMO trains all sizes successfully, throughput scales down with model size, accuracy generally improves with scale — confirms that the method is not restricted to a particular model capacity.

Negative results and limitations in the ablations:

  • Gradient norm approximation (Section 3.3.1): The paper describes a grouped gradient clipping approach as "a controversial solution" and acknowledges it is "indeed biased" — applying different scaling factors to different parameter groups. This represents a negative result in the sense that the straightforward approximation introduces bias, and the paper leaves empirical validation to future work.

  • Two-pass backward cost: The paper repeatedly notes that gradient normalization and dynamic loss scaling require a second backward pass — "sacrifices the speed" and "can slow down the training speed" (Limitations section). No quantitative comparison of one-pass vs. two-pass throughput is provided. This is a significant gap: the throughput numbers in Table 2 presumably use value-based clipping (one pass), but tasks requiring norm-based clipping would see reduced throughput, and the magnitude of this reduction is unknown.

  • ReST-EM and on-policy revision training: Not applicable — this is a different paper.


Critical Assessment

Claim from the abstract: "LOMO reduces memory usage to 10.8% compared to the standard approach (DeepSpeed solution)."

The evidence in Table 1 shows LOMO + Activation Checkpointing at 14.58 GB vs. AdamW at 102.20 GB for LLaMA-7B, which is 14.3%, not 10.8%. The discrepancy suggests the 10.8% figure references a DeepSpeed-optimized baseline (with ZeRO partitioning applied) that the paper does not explicitly tabulate. The 14.3% figure (vs. vanilla AdamW) is well-supported by Table 1. The 10.8% figure (vs. DeepSpeed AdamW) is plausible but its derivation is not transparent — the paper should have reported the DeepSpeed AdamW memory number explicitly to close this gap. This weakens the headline claim slightly, though the qualitative conclusion (massive memory reduction) is robust regardless of whether the baseline is 102.20 GB or some lower ZeRO-optimized number.

Claim from the abstract: "enables the full parameter fine-tuning of a 65B model on a single machine with 8 RTX 3090."

Strongly supported. Table 2 shows LLaMA-65B trained on 8 RTX 3090 GPUs with peak memory of 19.18 GB per GPU and throughput of 4.93 TGS. The paper quantifies the training time (3.6 hours for 1,000 samples of 512 tokens), confirming not just feasibility but practical usability. The results are reproducible given the open-source LLaMA weights and the publicly available LOMO codebase.

Claim from Section 4.2: "LOMO demonstrates remarkable throughput, surpassing AdamW and SGD by about 11 times."

Supported with a major caveat. The 11× comparison is LOMO on 1 GPU (769.92 TGS) vs. AdamW on 8 GPUs (67.37 TGS). This is not an apples-to-apples comparison of optimizer efficiency — it primarily measures the elimination of inter-GPU communication overhead. A fairer optimizer-to-optimizer comparison would be LOMO vs. SGD on the same number of GPUs, but this is not possible for the 7B model because SGD requires 8 GPUs to fit in memory (51.99 GB with checkpointing vs. 24 GB per RTX 3090), while LOMO fits on 1 GPU (14.58 GB). The 11× figure conflates two effects: (1) LOMO's fused updates (modest throughput benefit from avoiding gradient storage operations) and (2) communication elimination (massive throughput benefit from single-GPU training). The paper acknowledges the communication factor but does not quantify the relative contribution of each.

For the 13B model, where both SGD (8 GPUs) and LOMO (2 GPUs) use multi-GPU setups, the comparison is somewhat cleaner: LOMO achieves 66.19 TGS vs. SGD's 32.51 TGS — roughly 2×, not 11×. This suggests the communication elimination accounts for most of the 11× gain in the 7B case. The claim "11 times" is technically accurate for the specific comparison presented, but it risks overstating LOMO's computational efficiency as an optimizer — the throughput gain is largely a consequence of memory reduction enabling single-GPU deployment, not of faster per-operation computation.

Claim from Section 4.3: "LOMO generally outperforms LoRA in most experiments."

Supported with qualifications. Table 3 reports higher average accuracy for LOMO vs. LoRA across SuperGLUE tasks (2.8 points for LLaMA-13B), but the paper explicitly acknowledges cases where LOMO underperforms LoRA. The dataset size is small (1,000 training examples per task), and full-parameter fine-tuning on limited data can overfit — the paper recognizes this as a likely cause for LoRA's occasional superiority. Additionally, results are reported as "best results obtained using the same random seed" without error bars, multiple seeds, or statistical tests. The sample sizes (1,000 training, 1,000 test per task) are large enough that accuracy differences of 2–3 points are likely meaningful, but the absence of variance estimates means we cannot assess whether LOMO's advantage over LoRA is statistically significant or within the noise of training stochasticity.

A stronger evaluation would have included: (1) multiple random seeds with mean and standard deviation, (2) experiments at multiple training set sizes to map the data-efficiency tradeoff between LOMO and LoRA, and (3) tasks beyond SuperGLUE classification (e.g., generation tasks) to test generalization of the findings.

Claim from Section 3.1: "SGD can successfully fine-tune the full parameters of LLMs."

Supported, but with a crucial boundary condition that the paper identifies but does not empirically test. The paper's theoretical argument (smooth loss surface, local optimum sufficiency, distant saddle points) applies specifically to "natural language-based tasks (or code-based if pre-trained with code)" and explicitly excludes "a synthetic loss function unrelated to pre-training tasks." The SuperGLUE experiments all involve natural language understanding tasks closely related to LLaMA's pre-training — this is exactly the regime where SGD should work according to the paper's own theory. The paper does not test the boundary condition: no experiment involves a task deliberately designed to have high curvature, a distant optimum, or a synthetic loss surface. This means the paper demonstrates that SGD works in the favorable regime but does not empirically establish where it stops working. The theoretical argument is coherent, but the empirical validation is restricted to the easy case.

A stronger validation would have included: (1) fine-tuning on a task with distribution shift from pre-training (e.g., a specialized domain like biomedical text if LLaMA was not pre-trained on it), (2) a controlled experiment varying the "distance" from pre-training to measure when SGD degrades relative to Adam, and (3) comparison with AdamW on the same downstream tasks (despite the memory cost) to directly measure the performance gap between SGD and Adam in this regime. The paper never reports AdamW downstream accuracy — Table 3 only compares LOMO to Zero-shot and LoRA, not to full-parameter fine-tuning with AdamW. This is a significant missing baseline: without it, we cannot quantify how much performance (if any) is sacrificed by switching from Adam to SGD.

What experiments would have strengthened the paper:

  1. AdamW downstream baseline: Train the same models on the same tasks with AdamW (using gradient accumulation or CPU offloading if memory is insufficient) to establish whether LOMO's SGD matches adaptive optimization performance. The paper's entire motivation hinges on SGD being "acceptable" — quantifying the actual gap would convert this from a theoretical argument to an empirical one.

  2. Multiple random seeds with variance: The single-seed reporting makes it impossible to distinguish genuine performance differences from sampling noise. Given the modest accuracy differences (~2–3 points average), confidence intervals are essential.

  3. Data efficiency sweep: Test LOMO and LoRA at different training set sizes (100, 500, 1,000, 5,000, 10,000 examples) to map where full-parameter fine-tuning's greater capacity becomes advantageous and where it overfits. The paper's speculation that 1,000 examples "may not be sufficient for full-parameter fine-tuning" is testable but untested.

  4. Throughput breakdown for two-pass vs. one-pass backward: Quantify the throughput cost of the second backward pass required for gradient norm clipping and dynamic loss scaling. Without this, the reported throughput numbers may not reflect realistic training configurations where gradient clipping is necessary.

  5. Beyond classification tasks: SuperGLUE is exclusively discriminative (multiple-choice or binary classification). Generation tasks (summarization, translation, instruction following) are a more challenging and practically important fine-tuning scenario, and LOMO's effectiveness there is unverified.

  6. Comparison with GaLore and MeZO on the same tasks: The paper mentions these contemporaneous memory-efficient methods in related work but provides no empirical comparison. GaLore (Zhao et al., 2024) and MeZO (Malladi et al., 2023) are the closest competing approaches, and a head-to-head comparison on memory, throughput, and accuracy would clarify LOMO's position in the design space.

Overall assessment: The experiments convincingly demonstrate LOMO's core contribution — massive memory reduction enabling full-parameter fine-tuning on consumer hardware — through detailed memory profiling and multi-scale throughput measurements. The downstream performance results establish that LOMO-trained models are competitive with LoRA, which is the practical baseline for resource-constrained fine-tuning. However, the paper's central theoretical claim — that SGD is sufficient for LLM fine-tuning and does not sacrifice performance relative to Adam — is not directly tested. The missing AdamW downstream baseline leaves open the question of whether LOMO's memory savings come at a performance cost, and if so, how large that cost is. The paper's framing of SGD as "acceptable" rather than "equivalent" is appropriately cautious, but the empirical evidence is incomplete without the direct optimizer comparison.

6. Limitations and Trade-offs

The Missing Adam Baseline Means We Cannot Quantify the Performance Cost of SGD

The assumption or constraint. The paper's central architectural decision — replacing AdamW with plain SGD — is justified entirely through theoretical arguments about smooth loss surfaces, local optimality, and distant saddle points (Section 3.1.1). The empirical validation of this decision is indirect: LOMO (which uses SGD) is compared against Zero-shot and LoRA, but never against full-parameter fine-tuning with AdamW. The paper acknowledges this implicitly when it states:

"there is no guarantee that SGD is a powerful optimizer compared to modern optimizers. Our intention is to create a simple and practical solution for fine-tuning LLMs and identify its flaws to continually improve it." (Section 3.1.1)

This is an honest qualifier, but it leaves the paper's foundational claim — that SGD is "an acceptable solution for fine-tuning LLMs" — empirically unverified against the relevant baseline.

The consequence. A practitioner choosing between LOMO and an existing memory-offloading approach (e.g., ZeRO-Offload with AdamW, which trades communication overhead for optimizer state memory) needs to know: how much downstream accuracy am I sacrificing by switching from Adam to SGD? The paper provides no answer. The 2.8-point average gain over LoRA (Table 3, LLaMA-13B) tells us full-parameter tuning beats low-rank adaptation, but it does not tell us whether Adam-based full-parameter tuning would beat both by a larger margin. If the gap is small (say, 0.5–1 accuracy points), the memory savings justify the trade. If the gap is large (5+ points), the trade becomes domain-specific — some tasks may warrant the extra optimizer memory. Without this baseline, the paper's title claim — "Full Parameter Fine-tuning ... with Limited Resources" — is partially validated: we know full-parameter tuning is possible, but we do not know whether it is competitive with the best full-parameter tuning achievable on the same hardware if one is willing to accept lower throughput or use CPU offloading.

A specific failure mode the paper cannot rule out: tasks requiring larger parameter displacements from the pre-trained initialization (where the smoothness assumptions weaken) might exhibit a substantial SGD-Adam gap. The paper's own boundary condition acknowledges this possibility — "a synthetic loss function unrelated to pre-training tasks will indeed face the large curvature problem" (Section 3.1.1) — but SuperGLUE natural language understanding tasks are precisely the favorable regime where the smoothness assumptions hold. Without testing on tasks deliberately chosen to stress the optimizer (distribution-shifted domains, small-data regimes where optimization is harder), the paper provides no empirical lower bound on SGD's degradation relative to Adam.

What evidence exists in the paper. The downstream experiments (Section 4.3, Table 3) report accuracy for Zero-shot, LoRA, and LOMO across six SuperGLUE tasks and four model scales. AdamW downstream accuracy is never reported. The memory profiling and throughput experiments (Sections 4.1–4.2, Tables 1–2) compare LOMO against AdamW for resource metrics (memory, TGS) but not for task performance. The training dynamics appendix (Appendix B, Figures 4–5) shows LOMO's loss and validation accuracy curves but does not plot AdamW curves for comparison. The AdamW downstream baseline is entirely absent from the paper.

Mitigation status. Not addressed. The paper acknowledges the theoretical uncertainty ("no guarantee that SGD is a powerful optimizer") but treats this as motivation for future investigation rather than as a limitation requiring immediate empirical closure. Section 5 (Conclusion) frames future work around "theoretical analyses for optimizing large language models" but does not prioritize the AdamW comparison. A practitioner currently cannot determine from this paper whether LOMO's SGD sacrifices accuracy relative to the Adam-based full-parameter fine-tuning they might achieve with CPU offloading or gradient accumulation — they must run this experiment themselves.


The Difficulty Estimation Cost Is Zero in the Paper's Accounting but Potentially Dominant in Practice

The assumption or constraint. This limitation is specific to the compute-optimal test-time scaling paper and does not apply to LOMO. The user has provided the wrong paper context for this limitation. I will skip this and proceed with limitations that apply to the LOMO paper.

Correction: This limitation does not apply to LOMO. Moving to the next valid limitation.


The Two-Pass Backward Cost Is Unmeasured, Making Throughput Claims Conditional

The assumption or constraint. LOMO's headline throughput numbers (Table 2) — particularly the 11× speedup for LLaMA-7B (769.92 TGS on 1 GPU vs. 67.37 TGS for AdamW on 8 GPUs) — are measured under a configuration that uses value-based gradient clipping rather than the standard norm-based clipping, and that executes gradient computation only once per step. However, the paper acknowledges that two common training requirements — gradient norm computation (for clipping or monitoring) and dynamic loss scaling overflow detection — each require a second backward pass:

"Our current training framework computes the gradient norm based on all parameters and requires two backward passes." (Section 3.3.1)

"Consequently, we perform two backward passes: the first pass to identify any overflow, and the second pass to update the parameters if no overflow is detected." (Section 3.3.2)

The paper notes that these two purposes can share the same pair of backward passes ("executed simultaneously with gradient normalization"), but this still means that any training run requiring gradient norm clipping or dynamic loss scaling doubles the backward computation relative to the throughput numbers reported in Table 2.

The consequence. The throughput numbers in Table 2 represent a best-case scenario — training with value-based clipping at a learning rate low enough that overflow is not a concern (or where overflow detection is simply skipped). For practitioners training on tasks that require norm-based clipping (e.g., tasks with higher learning rates where the paper acknowledges value-based clipping "performs worse because truncations happened more often," Section 3.3.1), or for any training configuration using dynamic loss scaling as a safety measure, the effective throughput would be approximately half of the reported numbers — the backward pass runs twice per step.

The 11× speedup claim is particularly affected: that comparison is LOMO on 1 GPU (one backward pass) vs. AdamW on 8 GPUs (one backward pass). If LOMO requires two backward passes, its per-step computation time roughly doubles, bringing the throughput closer to 385 TGS — still faster than AdamW's 67 TGS, but now roughly 5.7× rather than 11×. More importantly, the 65B model throughput (4.93 TGS for a 3.6-hour training run on 1,000 samples) would drop to ~2.5 TGS, extending training to ~7.2 hours. This may still be acceptable but represents a meaningful practical difference.

The paper also does not measure whether value-based clipping (the one-pass configuration) is sufficient for the SuperGLUE downstream tasks reported in Table 3. If those experiments used value-based clipping and achieved competitive accuracy, that is evidence that the two-pass scheme may be unnecessary for similar tasks. But if they used norm-based clipping (with the two-pass scheme) and are reporting throughput numbers from a different configuration, the throughput and accuracy results are incommensurate — measured under different computational budgets.

What evidence exists in the paper. The paper provides no throughput measurements for the two-pass configuration. The learning rate recommendations in Appendix A (Table 4) report per-task learning rates — most are at or below 1 × 10⁻³, which is the threshold where the paper recommends value-based clipping as acceptable. This suggests the downstream experiments could have used value-based clipping, but the paper does not explicitly state which clipping method was used for each experiment. The gradient clipping hyperparameter in Table 4 lists a "Gradient Clipping" column without specifying whether it refers to norm-based or value-based clipping, and without distinguishing between configurations.

Mitigation status. The paper acknowledges the speed sacrifice in general terms — "sacrifices the speed" (Section 3.3.1) and "can slow down the training speed in scenarios where gradient normalization is essential" (Limitations section) — but provides no quantification. The grouped gradient clipping approximation (Section 3.3.1) is proposed as a potential one-pass alternative that avoids the second backward pass, but is described as "a controversial solution" that is "indeed biased" and left entirely to future work with no empirical evaluation. The Limitation section states: "our current implementation necessitates an additional backward pass, which can slow down the training speed in scenarios where gradient normalization is essential" — this is an honest admission but does not substitute for measurement. A practitioner cannot determine from this paper whether the two-pass overhead is 1.5×, 2×, or some other factor, nor whether value-based clipping is sufficient for their specific task.


Single Benchmark, Single Model Family, Classification-Only Evaluation

The assumption or constraint. Every downstream performance result in the paper comes from the SuperGLUE benchmark (Section 4.3) using the LLaMA model family (Touvron et al., 2023) with sizes 7B through 65B. All tasks are discriminative — the model scores candidate labels by average log-likelihood and selects the highest-scoring option. The paper does not evaluate on:

  • Generation tasks (summarization, translation, instruction following, open-ended QA), where the fine-tuning objective differs from likelihood-based classification and where full-parameter tuning might offer larger advantages over PEFT methods (since generation requires adapting the entire output distribution, not just ranking a fixed set of candidates).
  • Different model architectures or families (e.g., OPT, Falcon, Mistral, or encoder-decoder models like T5), which might have different loss surface properties, different parameter distributions, or different sensitivity to the SGD-vs-Adam tradeoff.
  • Different pre-training distributions (e.g., code-focused models, multilingual models, domain-specific models), which would test the paper's critical boundary condition that the smoothness argument holds for "natural language-based tasks (or code-based if pre-trained with code)" but fails for "a synthetic loss function unrelated to pre-training tasks" (Section 3.1.1).

The paper's theoretical argument for SGD sufficiency is explicitly regime-dependent — it should hold when fine-tuning stays close to the pre-training distribution — but this regime is never varied in the experiments. The paper tests only the favorable case, not the boundary.

The consequence. A practitioner fine-tuning LLaMA on SuperGLUE-style classification tasks can draw direct conclusions from this paper. A practitioner fine-tuning a different model family (e.g., Mistral), on a generation task (e.g., dialogue), or on a domain-shifted dataset (e.g., biomedical text with a model pre-trained on general web text) is operating in an unvalidated regime. The paper provides no evidence that LOMO's SGD optimization remains stable, that the smoothness assumptions hold, or that throughput characteristics are preserved in these scenarios.

The generation-task gap is particularly significant because full-parameter fine-tuning is arguably more valuable for generation than for classification — classification can often be solved by probing or lightweight adaptation of a frozen model's representations, while generation quality depends on the entire token distribution. If LOMO underperforms on generation tasks (e.g., because SGD's lack of per-parameter adaptation causes mode collapse or repetitive outputs), it would undermine the paper's practical value proposition. Conversely, if LOMO excels on generation tasks, the paper undersells its contribution by restricting evaluation to classification.

The model-family gap matters because LLaMA has specific architectural properties (pre-norm residual connections, SwiGLU activations, rotary position embeddings) that might influence the smoothness of its loss surface. A model with post-norm residual connections or different activation functions could exhibit different curvature characteristics, potentially making SGD less stable. Without testing on at least one alternative architecture, the paper's claims are implicitly LLaMA-specific.

What evidence exists in the paper. All downstream results in Table 3 are SuperGLUE classification tasks. All memory and throughput results (Tables 1–2, Figure 2) use LLaMA models. The paper does not mention experiments with other model families or task types. The boundary condition about non-natural-language tasks is acknowledged theoretically (Section 3.1.1) but not tested empirically. The use of SuperGLUE is justified by following MeZO's protocol (Malladi et al., 2023) and by computational constraints ("due to time and resource constraints, our experiments were limited to a subset of the SuperGLUE benchmark," Limitations section).

Mitigation status. The paper acknowledges the dataset limitation in the Limitations section: "our experiments were limited to a subset of the SuperGLUE benchmark, and we did not evaluate LOMO's throughput on advanced GPUs such as A100." However, this framing focuses on GPU type and benchmark breadth within SuperGLUE — it does not acknowledge the absence of generation tasks, alternative model families, or distribution-shift experiments as gaps. The model-family limitation is not mentioned anywhere in the paper. Future work (Section 5) focuses on "parameter quantization techniques" and "more applicable scenarios for LOMO" without specifically identifying generation or cross-model evaluation as priorities.

In practice, the paper's results should be interpreted as validated for LLaMA-family models on discriminative natural language understanding tasks with training data drawn from a distribution close to pre-training. Extrapolation to other models, tasks, or domains requires independent validation that the paper cannot provide.


Single Random Seed Without Variance Estimates Undermines Confidence in Performance Differences

The assumption or constraint. The paper reports downstream accuracy results (Table 3, Figure 3) based on "the best results obtained using the same random seed" (Section 4.3). This means that for each model size, task, and method (Zero-shot, LoRA, LOMO, LoRA+LOMO), a single training run was performed with a single random seed — the seed that controls data shuffling order, dropout mask generation, and other stochastic aspects of training. The paper does not report:

  • Mean and standard deviation across multiple random seeds.
  • Confidence intervals or statistical significance tests for pairwise method comparisons.
  • The specific random seed used, which would enable exact reproduction of the reported numbers.
  • Whether the same seed was used across methods (which would control for data order effects but not for other stochastic factors) or whether each method was tuned independently.

The paper attributes this to "limited computational resources" (Appendix A), which is understandable given the paper's focus on resource-constrained settings — running multiple seeds for 4 model scales × 6 tasks × 3 methods would multiply the already-substantial compute budget.

The consequence. The reported accuracy differences — particularly the 2.8-point average gain of LOMO over LoRA for LLaMA-13B — cannot be statistically distinguished from noise. With 1,000 test examples per task, the standard error of accuracy is approximately sqrt(p(1-p)/1000), where p is the accuracy. For p ≈ 0.7, this gives a standard error of roughly 1.4–1.5 percentage points. A 2.8-point difference across 6 tasks is larger than the per-task standard error, suggesting it may be a real effect, but without multiple seeds we cannot quantify how much of this difference is due to the method versus favorable data ordering or lucky initialization for LOMO relative to LoRA.

The "best results" reporting also introduces a selection bias concern: if the authors ran multiple configurations and reported the best result for each method, the single-seed reporting means we do not know whether the reported numbers reflect typical behavior or cherry-picked maxima. The paper states "the best results obtained using the same random seed" — this phrasing suggests a single run per configuration, not a sweep over seeds, but the ambiguity remains.

For the LOMO + LoRA combination results (Figure 3), the sample size is even more fragile: these experiments are reported for only two tasks (BoolQ and MultiRC) on a single model size (LLaMA-13B). A 1–2 point accuracy difference on two tasks with one seed provides almost no statistical evidence for the claim that the combination is "complementary" — the observed effect could easily reverse with a different seed.

What evidence exists in the paper. The paper explicitly states the single-seed protocol: "we report the best results obtained using the same random seed" (Section 4.3). Appendix A reiterates: "Due to limited computational resources, we report the highest results of experiments conducted with the same random seed." The training dynamics appendix (Figures 4–5) shows loss and accuracy curves for a single run, consistent with single-seed reporting. No variance estimates, error bars, or seed values appear anywhere in the paper.

Mitigation status. Partially addressed through transparency. The paper is honest about the single-seed limitation and the resource constraint that caused it. The Limitation section does not explicitly mention this as a limitation, but the acknowledgment in Appendix A serves as a de facto caveat. The paper does not attempt to mitigate the statistical weakness through alternative approaches that would not require additional training runs — for example, bootstrap confidence intervals on the test set predictions, or reporting per-task accuracy rather than (or in addition to) averages so readers can assess consistency across tasks.

A practitioner reading the 2.8-point average gain should interpret it as a point estimate from a single realization of a stochastic process. The direction of the effect (LOMO outperforming LoRA on average) is consistent with the theoretical argument that full-parameter tuning should beat low-rank adaptation given sufficient data, but the magnitude and statistical reliability are unknown. For high-stakes deployment decisions, this uncertainty matters — if the true average gain is 0.5 points rather than 2.8, the case for full-parameter tuning over LoRA is substantially weaker.


Hard Problems Remain Unsolved and the SGD-as-Sufficient Argument Is Never Stress-Tested

The assumption or constraint. The paper's theoretical justification for replacing Adam with SGD rests on three claims about the fine-tuning loss landscape: (1) large curvature is unlikely for natural language tasks, (2) a local optimum is sufficient, and (3) saddle points are distant from the pre-trained initialization (Section 3.1.1). These claims are supported by citation to prior work (Hao et al., 2019; Kawaguchi et al., 2019) and the implicit batch size analysis (Section 3.1.2), but they collectively describe a best-case scenario — fine-tuning on in-distribution natural language tasks with a well-pre-trained model.

The paper explicitly identifies the boundary of this scenario:

"Note that this holds only when we teach the LLMs natural language-based tasks (or code-based if pre-trained with code). A synthetic loss function unrelated to pre-training tasks will indeed face the large curvature problem." (Section 3.1.1)

This is a precise boundary condition, but the paper never empirically probes where this boundary lies. All SuperGLUE tasks are natural language understanding benchmarks closely related to LLaMA's pre-training corpus. There is no experiment that deliberately increases the distance between pre-training and fine-tuning — no domain-shifted dataset (e.g., biomedical literature, legal documents, low-resource languages), no task with a deliberately adversarial objective, and no controlled study that varies a "distribution shift" parameter and measures when SGD begins to fail relative to Adam.

The consequence. The paper validates LOMO in the regime where the theory predicts it should work, but provides no empirical guidance on where it stops working. A practitioner fine-tuning on a moderately out-of-distribution task — say, adapting a general-purpose LLM to a specialized technical domain — cannot determine from this paper whether their task is "natural language-based" enough to fall within LOMO's validated regime or "unrelated to pre-training tasks" enough to trigger the curvature problems the paper warns about. The boundary is drawn in principle but not located in practice.

This is analogous to a compass that points correctly at the North Pole but is never tested at the equator — it works at the known-good point but its reliability at intermediate distances is unknown. The paper's theoretical argument predicts that performance should degrade gradually as the task diverges from pre-training (since smoothness is a continuous property), but the shape of this degradation curve — linear, threshold, task-dependent — is entirely unexplored.

There is also a failure mode the paper cannot diagnose: if a practitioner uses LOMO and gets poor results, is it because (a) SGD is failing due to loss surface curvature (a fundamental limitation of the method for that task), or (b) the hyperparameters are suboptimal (a fixable configuration issue)? Without baselines comparing LOMO against AdamW on the same task, these two failure modes are indistinguishable. The practitioner would need to run an AdamW baseline themselves to determine whether the problem is the optimizer or the tuning configuration — exactly the resource-intensive process LOMO aims to avoid.

What evidence exists in the paper. SuperGLUE results (Table 3) are all in the favorable in-distribution regime. The paper does not report any experiments with domain-shifted or non-natural-language tasks. The limitation section does not mention this gap. The theoretical analysis (Section 3.1.1–3.1.2) identifies the boundary but treats its empirical characterization as beyond scope.

Mitigation status. Not addressed. The paper's future work section (Section 5) mentions "more applicable scenarios for LOMO" and "theoretical analyses for optimizing large language models" but does not prioritize mapping the SGD-failure boundary. The Limitation section focuses on implementation-level issues (two-pass backward speed, limited SuperGLUE subset, no A100 throughput tests) rather than this conceptual gap. A practitioner currently must treat LOMO as validated for in-distribution natural language tasks and assume additional risk for any task that deviates from this profile — the paper provides no tools or heuristics for assessing this risk before committing to training.

7. Implications and Future Directions

How This Work Changes the Landscape

LOMO changes the landscape by shifting the framing of the memory problem from partitioning to elimination. Before LOMO, the dominant strategy for fitting large-model training into limited GPU memory was to partition the problem: spread optimizer states, gradients, and parameters across multiple GPUs (ZeRO), offload to CPU or NVMe (ZeRO-Offload), or reduce the number of trainable parameters (PEFT). All of these strategies accept the fundamental memory profile of the training loop — forward pass, accumulate all gradients, update all parameters — as a given and work around it. LOMO challenges the profile itself by asking: must we store all gradients simultaneously? The answer, it turns out, is no — at least not for SGD fine-tuning.

This is a reframing of the training loop architecture, not merely an incremental memory optimization. The paper demonstrates that the two-phase training loop (all gradients → all updates) is a convention inherited from early deep learning frameworks, not a logical necessity. By fusing gradient computation with parameter updates in a single backward pass, LOMO eliminates gradient storage as a memory category — dropping gradient memory from O(P) to O(1), where P is the total parameter count. This is a categorical elimination (the memory category vanishes for practical purposes) rather than a proportional reduction (making it smaller by a constant factor).

The magnitude of this shift is substantiated by the numbers. As shown in Table 1 and Figure 2, optimizer states consume 73.7% of memory in standard AdamW training for LLaMA-7B, and gradients consume an additional 9.8%. LOMO eliminates the first category entirely (by switching to SGD, which has no momentum or variance buffers) and reduces the second to negligible size (by fusing updates, which means only one layer's gradient exists at any moment). The residual memory — 14.58 GB for a 7B model — is essentially the inference footprint plus FP32 master weights. This establishes a hard lower bound: training memory cannot drop below inference memory without compressing the model itself. The paper reaches this bound, demonstrating that the optimizer-side memory problem for full-parameter fine-tuning has been solved to its theoretical limit.

The work also reconciles a contradiction in the literature about SGD's viability for large models. Historically, SGD was the default optimizer for deep learning but was largely abandoned for transformer training because adaptive methods like Adam consistently outperformed it on smaller models and from-scratch training. Practitioners internalized the lesson that "SGD doesn't work well for transformers." Yet some researchers observed that SGD could fine-tune large pre-trained models reasonably well — an empirical puzzle without a clear explanation. The paper's three-part theoretical argument (smooth loss surfaces reduce curvature problems, local optima suffice for fine-tuning, saddle points are distant from pre-trained initializations) plus the implicit batch size analysis provide a mechanistic explanation for why SGD works in the fine-tuning regime despite failing during pre-training. The implicit batch size analysis (Section 3.1.2) is particularly valuable because it shows that under a smoothness assumption, sequential SGD updates on small batches approximate larger-batch updates — a concrete mechanism for the stability practitioners observed but couldn't explain.

This reframing redirects research attention toward several newly visible problems:

  • Verifier and optimizer robustness under aggressive optimization becomes less urgent for LOMO specifically (since LOMO uses exact first-order gradients, unlike MeZO's zeroth-order estimates or GaLore's low-rank approximations), but the broader question of when SGD is sufficient versus when adaptive optimization is necessary becomes newly tractable. The paper defines the boundary — natural language tasks close to the pre-training distribution — but doesn't map it empirically.

  • The inference-memory floor as a target means that further memory reductions must come from model compression (parameter quantization, pruning, distillation) rather than optimizer improvements. The paper explicitly flags this in Section 5: "one promising direction is the exploration of parameter quantization techniques, which could significantly reduce memory usage." This is not a generic future-work suggestion; it is a logical consequence of having reached the optimizer-side bound.

  • The backward-pass architecture as a design space is now open. LOMO shows that interleaving gradient computation with parameter updates is both correct (for SGD on smooth landscapes) and memory-efficient. Future optimizers or training frameworks might explore more sophisticated interleaving schedules — updating certain layers immediately while deferring others, or using partial gradient information to make early update decisions — that exploit the same temporal-ordering flexibility LOMO identifies.

Less visibly, LOMO changes the economics of who can participate in LLM research. The paper demonstrates that a 65B model can be full-parameter fine-tuned on 8 consumer-grade RTX 3090 GPUs in approximately 3.6 hours for 1,000 samples of 512 tokens each (Section 4.2). This is not a datacenter-scale resource — it's a workstation that a well-funded academic lab or small company can assemble. By lowering the hardware threshold from "8×80GB datacenter GPUs" to "8×24GB consumer GPUs," LOMO partially democratizes access to full-parameter fine-tuning. This is not a complete solution — the hardware cost is still non-trivial, and the 65B model requires all 8 GPUs — but it shifts the boundary of who can participate from "only industry labs with cluster-scale resources" to "any group that can assemble a multi-GPU workstation."

Follow-Up Research This Work Enables

Mapping the SGD-failure boundary across task distributions. The paper's theoretical argument for SGD sufficiency is explicitly regime-dependent: it holds for natural language tasks close to the pre-training distribution, and the paper warns that "a synthetic loss function unrelated to pre-training tasks will indeed face the large curvature problem" (Section 3.1.1). But the boundary between "close enough" and "too far" is entirely unmapped. A strong follow-up would systematically vary the distribution shift between pre-training and fine-tuning — for example, fine-tuning LLaMA on tasks from progressively more distant domains (general web text → Wikipedia → news → biomedical literature → legal documents → code → synthetic formal languages) and measuring both downstream accuracy and optimizer-specific metrics (gradient norm variance, loss surface curvature estimates, update-to-parameter ratio). The key measurement would be the SGD-Adam performance gap as a function of distribution distance, using AdamW with CPU offloading or gradient accumulation as the baseline when memory constraints prevent a direct in-GPU comparison. This would convert the paper's theoretical boundary condition into an empirical map that practitioners can use to assess whether their specific fine-tuning task falls within LOMO's safe operating envelope.

LOMO with generation tasks and instruction tuning. Every downstream evaluation in the paper is a SuperGLUE classification task scored by average log-likelihood of candidate labels. This is a narrow evaluation regime that sidesteps the central challenge of modern LLM fine-tuning: adapting models for open-ended generation (instruction following, dialogue, summarization). Full-parameter fine-tuning is arguably more consequential for generation than for classification — generation quality depends on the entire output token distribution, and PEFT methods' restricted capacity may impose a harder ceiling on generation than on discriminative tasks. A critical follow-up would fine-tune LLaMA models with LOMO on instruction-tuning datasets (e.g., Alpaca, Dolly, or the OpenAssistant corpus) and evaluate using both automated metrics (ROUGE, BERTScore) and human preference judgments (or LLM-as-judge with a strong evaluator model). The comparison should include: (1) LOMO vs. LoRA at matched memory budgets, (2) LOMO vs. AdamW full-parameter fine-tuning (using CPU offloading or gradient accumulation to fit AdamW within the available memory), and (3) LOMO + LoRA combination. The key question is whether SGD's lack of per-parameter adaptive learning rates causes mode collapse, repetitive outputs, or degraded fluency in generation — failure modes that classification accuracy would not reveal.

The two-pass backward overhead: measurement and mitigation. The paper acknowledges that gradient norm clipping and dynamic loss scaling each require a second backward pass but provides no throughput measurements for the two-pass configuration (Section 3.3.1–3.3.2). The headline 11× throughput improvement (LOMO on 1 GPU vs. AdamW on 8 GPUs for LLaMA-7B, Table 2) is measured with value-based clipping in a one-pass backward. A straightforward but practically essential follow-up would measure throughput for the two-pass configuration across model scales (7B–65B) and compare against: (1) one-pass LOMO with value-based clipping, (2) the grouped gradient clipping approximation described in Section 3.3.1, and (3) SGD with norm-based clipping and standard gradient storage (no fusion). The throughput numbers should be reported alongside downstream accuracy on tasks that genuinely require norm-based clipping (higher learning rates, tasks with observed gradient spikes). This would provide practitioners with the data needed to decide whether the two-pass overhead is acceptable for their use case, and whether the grouped clipping approximation — described as "a controversial solution" that is "indeed biased" — is worth the throughput savings.

LOMO across model architectures and pre-training objectives. The paper's experiments are restricted to the LLaMA model family, which has specific architectural properties (pre-norm residual connections, SwiGLU activations, rotary position embeddings). The loss surface smoothness that LOMO depends on might vary with architecture — post-norm transformers, models with different activation functions, or encoder-decoder architectures could exhibit different curvature characteristics. A systematic study would apply LOMO to a diverse set of architectures: LLaMA (pre-norm, decoder-only), Falcon (parallel attention), OPT (pre-norm, decoder-only but different training data), T5 (encoder-decoder), and Mistral (sliding window attention, architectural variants). For each architecture, measure: (1) memory reduction relative to AdamW, (2) throughput at comparable batch sizes, (3) downstream accuracy vs. AdamW on a fixed set of tasks (both classification and generation), and (4) training stability metrics (gradient norm trajectories, loss volatility). The goal is to determine whether LOMO's effectiveness is architecture-specific or a general property of large pre-trained transformers — the paper's theoretical arguments suggest generality for any smooth loss surface, but the empirical evidence is LLaMA-only.

LOMO with parameter quantization: pushing below the inference floor. The paper reaches the inference-memory floor for optimizer and gradient memory but acknowledges that parameters themselves remain a significant consumer (the FP32 master weights and FP16 working copies together account for roughly 42 GB for a 65B model — the majority of the remaining memory in Table 2). Section 5 identifies parameter quantization as "one promising direction" for further reduction. A concrete follow-up would integrate LOMO with 4-bit or 8-bit parameter quantization (e.g., QLoRA-style NF4 quantization, or INT8 quantization with stochastic rounding) and measure: (1) the combined memory footprint — can a 65B model be fine-tuned on 4 GPUs instead of 8? on 2 GPUs? — (2) the throughput impact of quantization/dequantization operations during the fused backward pass, and (3) the downstream accuracy relative to FP16 LOMO and FP16 AdamW. The key technical challenge is that the fused update param_fp32 -= lr * grad_fp32 requires FP32 arithmetic for the master weights; quantized training would need to maintain the master weights in FP32 but store and forward-propagate quantized versions, adding dequantization overhead at each layer's forward pass and quantization overhead after each update. Measuring whether this overhead negates LOMO's throughput advantages would determine whether quantization+LOMO is a practical path to single-GPU fine-tuning of 30B+ models.

Training from scratch with LOMO-style fused updates for large models. The paper's theoretical argument — that SGD suffices because fine-tuning starts from a smooth region of the loss landscape — explicitly excludes training from scratch. However, as model sizes continue to grow, even pre-training increasingly resembles fine-tuning from an initialization that is already in a relatively smooth region (random initialization followed by a few steps of warmup often produces parameters with reasonably well-behaved gradients). A provocative follow-up would test whether LOMO-style fused SGD can train a moderately large model (1B–3B parameters) from scratch on a language modeling corpus, compared against AdamW at matched compute budgets. The key measurement is not whether SGD matches AdamW's final loss — it almost certainly won't, given decades of evidence for Adam's superiority in from-scratch training — but how large the gap is and when it emerges. If the gap is small for the first few billion tokens of training (while the model is still in the early phases of representation learning), it suggests that LOMO-style optimizers might be useful for the early stages of very large model training before switching to Adam for fine-grained optimization. If the gap is immediately large, it confirms the paper's boundary condition and sharpens our understanding of exactly when adaptive optimization becomes necessary.

Practical Applications and Downstream Use Cases

Academic and small-company LLM research. The paper's headline result — full-parameter fine-tuning of a 65B model on 8 consumer GPUs (RTX 3090, 24GB each) — directly enables research groups with workstation-level hardware to participate in full-parameter adaptation studies that were previously gated behind datacenter-scale resources. A typical academic lab might have one or two GPU workstations with 4–8 consumer GPUs; before LOMO, such a lab could only use PEFT methods (LoRA, prefix-tuning) or zero-order optimizers (MeZO) for models above 7B. With LOMO, the same hardware can perform full-parameter fine-tuning on a 30B model with 4 GPUs (19.78 GB per GPU, Table 2) and a 65B model with 8 GPUs (19.18 GB per GPU). The practical training times are overnight-scale: ~3.6 hours for 1,000 samples of 512 tokens on the 65B model (Section 4.2). This means a research group can iterate on full-parameter fine-tuning experiments daily — testing different hyperparameters, data mixtures, or task formulations — rather than waiting for scarce cluster allocation or being forced to use PEFT methods that restrict the research questions they can ask.

Domain adaptation of open-weight LLMs in resource-constrained industries. Organizations with domain-specific text data (legal, medical, financial, technical support) but limited ML infrastructure can use LOMO to fully adapt open-weight models like LLaMA to their domain without investing in datacenter GPUs. The paper shows that the LOMO+LoRA combination (Section 4.3.2, Figure 3) provides additive benefits — full-parameter tuning of the base weights plus low-rank adaptation modules — which is particularly relevant for domain adaptation where both the base language understanding needs refinement and domain-specific patterns need to be learned. The combined memory footprint is not separately reported, but since LOMO and LoRA are orthogonal (LOMO reduces gradient and optimizer memory; LoRA reduces the number of trainable parameters), the additive memory cost of LoRA modules should be small relative to the LOMO savings. A legal tech startup with a single 8-GPU workstation could take a 65B open-weight model, fine-tune it fully on their corpus of legal documents using LOMO, and deploy a domain-adapted model without ever renting cloud GPU instances.

Cost-efficient batch inference with full-parameter fine-tuned models. For organizations running periodic fine-tuning pipelines — retraining models on fresh data weekly or monthly — LOMO's throughput advantages translate directly to cost savings. The 7B model with LOMO processes 769.92 tokens per GPU per second on a single consumer GPU (Table 2), compared to AdamW's 67.37 TGS on 8 GPUs. Even accounting for the two-pass backward if gradient norm clipping is needed (~385 TGS on a single GPU), LOMO processes roughly 5.7× more tokens per second than AdamW on 8 GPUs — meaning the fine-tuning phase of the pipeline completes in less than one-fifth the time on one-eighth the hardware. For a batch inference pipeline that fine-tunes on 10,000 examples of 512 tokens each (5.12 million tokens total), LOMO on 1 GPU would complete in roughly 3.7 hours (with one-pass backward) or 7.4 hours (with two-pass), compared to AdamW on 8 GPUs at roughly 10.6 hours. The savings in GPU-hours are dramatic: ~3.7–7.4 GPU-hours vs. ~85 GPU-hours for AdamW — a 11–23× reduction in total compute cost for the fine-tuning phase.

Prototyping full-parameter fine-tuning recipes before scaling to production hardware. A common workflow in industry is to prototype fine-tuning configurations (learning rates, data mixtures, number of epochs) on smaller models or smaller datasets before deploying at scale. LOMO makes this prototyping faster and cheaper: the 7B model fits on a single consumer GPU with LOMO, enabling rapid hyperparameter sweeps without multi-GPU orchestration. Once a configuration is validated, the same recipe can be applied to larger models using the appropriate number of GPUs (2 for 13B, 4 for 30B, 8 for 65B). The paper's demonstration that LOMO's memory and throughput characteristics scale predictably with model size (Table 2) means the prototyping-to-production transfer is reliable — a hyperparameter configuration that works well for LLaMA-7B on 1 GPU will likely work for LLaMA-65B on 8 GPUs, since the optimizer (SGD) and update mechanism (fused backward) are identical across scales.

When to Prefer This Method

The paper itself does not articulate a formal decision framework for choosing between LOMO, LoRA, AdamW-based full-parameter fine-tuning, and other memory-efficient methods. It demonstrates that LOMO can be used in resource-constrained settings and that it generally outperforms LoRA, but does not specify the conditions under which each approach is preferable. The paper's philosophy — expressed in the conclusion as "a simple and practical solution for fine-tuning LLMs and identify its flaws to continually improve it" — is one of progressive refinement rather than prescriptive decision rules. Constructing a "prefer X when Y" matrix would impose a formalism the paper does not provide and could misrepresent its intentionally exploratory stance.