Mixture of Experts in LLMs

Copyimport torch import torch.nn as nn import torch.nn.functional as F # A single Expert = a 2-layer feedforward network (same as a transformer FFN) class Expert(nn.Module): def __init__(self, d_model: int, d_ff: int): super().__init__() self.fc1 = nn.Linear(d_model, d_ff) # expand self.fc2 = nn.Linear(d_ff, d_model) # project back def forward(self, x: torch.Tensor) -> torch.Tensor: # x: [num_tokens, d_model] return self.fc2(F.relu(self.fc1(x))) # Try it out expert = Expert(d_model=512, d_ff=2048) tokens = torch.randn(4, 512) # batch of 4 tokens output = expert(tokens) print(f"Input: {tokens.shape}") # [4, 512] print(f"Output: {output.shape}") # [4, 512] — same shape!

Copyclass TopKRouter(nn.Module): """Routes each token to the top-k most relevant experts.""" def __init__(self, d_model: int, num_experts: int, top_k: int = 2): super().__init__() self.top_k = top_k self.num_experts = num_experts # Learned: maps token embedding → score per expert self.gate = nn.Linear(d_model, num_experts, bias=False) def forward(self, x: torch.Tensor): # x: [num_tokens, d_model] # Step A: raw routing scores for every expert logits = self.gate(x) # [num_tokens, num_experts] # Step B: softmax → probability distribution probs = F.softmax(logits, dim=-1) # [num_tokens, num_experts] # Step C: pick the top-k experts per token top_k_probs, top_k_idx = torch.topk(probs, self.top_k, dim=-1) # top_k_probs: [num_tokens, k] — how much weight per expert # top_k_idx: [num_tokens, k] — which expert indices # Step D: renormalize so the k weights sum to 1 top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True) return top_k_probs, top_k_idx, probs # probs used for load-balancing loss # Example: 4 tokens, 8 experts, select top-2 router = TopKRouter(d_model=512, num_experts=8, top_k=2) tokens = torch.randn(4, 512) weights, indices, all_probs = router(tokens) print("Expert indices per token:", indices) # e.g. tensor([[3, 6], [1, 7], [0, 2], [5, 3]]) print("Expert weights per token: ", weights.round(decimals=2)) # e.g. tensor([[0.62, 0.38], [0.71, 0.29], ...])

Copyclass MoELayer(nn.Module): """ Sparse Mixture of Experts Layer. Replaces a single FFN in a Transformer with N expert FFNs. Only k out of N experts activate per token (k << N). """ def __init__(self, d_model: int, d_ff: int, num_experts: int = 8, top_k: int = 2): super().__init__() self.top_k = top_k self.num_experts = num_experts # N independent expert networks self.experts = nn.ModuleList([ Expert(d_model, d_ff) for _ in range(num_experts) ]) # Router selects which experts process each token self.router = TopKRouter(d_model, num_experts, top_k) def forward(self, x: torch.Tensor): # x: [batch, seq_len, d_model] B, S, D = x.shape # Flatten batch and sequence dims — treat all tokens equally x_flat = x.reshape(-1, D) # [B*S, d_model] N = x_flat.shape[0] # total tokens # Get routing decisions weights, indices, all_probs = self.router(x_flat) # weights: [N, top_k] — how much each expert contributes # indices: [N, top_k] — which expert to use output = torch.zeros_like(x_flat) # [N, d_model] # For each expert slot (0 to top_k-1) for k in range(self.top_k): expert_ids = indices[:, k] # [N] which expert for slot k slot_weight = weights[:, k:k+1] # [N, 1] contribution weight # Group tokens by their assigned expert and run each expert once for e in range(self.num_experts): mask = (expert_ids == e) # which tokens go to expert e if not mask.any(): continue # skip unused experts (saves compute!) expert_out = self.experts[e](x_flat[mask]) # [n_e, D] output[mask] += slot_weight[mask] * expert_out # weighted add return output.reshape(B, S, D), all_probs # restore original shape # Test the full MoE layer moe = MoELayer(d_model=512, d_ff=2048, num_experts=8, top_k=2) x = torch.randn(2, 10, 512) # batch=2, seq_len=10 out, probs = moe(x) print(f"Input: {x.shape}") # torch.Size([2, 10, 512]) print(f"Output: {out.shape}") # torch.Size([2, 10, 512]) ✓ # Count parameters: 8 experts × (512×2048 + 2048×512) = ~33.6M expert params # But only 2/8 experts activate per token = 25% active compute total = sum(p.numel() for p in moe.parameters()) active = total * (moe.top_k / moe.num_experts) print(f"Total params: {total:,} | Active per token: {int(active):,}")

