AxMoE: System-Level Realism in MoE Research

• AxMoE is a cluster of mixture-of-experts approaches that address training efficiency, serving decoupling, and arithmetic approximations for real-world deployment.
• The design employs methodologies like step-time measurement, 3D sharding, and adaptive batching to balance compute, communication, and accuracy in model training and inference.
• Empirical studies across varied contexts demonstrate that AxMoE’s innovations significantly improve throughput and power efficiency compared to traditional dense models.

Searching arXiv for papers using the term "AxMoE" and closely related usage. AxMoE is a label that has been used in multiple, technically distinct contexts in mixture-of-experts research. In one usage, AxMoE denotes Apple’s Mixture-of-Experts LLM training recipe and implementation in the AxLearn framework, introduced to revisit MoE–dense speed–accuracy comparisons under step-time budgeting and Chinchilla-style compute-optimal training (Du et al., 2024). In a second usage, AxMoE is the title of a study on approximate multiplication in MoE DNN architectures, spanning Hard MoE, Soft MoE, and Cluster MoE under approximate-aware retraining and normalized inference-power analysis (Shende et al., 6 May 2026). A third, informal usage appears in asynchronous MoE serving: “AxMoE” is used as shorthand for asynchronous MoE serving, while AMoE is the concrete serving system that instantiates Asynchronous Expert Parallelism (AEP) through per-layer $\mu$-queuing, adaptive re-batching, a defragging scheduler, and two-phase P2P transport (Wang et al., 13 May 2025). This multiplicity of meanings makes AxMoE less a single method than a cluster of MoE research directions centered on training efficiency, serving efficiency, and arithmetic-efficient inference.

1. AxMoE as a training recipe for fair MoE–dense comparison

AxMoE is described as Apple’s Mixture-of-Experts LLM training recipe and implementation built in the AxLearn framework, designed expressly to revisit the speed–accuracy comparison between sparse MoE models and dense transformers under modern LLM training practice (Du et al., 2024). Its core goal is to put MoE–dense comparisons on a fair footing by measuring model complexity using step time rather than FLOPs or activated parameters, and by determining the total training compute under the Chinchilla compute-optimal settings.

The motivating critique is specific. Common practice had been to compare MoE to dense models at equal FLOPs or equal activated parameters and to train all models for the same number of tokens. For MoE layers, both measures miss the potentially large communication overhead, including all-to-all dispatch/combine and all-reduce, so MoE often enjoys more actual wall-clock training budget than dense under these comparisons (Du et al., 2024). AxMoE replaces FLOPs with step time as the complexity measure because wall-clock time per training step captures both compute and communication, including synchronization and interconnect effects.

The compute-optimal framing follows dense-optimized Chinchilla settings. Dense models are trained under the recommendation $T:N = 20:1$ tokens-to-parameters, with total training compute budget expressed as

$C_{\text{total}} = t_{\text{step}} \times S.$

AxMoE holds hardware and batch size fixed, measures $t_{\text{step}}$ for the dense baseline, and chooses MoE designs so that $$t_{\text{step}}(\text{MoE}) \times S(\text{MoE}) \approx C_{\text{total}}(\text{dense})$$ while using the same batch size and hardware (Du et al., 2024). This permits MoE systems to trade parameter count for tokens within an equal wall-clock budget.

A plausible implication is that AxMoE’s central intervention is methodological as much as architectural: it reframes the MoE–dense comparison around realized system cost rather than nominal sparsity measures.

2. Architectural design and sharding strategy in Apple’s AxMoE

The AxMoE architecture is based on a LLaMA2-style transformer with pre-normalization and RMSNorm, SwiGLU activations, and rotary positional embeddings; grouped-query attention is not used (Du et al., 2024). Sparse FFNs are inserted in an “Every-4” pattern, replacing the last FFN out of each group of four transformer layers with an MoE FFN. This placement is reported to balance added capacity with step-time impact and to give the best speed–accuracy trade-off.

Routing is token-level Top-K routing with $K=2$. The gate is

$p_i(x) = \operatorname{softmax}(W_g x)_i,$

and the top-2 experts are selected for each token; the second expert is chosen naively, which the paper found to work as well as GShard’s random tie-break (Du et al., 2024). The capacity factor is fixed at $C = 2.0$ for both training and evaluation. Load balancing is encouraged through a GShard-style auxiliary loss with coefficient $0.01$, and router stabilization uses a z-loss with coefficient $0.001$.

