PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Nuemaan Malik

arXiv:2607.19058v2cs.LGcs.AI

TL;DR

MoE training spends substantial memory on optimizer state because sparse activation does not reduce state per parameter. SkewAdam assigns tiered state to the backbone, experts, and router, reducing memory while matching or exceeding baseline perplexity.

  • Problem

    MoE optimizer state remains proportional to total parameters despite sparse activation, creating a major memory bottleneck for training.

  • Method

    SkewAdam assigns momentum and factored or exact second-moment state by parameter role, using fuller state for dense or sensitive populations and less for sparse experts.

  • Results

    1.29 GB of optimizer state, 2.6% of AdamW’s, reduces peak memory to 31.3 GB while achieving 108.4 validation perplexity and near-uniform routing.

  • Takeaways & Limitations

    SkewAdam’s tiered policy provides Adam-family quality at 2.6% of AdamW’s state, with the contribution attributed to memory rather than improved optimization.

  • Takeaways & Limitations

    The study uses a two-block MoE with mostly single-run configurations and does not validate the recipe over production-scale horizons with active weight decay.

Abstract

from arXiv · show

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 differ enough in size and gradient statistics that they should not receive the same state. Those populations are the dense backbone, the experts and the router. 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 or 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 reaches the same value while carrying twenty times the state, so the tiers buy memory rather than accuracy. Same-platform runs separate what does earn it. Removing momentum costs 31 perplexity points (tuned Adafactor, 139.7) and replacing the factored second moment and its update clipping with a full second moment costs 10 (tuned AdamW, 118.5), so neither tuned baseline reaches the untuned tiered policy. Where optimizer state lives, these results suggest, matters at least as much as how much of it there is.

1 Introduction

SkewAdam treats MoE parameter populations as distinct memory-allocation tiers rather than assigning uniform optimizer state. This yields 1.29 GB of optimizer state and 31.3 GB peak training memory, while achieving 108.4 validation perplexity in a controlled comparison.

  • Memory motivation: 1.29 GB of optimizer state is 2.6% of AdamW’s 50.55 GB on the same model.GB denotes GiB (2^30 bytes) throughout.
  • Tiered allocation: 2 MB of router state preserves an exact second moment for the gate, whose routing decisions depend on relative gradient magnitudes.The tier design keeps exact second-moment information where routing-sensitive statistics matter most.
  • Tiered allocation: SkewAdam assigns float32 momentum and factored second moments to the 5% dense backbone, but only factored second moments to the 95% experts.Experts process roughly 1/64 of tokens each under top-2 routing over 128 experts, making their factored-only state the largest savings.
  • Memory results: 31.3 GB of peak training memory follows from the tiered allocation policy and its closed-form memory model.The resulting footprint places training within the stated 40 GB accelerator budget.
  • Controlled comparison: 108.4 validation perplexity is achieved by SkewAdam in a controlled single-GPU comparison with identical initialization, data order, and bfloat16 update paths.Throughput remains within 1.5% of the fastest baseline.

2 Related work

Prior work provides memory-saving estimators and alternative update geometries, but applies them uniformly, leaving optimizer-state allocation across MoE parameter populations open. The paper also situates its method within MoE routing and low-precision training practices.

  • Memory-efficient optimizers: Adafactor supplies SkewAdam’s factored estimator and update-RMS clipping, but its uniform momentum and matrix-factorization recipes leave allocation across parameter populations open.The paper claims no novelty for either component and instead studies where optimizer state should live.
  • Sign- and geometry-based updates: Lion uses one sign-based momentum buffer, while Muon applies Newton–Schulz orthogonalization and routes non-matrix parameters to an Adam-style rule.Both methods remain uniform over hidden matrices, regardless of whether they belong to experts.
  • MoE training: MoE research has focused mainly on architectures and losses, including load balancing and router z-losses, rather than optimizer-state allocation.The cited systems span sparsely gated MoE, GShard, Switch Transformers, ST-MoE, Mixtral, and DeepSeekMoE.
  • Low-precision training: The study uses bfloat16 master weights with dithered stochastic rounding applied identically across every compared optimizer.This choice is motivated by the memory appeal of pure-bfloat16 training while addressing small-update loss from rounding.

3 SkewAdam: tiered state allocation

