Upload src/moe.py

src/moe.py

@@ -0,0 +1,163 @@
+"""
+Mixture of Experts (MoE) Implementation for nanoKimi
+
+This module implements the MoE layer used in Kimi-K2, which allows
+for efficient scaling by routing tokens to different expert networks.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MoELayer(nn.Module):
+    """
+    Mixture of Experts Layer
+    
+    Routes input tokens to different expert networks based on a learned gating function.
+    Only the top-k experts are activated for each token, making the computation sparse.
+    
+    Args:
+        n_embd: embedding dimension
+        num_experts: number of expert networks
+        expert_capacity: capacity of each expert (max tokens per expert)
+        top_k: number of experts to route each token to
+        dropout: dropout probability
+        bias: whether to use bias in linear layers
+    """
+    
+    def __init__(self, n_embd, num_experts=8, expert_capacity=32, top_k=2, dropout=0.0, bias=True):
+        super().__init__()
+        
+        self.n_embd = n_embd
+        self.num_experts = num_experts
+        self.expert_capacity = expert_capacity
+        self.top_k = top_k
+        
+        # Gating network - decides which experts to use
+        self.gate = nn.Linear(n_embd, num_experts, bias=bias)
+        
+        # Expert networks - simple FFN for each expert
+        self.experts = nn.ModuleList([
+            ExpertFFN(n_embd, dropout=dropout, bias=bias) 
+            for _ in range(num_experts)
+        ])
+        
+        # Load balancing loss coefficient
+        self.load_balance_loss_coef = 0.01
+        
+    def forward(self, x):
+        B, T, C = x.shape
+        
+        # Flatten to (B*T, C) for easier processing
+        x_flat = x.view(-1, C)
+        
+        # Compute gating scores
+        gate_logits = self.gate(x_flat)  # (B*T, num_experts)
+        gate_scores = F.softmax(gate_logits, dim=-1)
+        
+        # Select top-k experts for each token
+        top_k_scores, top_k_indices = torch.topk(gate_scores, self.top_k, dim=-1)
+        
+        # Normalize top-k scores
+        top_k_scores = top_k_scores / top_k_scores.sum(dim=-1, keepdim=True)
+        
+        # Initialize output
+        output = torch.zeros_like(x_flat)
+        
+        # Process each expert
+        for expert_idx in range(self.num_experts):
+            # Find tokens assigned to this expert
+            expert_mask = (top_k_indices == expert_idx).any(dim=-1)
+            
+            if expert_mask.sum() == 0:
+                continue
+                
+            # Get tokens for this expert
+            expert_tokens = x_flat[expert_mask]
+            
+            # Apply capacity constraint
+            if expert_tokens.size(0) > self.expert_capacity:
+                # Random sampling if too many tokens
+                perm = torch.randperm(expert_tokens.size(0))[:self.expert_capacity]
+                expert_tokens = expert_tokens[perm]
+                expert_mask_indices = torch.where(expert_mask)[0][perm]
+            else:
+                expert_mask_indices = torch.where(expert_mask)[0]
+            
+            # Process through expert
+            expert_output = self.experts[expert_idx](expert_tokens)
+            
+            # Weight by gating scores and add to output
+            for i, token_idx in enumerate(expert_mask_indices):
+                # Find which position in top_k this expert is for this token
+                expert_positions = (top_k_indices[token_idx] == expert_idx).nonzero(as_tuple=True)[0]
+                if len(expert_positions) > 0:
+                    weight = top_k_scores[token_idx, expert_positions[0]]
+                    output[token_idx] += weight * expert_output[i]
+        
+        # Reshape back to original shape
+        output = output.view(B, T, C)
+        
+        # Compute load balancing loss
+        load_balance_loss = self._compute_load_balance_loss(gate_scores)
+        
+        return output, load_balance_loss
+    
+    def _compute_load_balance_loss(self, gate_scores):
+        """
+        Compute load balancing loss to encourage equal usage of experts
+        """
+        # Compute the fraction of tokens routed to each expert
+        expert_usage = gate_scores.mean(dim=0)  # (num_experts,)
+        
+        # Target is uniform distribution
+        target_usage = torch.ones_like(expert_usage) / self.num_experts
+        
+        # L2 loss between actual and target usage
+        load_balance_loss = F.mse_loss(expert_usage, target_usage)
+        
+        return self.load_balance_loss_coef * load_balance_loss
+
+
+class ExpertFFN(nn.Module):
+    """
+    Expert Feed-Forward Network
+    
+    A simple two-layer MLP that serves as an expert in the MoE layer.
+    """
+    
+    def __init__(self, n_embd, dropout=0.0, bias=True):
+        super().__init__()
+        
+        # Typical GPT-style FFN with 4x expansion
+        self.fc1 = nn.Linear(n_embd, 4 * n_embd, bias=bias)
+        self.fc2 = nn.Linear(4 * n_embd, n_embd, bias=bias)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, x):
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+
+class StandardFFN(nn.Module):
+    """
+    Standard Feed-Forward Network for comparison with MoE
+    """
+    
+    def __init__(self, n_embd, dropout=0.0, bias=True):
+        super().__init__()
+        
+        self.fc1 = nn.Linear(n_embd, 4 * n_embd, bias=bias)
+        self.fc2 = nn.Linear(4 * n_embd, n_embd, bias=bias)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, x):
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x, 0.0  # Return 0 load balance loss for consistency