Mixture of Experts in LLMs — The 3am Router Collapse That Killed Our P99 Latency

import torch
import torch.nn as nn
import torch.nn.functional as F

class MoELayer(nn.Module):
    def __init__(self, d_model, num_experts=8, top_k=2, expert_capacity_factor=1.5, router_temperature=0.3):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.expert_capacity = None  # set per-batch
        self.expert_capacity_factor = expert_capacity_factor
        self.router_temperature = router_temperature
        
        # Experts: each is a simple FFN (2-layer MLP)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_model * 4),
                nn.GELU(),
                nn.Linear(d_model * 4, d_model)
            ) for _ in range(num_experts)
        ])
        
        # Router: single linear layer, no bias
        self.router = nn.Linear(d_model, num_experts, bias=False)
        
    def forward(self, x):
        # x: (batch, seq_len, d_model)
        batch, seq_len, d_model = x.shape
        
        # Router logits
        router_logits = self.router(x)  # (batch, seq_len, num_experts)
        
        # Apply temperature scaling to prevent logit saturation
        router_logits = router_logits / self.router_temperature
        
        # Softmax over experts
        router_weights = F.softmax(router_logits, dim=-1)  # (batch, seq_len, num_experts)
        
        # Top-k selection
        top_k_weights, top_k_indices = torch.topk(router_weights, self.top_k, dim=-1)
        # top_k_weights: (batch, seq_len, top_k), top_k_indices: (batch, seq_len, top_k)
        
        # Normalize top-k weights to sum to 1
        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)
        
        # Compute expert capacity: max tokens per expert
        # Capacity = (batch * seq_len * top_k) / num_experts * capacity_factor
        total_tokens = batch * seq_len
        self.expert_capacity = int((total_tokens * self.top_k) / self.num_experts * self.expert_capacity_factor)
        
        # Initialize output and token dropping counter
        output = torch.zeros_like(x)
        tokens_dropped = 0
        
        # For each expert, gather tokens assigned to it, process, and scatter back
        for expert_idx in range(self.num_experts):
            # Find tokens where this expert is in top-k
            # top_k_indices shape: (batch, seq_len, top_k)
            # We need to find all (batch, seq) pairs where top_k_indices[b, s, :] == expert_idx
            mask = (top_k_indices == expert_idx).any(dim=-1)  # (batch, seq_len)
            
            # Get the indices of these tokens
            token_indices = mask.nonzero(as_tuple=False)  # (N, 2) where N is number of tokens assigned to this expert
            
            if token_indices.size(0) == 0:
                continue
            
            # If tokens exceed capacity, drop the excess
            if token_indices.size(0) > self.expert_capacity:
                # Randomly select tokens to keep (or you could do first-come-first-serve)
                perm = torch.randperm(token_indices.size(0))
                token_indices = token_indices[perm[:self.expert_capacity]]
                tokens_dropped += token_indices.size(0) - self.expert_capacity
            
            # Gather the token embeddings
            selected_tokens = x[token_indices[:, 0], token_indices[:, 1]]  # (N, d_model)
            
            # Process through expert
            expert_output = self.experts[expert_idx](selected_tokens)  # (N, d_model)
            
            # Get the router weight for this expert for these tokens
            # router_weights shape: (batch, seq_len, num_experts)
            expert_weights = router_weights[token_indices[:, 0], token_indices[:, 1], expert_idx]  # (N,)
            
            # Weight the output
            expert_output = expert_output * expert_weights.unsqueeze(-1)  # (N, d_model)
            
            # Scatter back to output
            output[token_indices[:, 0], token_indices[:, 1]] += expert_output
        
        # Log token dropping rate (in production, use a proper logger)
        if tokens_dropped > 0:
            print(f"Warning: {tokens_dropped} tokens dropped ({(tokens_dropped / total_tokens) * 100:.2f}%)")
        
        return output

# Example usage
if __name__ == "__main__":
    batch, seq_len, d_model = 2, 4, 512
    x = torch.randn(batch, seq_len, d_model)
    moe = MoELayer(d_model, num_experts=8, top_k=2)
    output = moe(x)
    print(f"Input shape: {x.shape}, Output shape: {output.shape}")
    print(f"Expert capacity: {moe.expert_capacity}")

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from datasets import load_dataset
import tiktoken

