Research map / Architectures and efficiency / Mixture of experts
Mozart: Modularized and Efficient MoE Training on 3.5D Wafer-Scale Chiplet Architectures
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Mozart, an algorithm-hardware co-design that maps MoE-LLM post-training onto a 3.5D wafer-scale chiplet system using expert clustering/allocation, fine-grained streaming scheduling, and a 2.5D NoP-Tree interconnect with hierarchical memory.
Problem
Efficient post-training of sparse Mixture-of-Experts LLMs on chiplet hardware suffers from memory locality issues, all-to-all communication overhead, and poor compute utilization due to dynamic, uneven expert workloads.
Approach
Co-locate frequently co-activated experts on the same or adjacent chiplets (balanced by profiled workload) to reduce all-to-all replication, and stream experts and tokens to overlap DRAM weight loading with on-chip compute, on a 3D logic-on-memory, 2.5D NoP-Tree wafer-scale architecture.
Key results
- Mozart-C achieves 1.92x speedup on Qwen3-30B-A3B-Base, 2.37x on OLMoE-1B-7B-0924, and 2.17x on DeepSeek-MoE-16B-Base versus the unoptimized baseline (simulation).
- Communication-computation overlap is reported as the most critical optimization, giving 1.33x to 1.58x speedup alone on the baseline.
- Sequence length 512 speedup reaches 2.34x over baseline, while DRAM bandwidth (HBM2 vs SSD) changes the gains, with HBM2 enabling larger speedups.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add communication-computation overlap (ablation) | Baseline MoE training on chiplets · Qwen3-30B-A3B-Base | speedup: 1.33× acceleration | Helped |
| Add communication-computation overlap (ablation) | Baseline MoE training on chiplets · OLMoE-1B-7B-0924 | speedup: 1.58× | Helped |
| Add communication-computation overlap (ablation) | Baseline MoE training on chiplets · DeepSeek-MoE-16B-Base | speedup: 1.49× | Helped |
| Add efficient all-to-all communication (ablation) | Mozart-A · Qwen3-30B-A3B-Base | normalized latency: 0.73 -> 0.59 | Helped |
| Add efficient all-to-all communication (ablation) | Mozart-A · OLMoE-1B-7B-0924 | normalized latency: 0.63 -> 0.48 | Helped |
| Add efficient all-to-all communication (ablation) | Mozart-A · DeepSeek-MoE-16B-Base | normalized latency: 0.67 -> 0.56 | Helped |
| Add specialized expert layout on chiplets (ablation) | Mozart-B · Qwen3-30B-A3B-Base | normalized latency: 0.59 -> 0.52 | Helped |
| Add specialized expert layout on chiplets (ablation) | Mozart-B · OLMoE-1B-7B-0924 | normalized latency: 0.48 -> 0.422 | Helped |
| Add specialized expert layout on chiplets (ablation) | Mozart-B · DeepSeek-MoE-16B-Base | normalized latency: 0.56 -> 0.46 | Helped |
| Use Mozart-C instead of the unoptimized baseline (ablation) | MoE training on chiplets · sequence length 512 | speedup: 2.34× | Helped |
| Use Mozart-C instead of the unoptimized baseline (ablation) | MoE training on chiplets · sequence length 128 | speedup: 1.47× | Helped |
Surprising findings
- Training under the proposed architecture is memory-bound rather than compute-bound, limited by sequential MoE weight loading.
- Under SSD-level DRAM bandwidth, expert weight streaming dominates latency and the co-design gains are constrained.
Limitations
- Attention is assigned to a single chiplet, which may limit latency due to limited resources.
- Switches can become performance bottlenecks under high communication demand.
- Results come from cycle-accurate simulation and Verilog synthesis rather than physical hardware.
Experimental setup
Datasets: Alpaca
Benchmarks: Alpaca
Metrics: end-to-end latency energy normalized latency all-to-all communication complexity C_T
Baselines: Mozart baseline without optimizations Mozart-A Mozart-B MegaBlocks
Builds on: expert parallelism farthest point sampling FRED Tutel
Closest papers by meaning
- Oracle-MoE: Locality-preserving Routing in the Oracle Space for Memory-constrained Large Language Model Inference (ICML 2025 · similarity 0.86)
- MoE-CAP: Benchmarking Cost, Accuracy and Performance of Sparse Mixture-of-Experts Systems (NeurIPS 2025 · similarity 0.86)
- BigMac: A Communication-Efficient Mixture-of-Experts Model Structure for Fast Training and Inference (AAAI 2025 · similarity 0.86)
- FlowMoE: A Scalable Pipeline Scheduling Framework for Distributed Mixture-of-Experts Training (NeurIPS 2025 · similarity 0.86)
- MoE++: Accelerating Mixture-of-Experts Methods with Zero-Computation Experts (ICLR 2025 · similarity 0.85)
- Mixture Compressor for Mixture-of-Experts LLMs Gains More (ICLR 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.