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

By Shuqing Luo ¹, Han Ye  ², Pingzhi Li ¹, Jiayin Qin ², Jie Peng ¹, Yang (Katie) Zhao ², Yu (Kevin) Cao ², Tianlong Chen ^1
1 University of North Carolina at Chapel Hill 
² University of Minnesota - Twin Cities

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 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, 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, 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.

To read the full article, click here