Load balancing keeps sparse mixture-of-experts routing from sending most tokens to a few experts while leaving others idle. Routing strategies such as auxiliary losses and expert-choice selection reduce expert collapse, but expert-choice routing can leave some tokens uncovered. section 3.5.1
A minimal top-k gate computes expert logits, selects the highest-scoring experts, and normalizes their scores. GShard-style top-2 gating extends this with capacity settings, load-balancing logic, and optional random routing. section 3.5.1
Listing 3.9 GShard-style top-2 gating with load balancing and random routing
class GshardGate(NativeGate):
def __init__(self, d_model, num_expert, world_size, top_k=2,
capacity=(1.2, 2.4), random_routing=True, gate_bias=True):
assert top_k == 2, "top_k should be 2 in gshard"
super().__init__(d_model, num_expert, world_size,
top_k=2, gate_bias=gate_bias)
self.capacity = capacity
self.random_routing = random_routing
def forward(self, x):
naive_outs = super().forward(x, return_all_scores=True)
topk_idx, topk_val, gate_score = naive_outs
s = gate_score.shape[0] #1
top1_idx = topk_idx.view(-1, self.top_k)[:, 0]
c_e = torch.scatter_add(
torch.zeros(self.tot_expert, device=top1_idx.device),
0,
top1_idx,
torch.ones_like(top1_idx, dtype=torch.float32),
) / s
m_e = torch.mean(F.softmax(gate_score, dim=1), dim=0)
loss = torch.mean(c_e * m_e) * (self.num_expert ** 2)
self.set_loss(loss)
cap_rate = self.capacity
[0 if self.training else 1] #2
capacity = math.ceil(cap_rate * x.shape[0])
capacity = capacity * self.top_k / (self.world_size *
self.num_expert)
capacity = torch.ones(self.num_expert * self.world_size,
dtype=torch.int32,
device=topk_idx.device) * capacity
topk_idx = fmoe_native_prune_gate_by_capacity(topk_idx,
capacity,
self.num_expert,
self.world_size)
if self.random_routing: #3
rand_routing_prob =
torch.rand(gate_score.size(0), device=x.device)
mask = (2 * topk_val[:, 1] < rand_routing_prob)
topk_idx[:, 1].masked_fill(mask, -1)
return topk_idx, topk_val
Use the GShard-style implementation to trace top-2 selection and the first steps of distributing tokens across experts. section 3.5.1