Copydef auxiliary_load_balancing_loss( router_probs: torch.Tensor, # [N, num_experts] — softmax output expert_indices: torch.Tensor, # [N, top_k] — selected expert IDs num_experts: int ) -> torch.Tensor: """ Switch Transformer / Mixtral style auxiliary load-balancing loss. Encourages all experts to receive roughly equal token counts. Loss = num_experts * sum_i(f_i * P_i) f_i = fraction of tokens routed to expert i P_i = mean router probability for expert i Both terms should be ~1/num_experts for uniform distribution. """ N = router_probs.shape[0] # f_i: fraction of tokens dispatched to each expert one_hot = F.one_hot(expert_indices, num_experts).float() # [N, k, E] tokens_per_expert = one_hot.sum(dim=[0, 1]) / (N * expert_indices.shape[1]) # Shape: [num_experts] — ideally all equal to 1/num_experts # P_i: mean router probability assigned to each expert mean_probs = router_probs.mean(dim=0) # Shape: [num_experts] — ideally all equal to 1/num_experts # Dot product × num_experts: 1.0 when perfectly balanced loss = num_experts * (tokens_per_expert * mean_probs).sum() return loss # ------ How to use it in your training loop ------ moe = MoELayer(d_model=512, d_ff=2048, num_experts=8, top_k=2) x = torch.randn(2, 10, 512) out, all_probs = moe(x) _, indices, _ = moe.router(x.reshape(-1, 512)) # Main task loss (e.g. cross-entropy for language modelling) task_loss = torch.tensor(2.5) # placeholder # Add a small load-balancing penalty (coefficient α = 0.01) balance_loss = auxiliary_load_balancing_loss(all_probs, indices, num_experts=8) total_loss = task_loss + 0.01 * balance_loss print(f"Task loss: {task_loss.item():.4f}") print(f"Balance loss: {balance_loss.item():.4f}") # ideally close to 1.0 print(f"Total loss: {total_loss.item():.4f}")

Copy# Install: pip install transformers accelerate bitsandbytes from transformers import AutoModelForCausalLM, AutoTokenizer import torch # ── Option A: Mixtral 8x7B (open weights, Apache-2.0) ────────────────────── model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1" tokenizer = AutoTokenizer.from_pretrained(model_name) # Load with 4-bit quantization: reduces VRAM from ~90 GB → ~24 GB model = AutoModelForCausalLM.from_pretrained( model_name, device_map="auto", # spread across available GPUs/CPU torch_dtype=torch.float16, load_in_4bit=True, # requires bitsandbytes ) # Prepare a chat-formatted prompt messages = [ {"role": "user", "content": "Explain Mixture of Experts in three sentences."} ] input_ids = tokenizer.apply_chat_template( messages, return_tensors="pt", add_generation_prompt=True ).to(model.device) # Generate a response with torch.no_grad(): output_ids = model.generate( input_ids, max_new_tokens=256, temperature=0.7, do_sample=True, ) response = tokenizer.decode( output_ids[0][input_ids.shape[1]:], skip_special_tokens=True ) print(response) # ── Option B: Inspect which experts were activated ───────────────────────── # Access the router logits for interpretability outputs = model(input_ids, output_router_logits=True) router_logits = outputs.router_logits # tuple of tensors per MoE layer for layer_idx, logits in enumerate(router_logits): top2 = logits.topk(2, dim=-1).indices # [seq, 2] print(f"Layer {layer_idx}: top-2 experts used → {top2[0].tolist()}")