Where Should Optimizer State Live? Tiered State Allocation for Memory-Efficient Mixture-of-Experts Training

• URL: http://arxiv.org/abs/2607.19058v1
• Date: Tue, 21 Jul 2026 12:47:15 GMT
• Title: Where Should Optimizer State Live? Tiered State Allocation for Memory-Efficient Mixture-of-Experts Training
• Abstract summary: We study SkewAdam, an estimator built on the observation that the three parameter populations of an MoE differ enough in size and gradient statistics that they should not receive the same state.<n>SkewAdam reaches perplexity 108.4, ahead of AdamW (126.8), Muon (120.2), and Lion (393.7), and router load balance settles to within 1% of its uniform floor.
• Score: 0.0
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Optimizer state is the largest single line item in the memory budget of mixture-of-experts (MoE) training: on a 6.78B-parameter MoE language model, AdamW keeps 50.6 GB of first and second moments to update 12.6 GB of bfloat16 weights. We study SkewAdam, an optimizer built on the observation that the three parameter populations of an MoE - the dense backbone, the experts, and the router - differ enough in size and gradient statistics that they should not receive the same state. SkewAdam keeps float32 momentum plus a factored second moment for the backbone (5% of parameters), a factored second moment alone for the experts (95%), and an exact second moment for the router (<0.01%). The resulting state occupies 1.29 GB, 2.6% of AdamW's, and peak training memory falls from 81.4 GB to 31.3 GB, within the budget of a 40 GB accelerator. In a controlled comparison from identical initializations over 82M tokens, SkewAdam reaches validation perplexity 108.4, ahead of AdamW (126.8), Muon (120.2), and Lion (393.7), and settles router load balance to within 1% of its uniform floor. The allocation is not what earns that perplexity: a tier ablation matches it with twenty times the state, and Adafactor, which shares the factored estimator but drops momentum, plateaus 40 points behind. The tiers buy memory at no cost to accuracy; the accuracy comes from keeping momentum, which a uniform optimizer shares too. Sweeping the baselines' learning rates narrows but does not close the gap: the best tuned AdamW reaches 118.5, tuned Adafactor 139.7. Where optimizer state lives, these results suggest, matters at least as much as how much of it there is.

Related papers

• Scaling Post-Training Ternarisation to Qwen3-8B Capability Retention, Reproduction, Lossless Packing, and Packed Execution [0.0]
  We study a scale-up of an aggressive post-training conversion pipeline from Qwen3-4B to Qwen3-8B.<n>The conversion uses KOTMS rotation, E2M-ATQ adaptive ternarisation, and GPTQ-style error compensation in a weight-only A16 configuration.
  arXiv  Detail & Related papers  (2026-09-08T06:21:44Z)
• Post-Training Ternarization of Qwen3-4B Capability, Effective Bit Budget, Storage Compression, and Deployment [0.0]
  Ultra-low-bit language models can reduce storage and memory bandwidth, but a nominal "1.58-bit" label does not fully describe the stored representation, retained capability, or runtime behavior.<n>We study an end-to-end post-training conversion of Qwen, an instruction-tuned 4B- parameter model, using KOTMS rotation, E2M-ATQ ternarization, and GPTQ-style error compensation from TWLA.<n>We evaluate effective bit accounting, task capability retention, perplexity, calibration sensitivity, checkpoint composition, and deployment behavior.
  arXiv  Detail & Related papers  (2026-09-02T00:42:54Z)
• When Does Trace-Driven Evaluation Mislead MoE Expert Caching? Replay Semantics, Workload Contamination, and Operating Regimes [3.471896896439183]
  Mixture-of-Experts (MoE) models have outgrown accelerator memory.<n>Policy that raised the hit rate would cut expert traffic per token.<n>We evaluate three evaluation axes that change conclusions, not just numbers.
  arXiv  Detail & Related papers  (2026-08-08T04:29:20Z)
• Variable Bit-width Quantization: Learning Per-Group Precision for "Bigger-but-Smaller" Language Models [0.0]
  We introduce Variable Bit-width Quantization (VBQ)<n>Training-time method in which each contiguous group of 64 weights learns its own resolution from 1,2,4,8 bits.<n>As quality-per-byte, VBQ is 3.9-8.4x more efficient than FP16.
  arXiv  Detail & Related papers  (2026-07-03T02:42:04Z)
