🤖 machine learning

Runtime-Orchestrated Second-Order Optimization for Scalable LLM Training

The paper introduces Asteria, a runtime system that enables practical, scalable second-order optimization for large language models by dynamically distributing optimizer states across heterogeneous memory hierarchies and decoupling expensive preconditioner computations from the critical GPU training path to reduce overhead and accelerate convergence.

Original authors: Yishun Lu, Junhao Zhang, Zeyu Yang, Wes Armour

Published 2026-05-18

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yishun Lu, Junhao Zhang, Zeyu Yang, Wes Armour

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

Imagine you are trying to teach a giant, super-smart robot (a Large Language Model) to write stories. To do this, you need a "teacher" (an optimizer) that looks at the robot's mistakes and tells it how to improve.

Currently, most teachers use a simple method called AdamW. It's like a student who checks their homework, sees a few wrong answers, and makes small, quick corrections. It's fast and works well, but it takes a lot of practice (millions of examples) to get really good.

There is a smarter, more advanced teacher method called Second-Order Optimization (like SOAP or Shampoo). This teacher doesn't just look at what was wrong; it analyzes the shape of the mistakes to understand the problem deeply. It learns much faster and needs fewer examples. However, there's a huge catch: this smart teacher is incredibly heavy. It requires massive amounts of memory and computing power to do its complex math, often causing the computer to crash or run so slowly that the simple teacher actually finishes the job first.

Enter Asteria: The "Smart Logistics Manager"

The researchers at Oxford built a new system called Asteria. Think of Asteria not as a new teacher, but as a brilliant logistics manager who reorganizes the entire classroom so the smart teacher can finally do their job without breaking the school.

Here is how Asteria solves the three biggest problems, using simple analogies:

1. The "Furniture in a Tiny Room" Problem (Memory)

The Issue: The smart teacher needs to keep huge, complex charts (matrices) in their head (the GPU memory). On a standard computer, there isn't enough space. It's like trying to fit a king-sized bed, a dining table, and a wardrobe into a studio apartment. The room gets full, and the teacher has to quit.
The Asteria Solution: Asteria acts like a folding furniture expert. It realizes that the teacher doesn't need to hold everything in their hands at once.

It keeps the most active parts of the charts on the GPU (the desk).
It moves the heavy, less-frequently-used parts to the CPU (the hallway) or even to a hard drive (the garage).
Crucially, it only brings the heavy parts back to the desk exactly when they are needed. This allows the smart teacher to work on a tiny laptop just as well as on a supercomputer.

2. The "Traffic Jam" Problem (Speed)

The Issue: Every few steps, the smart teacher has to do a massive, complex calculation (like solving a giant puzzle) to update their charts. This takes a long time and stops the whole class from working. It's like a delivery truck blocking the entire highway while it unloads a single package. The GPU (the computer's brain) sits idle, waiting.
The Asteria Solution: Asteria introduces a parallel assembly line.

Instead of stopping the main class to do the heavy math, Asteria sends the "heavy lifting" to a team of workers in the back office (the CPU).
While the main class (the GPU) keeps teaching the robot, the back office workers are quietly solving the complex puzzles in the background.
By the time the class needs the new charts, the back office has already finished them and slid them over. The main class never has to stop. The "traffic jam" disappears.

3. The "Group Chat" Problem (Distributed Training)

The Issue: When training with many computers at once, everyone usually has to stop and agree on the exact same information before moving to the next step. This is like a group chat where everyone has to wait for the slowest person to reply before anyone can type again.
The Asteria Solution: Asteria uses a "Good Enough" policy.

Instead of waiting for everyone to be perfectly synchronized, it allows computers to work with slightly "stale" (slightly outdated) information for a short time.
It only sends updates when they are truly necessary.
This keeps the group moving fast, like a team that trusts each other to keep working while catching up on updates later, rather than stopping the whole project every five minutes.

The Results: What Did They Find?

