Mozart: Modularized and Efficient MoE Training on 3.5D Wafer-Scale Chiplet Architectures

Tianlong Chen ( University of North Carolina at Chapel Hill) · Yang Zhao (Harbin Institute of Technology) · Shuqing Luo (School of Electronic and Computer Engineering, Peking University Shenzhen Graduate School) · Ye Han (Vanderbilt University) · Pingzhi Li (UNC-Chapel Hill) · Jiayin Qin (University of Minnesota - Twin Cities) · Jie Peng (University of Science and Technology of China) · Yu Cao (University of California, Berkeley)
3.5d wafer-scale chiplet architecturesalgorithm-hardware co-designcommunication-computation overlapefficiency gainsexpert allocation strategyfine-grained schedulingheterogeneous moduleshierarchical memory structuremixture-of-expertsmodularized computationnop-tree topologyon-package communicationparallelizationresource utilizationstreaming tokens

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.