• Slimmable ConvNeXt: Width-Adaptive Inference for Efficient Multi-Device Deployment [4.976815699476327]
  We present Slimmable ConvNeXt, which shows that ConvNeXt's modern design, specifically LayerNorm and inverted bottlenecks, makes it particularly suited for channel-width slimming.<n>On ImageNet-1k, Slimmable ConvNeXt-T with 3works achieves 80.8% top-1 accuracy at 4.5 GMACs and 77.4% at 1.2 GMACs, trained from scratch for 600 epochs.
  arXiv  Detail & Related papers  (2026-05-21T16:20:18Z)
• ECO: Quantized Training without Full-Precision Master Weights [58.97082407934466]
  Error-Compensating (ECO) eliminates master weights by applying updates directly to quantized parameters.<n>We show that ECO converges to a constant-radius neighborhood of the optimum, while naive master-weight removal can incur an error that is inversely proportional to the learning rate.
  arXiv  Detail & Related papers  (2026-01-29T18:35:01Z)
• Uncertainty Makes It Stable: Curiosity-Driven Quantized Mixture-of-Experts [6.221156050218661]
  We present a curiosity-driven quantized Mixture-of-Experts framework for deep neural networks on resource-constrained devices.<n>Our 4-bit quantization maintains 99.9 percent of 16-bit accuracy (0.858 vs 0.859 F1) with 4x compression and 41 percent energy savings.<n>Our information-theoretic routing demonstrates that adaptive quantization yields accurate (0.858 F1, 1.2M params), energy-efficient (3.87 F1/mJ), and predictable edge models.
  arXiv  Detail & Related papers  (2025-11-13T15:32:41Z)
• Fantastic Pretraining Optimizers and Where to Find Them [59.56075036649332]
  AdamW has long been the dominant gradients in language model pretraining.<n>Speedup of matrix-based matrices is inversely proportional to model scale.
  arXiv  Detail & Related papers  (2025-09-02T07:43:22Z)
• Foxtsage vs. Adam: Revolution or Evolution in Optimization? [4.939986309170004]
  This study introduces Foxtsage, a novel hybrid optimisation approach that integrates the Hybrid FOX-TSA with Gradient Descent for training Multi-Layer Perceptron models.<n> Experimental results demonstrate that Foxtsage achieves a 42.03% reduction in loss mean (Foxtsage: 9.508, Adam: 16.402) and a 42.19% improvement in loss standard deviation (Foxtsage: 20.86, Adam: 36.085)
  arXiv  Detail & Related papers  (2024-12-20T03:20:58Z)
• 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)
• Adapprox: Adaptive Approximation in Adam Optimization via Randomized Low-Rank Matrices [24.319712013824876]
  Adapprox is a novel approach that employs randomized low-rank matrix approximation for a more accurate approximation of Adam's second moment. In GPT-2 training and downstream tasks, Adapprox surpasses AdamW by achieving 34.5% to 49.9% memory savings.
  arXiv  Detail & Related papers  (2024-03-22T05:23:31Z)
• GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection [133.45193150403537]
  Training Large Language Models (LLMs) presents significant memory challenges due to the growing size of weights and GPU states. In this work, we propose Gradient Low-Rank Projection (GaLore) as a memory-efficient training strategy. Our 8-bit GaLore further reduces memory by up to 82.5% and total training memory by 63.3%, compared to a BF16 baseline.
  arXiv  Detail & Related papers  (2024-03-06T07:29:57Z)
• Alexa Teacher Model: Pretraining and Distilling Multi-Billion-Parameter Encoders for Natural Language Understanding Systems [63.713297451300086]
  We present results from a large-scale experiment on pretraining encoders with non-embedding parameter counts ranging from 700M to 9.3B. Their subsequent distillation into smaller models ranging from 17M-170M parameters, and their application to the Natural Language Understanding (NLU) component of a virtual assistant system.
  arXiv  Detail & Related papers  (2022-06-15T20:44:23Z)
• Non-Parametric Adaptive Network Pruning [125.4414216272874]
  We introduce non-parametric modeling to simplify the algorithm design. Inspired by the face recognition community, we use a message passing algorithm to obtain an adaptive number of exemplars. EPruner breaks the dependency on the training data in determining the "important" filters.
  arXiv  Detail & Related papers  (2021-01-20T06:18:38Z)

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.