SkewAdam allocates optimizer state by parameter role, using richer state for the dense backbone and cheaper factored or exact second moments elsewhere. Factoring reduces matrix-state storage structurally, while clipped float32 updates and low-precision writes complete the implementation, with weight decay effectively inactive.

  • Factored second moments: 32 KB rather than 64 MB stores the second moment for a 4096 × 4096 expert matrix under factored estimation.Storage falls from nm to n + m floats per matrix.
  • The tiers: 1.29 GB of state results from assigning float32 momentum plus factored second moments to the backbone, factored second moments to experts, and full moments to routers.The backbone is 5% of parameters and experts are 95%; the router path uses a full second moment.
  • Update clipping and low-precision updates: Unit-root-mean-square update clipping and float32 computation with dithered bfloat16 writes limit update magnitude while approximating unbiased stochastic rounding.The clipping threshold is 1, and uniform one-ULP noise is added before casting.
  • Limitation: 1.5 × 10−5 relative weight decay is more than two orders of magnitude below the bfloat16 ULP of 2−7, making decay a no-op in all runs.This occurs with η = 3 × 10−4 and λ = 0.05, leaving comparisons effectively unregularized.

4 Experimental setup

The experiments use a 6,784M-parameter decoder-only MoE transformer with 128 experts and top-2 routing, trained for 10,000 steps on a document-disjoint OpenWebText split. All optimizers share initialization, batches, precision, and validation batches, while the study also measures implementation limits affecting quantized state.

  • Model: 6,784M parameters comprise a two-block decoder-only transformer with a dense SwiGLU block followed by a 128-expert top-2-routed MoE block.The model has width 4096 throughout, with about 440M active parameters per token.
  • Data and protocol: 10,000 steps use batch size 64 × 128 tokens, yielding 81.9M tokens in one epoch without batch repeats.OpenWebText documents are hashed into a 95/5 train/validation split with no shared documents.
  • Data and protocol: 64 held-out validation batches, approximately 0.5M tokens, are reused for every optimizer.All four optimizers start from one shared initialization and consume identical batches in identical order under bfloat16 autocast.
  • Implementation constraints: 231 elements is the measured failure boundary for current 8-bit optimizer kernels, while factored state crosses it using native PyTorch operations with 64-bit indexing.On an A100, the 8-bit step succeeds at 231 −1 elements and kills the process at 231.
  • Evaluation overview: 14.5% below AdamW is SkewAdam’s final convergence margin, and only SkewAdam clears the 40 GB memory line.AdamW and Muon lead through step 3,000 before SkewAdam overtakes both by step 4,000.

5 Results

SkewAdam reduces optimizer-state and peak-memory costs while achieving the best reported validation perplexity, despite tier ablations showing that the allocation primarily saves memory rather than accuracy. Its advantage over tuned baselines is tied to retaining momentum and using factored second-moment updates with clipping, while routing remains near uniformly balanced.

  • Memory and throughput: 1.29 GB of optimizer state lowers SkewAdam’s peak memory to 31.3 GB, below the 40 GB accelerator budget, versus AdamW’s 81.4 GB peak.SkewAdam sustains 5,000 tokens/s, 6.6% above AdamW and 1.5% below Lion.
  • Convergence: 108.4 validation perplexity lets SkewAdam finish ahead of Muon at 120.2, AdamW at 126.8, and Lion at 393.7 after initially slower convergence.SkewAdam passes AdamW and Muon by step 4,000.
  • Load balancing: Within 1% of the uniform-routing floor, SkewAdam and AdamW maintain balanced routing from step 4,000 onward, while Muon ends 22% above it.The floor is α = 0.05; final values are 0.0505 for SkewAdam and 0.0502 for AdamW.
  • Which tier does the work?: 108.2–108.9 validation perplexity and about 0.050 load balance remain flat across tier variants, while optimizer state changes twentyfold.Restoring expert momentum costs 24 GB and changes perplexity by 0.2; factoring the router changes neither perplexity nor balance.
  • Tuning the baselines: 118.5 ± 0.5 for tuned AdamW and 139.7 for tuned Adafactor still exceed SkewAdam’s 108.4, whose learning rate remains untuned.AdamW improves from 126.8 and Adafactor from 149.5 when tuned to 10−4.

6 Limitations

The evaluation uses a shallow two-block MoE, mostly one expert bank, short training runs, and limited tuning, so deeper routing behavior and production-horizon effects remain untested. Weight decay was inert in all runs and must be reintroduced for production scaling.

  • Architectural scope: 95% of parameters occupy one expert bank in the two-block model, leaving tiered allocation across many composing MoE layers untested.The architecture was chosen as a stress test but does not establish behavior for deeper routing compositions.
  • Evaluation scope: Most configurations use one run, although key Table 4 sweeps used repeated seeds and SkewAdam was replicated on three GPUs.SkewAdam itself was not tuned, AdamW was not probed below 10−4, and Lion and Muon retained single untuned rates.
  • Evaluation scope: 82M-token training horizons and 128-token contexts are small, limiting evidence about longer-horizon behavior.The paper notes that production-scale use would require evaluation beyond these settings.
  • Production scaling: Weight decay was inert in all runs and must be reintroduced for production scaling.The decay term should be fused into the float32 update before the stochastically rounded bfloat16 cast so it survives the update.

