Occult: Optimizing Collaborative Communication across Experts for Accelerated Parallel MoE Training and Inference

TL;DR

Occult introduces a collaborative communication framework for mixture-of-experts (MoE) that treats cross-expert all-to-all as a controllable resource. By redefining routing with intra- and inter-collaboration, and by engineering a sparse-matrix-multiplication pipeline with BRIM data structures and 3 token states, it reduces communication overhead while preserving accuracy. The approach combines expert-placement rescheduling and two pruning schemes—routing-score-based and similarity-based—to minimize the all-to-all budget $C_{\mathcal{T}}$ while maintaining performance. Empirical results across three frontier MoE-LLMs show substantial speedups in prefilling, decoding, and training on multi-GPU hardware, often surpassing state-of-the-art frameworks. Overall, Occult demonstrates a scalable, cross-layer solution to accelerate MoE workloads with practical implications for deploying large MoE-LLMs.

Abstract

Mixture-of-experts (MoE) architectures could achieve impressive computational efficiency with expert parallelism, which relies heavily on all-to-all communication across devices. Unfortunately, such communication overhead typically constitutes a significant portion of the total runtime, hampering the scalability of distributed training and inference for modern MoE models (consuming over $40\%$ runtime in large-scale training). In this paper, we first define collaborative communication to illustrate this intrinsic limitation, and then propose system- and algorithm-level innovations to reduce communication costs. Specifically, given a pair of experts co-activated by one token, we call them "collaborated", which comprises $2$ cases as intra- and inter-collaboration, depending on whether they are kept on the same device. Our pilot investigations reveal that augmenting the proportion of intra-collaboration can accelerate expert parallelism at scale. It motivates us to strategically optimize collaborative communication for accelerated MoE training and inference, dubbed Occult. Our designs are capable of either delivering exact results with reduced communication cost or controllably minimizing the cost with collaboration pruning, materialized by modified fine-tuning. Comprehensive experiments on various MoE-LLMs demonstrate that Occult can be faster than popular state-of-the-art inference or training frameworks (more than $1.5\times$ speed up across multiple tasks and models) with comparable or superior quality compared to the standard fine-tuning. Code is available at $\href{https://github.com/UNITES-Lab/Occult}{https://github.com/UNITES-Lab/Occult}$.

Figures (12)

• Figure 1: Latency Comparison with Multiple Models & Tasks.Occult can accelerate training & inference for modern MoE-LLMs on communication-intensive tasks.
• Figure 2: MoE workflows with $3$ all-to-all communication strategies. We take $2$ devices ($\texttt{D}_0$ and $\texttt{D}_1$) for expert parallel, and $\texttt{D}_{i}^{j}$ denotes tokens from device $i$ dispatched to device $j$.
• Figure 3: Empirical correlation between the expert amount in the maximal connected sub-graph and token amount in a mini-batch for top-$2$, $4$, $6$, and $8$ routing. We examine the last MoE layer for each model with the same inputs.
• Figure 4: MoE workflow with proposed bidirectional re-index-guided all-to-all communication and expert computing. We take the forward process for illustration, where token tensors are stated as ORI, SFD and EPD, compactly stored on HBM and guided by $\textit{BRIM}$s.
• Figure 5: Functions of $2$$\textit{BRIM}s$. Both $\textit{BRIM}_0$ and $\textit{BRIM}_1$ are 2D matrices. $\textit{BRIM}_0$'s dimensions are defined by the number of devices for expert parallelism (first dimension) and the token count in the ORI tensor (second dimension). $\textit{BRIM}_1$'s dimensions are defined by the number of local experts per device (first dimension) and the token count in the SFD tensor (second dimension).