Communication Optimization for MoE Expert Parallelism in Distributed Training

Introduction Mixture of Experts (MoE) is a sparse neural network architecture that significantly scales model parameters without proportionally increasing computational cost. In a MoE architecture, each input token activates only a small number of expert sub-networks (typically 1 to 2) for processing, enabling efficient utilization of compute resources. Expert Parallelism (EP) is a key parallelism strategy for MoE training. By distributing different experts across multiple compute nodes (e.g., GPUs), EP achieves parallel computation across nodes, significantly reducing per-node load and improving training throughput. However, EP relies heavily on All-to-All collective communication to dispatch tokens to the correct experts and gather results. Experiments have shown that All-to-All communication accounts for over 34.1% of total training time , constituting the primary scalability bottleneck. Furthermore, during the backward pass, All-to-All communication contends for bandwidth with AllReduce communication used for gradient aggregation in data parallelism. This document proposes two mechanisms to address these challenges: (1) Adaptive Communication Mode Selection: Dynamically selects between Data-Centric and Expert-Centric All-to-All patterns based on model and cluster parameters. (2) Priority-Based Co-Scheduling: Coordinates All-to-All and AllReduce communications through priority scheduling, chunked transmission, and in-network aggregation to eliminate bandwidth contention.

Requirements Language The key words "MUST", "MUST NOT", "REQUIRED", "SHALL", "SHALL NOT", "SHOULD", "SHOULD NOT", "RECOMMENDED", "NOT RECOMMENDED", "MAY", and "OPTIONAL" in this document are to be interpreted as described in BCP 14 when, and only when, they appear in all capitals, as shown here.

Terminology

MoE (Mixture of Experts):: A sparse neural network architecture where each input activates only a subset of expert sub-networks.
EP (Expert Parallelism):: A parallelism strategy that distributes different experts across multiple compute nodes.
All-to-All:: A collective communication primitive where each node sends distinct data to every other node.
AllReduce:: A collective communication primitive that performs a reduction operation across all nodes and distributes the result.
Token:: The smallest unit of input data processed by the model.
Worker:: A single compute unit participating in training, typically corresponding to one GPU.
H:: The hidden dimension size of the model.
T:: Total number of tokens to process, determined by Batch Size (B) and Sequence Length (S).
E:: Total number of experts in the MoE model.
Top-K:: The number of experts selected per token by the MoE router.
GDR (GPU Direct RDMA):: A technology enabling direct memory access between GPUs across a network without CPU involvement.
NVLink:: A high-speed interconnect bus between NVIDIA GPUs within a node.
NVSwitch:: A high-speed switching chip for inter-GPU communication within a node.
ToR Switch:: Top-of-Rack switch, used for inter-node communication.
In-Network Aggregation:: The capability of performing reduction operations on network switching devices.

Problem Statement

All-to-All Communication Bottleneck MoE models achieve computational efficiency through dynamic expert activation. However, the All-to-All communication required for Expert Parallelism constitutes the primary bottleneck in scaling distributed training. Empirical studies demonstrate that All-to-All communication accounts for over 34.1% of total training time. The All-to-All communication exhibits a "barrel effect": overall communication latency is determined by the slowest node. Since processing capacity and token count vary across nodes, individual slow nodes can degrade overall training throughput.

Static Communication Mode Limitation Existing implementations typically employ a fixed All-to-All communication pattern, either exchanging data or exchanging expert parameters. However, the communication cost of each pattern varies significantly under different model configurations and cluster topologies. A static approach cannot adapt to all scenarios.

All-to-All and AllReduce Bandwidth Contention During the backward pass of training, All-to-All communication (for EP gradient synchronization) and AllReduce communication (for data-parallel gradient aggregation) share physical link bandwidth. Concurrent execution of both communication types leads to bandwidth contention, causing significant increases in overall communication latency and degraded training performance.

Adaptive Communication Mode Selection

Overview This mechanism dynamically selects the optimal All-to-All communication mode by comparing expert parameter size against token data size. Data-Centric Mode: When expert parameter size is smaller than token data size, workers dynamically exchange expert parameters while keeping data stationary. Suitable for scenarios with smaller experts and larger data volumes. Expert-Centric Mode: When expert parameter size is larger than token data size, workers exchange token data while keeping expert parameters stationary. Suitable for scenarios with larger experts and smaller data volumes.

