Runtime-Orchestrated Second-Order Optimization for Scalable LLM Training
- URL: http://arxiv.org/abs/2605.16184v1
- Date: Fri, 15 May 2026 17:03:55 GMT
- Title: Runtime-Orchestrated Second-Order Optimization for Scalable LLM Training
- Abstract summary: We introduce textAsteria, a runtime system that separates second-order optimization logic from the critical GPU training path.<n>Asteria dynamically distributes state across GPU memory, CPU memory, and optional storage according to architectural constraints.<n>We evaluate Asteria on both memory-constrained and distributed training settings.
- Score: 4.950833328317384
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: 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 \textbf{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.
Related papers
- Convex Optimization for Alignment and Preference Learning on a Single GPU [52.997197698288936]
Fine-tuning large language models to align with human preferences has driven the success of systems such as Gemini and ChatGPT.<n>Direct Preference Optimization (DPO) offers a simpler alternative but has limitations such as inconsistent ranking accuracy and high dependence on GPU resources.<n>We propose the Convex Optimization for Alignment and Preference Learning Algorithm (COALA), a novel lightweight strategy with strong theoretical guarantees.
arXiv Detail & Related papers (2026-05-22T05:25:00Z) - AdaFuse: Accelerating Dynamic Adapter Inference via Token-Level Pre-Gating and Fused Kernel Optimization [84.25316984309725]
Integration of dynamic, sparse structures with parameter-efficient adapters (e.g., LoRA) is a powerful technique for enhancing Large Language Models (LLMs)<n>Despite minimal increases in computational load, the computation latency often skyrockets, leading to decoding speeds slowing by over 2.5 times.<n>AdaFuse is a framework built on a tight co-design between the algorithm and the underlying hardware system to enable efficient dynamic adapter execution.
arXiv Detail & Related papers (2026-03-12T12:46:42Z) - AutoQRA: Joint Optimization of Mixed-Precision Quantization and Low-rank Adapters for Efficient LLM Fine-Tuning [23.59600455731982]
We propose a joint optimization framework that simultaneously optimize the bit-width and LoRA rank configuration for each layer during the mixed quantized fine-tuning process.<n>Experiments show that AutoQRA achieves performance close to full-precision fine-tuning with a memory footprint comparable to uniform 4-bit methods.
arXiv Detail & Related papers (2026-02-25T07:18:08Z) - PT$^2$-LLM: Post-Training Ternarization for Large Language Models [52.4629647715623]
Large Language Models (LLMs) have shown impressive capabilities across diverse tasks, but their large memory and compute demands hinder deployment.<n>We propose PT$2$-LLM, a post-training ternarization framework tailored for LLMs.<n>At its core is an Asymmetric Ternary Quantizer equipped with a two-stage refinement pipeline.
arXiv Detail & Related papers (2025-09-27T03:01:48Z) - CSGO: Generalized Optimization for Cold Start in Wireless Collaborative Edge LLM Systems [62.24576366776727]
We propose a latency-aware scheduling framework to minimize total inference latency.<n>We show that the proposed method significantly reduces cold-start latency compared to baseline strategies.
arXiv Detail & Related papers (2025-08-15T07:49:22Z) - DistZO2: High-Throughput and Memory-Efficient Zeroth-Order Fine-tuning LLMs with Distributed Parallel Computing [4.589472292598182]
Fine-tuning large language models (LLMs) remains resource-intensive due to their sheer scale.<n>We present DistZO2, a memory-efficient framework for distributed zeroth-order fine-tuning of LLMs.
arXiv Detail & Related papers (2025-07-03T22:53:34Z) - A Minimalist Optimizer Design for LLM Pretraining [31.996047271119156]
Training large language models typically relies on adaptives such as Adam.<n>Recent works such as GaLore Fira, and APOLLO have proposed state-compressed variants to reduce memory consumption.<n>In this work, we investigate what is the minimal amount of state that is truly necessary to retain state-of-the-art performance in LLM pretraining.
arXiv Detail & Related papers (2025-06-20T00:10:35Z) - APOLLO: SGD-like Memory, AdamW-level Performance [61.53444035835778]
Large language models (LLMs) are notoriously memory-intensive during training.<n>Various memory-efficient Scals have been proposed to reduce memory usage.<n>They face critical challenges: (i) costly SVD operations; (ii) significant performance trade-offs compared to AdamW; and (iii) still substantial memory overhead to maintain competitive performance.
arXiv Detail & Related papers (2024-12-06T18:55:34Z) - Simultaneous Computation and Memory Efficient Zeroth-Order Optimizer for Fine-Tuning Large Language Models [33.911521719528686]
Fine-tuning is powerful for adapting large language models to downstream tasks, but it often results in huge memory usages.
A promising approach is using Zeroth-Order (ZO) gradients, which estimates to replace First-Order (FO) gradients.
We introduce a novel layer-wise sparse computation and memory efficient ZO, named LeZO.
arXiv Detail & Related papers (2024-10-13T12:47:37Z) - Bypass Back-propagation: Optimization-based Structural Pruning for Large Language Models via Policy Gradient [57.9629676017527]
We propose an optimization-based structural pruning that learns the pruning masks in a probabilistic space directly by optimizing the loss of the pruned model.<n>We achieve this by learning an underlying Bernoulli distribution to sample binary pruning masks.<n>Experiments conducted on LLaMA, LLaMA-2, LLaMA-3, Vicuna, and Mistral models demonstrate the promising performance of our method in efficiency and effectiveness.
arXiv Detail & Related papers (2024-06-15T09:31:03Z) - ACCO: Accumulate While You Communicate for Communication-Overlapped Sharded LLM Training [22.940404796500985]
We propose a memory-efficient optimization algorithm for distributed training LLMs.<n>By synchronizing delayed gradients while computing new ones, ACCO reduces GPU idle time and supports heterogeneous hardware.<n>Compared to ZeRO-1, our approach is significantly faster and scales effectively across heterogeneous hardware.
arXiv Detail & Related papers (2024-06-03T08:23:45Z) - AdaLomo: Low-memory Optimization with Adaptive Learning Rate [59.64965955386855]
We introduce low-memory optimization with adaptive learning rate (AdaLomo) for large language models.
AdaLomo results on par with AdamW, while significantly reducing memory requirements, thereby lowering the hardware barrier to training large language models.
arXiv Detail & Related papers (2023-10-16T09:04:28Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.