上一代的问题：Sparse MoE (Shazeer 2017) used learned gating with RNNs, but no one had made MoE work with Transformers at scale. Dense Transformer scaling hit the wall — going from 2.3B to 600B was impractical due to O(L2) attention cost and memory bandwidth limits.
这代改了什么：First practical MoE Transformer. Replace every other FFN with a Top-2 gated expert layer. Use distributed sharding across TPU devices with All-to-All communication. Add auxiliary load balancing loss to ensure uniform expert utilization.
效果：600B MoE trained in 4 days on 2048 TPU v3 cores (22.4 TPU core-years). Quality far exceeded dense baseline. Direct predecessor of Switch Transformer and all modern MoE LLMs.
GShard takes a standard Transformer and replaces every other FFN with a Mixture-of-Experts layer. The core idea: instead of one dense FFN that processes every token with the same parameters, we have E expert FFNs (FFN1 ... FFNE) and a routing mechanism that sends each token to the most relevant experts.
Each token is routed to exactly 2 experts (Top-2 routing), and the outputs are weighted by the gate values:
To ensure balanced computation across devices, each expert has a fixed capacity:
where N is the total number of tokens per batch and E is the number of experts. Tokens that overflow (routed to a full expert) skip the MoE layer via a residual connection, passing through unchanged.
GShard explores two gating strategies:
GShard's key engineering contribution: how to distribute experts across a TPU pod (2048 cores) while keeping communication efficient.
Each expert is sharded across multiple devices. The standard configuration places one expert replica per group of devices. The communication pattern is All-to-All: tokens must be sent from their original device to the device hosting their chosen expert(s).
# GShard distributed MoE: token routing
# D = total devices, E = total experts
# Phase 1: Gate (local computation)
for each token x on device d:
scores = softmax(top2(x @ W_g)) # E-dimensional gate vector
expert_ids = argsort(scores)[:2] # pick top 2 experts
# Phase 2: All-to-All communication
send(token_data, expert_ids) # token → expert's device
receive(tokens, other_devices) # expert receives its tokens
# Phase 3: Expert computation (local FFN)
for each assigned token:
output += weight * FFN_expert(token)
# Phase 4: All-to-All return
send(output, original_device) # result → token's deviceGShard achieves sub-linear communication scaling: O(√D) for All-to-All on a 2D TPU mesh. This is crucial — if communication scaled linearly with D, adding more devices would not help.
GShard uses the Single Program, Multiple Data (SPMD) compiler strategy — one program runs on all devices, with three annotation APIs to control distribution:
A critical innovation: without explicit load balancing, the gate network would quickly learn to route all tokens to the same few experts, defeating the purpose of MoE. GShard introduces an auxiliary loss:
where:
The auxiliary loss is added to the main training loss with a small coefficient (typically 0.01). The term me makes the loss differentiable — without it, the routing decision (ce) is a discrete choice that blocks gradient flow.
GShard trained a 600B parameter MoE model (dense: 6.7B parameters per layer, 64 experts) for multilingual translation (100 languages to English):
| Metric | Dense Baseline | GShard MoE (600B) | Improvement |
|---|---|---|---|
| Average BLEU (100-En) | 36.9 | 44.3 | +13.5 (36.6%) |
| Training Cost | 235.5 TPU core-years | 22.4 TPU core-years | 10.5x more efficient |
| Training Time | estimated months | 4 days | on 2048 TPU v3 cores |
| Total Parameters | ~6.7B per layer | ~600B total | ~90x more capacity |
| Active Parameters | ~6.7B per layer | ~7.8B per layer | ~1.2x compute per token |
GShard attempted a 1 trillion parameter configuration but encountered numerical stability issues. The model would diverge during training, suggesting scaling laws for MoE training are more complex than increasing parameters alone. This foreshadowed the need for careful initialization and training stability techniques used in later MoE models.
GShard demonstrated three critical things simultaneously:
This is the fundamental MoE trade-off: you get the representational capacity of a 600B model (because each expert specializes in different domains/patterns), but you only pay the compute cost of a 7.8B model per forward pass.
上一篇：Outrageously Large Neural Networks: The Sparsely-Gated MoE Layer (Shazeer et al., 2017) — introduced learned gating for MoE with RNNs. GShard adapts this to Transformers and solves the distributed training problem.
下一篇：Switch Transformer (Fedus et al., 2021) — simplifies Top-2 to Top-1 routing, simplifies the distributed communication, and scales to 1.6T parameters. GShard's auxiliary loss and expert capacity mechanisms are inherited.
引用/采用：Every modern MoE LLM builds on GShard's foundation:
GShard is a systems paper disguised as a modeling paper. Its most impactful contribution is not the gating mechanism or the auxiliary loss (both adapted from prior work), but the engineering blueprint for distributed MoE: how to shard experts across devices, how to communicate tokens efficiently, and how to keep load balanced at 600B+ scale.
The numbers tell a clear story: training a 600B dense model would have cost 235+ TPU core-years and likely been impossible in practice. GShard's MoE achieved better quality with 10x less compute. This efficiency lesson is why every major LLM provider now has an MoE variant.
The 1T parameter failure is also instructive: scaling MoE isn't just about adding experts. The interaction between the routing mechanism and the optimizer creates new failure modes that don't exist in dense training. This is a recurring theme in deep learning — every new architecture class brings its own unique training pathologies.