Mode Selection Criteria The collective communication library SHOULD automatically select the communication mode through the following procedure: Step 1 - Compute Sizes: Expert parameter size S_expert = 8 * H^2; Token data size S_data = 2 * H * T (where T = B * S * K). Step 2 - Detect Cluster Topology: The library MUST automatically detect node configuration, GPU type, network topology, and inter-node interconnect bandwidth. Step 3 - Compute Cross-Node Communication Volume: Data-Centric volume C_dc = 8 * H * E * m * (n - 1); Expert-Centric volume C_ec = 2 * m * H * T * (n - 1) / n (E = expert count, n = node count, m = workers per node). Step 4 - Select Mode: If C_dc < C_ec, select Data-Centric mode. Otherwise, select Expert-Centric mode.

Data-Centric Mode Execution When Data-Centric mode is selected, the following procedure MUST be executed: (1) Each worker fetches remote expert parameters to local node CPU memory via GDR NIC; fetch targets SHOULD be staggered across workers to avoid link contention. (2) The fetch operations are coordinated by the collective communication library or scheduled by in-network compute switches. (3) Intra-node expert transfer between workers is performed via NVLink, scheduled by the collective communication library or NVSwitch. (4) Expert parameters are loaded into workers (GPUs) via the PCIe bus in parallel. (5) Inter-node Ethernet transfers and intra-node PCIe/NVLink transfers SHOULD proceed concurrently to achieve pipelining.

All-to-All and AllReduce Co-Scheduling

Overview During the backward pass, this mechanism eliminates bandwidth contention between All-to-All and AllReduce communications through: Chunked Transmission (AllReduce gradient data is split into fixed-size chunks); Priority Scheduling (All-to-All traffic is assigned high priority; AllReduce chunks are assigned low priority); Gap Filling (AllReduce chunks are transmitted during idle intervals between All-to-All bursts); In-Network Aggregation (AllReduce traffic passing through ToR switches and NVSwitch is accelerated by performing aggregation operations on network devices).

Scheduling Procedure The collective communication library MUST execute the following scheduling rules: (1) AllReduce gradient data MUST be split into fixed-length chunks; chunk size MAY be statically configured or dynamically adjusted; each chunk is assigned low priority and enqueued. (2) All-to-All gradient synchronization data MUST be assigned high priority and enqueued. (3) The library MUST schedule the send queue in priority order, transmitting All-to-All data first and filling idle gaps with AllReduce chunks. (4) The library MAY monitor or predict contention in real time, dynamically adjusting chunk size to optimize throughput. (5) AllReduce traffic traversing ToR switches and NVSwitch SHOULD leverage in-network aggregation for acceleration.

Communication Hierarchy This mechanism follows an "intra-node first, inter-node second" communication principle. Intra-node communication (typical bandwidth: 600 GB/s via NVLink) completes local All-to-All and AllReduce operations first. Inter-node communication (typical bandwidth: 200 Gbps via Ethernet/InfiniBand) then completes cross-node synchronization. For AllReduce, inter-node traffic is aggregated via ToR switches, and intra-node traffic is aggregated via NVSwitch.

Protocol Operations

Collective Communication Library Extensions To support this mechanism, the collective communication library SHOULD provide the following extension interfaces: Mode Query Interface (returns the currently recommended communication mode); Topology Detection Interface (returns cluster physical topology information); Chunk Size Configuration Interface (sets AllReduce chunk size); Priority Configuration Interface (sets priority levels for different communication task types); In-Network Aggregation Control Interface (enables or disables in-network aggregation).

Mode Re-evaluation Triggers Communication mode selection is performed during training initialization. The library MUST re-evaluate mode selection upon: training initialization; cluster topology changes (e.g., node addition or removal); model configuration changes (e.g., changes to expert count, hidden dimension, or other relevant parameters).

Compatibility This mechanism is designed as an incremental extension to existing collective communication libraries (e.g., NCCL, HCCL). It does not alter existing All-to-All and AllReduce semantic interfaces. Mode selection and scheduling strategies are implemented as internal optimizations, supporting seamless integration with existing training frameworks (e.g., PyTorch, MindSpore).

Security Considerations This mechanism involves communication optimization across multiple nodes and the following security aspects MUST be considered. Cluster Topology Protection: Auto-detected cluster topology information constitutes sensitive infrastructure data; the library MUST ensure this information is not accessible to unauthorized parties. Communication Integrity: In Data-Centric mode, expert parameters are transmitted over the network; implementations SHOULD ensure the integrity of parameter data during transmission to prevent man-in-the-middle tampering that could compromise training results. Denial of Service Protection: Misuse of in-network aggregation capabilities could exhaust switch resources; network devices SHOULD implement rate limiting and resource isolation for in-network aggregation tasks. Priority Abuse Prevention: Improper use of the priority mechanism could cause persistent starvation of low-priority traffic; the library SHOULD enforce minimum guaranteed bandwidth to ensure AllReduce communication is not indefinitely deferred.

IANA Considerations This document has no IANA actions.