7 Conclusion

SkewAdam allocates optimizer state by MoE parameter population, reducing memory substantially while preserving the perplexity and routing balance of a uniform optimizer with twenty times more state. The contribution is memory efficiency rather than improved optimizer quality.

  • Memory-efficient allocation: 97.4% lower optimizer state and 61% lower peak training memory result from dense-backbone Adam state, factored expert variance, and an exact router second moment.These reductions are reported for a 6.78B-parameter MoE.
  • Ablation: The tier ablation reaches the same perplexity and routing balance as a uniform optimizer carrying twenty times the state.The additional state therefore does not improve these outcomes in the reported comparison.
  • Contribution: 2.6% of Adam’s state delivers Adam-family quality, making memory reduction—not a better optimizer—the contribution.The conclusion attributes the gain to where optimizer state is allocated.

A Hyperparameters · B Memory accounting

The configuration specifies the released training-script settings for Lion, Muon, GaLore-style updates, Adam, and weight decay. Memory accounting gives the tiered optimizer-state footprint and a 6,784.3M-parameter model breakdown across backbone, experts, and router.

  • A Hyperparameters: The released training script provides the complete optimizer configuration.Table 5 identifies the configuration as complete and notes that all values are used in the released training script.
  • A Hyperparameters: Lion uses η = 1 × 10−4 with β = (0.9, 0.99).
  • A Hyperparameters: Muon uses η = 0.02, momentum 0.95, and 3 NS steps, with internal Adam η = 10−3.
  • A Hyperparameters: Rank 128, η = 3 × 10−4, β = (0.9, 0.999), and weight decay 0.05 define the GaLore-style and full-Adam settings.Weight decay is 0 on the router; the GaLore-style update is inert under bf16 rounding.
  • B Memory accounting: The dense-backbone momentum buffer is, to first order, the entire SkewAdam optimizer-state footprint.
  • B Memory accounting: 6,784.3M total parameters comprise a 341.4M backbone, 6,442.5M experts, and a 0.52M router.The backbone includes token embeddings, positions, attention, dense FFN, and norms; the embedding is tied with the output head.

C Convergence detail

Convergence comparisons show optimizer leadership changing over training: Muon leads early, while SkewAdam leads later. In the H100 follow-up, Adafactor leads initially before SkewAdam takes over.

  • Optimizer convergence: Muon leads validation perplexity for the first three thousand steps, while SkewAdam leads from step four thousand on.
  • H100 follow-up: In the H100 follow-up, Adafactor leads validation perplexity for the first two thousand steps and SkewAdam thereafter.
  • Training dynamics: Training-loss traces are single-batch measurements at evaluation steps, producing visible variance alongside the reported throughput.

D Zero-shot evaluation detail

Zero-shot evaluation found no meaningful optimizer differences: no pairwise difference exceeded two standard errors on any task. Table 9 reports accuracy-based metrics for PIQA, WinoGrande, HellaSwag, and ARC-Challenge, with task-specific chance levels.

  • Zero-shot evaluation detail: Two standard errors: no pairwise difference between optimizers exceeded this threshold on any task.Table 9 expands the zero-shot summary in Section 5.
  • Zero-shot evaluation detail: 50% chance applies to PIQA and WinoGrande, while 25% applies to HellaSwag and ARC-Challenge.PIQA and WinoGrande use accuracy; HellaSwag and ARC-Challenge use length-normalized accuracy.

E Reproducibility

The reported comparisons use released, reproducible trainers with shared initialization and cached batches for the four main optimizers. Follow-up experiments preserve the validation split across hardware and document their scripts, metrics, and logs.

  • Main comparison: The main comparison ran on a single NVIDIA H200 (141 GB) using the released single-file trainer and evaluator, with code, logs, and figures available online.The environment uses torch, numpy, transformers, bitsandbytes, datasets, and lm_eval.
  • Main comparison: One invocation trained all four optimizers in sequence, while one initialization and cached batch sequence were reseeded and reused for matched conditions.The optimizer sequence was SkewAdam, Adam, Lion, and Muon.
  • Follow-up experiments: The Table 2 follow-up ran on an NVIDIA H100 NVL 47 GB MIG slice under the same protocol, using a canonical parquet mirror with an unchanged document-hash split.Metrics and training logs are kept under runs/h100/.
  • Follow-up experiments: The tier ablation and learning-rate sweeps ran on an MI300X using dedicated experiment scripts, with per-run metrics stored under runs/amd-.The scripts were experiments/tier-ablation/run_ablation.py and experiments/lr-sweep/run_sweep.py / run_fix.py.
Loading 2607.19058v2…