A defining systems contribution is 3D sharding over Data, Expert, and Model axes, using GShard’s MoE layer implementation atop GSPMD, Praxis’s Einsum-based routing/dispatch/combine, and Mesh-TensorFlow’s “outer batch” trick (Du et al., 2024). The sharding policy is carefully partitioned: the Data axis is pure data parallel, the Expert axis shards experts across cores with a design goal of at most one expert per core, and the Model axis shards attention heads and FFN hidden dimensions. MoE dispatch and combine use all-to-all collectives across the Expert-parallel group, while dense model-parallel shards use all-reduce and all-gather on the Model axis.

The step time is decomposed as

$T:N = 20:1$0

The design objective is to keep the extra MoE communication overhead modest. Empirically, AxMoE keeps dense-to-MoE step-time increases within a healthy range, often at or below $T:N = 20:1$1 across scales, and nearly flat while scaling experts up to roughly 256 experts (Du et al., 2024).

Aspect	Apple AxMoE specification
Base transformer	LLaMA2-style, pre-normalization, RMSNorm, SwiGLU, RoPE
Sparse-layer placement	Every-4 sparse layers
Routing	Top-2 token routing
Capacity factor	$T:N = 20:1$2
Load balancing	Auxiliary loss coefficient $T:N = 20:1$3
Router stabilization	z-loss coefficient $T:N = 20:1$4
Sharding axes	Data, Expert, Model

These choices are tied directly to the paper’s fairness argument. The 3D sharding strategy is not presented as an incidental implementation detail, but as the mechanism that prevents communication costs from invalidating the wall-clock comparison.

3. Empirical results of the training-oriented AxMoE

AxMoE is evaluated across three dense baselines—$T:N = 20:1$5, $T:N = 20:1$6, and $T:N = 20:1$7—and corresponding MoE alternatives trained under matched wall-clock budgets on TPU-v4 and TPU-v5e systems (Du et al., 2024). The training corpus is described as a mixture similar to LLaMA2, and evaluation spans Core English tasks, MMLU 5-shot, and GSM8K 8-shot.

The reported outcome is that MoE dominates dense on speed–accuracy at all three scales. At the $T:N = 20:1$8 budget on 512 TPU-v4 cores, dense $T:N = 20:1$9 reaches $C_{\text{total}} = t_{\text{step}} \times S.$0 CoreEN (0S) with $C_{\text{total}} = t_{\text{step}} \times S.$1, while $C_{\text{total}} = t_{\text{step}} \times S.$2 reaches $C_{\text{total}} = t_{\text{step}} \times S.$3 with $C_{\text{total}} = t_{\text{step}} \times S.$4 and $C_{\text{total}} = t_{\text{step}} \times S.$5 reaches $C_{\text{total}} = t_{\text{step}} \times S.$6 with $C_{\text{total}} = t_{\text{step}} \times S.$7 (Du et al., 2024). At the $C_{\text{total}} = t_{\text{step}} \times S.$8 budget on 1024 TPU-v4 cores, dense $C_{\text{total}} = t_{\text{step}} \times S.$9 achieves $t_{\text{step}}$0 CoreEN (all), MMLU $t_{\text{step}}$1, GSM8K $t_{\text{step}}$2, with $t_{\text{step}}$3, while $t_{\text{step}}$4 improves to $t_{\text{step}}$5, MMLU $t_{\text{step}}$6, GSM8K $t_{\text{step}}$7, with $t_{\text{step}}$8 (Du et al., 2024). At the $t_{\text{step}}$9 budget, dense $$t_{\text{step}}(\text{MoE}) \times S(\text{MoE}) \approx C_{\text{total}}(\text{dense})$$0 obtains $$t_{\text{step}}(\text{MoE}) \times S(\text{MoE}) \approx C_{\text{total}}(\text{dense})$$1 CoreEN (all), MMLU $$t_{\text{step}}(\text{MoE}) \times S(\text{MoE}) \approx C_{\text{total}}(\text{dense})$$2, GSM8K $$t_{\text{step}}(\text{MoE}) \times S(\text{MoE}) \approx C_{\text{total}}(\text{dense})$$3, whereas $$t_{\text{step}}(\text{MoE}) \times S(\text{MoE}) \approx C_{\text{total}}(\text{dense})$$4 reaches $$t_{\text{step}}(\text{MoE}) \times S(\text{MoE}) \approx C_{\text{total}}(\text{dense})$$5, MMLU $$t_{\text{step}}(\text{MoE}) \times S(\text{MoE}) \approx C_{\text{total}}(\text{dense})$$6, GSM8K $$t_{\text{step}}(\text{MoE}) \times S(\text{MoE}) \approx C_{\text{total}}(\text{dense})$$7, at $$t_{\text{step}}(\text{MoE}) \times S(\text{MoE}) \approx C_{\text{total}}(\text{dense})$$8 (Du et al., 2024).

