# Cobalt: Leveraging Expert Co-activation for Efficient Distributed MoE Training > Cobalt exploits expert co-activation to co-locate frequently paired experts, cutting cross-node traffic up to 99% and achieving up to 2.41x speedup in distributed MoE training. - **Source:** [arXiv](https://arxiv.org/abs/2609.36959) - **Published:** 2026-10-03 - **Permalink:** https://picx.dev/p/NmBSPA - **Whiteboard:** https://picx.dev/p/NmBSPA/image ## Summary ## Summary (Overview) - **Cobalt** is a novel distributed training framework for Mixture-of-Experts (MoE) models that leverages the observation that many expert pairs are *frequently co-activated* by the same tokens. - It introduces a **two-stage expert layout planner** that (1) periodically performs a global re-layout of experts based on co-activation and workload statistics, and (2) performs per-step intra-node replica rebalancing to ensure workload balance. - A **communication-aware task assignment** algorithm routes tokens to the fewest remote nodes possible, significantly reducing cross-node traffic. - Experiments on 32 NVIDIA B200 GPUs show Cobalt achieves **up to 1.53–2.41× speedup** (1.28–1.89× on average) over existing frameworks, while reducing cross-node token transfer volume by **75.74%–99.26%**. - This is the first work to exploit expert co-activation for communication reduction *and* workload balance in distributed MoE training. ## Introduction and Theoretical Foundation **Background & Motivation:** MoE architectures scale LLMs by routing each token to a small subset of experts, keeping per-token computation nearly constant. Training large MoE models relies on Expert Parallelism (EP), which distributes expert replicas across GPUs and uses all-to-all communication. EP is bottlenecked by two factors: 1. **Cross-node token transfers**: Inter-node bandwidth (e.g., 400 GB/s InfiniBand) is significantly lower than intra-node bandwidth (e.g., 1.8 TB/s NVLink). Communication can account for 58%–83% of training time. 2. **Skewed expert workloads**: Uneven token-to-expert assignments lead to imbalanced computation across GPUs. **Key Observation (Expert Co-activation):** The authors empirically track the set of active tokens for each expert and compute the **Ochiai similarity** between expert pairs: $$\text{Ochiai}(A, B) = \frac{|A \cap B|}{\sqrt{|A| \cdot |B|}}$$ where $A$ and $B$ are the sets of tokens activating experts $A$ and $B$. During training of GLM-4.5-Air, the average Ochiai similarity rises from **0.41 to 0.89** (Layer 35) and from **0.28 to 0.86** (Layer 41). This indicates that many expert pairs are co-activated by the same tokens, and this co-activation strengthens over training. Critically, when two selected experts reside on the same remote node, node-level token deduplication means the token is only sent cross-node *once*, so co-locating co-activated experts can dramatically reduce cross-node traffic. **Problem Formulation:** The paper formalizes the joint optimization of expert layout and task assignment. The communication cost is modeled as the total number of cross-node token transfers: $$V = \sum_{t=1}^{T} \left| \bigcup_{e \in D_t} \{N(r_t^e)\} \setminus \{N(s_t)\} \right|$$ where $T$ is the number of tokens, $N(r)$ is the node containing rank $r$, $D_t$ is the set of selected experts for token $t$, $s_t$ is the source rank, and $r_t^e$ is the rank assigned to token-expert task $(t,e)$. The full optimization problem is: $$\arg\min_{\{P_r\}, \{r_t^e\}} V$$ $$\text{s.t.} \sum_r \mathbb{I}[e \in P_r] \geq 1, \forall e, \quad |P_r| \leq (1+\rho)E/G, \forall r$$ $$e \in P_{r_t^e}, \forall t, e \in D_t, \quad \sum_{e \in P_r} w_r^e \leq (1+\epsilon)KT/G, \forall r$$ Here, $P_r$ is the set of experts with a replica on rank $r$, $\rho$ is the replica budget, $\epsilon$ is the workload imbalance tolerance, $E$ is the number of experts, $G$ is the number of ranks, and $K$ is the number of selected experts per token. ## Methodology **System Architecture:** Cobalt separates expert layout planning from task assignment. ### 1. Two-Stage Expert Layout Planner **Stage 1: Periodic Global Planning.** Every 10 training steps, a global re-layout is computed using historical statistics: - $\hat{w}_e$: EMA of expert $e$'s assigned computation task count - $\hat{c}_{ij}$: EMA of the number of tokens selecting both experts $i$ and $j$ The planning problem is reformulated to maximize total co-activation under workload constraints: $$\arg\max_{\{P_r\}, \{w'_{er}\}} \sum_n \sum_{i