The researchers tested Asteria on real hardware, including a powerful single computer (Nvidia DGX Spark) and a large cluster of computers (Nvidia GH200).

It works on small machines: They were able to train a large language model (1 billion parameters) on a single machine that previously couldn't handle the memory requirements of this smart teacher.
It's faster in real time: Because it hides the "traffic jams," the training finishes in less actual clock time, even though the math is harder.
It saves energy: By keeping the computer's brain (GPU) busy and not letting it sit idle waiting for calculations, Asteria actually uses less total energy to achieve the same result compared to the old, clunky methods.
It doesn't lose quality: The model learns just as well as it would with the "native" (unoptimized) smart teacher, but it gets there much faster and cheaper.

In Summary:
Asteria doesn't invent a new math formula. Instead, it rethinks how the computer runs that math. It separates the heavy lifting from the main work, uses all available storage space wisely, and lets computers work asynchronously. This turns a "theoretically great but practically impossible" method into a practical tool for training the next generation of AI.

Technical Summary: Asteria – Runtime-Orchestrated Second-Order Optimization for Scalable LLM Training

1. Problem Statement

Second-order optimization methods (e.g., Shampoo, SOAP, KL-Shampoo) offer superior statistical efficiency and faster convergence for Large Language Model (LLM) training compared to first-order optimizers like AdamW. However, their practical deployment at scale is hindered by severe system-level barriers:

Memory Constraints: Second-order methods require storing large covariance matrices ( $L$ and $R$ ) and performing dense eigendecompositions, leading to an $O(N^2)$ memory footprint and $O(N^3)$ computational cost. In memory-constrained environments (e.g., single-node workstations with Unified Memory Architectures), these states exhaust physical memory, causing Out-of-Memory (OOM) failures or forcing reductions in model size and batch size.
Latency Spikes: The periodic computation of inverse roots (e.g., $L^{-1/4}$ ) involves cubic-cost matrix operations. In existing implementations, these synchronous updates block the GPU training pipeline, causing severe periodic latency spikes that disrupt the fine-grained overlap between computation and communication essential for efficient distributed training (e.g., FSDP).
Synchronization Overhead: Current distributed second-order methods rely on rigid, synchronous global state updates (AllGather) across all ranks. This couples fast intra-node communication with slow inter-node synchronization, preventing the system from exploiting asynchronous, topology-aware staleness bounds.

Existing solutions, such as DTensor-based sharding, fail to address these issues in resource-constrained single-node settings or when the communication overhead of global synchronization outweighs the benefits.

2. Methodology: The Asteria Runtime

Asteria is a hardware-software co-designed runtime system built on top of PyTorch FSDP. It decouples second-order optimization logic from the critical GPU training path through three core mechanisms:

A. Architecture-Adaptive Asymmetric Memory Tiering

Unlike first-order optimizers that treat state as homogeneous vectors, Asteria recognizes the asymmetric lifecycle of second-order states:

GPU: Stores factor_matrices (Kronecker factors) which are updated frequently during training.
Host (CPU/UVM): Stores inv_factor_matrices (inverse factors). These are maintained in CPU-accessible memory to avoid GPU memory pressure.
NVMe (Optional): Provides a staging tier for capacity-limited runs, allowing states to be swapped out via serialized computing or JIT paging.
Mechanism: The runtime captures host-side snapshots of factor matrices, dispatches asynchronous CPU workers to compute inverse roots, and writes results back to host memory. When needed, the GPU accesses these via Unified Virtual Memory (UVM) views or explicit prefetching, avoiding redundant host-to-device copies.

B. Hook-Orchestrated Shadow-State Pipeline

To eliminate latency spikes, Asteria removes the $O(N^3)$ inverse-root computation from the critical GPU path:

