A working MoE layer in PyTorch

pip install torch

import time, torch, torch.nn as nn, torch.nn.functional as F

class SwiGLU(nn.Module):
    def __init__(self, d, ff):
        super().__init__()
        self.gate = nn.Linear(d, ff, bias=False)
        self.up   = nn.Linear(d, ff, bias=False)
        self.down = nn.Linear(ff, d, bias=False)
    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

class MoELayer(nn.Module):
    """Routed experts + one always-on shared expert, with sorted dispatch."""
    def __init__(self, d=64, ff=128, n_exp=8, top_k=2, n_shared=1, aux_alpha=0.01):
        super().__init__()
        self.k, self.n, self.alpha = top_k, n_exp, aux_alpha
        self.router  = nn.Linear(d, n_exp, bias=False)
        self.experts = nn.ModuleList(SwiGLU(d, ff) for _ in range(n_exp))
        self.shared  = nn.ModuleList(SwiGLU(d, ff) for _ in range(n_shared))
        self.aux_loss = torch.tensor(0.0)

    def forward(self, x):
        shape = x.shape
        x = x.reshape(-1, shape[-1])
        T = x.shape[0]

        logits = self.router(x)
        probs  = F.softmax(logits.float(), dim=-1).to(x.dtype)   # float32 router, as in practice
        w, idx = probs.topk(self.k, dim=-1)
        w = w / w.sum(-1, keepdim=True)

        flat = idx.reshape(-1)                       # (T*k,)
        order = flat.argsort()                       # group assignments by expert
        counts = torch.bincount(flat, minlength=self.n)
        token_of = order // self.k                   # which token each slot belongs to
        weight_of = w.reshape(-1)[order]

        xs = x[token_of]                             # tokens laid out expert by expert
        ys = torch.empty_like(xs)
        start = 0
        for e in range(self.n):                      # contiguous slice, no masking
            end = start + int(counts[e])
            if end > start:
                ys[start:end] = self.experts[e](xs[start:end])
            start = end
        ys = ys * weight_of[:, None]

        out = torch.zeros_like(x).index_add_(0, token_of, ys)
        for s in self.shared:                        # the shared expert sees every token
            out = out + s(x)

        f = counts.float() / (T * self.k)            # Switch load-balancing loss
        P = probs.mean(0)
        self.aux_loss = self.alpha * self.n * (f * P).sum()
        return out.reshape(shape), counts

    def naive_forward(self, x):
        """Reference implementation. Same maths, easy to check by eye, slow."""
        shape = x.shape; x = x.reshape(-1, shape[-1])
        probs = F.softmax(self.router(x).float(), dim=-1).to(x.dtype)
        w, idx = probs.topk(self.k, dim=-1); w = w / w.sum(-1, keepdim=True)
        out = torch.zeros_like(x)
        for s in range(self.k):
            for e in range(self.n):
                hit = idx[:, s] == e
                if hit.any():
                    out[hit] += w[hit, s, None] * self.experts[e](x[hit])
        for s in self.shared:
            out = out + s(x)
        return out.reshape(shape)

torch.manual_seed(0)
moe = MoELayer()
x = torch.randn(4, 32, 64)                          # batch 4, 32 tokens, dim 64

fast, counts = moe(x)
slow = moe.naive_forward(x)
print("sorted dispatch matches the reference:", torch.allclose(fast, slow, atol=1e-5))
print("max difference:", (fast - slow).abs().max().item())
print("tokens per expert:", counts.tolist(), " (128 tokens x top-2 = 256 assignments)")
print("aux loss:", round(float(moe.aux_loss), 6))
per_expert = sum(p.numel() for p in moe.experts[0].parameters())
total_p  = sum(p.numel() for p in moe.parameters())
active_p = (moe.k + len(moe.shared)) * per_expert + moe.router.weight.numel()
print(f"parameters: {total_p:,} total | {active_p:,} active per token "
      f"({active_p/total_p:.1%})")

big = torch.randn(2048, 64)
for name, fn in (("naive loop", moe.naive_forward), ("sorted dispatch", lambda z: moe(z)[0])):
    with torch.no_grad():
        fn(big); t0 = time.perf_counter()
        for _ in range(20): fn(big)
        print(f"  {name:<18}{(time.perf_counter()-t0)/20*1000:7.2f} ms per call (CPU)")

print("\ntraining it for 300 steps on a toy task")
torch.manual_seed(1)
model = MoELayer(d=32, ff=64, n_exp=8, top_k=2)
opt = torch.optim.Adam(model.parameters(), lr=3e-3)
gen = torch.Generator().manual_seed(2)
W = torch.randn(32, 32)
for step in range(301):
    xb = torch.randn(256, 32, generator=gen)
    yb = torch.tanh(xb @ W)
    out, counts = model(xb)
    loss = F.mse_loss(out, yb) + model.aux_loss
    loss.backward(); opt.step(); opt.zero_grad()
    if step % 100 == 0:
        share = counts.float() / counts.sum()
        print(f"  step {step:>3}  loss {loss.item():.4f}  "
              f"busiest expert {share.max():.1%}  quietest {share.min():.1%}")

class Block(nn.Module):
    def __init__(self, d, n_heads, **moe_kwargs):
        super().__init__()
        self.n1, self.n2 = nn.RMSNorm(d), nn.RMSNorm(d)
        self.attn = nn.MultiheadAttention(d, n_heads, batch_first=True)
        self.moe = MoELayer(d=d, **moe_kwargs)

    def forward(self, x):
        h = self.n1(x)
        x = x + self.attn(h, h, h, need_weights=False)[0]
        y, _ = self.moe(self.n2(x))
        return x + y

aux = sum(b.moe.aux_loss for b in model.blocks)
loss = task_loss + aux