File size: 7,641 Bytes
f90f02b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 | """
Mixture of Experts (MoE) implementation for Transformer architectures.
Implements the 5 required FFN variants:
1. Standard 2-layer MLP (dense baseline)
2. MoE with 4 experts, 1 active at a time (top-1 routing)
3. MoE with 4 experts, 2 active at a time (top-2 routing)
4. MoE with 1 shared expert (always active) + 3 routed experts (1 active at a time, 2 total active)
5. MoE with 4 experts, 2 active at a time, matched active parameters to dense baseline (1)
"""
import math
from typing import Dict, Tuple, Optional, List
import torch
import torch.nn as nn
import torch.nn.functional as F
from .config import TransformerConfig
class MLP(nn.Module):
"""
Standard 2-layer Feed-Forward Network (Dense Baseline).
Input: (B, N_ctx, D_model)
Output: (B, N_ctx, D_model)
"""
def __init__(self, config: TransformerConfig):
super().__init__()
self.D_model = config.D_model
self.D_ff = config.D_ff
self.w1 = nn.Linear(config.D_model, config.D_ff, bias=True)
self.act = nn.GELU()
self.dropout = nn.Dropout(config.dropout)
self.w2 = nn.Linear(config.D_ff, config.D_model, bias=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x: (B, N_ctx, D_model)
h = self.act(self.w1(x))
h = self.dropout(h)
out = self.dropout(self.w2(h))
return out
class Expert(nn.Module):
"""
Individual Expert Feed-Forward Network.
Input: (..., D_model)
Output: (..., D_model)
"""
def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
super().__init__()
self.d_model = d_model
self.d_ff = d_ff
self.w1 = nn.Linear(d_model, d_ff, bias=True)
self.act = nn.GELU()
self.dropout = nn.Dropout(dropout)
self.w2 = nn.Linear(d_ff, d_model, bias=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
h = self.act(self.w1(x))
h = self.dropout(h)
return self.dropout(self.w2(h))
class MoE(nn.Module):
"""
Mixture-of-Experts Module with drop-in replacement support for MLP.
Supports Top-1, Top-2, Shared+Routed, and Matched-Active configurations.
"""
def __init__(self, config: TransformerConfig):
super().__init__()
self.config = config
self.D_model = config.D_model
self.ablation_variant = config.ablation_variant
self.aux_loss_coef = config.aux_loss_coef
# Determine architecture based on ablation variant
if self.ablation_variant == "shared_routed":
# 1 shared expert + 3 routed experts
self.num_routed_experts = 3
self.shared_expert = Expert(config.D_model, config.expert_d_ff, config.dropout)
self.routed_experts = nn.ModuleList([
Expert(config.D_model, config.expert_d_ff, config.dropout)
for _ in range(self.num_routed_experts)
])
self.router = nn.Linear(config.D_model, self.num_routed_experts, bias=False)
self.top_k = 1
else:
# top1, top2, matched_active
self.num_routed_experts = config.num_experts
self.shared_expert = None
self.routed_experts = nn.ModuleList([
Expert(config.D_model, config.expert_d_ff, config.dropout)
for _ in range(self.num_routed_experts)
])
self.router = nn.Linear(config.D_model, self.num_routed_experts, bias=False)
self.top_k = config.num_active_experts
# Running telemetry for expert utilization and language analysis
self.register_buffer(
"expert_selection_counts",
torch.zeros(self.num_routed_experts, dtype=torch.int64),
persistent=False
)
self.last_expert_indices: Optional[torch.Tensor] = None
self.last_aux_loss: torch.Tensor = torch.tensor(0.0)
def get_active_parameters(self) -> int:
"""Computes active parameters per token for this layer."""
router_params = sum(p.numel() for p in self.router.parameters())
if self.ablation_variant == "shared_routed":
shared_params = sum(p.numel() for p in self.shared_expert.parameters())
routed_params = sum(p.numel() for p in self.routed_experts[0].parameters()) * self.top_k
return router_params + shared_params + routed_params
else:
routed_params = sum(p.numel() for p in self.routed_experts[0].parameters()) * self.top_k
return router_params + routed_params
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Forward pass for MoE layer.
Args:
x: Input tensor of shape (B, N_ctx, D_model)
Returns:
out: Output tensor of shape (B, N_ctx, D_model)
"""
batch_size, seq_len, d_model = x.shape
x_flat = x.reshape(-1, d_model) # (N, D) where N = B * N_ctx
num_tokens = x_flat.size(0)
# 1. Router forward
router_logits = self.router(x_flat) # (N, num_routed_experts)
router_probs = F.softmax(router_logits, dim=-1) # (N, num_routed_experts)
# 2. Expert selection
if self.top_k == 1:
top_weights, top_indices = router_probs.max(dim=-1, keepdim=True) # (N, 1)
else:
top_weights, top_indices = torch.topk(router_probs, k=self.top_k, dim=-1) # (N, top_k)
# Re-normalize weights among selected top-k experts
top_weights = top_weights / (top_weights.sum(dim=-1, keepdim=True) + 1e-9)
# Cache routing decisions for downstream language heatmap extraction
self.last_expert_indices = top_indices.detach()
# Update cumulative selection counts
with torch.no_grad():
for i in range(self.num_routed_experts):
self.expert_selection_counts[i] += (top_indices == i).sum().item()
# 3. Compute auxiliary load-balancing loss (Switch / GShard formulation)
# L_aux = alpha * num_experts * sum_i(f_i * P_i)
# where f_i is fraction of tokens dispatched to expert i, P_i is average routing prob
num_experts = self.num_routed_experts
expert_mask = F.one_hot(top_indices, num_classes=num_experts) # (N, top_k, num_experts)
tokens_per_expert = expert_mask.sum(dim=1).sum(dim=0).float() # (num_experts,)
f_i = tokens_per_expert / (num_tokens * self.top_k + 1e-9)
P_i = router_probs.mean(dim=0) # (num_experts,)
self.last_aux_loss = self.aux_loss_coef * (num_experts * torch.sum(f_i * P_i))
# 4. Dispatch tokens through routed experts
routed_out = torch.zeros_like(x_flat)
for expert_idx in range(num_experts):
# Check which tokens selected this expert
# top_indices shape: (N, top_k)
token_matches = (top_indices == expert_idx)
if not token_matches.any():
continue
token_indices, k_positions = torch.where(token_matches)
selected_tokens = x_flat[token_indices] # (M, D)
weights = top_weights[token_indices, k_positions].unsqueeze(-1) # (M, 1)
expert_output = self.routed_experts[expert_idx](selected_tokens)
routed_out.index_add_(0, token_indices, expert_output * weights)
# 5. Add shared expert output if applicable
if self.shared_expert is not None:
shared_out = self.shared_expert(x_flat)
total_out = shared_out + routed_out
else:
total_out = routed_out
return total_out.reshape(batch_size, seq_len, d_model)
|