# Reuse the MoELayer from above
from moe_layer_production import MoELayer

class TransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads, num_experts, top_k, dropout=0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, num_heads, dropout=dropout, batch_first=True)
        self.ln2 = nn.LayerNorm(d_model)
        self.moe = MoELayer(d_model, num_experts=num_experts, top_k=top_k)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, attn_mask=None):
        # Self-attention with residual
        x = x + self.dropout(self.attn(self.ln1(x), self.ln1(x), self.ln1(x), attn_mask=attn_mask)[0])
        # MoE with residual
        x = x + self.dropout(self.moe(self.ln2(x)))
        return x

class MoETransformer(nn.Module):
    def __init__(self, vocab_size, d_model=256, num_heads=8, num_layers=6, num_experts=8, top_k=2, max_seq_len=512):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.pos_embedding = nn.Embedding(max_seq_len, d_model)
        self.blocks = nn.ModuleList([
            TransformerBlock(d_model, num_heads, num_experts, top_k)
            for _ in range(num_layers)
        ])
        self.ln_f = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        
    def forward(self, input_ids, labels=None):
        batch, seq_len = input_ids.shape
        # Token + position embeddings
        x = self.token_embedding(input_ids) + self.pos_embedding(torch.arange(seq_len, device=input_ids.device))
        # Causal mask
        attn_mask = torch.triu(torch.ones(seq_len, seq_len, device=input_ids.device) * float('-inf'), diagonal=1)
        # Pass through blocks
        for block in self.blocks:
            x = block(x, attn_mask=attn_mask)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
            return loss, logits
        return logits

# Training setup
class WikiTextDataset(Dataset):
    def __init__(self, split='train', max_seq_len=512):
        dataset = load_dataset('wikitext', 'wikitext-2-raw-v1', split=split)
        self.enc = tiktoken.get_encoding('gpt2')
        self.max_seq_len = max_seq_len
        # Tokenize all text
        self.tokens = []
        for example in dataset:
            tokens = self.enc.encode(example['text'])
            self.tokens.extend(tokens)
        # Split into chunks
        self.chunks = [self.tokens[i:i+max_seq_len] for i in range(0, len(self.tokens)-max_seq_len, max_seq_len)]
        
    def __len__(self):
        return len(self.chunks)
    
    def __getitem__(self, idx):
        chunk = self.chunks[idx]
        # Pad if necessary
        if len(chunk) < self.max_seq_len:
            chunk = chunk + [self.enc.eot_token] * (self.max_seq_len - len(chunk))
        return torch.tensor(chunk[:self.max_seq_len])

if __name__ == "__main__":
    # Hyperparams
    vocab_size = 50257  # GPT-2 vocab size
    d_model = 256
    num_heads = 8
    num_layers = 6
    num_experts = 8
    top_k = 2
    batch_size = 4
    max_seq_len = 128
    lr = 3e-4
    
    # Model
    model = MoETransformer(vocab_size, d_model, num_heads, num_layers, num_experts, top_k, max_seq_len)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    
    # Data
    dataset = WikiTextDataset(split='train', max_seq_len=max_seq_len)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    
    # Training loop (just 10 steps for demo)
    model.train()
    for step, batch in enumerate(dataloader):
        if step >= 10:
            break
        input_ids = batch
        loss, _ = model(input_ids, labels=input_ids)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad()
        print(f"Step {step}, Loss: {loss.item():.4f}")
    
    # Save model
    torch.save(model.state_dict(), 'moe_transformer.pt')
    print("Model saved.")

import torch
import time
from moe_layer_production import MoELayer

# Benchmark dense vs MoE
# Dense: single FFN with 4x hidden dimension
# MoE: 8 experts, top-2, each with 4x hidden dimension (same total FLOPs per token)

def benchmark_layer(layer, x, num_runs=100):
    # Warmup
    for _ in range(10):
        _ = layer(x)
    torch.cuda.synchronize()
    
    start = time.time()
    for _ in range(num_runs):
        _ = layer(x)
    torch.cuda.synchronize()
    end = time.time()
    
    return (end - start) / num_runs * 1000  # ms