The sharding study is equally central. For $$t_{\text{step}}(\text{MoE}) \times S(\text{MoE}) \approx C_{\text{total}}(\text{dense})$$9 on 256 TPU-v5e cores, naive 2D sharding yields $K=2$0 total parameters and $K=2$1, 2D sharding with padding yields $K=2$2 parameters and $K=2$3, while AxMoE 3D sharding retains $K=2$4 parameters and reaches $K=2$5 (Du et al., 2024). This result is used to argue that 3D sharding attains the speed of padded 2D without parameter inflation.

A plausible implication is that AxMoE’s reported advantage depends on the co-design of architecture, routing, and sharding; the fairness claim would weaken if MoE were evaluated with suboptimal partitioning.

4. AxMoE in asynchronous MoE serving: the AMoE and AEP connection

A different usage of the term appears in serving. In the asynchronous serving paper, “AxMoE” is used informally to denote asynchronous MoE serving, while the technical approach is Asynchronous Expert Parallelism (AEP) and the implemented system is AMoE (Wang et al., 13 May 2025). AMoE operationalizes AEP through one runtime per GPU, a central coordinator, per-layer $K=2$6-queues, adaptive re-batching, a defragging scheduler, and a two-phase P2P communication substrate.

The targeted problems are expert load skew, GPU idling under barriers, and small-batch inefficiency (Wang et al., 13 May 2025). Traditional synchronous expert parallelism collects all tokens for a layer, executes experts in parallel, performs all-to-all communication, and only then allows the model to advance. Under skewed expert loads, the hottest expert becomes a straggler, forcing all GPUs to idle and leaving cold experts with small, memory-bandwidth-limited microbatches.

AEP removes barrier-style synchronization around expert layers. Each runtime maintains per-layer $K=2$7-queues holding token metadata, while activations remain on GPU. When a GPU becomes available, the scheduler selects a layer to run and forms a contiguous batch by draining that layer’s queue. For top-$K=2$8, duplicated expert outputs are later rejoined through a rendezvous keyed by $K=2$9 before enqueuing the merged token to the next layer (Wang et al., 13 May 2025). Attention and expert layers are disaggregated, with attention on its own data-parallel ranks and experts on expert-parallel ranks.

The queueing-theoretic view is explicit. For expert $p_i(x) = \operatorname{softmax}(W_g x)_i,$0, with arrival rate $p_i(x) = \operatorname{softmax}(W_g x)_i,$1, service rate $p_i(x) = \operatorname{softmax}(W_g x)_i,$2, and utilization $p_i(x) = \operatorname{softmax}(W_g x)_i,$3, an $p_i(x) = \operatorname{softmax}(W_g x)_i,$4 approximation gives

$p_i(x) = \operatorname{softmax}(W_g x)_i,$5

with per-expert latency approximated as

$p_i(x) = \operatorname{softmax}(W_g x)_i,$6

AEP is described as reducing $p_i(x) = \operatorname{softmax}(W_g x)_i,$7-induced stalls by decoupling layers, increasing effective $p_i(x) = \operatorname{softmax}(W_g x)_i,$8 through adaptive batching, and avoiding synchronized all-to-all paths dominated by the hottest expert’s completion (Wang et al., 13 May 2025).

AMoE also replaces synchronous all-to-all with a two-phase protocol: control-plane metadata via ZeroMQ on CPU and data-plane transfer via NCCL P2P on GPU. The scheduler scores each layer by local queue occupancy plus a decayed lookahead over subsequent blocks, with time complexity per decision of $p_i(x) = \operatorname{softmax}(W_g x)_i,$9 (Wang et al., 13 May 2025).

5. Performance characteristics of asynchronous AxMoE/AMoE serving

AMoE is evaluated against SGLang as a baseline expert-parallel system, on both single-node and multi-node A100 platforms, focusing exclusively on decoding after removing prefill in both systems (Wang et al., 13 May 2025). The base model is Mixtral $C = 2.0$0 with 32 decoding blocks and 8 experts per block, using Top-1 and Top-2 routing with exponential routing skew induced from Dolly traces.

The reported single-node result is that AMoE achieves up to $C = 2.0$1 higher throughput than SGLang under Top-1 routing, specifically $C = 2.0$2 on short workloads, $C = 2.0$3 on medium, and $C = 2.0$4 on reasonable workloads (Wang et al., 13 May 2025). Under Top-2, gains persist but are smaller because the merge acts as a partial barrier. In a scaled multi-node configuration with 16 experts across 16 GPUs, AMoE scales nearly linearly, whereas SGLang shows no throughput improvement when doubling GPUs under skew (Wang et al., 13 May 2025).

The throughput contrast is summarized with idealized models. For synchronous expert parallelism,

$C = 2.0$5

where $C = 2.0$6 is the end-to-end time for expert $C = 2.0$7 in a step. For AEP,

