Title: Mozart: Modularized and Efficient MoE Training on 3.5D Wafer-Scale Chiplet Architectures
Venue: NeurIPS 2025
Accepted: yes
Keywords: MoE

Chiplet Architecture

Post-Training Efficiency
Abstract: Mixture-of-Experts (MoE) architecture offers enhanced efficiency for Large Language Models (LLMs) with modularized computation, yet its inherent sparsity poses significant hardware deployment challenges, including memory locality issues, communication overhead, and inefficient computing resource utilization. Inspired by the modular organization of the human brain, we propose $\texttt{Mozart}$, a novel algorithm-hardware co-design framework tailored for efficient training of MoE-based LLMs on 3.5D wafer-scale chiplet architectures. On the algorithm side, $\texttt{Mozart}$ exploits the inherent modularity of chiplets and introduces: 
($1$) an expert allocation strategy that enables efficient on-package all-to-all communication, and ($2$) a fine-grained scheduling mechanism that improves communication-computation overlap through streaming tokens and experts. On the architecture side, $\texttt{Mozart}$ adaptively co-locates heterogeneous modules on specialized chiplets with a 2.5D NoP-Tree topology and hierarchical memory structure.
Evaluation across three popular MoE models demonstrates significant efficiency gains, enabling more effective parallelization and resource utilization for large-scale modularized MoE-LLMs.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nMozart is a novel and comprehensive hardware-software co-design framework for training MoE models on 3.5D wafer-scale chiplet architectures. On the software part, Mozart optimizes the expert placement to reduce inter-chiplet communication, and it uses streaming to improve the communication-compute overlap. On the hardware part, Mozart can place heterogeneous components on each chiplet, and it connects them via a 2.5D NoP-Tree.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\n\n1. This paper focuses on popular topics like MoE, wafer-scale, and 3D architecture.\n\n2. This paper proposes a comprehensive solution for placing the experts based on the hardware architecture.\n\n3. The evaluation shows impressive performance benefits.\n\nWeakness:\n\n1. The compute-communication overlap seems similar to existing works like DualPipe from DeepSeek and TileLink from ByteDance.\n\n2. This paper did not discuss that, as model sizes scale quickly, what will happen if the size of an expert exceeds the capacity of one chiplet group. Specifically, when each expert is mapped to multiple chiplet groups, how will you handle the parallelism?\n\nQUESTIONS:\nOverall, I enjoy reading this paper and have learned a lot!\n\nSee weakness 1, the authors can discuss more on why the proposed overlap scheme is better than existing works for this hardware.\n\nTo support larger models, how is the proposed solution compatible with existing distributed training optimizations like model parallelism?\n\nLIMITATIONS:\nThe paper is great! I can only observe one limitation.\n\nSee weakness 2, as LLMs scale quickly, how does a Mozart hardware manufactured today support future larger models?\nIf we support that by adapting model parallelism on top of your design, will your performance benefits degrade?\nAnd if so, by how much?\nDiscussing these questions would help address the limitations on hardware scalability.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-06-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a 3.5-D hardware architecture that leverages new in-package semiconductor technologies to tightly integrate small chiplets at a large scale, connected together using a hierarchical network on-chip. It is labeled as 3.5-D (instead of 2.5-D) because each chiplet has 3D integration of SRAM (memory) on top of a logic die as well. On the algorithmic side, the paper presents a number of architecture-specific optimizations, specifically to optimize the post-training fine-tuning of mixture-of-expert architectures such as Deepseek, Olmo, and others. The proposed optimizations are: (1) clustering experts and pins them to the suitable chiplet in the system hierarchy to minimize overall communication, and (2) pipelines memory access to hide computation latency.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: the presented architecture is very interesting and combines many new ideas to tackle large-model fine-tuning. Furthermore, there are RTL simulated areas with details down to the pitch of TSVs suggesting a realistic hardware architecture. Although the focus of the paper is on the mapping of MOE model mapping onto the device. On that front, the presented ideas are interesting and logical, but:\n\nWeaknesses:\n1- There are no comparisons, even high-level, rough comparisons to any available hardware. This makes it very difficult to appreciate the presented hardware architecture, its uitility, or efficiency compared to available solutions such as GPUs.\n2- Validation: Barring tape-out, it is conventional in the hardware space to map a small design onto an FPGA or a cycle-accurate simulator to validate the cycle-accurate assumptions made in the paper.\n3- MOE mapping methods are somewhat straightforward. If I assume the novelty of the paper is on the hardware side, I find that this is not the right venue for presentation, if the novelty is in specifically mapping MOE to this new hardware architecture, I find that the presented methods are practically sound and logical but somewhat straightforward.\n\nQUESTIONS:\nThe main missing component of this paper is to compare to any existing hardware solution to enable some understanding of the impact of this work. Are you able to compare, for example, to 8xA100?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces Mozart, a software-hardware co-design framework for the efficient post-training of MOE LLMs. This work is motivated by the hardware challenges posed by the sparse and modular nature of MoE models especially memory locality and communication overhead. The proposed strategy tries to addresses expert routing and communication overheads in MoE-style LLMs, On the algorithmic side, Mozart implements an expert allocation strategy that clusters experts based on their co-activation patterns to minimize inter-chiplet communication and balance workloads. It also implements a fine-grained scheduling mechanism to improve the overlap between communication and computation. On the hardware side, Mozart designs a specialized 3.5D wafer-scale chiplet architecture with a hierarchical memory system and a 2.5D NoP-Tree interconnect topology. Using cycle-accurate simulations, the paper shows that Mozart achieves speedups (over 1.9x) across three popular MoE models compared to baseline designs.\n\nDespite its specialization to MoE-style LLMs and modular hardware assumptions, the paper presents valuable design principles for future accelerators. Its latency-aware expert scheduling, communication modeling, and hardware-algorithm co-design offer practical insights for chiplet-based and sparse-aware systems. These contributions are particularly timely as AI hardware trends shift toward modular packaging, inter-die communication optimization, and efficient support for dynamic, sparse workloads. While generalization beyond MoEs is limited, the paper nonetheless informs a growing class of AI accelerators targeting scalable, fine-grained model execution.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The paper is well written and organized. \n- The authors address a timely topic.  MoE architectures are increasingly being adopted to scale LLMs to trillion-parameters. Hence addressing their communication overhead and resource utilization bottleneck is an important and practical research direction.\n- The concept of an algorithm-hardware co-design frameworks specifically designed for MoE models on emerging 3.5D wafer-scale architectures is a strong contribution. \n- The work is well motivated by a clear analysis of the system-level challenges. \n- The of a cycle-accurate simulator that has been validated against Verilog simulations is also a strong point. \n- The emprirical evaluation is generally well explained and done.\n\nWeaknesses/Areas of improvements:\n- The hardware evaluation is entirely simulation-based. While this is a critical first step in any new architecture research, the performance evaluations and gains depend heavily on the realization of the 3.5D wafer-scale system. I recommend that the authors include a detailed discussion about the assumptions made in the simulator and if there are any challenges to translate their design and simulation into a real physical hardware/chip.  \n- This statement is partially correct: “Prior works tend to neglect wafer-scale integration [7, 19, 42] and largely adopt coarse-grained, static workload partitioning strategies [29].”. Wafer-scale integration (e.g., Cerebras WSE) is rare and only done by a few companies due to extreme cost, yield, and packaging challenges. So although most existing work do not consider wafer-scal, this is rather a of practicality than oversight. Coarse-grained, static partitioning is common in chiplet designs: models are statically sharded across chiplets based on layers or tensor slices. That said, some recent works that are still evolving (e.g., FlexGen, Megatron-DeepSpeed) do explore dynamic routing or load-balancing. ”\n- The statement in the paper: “Prior works lack system-level coordination and optimization for MoE” uses reference 22, Future of Memory (Liu et al., IEDM 2024). This reference is not a MoE-LLM workload partitioning paper strategy. This work advocates for tighter integration of diverse memory systems but does not propose specific coordination or workload tiling strategies for MoE-LLMs or chiplet systems. So, this citation does not directly support the claim about workload tiling or coordination strategies for MoE.\n- Some key relevant works are missing from the related work discussion:\n   - GShard (Lepikhin et al., ICLR 2021): A foundational MoE system that introduced routing + automatic sharding across devices. Crucial for historical context. This is the first large-scale MoE at Google.\n   - BASE Layers (Lewkowycz et al., NeurIPS 2022): A hybrid sparse-dense expert design that maintains performance while reducing memory. This should be compared and contrasted Mozart’s memory-aware scheduling.\n    - MoEfication (Zhou et al., ICML 2022): Shows conversion of dense models to MoE post-training. This is related to the post-training MoE optimization approaches as proposed by Mozart.\n   - TaskMoE / Expert Tuning (Zhang et al., 2022): Specialization for multi-task models. \n   - Tutel & DeepSpeed-MoE need to be more discussed, especially in comparison to Mozart’s routing, scheduling, and clustering.\n- Suggestions for improving the related work section:\n   - Cite GShard and discuss how Mozart differs in routing granularity, post-training design, and hardware assumptions.\n   - Include Cerebras WSE to contrast wafer-scale monolithic vs. modular chiplet-based wafer-scale.\n   - Add a table or paragraph contrasting Mozart with FRED, Cambricon-LLM, Maestro, etc., clarifying MoE-awareness and system-level scheduling gaps.\n   - Consider adding a  short discussion of commercial accelerators (AMD MI300, Intel PV) for practical relevance.\n   - Cite MoEfication to better support the post-training and fine-grained modularity motivation.\n- The clustering method used in Equation 5 is based on the k-clustering with pairwise distances.This algorithm is computationally expensive (O(n²)), and no mention is made of approximation methods (e.g., hierarchical clustering, METIS). \n- The mathematical formulation in Equation 6 does not provide constraints on the number of experts per token (e.g., top-2) of the token capacity per expert (e.g., max tokens per batch). This needs to be clarified by the authors. \n- In Equation 7, runtime score for scheduling, how is the load defined? Is it measured in flops, tokens, latency? What is the utilization measured against? (e.g., theoretical peak). The formulation remains abstract and hard to reproduce without further clarification by the authors. \n- The proposed co-design strategy effectively addresses expert routing and communication overheads in MoE-style LLMs, however, it is highly specialized to this sparse model family. The system assumes hard top-k expert gating and token-level dispatch, which do not generalize to dense autoregressive models (e.g., GPT-3, LLaMA), diffusion models, or state-space architectures (e.g., Mamba) that lack modular sparse execution patterns. Additionally, the hardware assumptions—modular chiplet-based designs with a Network-on-Package interconnect—limit applicability to a narrow class of accelerators. Real-world systems such as monolithic GPUs, wafer-scale engines (e.g., Cerebras), and mobile AI chips follow vastly different architectural and memory paradigms. The lack of runtime adaptability in expert scheduling and coarse treatment of memory hierarchies further constrain its deployability. A discussion on these generalization limitations is absent.\n- Regarding the empirical evaluation. There is a lack of baselines beyong static scheduling. The main baseline is a static MoE routing strategy or uniform tiling. It would be good to add comparisons against DeepSpeed-MoE or Tutel routing strategies and Existing expert placement heuristics (e.g., token-balanced sharding, load-aware dispatch). The existing evaluation is not enough to claim that Mozart significantly outperforms state-of-the-art MoE runtimes.\n- Limited System Scaling Study. The paper does not analyze how Mozart behaves under increasing number of chiplets, varying interconnect topologies (e.g., mesh vs. ring) and realistic memory capacity constraints. Such scaling or sensitivity analysis can showcase if the benefits hold at larger or smaller deployment scales.\n- The paper’s results report mean latency or throughput, but does not show variance or standard deviation, confidence intervals, and sensitivity to input token distribution (e.g., random vs. clustered token-to-expert assignment). \n- While latency and utilization are discussed, the paper completely omits energy or area modeling. Given the focus on chiplet hardware, this is a notable omission—MoE communication patterns can be power-intensive.\n\nQUESTIONS:\n1.\tThe expert clustering approach (Algorithm 1) is based on a greedy, farthest-point-sampling method. Have you investigated or considered alternative graph clustering algorithms (e.g., spectral clustering, METIS) on the expert co-activation graph especially that the used algorithm is computationally intensive? It would be good to know if other methods could yield a better workload balance or further reduce communication volume.\n2.\tCan you elaborate on how the system's performance and communication patterns would be affected if you were to apply data or tensor parallelism across multiple chiplets for the attention computation? How would this interact with the NoP-Tree topology? Would it be compatible with existing TP approaches?\n3.\tThe expert allocation strategy is based on a static analysis of expert co-activation patterns from an instruction tuning dataset. How stable are these patterns during a long pre-training run? Would a dynamic approach, where expert clusters are periodically re-evaluated and re-allocated, be necessary or beneficial?\n4.\tAll-to-all communication is identified as a key bottleneck. Your approach reduces the data volume by co-locating experts. How does this compare to other techniques for optimizing all-to-all, such as using hierarchical communication algorithms or topology-aware communication libraries?\n5.\tYour proposed hardware–software co-design is tightly coupled to the modular and sparse nature of MoE architectures. I'm not expecting a complete solution here, but could you share how you might begin to extend or adapt your approach to support a broader range of model types—such as dense LLMs, hybrid MoE + dense architectures, or even non-transformer models like state-space models (e.g., Mamba)? In particular, how might your expert clustering, routing strategies, or inter-chiplet communication optimizations evolve (or be abstracted) to support models with different compute and dataflow characteristics?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-07"}