Shadow Stream: A low-priority CUDA stream handles GPU-side state staging (prefetching and paging) without blocking the primary training stream.
Asynchronous Execution: Expensive inverse-root updates are executed asynchronously on CPU worker threads.
Hook Integration: Lightweight PyTorch hooks (register_forward_hook, register_full_backward_pre_hook) act as scheduling signals. They trigger the draining of ready state transfers and prefetching of soon-to-be-consumed inverse factors, ensuring the GPU has fresh preconditioners without stalling the main training loop.

C. Topology-Aware Bounded-Staleness Coherence

For distributed training, Asteria replaces rigid global synchronization with a selective coherence protocol:

Staleness Tracking: A CoherenceRegistry tracks the freshness of each preconditioner block. Only blocks exceeding a configured staleness budget are synchronized.
Hierarchical Synchronization: Synchronization occurs in two stages: first within node-local process groups, then across node representatives. This minimizes cross-node traffic and avoids unnecessary CPU-to-GPU round trips by keeping communication on host-accessible memory.
Bounded Barrier: The runtime enforces a local barrier ensuring that if a preconditioner version becomes too stale while a CPU job is in flight, the optimizer waits for completion, balancing freshness with throughput.

3. Key Contributions

Runtime Decoupling: Asteria demonstrates that second-order optimization can be made practical not by simplifying the algorithm, but by rethinking state management, background computation, and synchronization at the runtime level.
Memory Scalability: It enables second-order training on single-node, memory-constrained hardware (e.g., a single GB10 GPU with 128GB unified memory) where traditional sharding is impossible.
Latency Hiding: By moving inverse-root computations to the host and using a shadow pipeline, Asteria flattens periodic latency spikes, achieving a smooth hardware execution profile comparable to first-order methods.
Energy Efficiency: The system improves the energy-loss tradeoff, achieving more loss reduction per unit of energy by reducing GPU idle time and optimizing the overlap of host-side computation.

4. Experimental Results

The authors evaluated Asteria on OLMo-family models (660M, 1B, 7B parameters) using the OLMo training framework on two platforms:

Single-Node (DGX Spark, GB10 GPU):
- Throughput: Asteria-SOAP and Asteria-KL-Shampoo eliminated the periodic step-time spikes seen in native implementations (reducing peak step time from ~96s to ~66.5s).
- Energy: While native second-order methods consumed ~17-20% more total SoC energy than AdamW, Asteria reduced this gap significantly (e.g., Asteria-KL-Shampoo consumed only ~7.9% more energy than AdamW).
- Efficiency: Asteria-KL-Shampoo achieved the highest normalized loss-reduction efficiency (3.3010), outperforming both native second-order methods and AdamW.
Multi-Node (GH200 Cluster):
- Convergence: Asteria variants maintained convergence trajectories close to their native second-order counterparts, significantly outperforming AdamW in wall-clock time to reach target loss levels.
- Staleness Tolerance: Experiments varying the staleness budget ( $S$ ) showed that increasing $S$ from 1 to 5 substantially reduced training time with negligible impact on final model quality.
- Strong Scaling: On 7B models scaling from 2 to 16 nodes, Asteria achieved higher speedup and lower per-step time compared to native baselines, proving effective even as the relative cost of state maintenance increases.

5. Significance and Claims

The paper argues that the primary obstacle to practical second-order LLM training is not the mathematical cost of curvature-aware updates, but the systems mismatch between second-order state maintenance and modern training stacks.

Reframing the Problem: The authors claim that second-order optimization should be viewed as a co-design problem involving algorithms, runtime scheduling, and heterogeneous memory systems.
Feasibility: They demonstrate that large aggregated GPU memory is not a fundamental prerequisite for second-order training; with lifecycle-aware placement and staged access, these methods can run on tight memory budgets.
Practicality: By treating preconditioner freshness and communication as runtime-managed resources rather than rigid global states, Asteria narrows the gap between the statistical promise of second-order optimization and its real-world deployability.

The paper concludes that Asteria makes second-order optimization a viable systems choice for LLM training, offering better convergence and energy efficiency without requiring algorithmic simplification.