# Runtime-Orchestrated Second-Order Optimization for Scalable LLM Training

arXiv: https://arxiv.org/abs/2605.16184 · HTML: https://mlsanity.com/paper/2605.16184/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 26 references that are in the map.
- First posted on arXiv 15 May 2026.

## Builds on (most cited first, 21 of 21)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Local SGD Converges Fast and Communicates Little](https://mlsanity.com/paper/1805.09767/) (2018)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [2 OLMo 2 Furious](https://mlsanity.com/paper/2501.00656/) (2024)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale](https://mlsanity.com/paper/2309.06497/) (2023)
- [ASDL: A Unified Interface for Gradient Preconditioning in PyTorch](https://mlsanity.com/paper/2305.04684/) (2023)
- [GossipGraD: Scalable Deep Learning using Gossip Communication based Asynchronous Gradient Descent](https://mlsanity.com/paper/1803.05880/) (2018)
- [Understanding and Improving Shampoo and SOAP via Kullback-Leibler Minimization](https://mlsanity.com/paper/2509.03378/) (2025)
- [PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management](https://mlsanity.com/paper/2108.05818/) (2021)
- [Part-time Power Measurements: nvidia-smi's Lack of Attention](https://mlsanity.com/paper/2312.02741/) (2023)
- [Clarifying Shampoo: Adapting Spectral Descent to Stochasticity and the Parameter Trajectory](https://mlsanity.com/paper/2602.09314/) (2026)
- [Deep Optimizer States: Towards Scalable Training of Transformer Models Using Interleaved Offloading](https://mlsanity.com/paper/2410.21316/) (2024)
- [Beyond the Mean: Fisher-Orthogonal Projection for Natural Gradient Descent in Large Batch Training](https://mlsanity.com/paper/2508.13898/) (2025)

## Built on by (most cited first, 1 of 1)

- [Detecting Contaminated Code-Generation Prompt Batches via Influence Functions](https://mlsanity.com/paper/2608.14303/) (2026)

## Abstract (arXiv)

Second-order methods offer an attractive path toward more sample-efficient LLM training, but their practical use is often blocked by the systems cost of maintaining and updating large matrix-based optimizer states. We introduce Asteria, a runtime system designed to remove this bottleneck by separating second-order optimization logic from the critical GPU training path. Rather than keeping all preconditioner state on the accelerator, Asteria dynamically distributes optimizer state across GPU memory, CPU memory, and optional NVMe storage according to architectural constraints and runtime pressure. It further uses training hooks to prepare shadow states in advance, allowing expensive inverse-root computations to proceed asynchronously on the host while GPU computation continues. For distributed training, Asteria employs a bounded-staleness protocol that limits synchronization frequency while preserving optimizer effectiveness through topology-aware coordination. We evaluate Asteria on both memory-constrained and distributed training settings. On a DGX Spark platform with a single GB10 GPU and 128GB unified memory, Asteria supports second-order training for a 1B-parameter language model. On multi-node GH200 systems, it lowers visible optimizer overhead, reduces recurring latency spikes, accelerates convergence in wall-clock time, and maintains the optimization advantages of SOAP and KL-Shampoo in a 7B-parameter language model. Our results suggest that second-order LLM training can be made practical not by simplifying the optimizer alone, but by rethinking how optimizer state, background computation, and distributed synchronization are managed at the runtime level.