src/part1/moe.py · Chiru20/ANLP

Upload ANLP Assignment 2 Model Card, Configs, Report, and Evaluation Deliverables

f90f02b 4 days ago

7.64 kB

"""
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)

	"""
	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)