import torch
import torch.nn as nn
import torch.nn.functional as F
class Expert(nn.Module):
def __init__(self, dim, intermediate_dim):
super().__init__()
self.gate_proj = nn.Linear(dim, intermediate_dim)
self.up_proj = nn.Linear(dim, intermediate_dim)
self.down_proj = nn.Linear(intermediate_dim, dim)
self.act = nn.SiLU()
def forward(self, x):
gate = self.gate_proj(x)
up = self.up_proj(x)
swish = self.act(gate)
output = self.down_proj(swish * up)
return output
class MoELayer(nn.Module):
def __init__(self, dim, intermediate_dim, num_experts, top_k=2):
super().__init__()
self.num_experts = num_experts
self.top_k = top_k
self.dim = dim
# Create expert networks
self.experts = nn.ModuleList([
Expert(dim, intermediate_dim) for _ in range(num_experts)
])
self.router = nn.Linear(dim, num_experts)
def forward(self, hidden_states):
batch_size, seq_len, hidden_dim = hidden_states.shape
# Reshape for expert processing, the compute routing probabilities
hidden_states_reshaped = hidden_states.view(–1, hidden_dim)
router_logits = self.router(hidden_states_reshaped) # (batch_size * seq_len, num_experts)
routing_probs = F.softmax(router_logits, dim=–1)
# Select top-k experts, and scale the probabilities to sum to 1
# output shape: (batch_size * seq_len, k)
top_k_probs, top_k_indices = torch.topk(routing_probs, self.top_k, dim=–1)
top_k_probs = top_k_probs / top_k_probs.sum(dim=–1, keepdim=True)
# Process through selected experts
output = []
for i in range(self.top_k):
expert_idx = top_k_indices[:, i]
expert_probs = top_k_probs[:, i]
# Process each vector in the batch and sequence with the selected expert
expert_output = torch.stack([
self.experts[j](hidden_states_reshaped[j])
for j in expert_idx
], dim=0)
# Weighted sum by routing probability
output.sum(expert_probs.unsqueeze(–1) * expert_output)
# Reshape back to original shape
output = sum(output).view(batch_size, seq_len, hidden_dim)
return output
class MoETransformerLayer(nn.Module):
def __init__(self, dim, intermediate_dim, num_experts, top_k=2, num_heads=8):
super().__init__()
self.attention = nn.MultiheadAttention(dim, num_heads, batch_first=True)
self.moe = MoELayer(dim, intermediate_dim, num_experts, top_k)
self.norm1 = nn.RMSNorm(dim)
self.norm2 = nn.RMSNorm(dim)
def forward(self, x):
# Attention sublayer
input_x = x
x = self.norm1(x)
attn_output, _ = self.attention(x, x, x)
input_x = input_x + attn_output
# MoE sublayer
x = self.norm2(input_x)
moe_output = self.moe(x)
return input_x + moe_output