$C = 2.0$8

These formulas are presented as idealized but intended to capture the shift from straggler-limited progress to the sum of independent services (Wang et al., 13 May 2025).

The paper attributes the gain to two mechanisms: adaptive re-batching increases expert arithmetic intensity, and barrier removal prevents hot experts from stalling cold ones. On A100, expert-layer throughput is observed to increase almost linearly up to batch size $C = 2.0$9, which makes re-batching especially important for cold experts (Wang et al., 13 May 2025). The defragging scheduler is reported to outperform MTFS and avoid FLFS livelock.

This suggests that the serving-oriented AxMoE usage is primarily about runtime decoupling and scheduling under skew, rather than about model architecture or training-time fairness.

6. AxMoE as a study of approximate multipliers in MoE DNNs

The 2026 paper uses AxMoE as the title for the first study of the impact of approximate multiplication on MoE DNN architectures (Shende et al., 6 May 2026). Its focus is edge inference, where approximate computing and MoE are treated as orthogonal efficiency mechanisms whose interaction had not previously been studied. The study evaluates Hard MoE, Soft MoE, and Cluster MoE against dense baselines across ResNet-20, VGG11_bn, VGG19_bn, and ViT-Small, using eight 8-bit signed multipliers from EvoApproxLib, including one exact baseline.

The formal definitions are explicit. With input $0.01$0, experts $0.01$1, and gate

$0.01$2

Soft MoE computes

$0.01$3

and Hard MoE computes

$0.01$4

Cluster MoE instead routes the entire input with a standalone gateway:

$0.01$5

The gate adds negligible MAC overhead and parameters per MoE layer, while Cluster MoE incurs gateway overhead of approximately $0.01$6 MACs for CNNs and approximately $0.01$7 MACs for ViT-Small (Shende et al., 6 May 2026).

Approximate multiplication is confined to expert computation. For CNNs it is applied to Conv2d layers, while Linear and BatchNorm remain exact; for ViT-Small it is applied to all Linear layers in transformer blocks, with patch embedding and classifier head exact (Shende et al., 6 May 2026). The study distinguishes moderate approximation regimes—KV8, KV9, KVM, KVP, L2J—from aggressive regimes KVA and L2L.

Normalized inference power is defined as

$0.01$8

Here, $0.01$9 is the effective MACs of the dense exact-multiplier reference, and $0.001$0 is the approximate MAC fraction (Shende et al., 6 May 2026). This metric combines routing-induced compute overhead with multiplier-level per-operation power savings.

Approximate-aware retraining lasts five epochs with SGD, freezes the gates, and updates only expert weights (Shende et al., 6 May 2026). LUT-based approximate multiplication incurs roughly $0.001$1–$0.001$2 wall-clock overhead per epoch relative to standard training.

7. Comparative interpretation and recurring themes

The approximate-multiplier AxMoE paper reports architecture-dependent robustness patterns. Without retraining, Dense is the most resilient topology across all CNN architectures, whereas on ViT-Small all topologies degrade at nearly identical rates because self-attention projections are approximated identically across variants (Shende et al., 6 May 2026). After retraining, ResNet-20 fully recovers across the entire multiplier range; VGG architectures recover at moderate multipliers but fail irreversibly at aggressive ones for all topologies except Cluster MoE on VGG11_bn; and on ViT-Small, Hard MoE outperforms Dense under aggressive approximation at equal normalized inference cost (Shende et al., 6 May 2026). Dense+L2J is reported as Pareto-optimal across CNNs, while Hard MoE with $0.001$3 is Pareto-optimal on ViT-Small under aggressive multipliers.

Across the three usages of AxMoE, a common theme is that nominal sparsity metrics are insufficient to characterize real performance. Apple’s AxMoE argues that FLOPs and activated parameters hide communication overhead and therefore advocates step time as the relevant complexity measure (Du et al., 2024). Asynchronous AxMoE/AMoE argues that synchronous expert-parallel abstractions hide straggler and batching pathologies, making queueing and scheduling central to throughput (Wang et al., 13 May 2025). The approximate-multiplier AxMoE paper argues that topology alone does not determine efficiency, since multiplier characteristics, approximate MAC fraction, and retraining recovery jointly shape the accuracy–power frontier (Shende et al., 6 May 2026).

A common misconception would be to treat AxMoE as a single standardized method. The published record instead shows three different uses: a training recipe and implementation in AxLearn, an informal shorthand for asynchronous MoE serving instantiated by AMoE, and a framework for studying approximate arithmetic in MoE DNNs. What unifies them is not a shared codebase or model family, but a shared concern with system-level realism in MoE evaluation: wall-clock training cost, barrier-free serving efficiency, and end-to-end accuracy–power behavior under approximate hardware.