if __name__ == "__main__":
    d_model = 1024
    batch_size = 32
    seq_len = 128
    x = torch.randn(batch_size, seq_len, d_model).cuda()
    
    # Dense layer
    dense = nn.Sequential(
        nn.Linear(d_model, d_model * 4),
        nn.GELU(),
        nn.Linear(d_model * 4, d_model)
    ).cuda()
    
    # MoE layer
    moe = MoELayer(d_model, num_experts=8, top_k=2, router_temperature=0.3).cuda()
    
    # Benchmark
    dense_time = benchmark_layer(dense, x)
    moe_time = benchmark_layer(moe, x)
    
    print(f"Dense layer: {dense_time:.2f} ms")
    print(f"MoE layer: {moe_time:.2f} ms")
    print(f"MoE overhead: {(moe_time / dense_time - 1) * 100:.1f}%")
    
    # With small batch
    x_small = torch.randn(4, seq_len, d_model).cuda()
    dense_small = benchmark_layer(dense, x_small)
    moe_small = benchmark_layer(moe, x_small)
    print(f"\nSmall batch (batch=4):")
    print(f"Dense: {dense_small:.2f} ms, MoE: {moe_small:.2f} ms")

import torch
import torch.distributed as dist

# Simulate expert parallelism with all-to-all communication
# Assume we have 4 GPUs, each with 2 experts
# This is a simplified version of what frameworks like Megatron-LM do

def all_to_all_expert_routing(token_embeddings, expert_assignments, num_experts, world_size):
    """
    token_embeddings: (batch, seq_len, d_model) on this GPU
    expert_assignments: (batch, seq_len, top_k) - which experts each token is assigned to
    num_experts: total number of experts across all GPUs
    world_size: number of GPUs
    """
    # Step 1: For each token, determine which GPU holds its assigned expert
    # experts_per_gpu = num_experts // world_size
    experts_per_gpu = num_experts // world_size
    
    # Step 2: Build send buffers: for each GPU, collect tokens that need to go there
    send_buffers = [[] for _ in range(world_size)]
    for b in range(token_embeddings.size(0)):
        for s in range(token_embeddings.size(1)):
            for k in range(expert_assignments.size(-1)):
                expert_idx = expert_assignments[b, s, k].item()
                target_gpu = expert_idx // experts_per_gpu
                send_buffers[target_gpu].append(token_embeddings[b, s].unsqueeze(0))
    
    # Step 3: All-to-all send/receive
    # In practice, you'd use torch.distributed.all_to_all or a custom communication primitive
    # For this demo, we just simulate the communication cost
    import time
    time.sleep(0.01)  # Simulate 10ms communication
    
    # Step 4: Process tokens on local experts
    # (Assume we have local experts stored in a list)
    local_experts = [None] * experts_per_gpu  # Placeholder
    local_outputs = []
    for tokens in send_buffers[dist.get_rank()]:
        # Process through the appropriate local expert
        # This is where the actual expert computation happens
        local_outputs.append(tokens)  # Placeholder
    
    # Step 5: All-to-all send results back
    time.sleep(0.01)  # Simulate 10ms communication
    
    # Step 6: Aggregate outputs
    # (In practice, you'd sum weighted outputs)
    return torch.cat(local_outputs, dim=0)

if __name__ == "__main__":
    # This is a conceptual example; requires torch.distributed to run
    print("Expert parallelism adds significant communication overhead.")
    print("Benchmark your specific network topology before deploying.")

import torch
import numpy as np

# Production monitoring function
def monitor_expert_utilization(router_weights, num_experts, log_every=100):
    """
    router_weights: (batch, seq_len, num_experts) - softmax output
    Logs expert utilization histogram and alerts if any expert is underused.
    """
    # Count tokens assigned to each expert (based on max weight)
    expert_assignments = router_weights.argmax(dim=-1)  # (batch, seq_len)
    utilization = torch.bincount(expert_assignments.flatten(), minlength=num_experts).float()
    utilization = utilization / utilization.sum()  # Normalize to percentages
    
    # Log
    print(f"Expert utilization: {utilization.tolist()}")
    
    # Alert if any expert has <5% utilization
    if (utilization < 0.05).any():
        underused = (utilization < 0.05).nonzero(as_tuple=True)[0].tolist()
        print(f"WARNING: Experts {underused} have less than 5% utilization!")
        # In production, send to alerting system (e.g., PagerDuty)
        # send_alert(f"MoE router collapse detected: experts {underused} underused")
    
    return utilization

