Brief Summary: Switch Transformers

Full title: Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity Authors: William Fedus, Barret Zoph, Noam Shazeer (Google) Year: 2022 (JMLR 23, submitted 8/21, published 4/22; arXiv version 2021) Venue: Journal of Machine Learning Research 23 (2022)

Problem

Dense transformer models use the same parameters for every input token, which is computationally expensive at scale. Mixture-of-Experts (MoE) models can increase parameter count without proportionally increasing FLOPs per token, but prior MoE implementations suffered from training instability, high communication cost (all-to-all across devices), and complexity (top-k routing with k≥2). Scaling to trillions of parameters remained impractical.

Core Insight

Routing each token to exactly one expert (k=1, "Switch Routing") rather than top-k experts simplifies routing computation, reduces communication cost (expert capacity can be halved), and empirically matches or exceeds the quality of top-2 routing. The parameter count (and thus model quality via scaling laws) can be increased by adding more experts, while keeping the FLOPs per token constant — a fourth scaling axis orthogonal to model depth, width, and training compute.

Method

• Switch Layer: Replaces the FFN in alternating transformer layers. A learned router (W_r: linear projection + softmax) assigns each token to one of N experts (N = number of devices in expert-parallel dimension). Each expert is a standard two-layer FFN with its own independent weights.
• Expert Capacity: Fixed buffer size per expert = (tokens_per_batch / num_experts) × capacity_factor. Overflow tokens are dropped (passed through residual). Capacity factor ≥1.0 reduces drops.
• Load Balancing Loss: Auxiliary loss = α·N·∑(f_i · P_i) where f_i = fraction of tokens dispatched to expert i, P_i = mean router probability for expert i. This encourages uniform routing. α = 10⁻² used throughout.
• Selective Precision: Router computation cast to float32 while model body stays bfloat16, preventing divergence without the full float32 communication cost overhead.
• Reduced Initialization Scale: Weight initialization standard deviation reduced by 10× (σ = √(s/n) with s=0.1) greatly stabilizes training and reduces variance across runs.
• Expert Dropout: During fine-tuning, increase dropout within expert FFN layers (rate=0.4) while keeping non-expert layers at standard dropout (rate=0.1) to prevent overfitting on small downstream tasks.
• Distributed Communication: Expert routing requires two all-to-all operations per Switch layer (once to send tokens to experts, once to return results). This is the dominant communication cost and scales with E × C × d_model.

Key Results

• Switch-Base (64 experts, 7B parameters) reaches the same pre-training quality as T5-Base (223M parameters) in 1/7th the wall-clock time using the same computational budget (32 TPUv3 cores).
• Switch-Base outperforms T5-Large (3.5× more FLOPs per token) by 2.5× on wall-clock basis.
• Switch-C (1.6T parameters, 2048 experts) is 4× faster than T5-XXL to fixed pre-training perplexity.
• Switch-XXL (395B parameters) achieves SQuAD 89.7 (vs. T5 SOTA 91.3), SuperGLUE 87.5 (vs. T5 SOTA 89.3), ANLI 65.7 (vs. prior best 49.4).
• Multilingual: mSwitch-Base improves over mT5-Base on all 101 languages, with mean 5× step speedup; 91% of languages achieve ≥4× speedup.
• Distillation: ~30% of sparse model quality gain preserved when distilling into a 223M dense model (99% parameter reduction).

Limitations

• Expert parallelism requires all-to-all communication which has high latency on bandwidth-limited inter-node links. Switch Transformers are designed for TPU pods (high-bandwidth interconnect); GPU clusters with slower inter-node links see more degradation.
• Training instability remains at large scale (Switch-XXL is sometimes unstable; Switch-C is stable but uses only expert parallelism, not model parallelism, so its per-token FLOPs are lower).
• Token dropping when experts overflow (capacity factor=1.0 drops ~1% of tokens) introduces a training/inference inconsistency.
• Fine-tuning gap: large sparse models do not always outperform dense counterparts on reasoning-heavy downstream tasks, despite superior pre-training perplexity. The largest Switch models (T5-XXL scale) do not yet match T5-XXL on all SuperGLUE tasks.
• No GPU-optimized sparse kernel support at time of writing; relies on dense matrix operations with padding/masking.

Relevance to DynamICCL

Switch Transformers introduce a new collective communication type — all-to-all — into the training loop, which DynamICCL must handle alongside AllReduce/ReduceScatter/AllGather.

• All-to-all dominates Switch layer cost: Each Switch layer requires 2 all-to-all collectives (token dispatch + result combine), each of size E × C × d_model per device. For 128 experts and capacity factor 1.0, this is a large, irregular, many-to-many message pattern very different from the bandwidth-optimal ring AllReduce.
• NCCL All-to-All tuning: NCCL's ncclAllToAll collective (or AllToAllv for irregular counts) is not tuned by the standard NCCL tuner plugin. DynamICCL would need to extend its Config Agent action space to cover all-to-all collectives.
• Expert load imbalance → congestion: Unbalanced routing causes some devices to receive far more tokens than others, creating asymmetric network load — a natural congestion trigger for DynamICCL's LSTM+CUSUM Trigger Agent.
• Data-parallel gradient AllReduce still present: Switch Transformer training combines expert-parallel all-to-all with data-parallel gradient AllReduce, creating simultaneous multi-collective traffic. DynamICCL must avoid assigning NCCL channels in ways that create bandwidth competition between the two.
• Low-bandwidth relevance: On Chameleon Cloud 1 GbE, the all-to-all communication cost of even a small Switch model would be prohibitive, making DynamICCL's ability to detect and react to this congestion especially important.