# Example
if __name__ == "__main__":
    # Simulate router weights where expert 0 gets 90% of tokens
    router_weights = torch.zeros(2, 10, 8)
    router_weights[:, :, 0] = 0.9
    router_weights[:, :, 1:] = 0.1 / 7
    
    monitor_expert_utilization(router_weights, 8)
    # Output: Expert utilization: [0.9, 0.014, 0.014, ...] -> alert

import torch
import time

# Simplified comparison of different architectures

def benchmark_model(model, x, num_runs=50):
    for _ in range(10):
        _ = model(x)
    torch.cuda.synchronize()
    start = time.time()
    for _ in range(num_runs):
        _ = model(x)
    torch.cuda.synchronize()
    return (time.time() - start) / num_runs * 1000

if __name__ == "__main__":
    d_model = 1024
    batch_size = 16
    seq_len = 256
    x = torch.randn(batch_size, seq_len, d_model).cuda()
    
    # Dense
    dense = nn.Sequential(
        nn.Linear(d_model, d_model * 4),
        nn.GELU(),
        nn.Linear(d_model * 4, d_model)
    ).cuda()
    
    # MoE (8 experts, top-2)
    from moe_layer_production import MoELayer
    moe = MoELayer(d_model, num_experts=8, top_k=2).cuda()
    
    # Mixture of Attention (simplified: multiple attention heads with routing)
    # This is a placeholder — real MoA is more complex
    moa = nn.MultiheadAttention(d_model, num_heads=8, batch_first=True).cuda()
    
    print(f"Dense: {benchmark_model(dense, x):.2f} ms")
    print(f"MoE: {benchmark_model(moe, x):.2f} ms")
    print(f"MoA (placeholder): {benchmark_model(moa, x):.2f} ms")
    print("\nNote: MoE is slower per-token but allows larger total model size.")

import torch
import numpy as np
from collections import deque

class MoEMonitor:
    def __init__(self, num_experts, alert_threshold=0.05, window_size=100):
        self.num_experts = num_experts
        self.alert_threshold = alert_threshold
        self.utilization_history = deque(maxlen=window_size)
        self.router_logit_std_history = deque(maxlen=window_size)
        self.token_drop_rate_history = deque(maxlen=window_size)
        
    def log_batch(self, router_weights, tokens_dropped, total_tokens):
        # Utilization
        expert_assignments = router_weights.argmax(dim=-1)
        utilization = torch.bincount(expert_assignments.flatten(), minlength=self.num_experts).float()
        utilization = utilization / utilization.sum()
        self.utilization_history.append(utilization.cpu().numpy())
        
        # Router logit std (approximate from weights)
        # In practice, log the actual logits before softmax
        self.router_logit_std_history.append(router_weights.std().item())
        
        # Token drop rate
        drop_rate = tokens_dropped / total_tokens if total_tokens > 0 else 0
        self.token_drop_rate_history.append(drop_rate)
        
        # Check alerts
        alerts = []
        if (utilization < self.alert_threshold).any():
            underused = (utilization < self.alert_threshold).nonzero(as_tuple=True)[0].tolist()
            alerts.append(f"Experts {underused} underused (utilization < {self.alert_threshold*100}%)")
        if drop_rate > 0.01:
            alerts.append(f"Token drop rate {drop_rate*100:.2f}% > 1%")
        if len(self.router_logit_std_history) > 10:
            avg_std = np.mean(self.router_logit_std_history)
            if router_weights.std() > 5 * avg_std:
                alerts.append(f"Router logit std dev spike: {router_weights.std():.4f} vs avg {avg_std:.4f}")
        
        return alerts
    
    def get_summary(self):
        if not self.utilization_history:
            return {}
        avg_utilization = np.mean(self.utilization_history, axis=0)
        return {
            "avg_utilization": avg_utilization.tolist(),
            "avg_router_std": np.mean(self.router_logit_std_history),
            "avg_drop_rate": np.mean(self.token_drop_rate_history)
        }

# Example usage
if __name__ == "__main__":
    monitor = MoEMonitor(num_experts=8)
    # Simulate a batch
    router_weights = torch.randn(2, 10, 8).softmax(dim=-1)
    alerts = monitor.log_batch(router_weights, tokens_dropped=5, total_tokens=20)
    print("Alerts:", alerts)
    print("Summary:", monitor.get_summary())