src/part1/model.py · Chiru20/ANLP

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

f90f02b 4 days ago

7.71 kB

"""
Decoder-only Transformer implementation with drop-in MoE support.
Follows Appendix Listing 1 and requirements of Assignment 2.
"""

import math
from typing import Optional, Dict, Any, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F

from .config import TransformerConfig
from .moe import MoE, MLP


class CausalSelfAttention(nn.Module):
    """
    Multi-head causal self-attention mechanism utilizing scaled_dot_product_attention.
    """
    def __init__(self, config: TransformerConfig):
        super().__init__()
        assert config.D_model % config.N_heads == 0
        self.D_model = config.D_model
        self.N_heads = config.N_heads
        self.head_dim = config.D_model // config.N_heads
        self.dropout = config.dropout

        # Combined Q, K, V projection
        self.c_attn = nn.Linear(config.D_model, 3 * config.D_model, bias=True)
        # Output projection
        self.c_proj = nn.Linear(config.D_model, config.D_model, bias=True)
        self.resid_dropout = nn.Dropout(config.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape  # Batch size, sequence length, hidden dimension

        # Compute queries, keys, and values
        qkv = self.c_attn(x)
        q, k, v = qkv.split(self.D_model, dim=2)

        # Reshape for multi-head attention: (B, N_heads, T, head_dim)
        q = q.view(B, T, self.N_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.N_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.N_heads, self.head_dim).transpose(1, 2)

        # PyTorch 2.0+ scaled_dot_product_attention with causal masking
        y = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=None,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=True
        )

        # Re-assemble head outputs: (B, T, C)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        out = self.resid_dropout(self.c_proj(y))
        return out


class TransformerBlock(nn.Module):
    """
    Standard Transformer Decoder Block containing:
    - LayerNorm + Causal Attention + Residual
    - LayerNorm + FFN (MLP or MoE) + Residual
    """
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.D_model)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.D_model)
        self.ffn = MoE(config) if config.use_moe else MLP(config)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln_1(x))
        x = x + self.ffn(self.ln_2(x))
        return x


class Transformer(nn.Module):
    """
    Decoder-only Transformer Language Model.
    Supports drop-in MoE or standard MLP feed-forward networks.
    """
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.config = config

        self.wte = nn.Embedding(config.N_vocab, config.D_model)
        self.wpe = nn.Embedding(config.N_ctx, config.D_model)
        self.drop = nn.Dropout(config.dropout)

        self.blocks = nn.ModuleList([
            TransformerBlock(config) for _ in range(config.N_layers)
        ])
        self.ln_f = nn.LayerNorm(config.D_model)
        self.head = nn.Linear(config.D_model, config.N_vocab, bias=False)

        # Tie weights between token embeddings and output projection
        self.head.weight = self.wte.weight

        # Initialize weights
        self.apply(self._init_weights)

    def _init_weights(self, module: nn.Module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.LayerNorm):
            nn.init.ones_(module.weight)
            nn.init.zeros_(module.bias)

    def get_aux_loss(self) -> torch.Tensor:
        """Sums auxiliary load-balancing losses across all MoE layers."""
        if not self.config.use_moe:
            return torch.tensor(0.0, device=next(self.parameters()).device)

        total_aux = torch.tensor(0.0, device=next(self.parameters()).device)
        for block in self.blocks:
            if isinstance(block.ffn, MoE):
                total_aux = total_aux + block.ffn.last_aux_loss.to(total_aux.device)
        return total_aux

    def get_active_parameters(self) -> int:
        """
        Computes the number of parameters actively participating in forward computation per token.
        For dense models, active_parameters == total_parameters.
        For MoE models, only router and active experts count.
        """
        active = 0
        # Embeddings
        active += self.wte.weight.numel()
        active += self.wpe.weight.numel()
        # Final norm and tied lm_head
        active += sum(p.numel() for p in self.ln_f.parameters())
        # (self.head shares weights with self.wte, so not double-counted)

        for block in self.blocks:
            active += sum(p.numel() for p in block.ln_1.parameters())
            active += sum(p.numel() for p in block.attn.parameters())
            active += sum(p.numel() for p in block.ln_2.parameters())
            if isinstance(block.ffn, MoE):
                active += block.ffn.get_active_parameters()
            else:
                active += sum(p.numel() for p in block.ffn.parameters())

        return active

    def forward(
        self,
        idx: torch.Tensor,
        targets: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]:
        """
        Forward pass for Transformer.
        Args:
            idx: Tensor of token indices with shape (B, T)
            targets: Optional ground-truth target tokens with shape (B, T)
        Returns:
            logits: (B, T, N_vocab)
            loss: Cross-entropy loss + aux_loss (if targets provided)
            aux_loss: Load balancing auxiliary loss
        """
        device = idx.device
        B, T = idx.shape
        assert T <= self.config.N_ctx, f"Sequence length {T} exceeds context window {self.config.N_ctx}"

        pos = torch.arange(0, T, dtype=torch.long, device=device).unsqueeze(0)  # (1, T)

        # Token + positional embeddings
        tok_emb = self.wte(idx)       # (B, T, D_model)
        pos_emb = self.wpe(pos)       # (1, T, D_model)
        x = self.drop(tok_emb + pos_emb)

        # Pass through Transformer blocks
        for block in self.blocks:
            x = block(x)

        x = self.ln_f(x)
        logits = self.head(x)  # (B, T, N_vocab)

        loss = None
        aux_loss = self.get_aux_loss()

        if targets is not None:
            # Shift tokens for next-token prediction
            if targets.shape[1] == logits.shape[1]:
                # Standard next-token prediction: position t predicts token t+1
                shift_logits = logits[:, :-1, :].contiguous().view(-1, self.config.N_vocab)
                shift_targets = targets[:, 1:].contiguous().view(-1)
            elif targets.shape[1] == logits.shape[1] - 1:
                # Targets are already pre-shifted by 1 token
                shift_logits = logits[:, :-1, :].contiguous().view(-1, self.config.N_vocab)
                shift_targets = targets.contiguous().view(-1)
            else:
                shift_logits = logits.contiguous().view(-1, self.config.N_vocab)
                shift_targets = targets.contiguous().view(-1)

            ce_loss = F.cross_entropy(shift_logits, shift_targets, ignore_index=-100)
            loss = ce_loss + aux_loss

        return logits, loss, aux_loss

	"""
	Decoder-only Transformer implementation with drop-in MoE support.
	Follows Appendix Listing 1 and requirements of Assignment 2.
	"""

	import math
	from typing import Optional, Dict, Any, Tuple
	import torch
	import torch.nn as nn
	import torch.nn.functional as F

	from .config import TransformerConfig
	from .moe import MoE, MLP


	class CausalSelfAttention(nn.Module):
	"""
	Multi-head causal self-attention mechanism utilizing scaled_dot_product_attention.
	"""
	def __init__(self, config: TransformerConfig):
	super().__init__()
	assert config.D_model % config.N_heads == 0
	self.D_model = config.D_model
	self.N_heads = config.N_heads
	self.head_dim = config.D_model // config.N_heads
	self.dropout = config.dropout

	# Combined Q, K, V projection
	self.c_attn = nn.Linear(config.D_model, 3 * config.D_model, bias=True)
	# Output projection
	self.c_proj = nn.Linear(config.D_model, config.D_model, bias=True)
	self.resid_dropout = nn.Dropout(config.dropout)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	B, T, C = x.shape # Batch size, sequence length, hidden dimension

	# Compute queries, keys, and values
	qkv = self.c_attn(x)
	q, k, v = qkv.split(self.D_model, dim=2)

	# Reshape for multi-head attention: (B, N_heads, T, head_dim)
	q = q.view(B, T, self.N_heads, self.head_dim).transpose(1, 2)
	k = k.view(B, T, self.N_heads, self.head_dim).transpose(1, 2)
	v = v.view(B, T, self.N_heads, self.head_dim).transpose(1, 2)

	# PyTorch 2.0+ scaled_dot_product_attention with causal masking
	y = F.scaled_dot_product_attention(
	q, k, v,
	attn_mask=None,
	dropout_p=self.dropout if self.training else 0.0,
	is_causal=True
	)

	# Re-assemble head outputs: (B, T, C)
	y = y.transpose(1, 2).contiguous().view(B, T, C)
	out = self.resid_dropout(self.c_proj(y))
	return out


	class TransformerBlock(nn.Module):
	"""
	Standard Transformer Decoder Block containing:
	- LayerNorm + Causal Attention + Residual
	- LayerNorm + FFN (MLP or MoE) + Residual
	"""
	def __init__(self, config: TransformerConfig):
	super().__init__()
	self.ln_1 = nn.LayerNorm(config.D_model)
	self.attn = CausalSelfAttention(config)
	self.ln_2 = nn.LayerNorm(config.D_model)
	self.ffn = MoE(config) if config.use_moe else MLP(config)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	x = x + self.attn(self.ln_1(x))
	x = x + self.ffn(self.ln_2(x))
	return x


	class Transformer(nn.Module):
	"""
	Decoder-only Transformer Language Model.
	Supports drop-in MoE or standard MLP feed-forward networks.
	"""
	def __init__(self, config: TransformerConfig):
	super().__init__()
	self.config = config

	self.wte = nn.Embedding(config.N_vocab, config.D_model)
	self.wpe = nn.Embedding(config.N_ctx, config.D_model)
	self.drop = nn.Dropout(config.dropout)

	self.blocks = nn.ModuleList([
	TransformerBlock(config) for _ in range(config.N_layers)
	])
	self.ln_f = nn.LayerNorm(config.D_model)
	self.head = nn.Linear(config.D_model, config.N_vocab, bias=False)

	# Tie weights between token embeddings and output projection
	self.head.weight = self.wte.weight

	# Initialize weights
	self.apply(self._init_weights)

	def _init_weights(self, module: nn.Module):
	if isinstance(module, (nn.Linear, nn.Embedding)):
	nn.init.normal_(module.weight, mean=0.0, std=0.02)
	if isinstance(module, nn.Linear) and module.bias is not None:
	nn.init.zeros_(module.bias)
	elif isinstance(module, nn.LayerNorm):
	nn.init.ones_(module.weight)
	nn.init.zeros_(module.bias)

	def get_aux_loss(self) -> torch.Tensor:
	"""Sums auxiliary load-balancing losses across all MoE layers."""
	if not self.config.use_moe:
	return torch.tensor(0.0, device=next(self.parameters()).device)

	total_aux = torch.tensor(0.0, device=next(self.parameters()).device)
	for block in self.blocks:
	if isinstance(block.ffn, MoE):
	total_aux = total_aux + block.ffn.last_aux_loss.to(total_aux.device)
	return total_aux

	def get_active_parameters(self) -> int:
	"""
	Computes the number of parameters actively participating in forward computation per token.
	For dense models, active_parameters == total_parameters.
	For MoE models, only router and active experts count.
	"""
	active = 0
	# Embeddings
	active += self.wte.weight.numel()
	active += self.wpe.weight.numel()
	# Final norm and tied lm_head
	active += sum(p.numel() for p in self.ln_f.parameters())
	# (self.head shares weights with self.wte, so not double-counted)

	for block in self.blocks:
	active += sum(p.numel() for p in block.ln_1.parameters())
	active += sum(p.numel() for p in block.attn.parameters())
	active += sum(p.numel() for p in block.ln_2.parameters())
	if isinstance(block.ffn, MoE):
	active += block.ffn.get_active_parameters()
	else:
	active += sum(p.numel() for p in block.ffn.parameters())

	return active

	def forward(
	self,
	idx: torch.Tensor,
	targets: Optional[torch.Tensor] = None
	) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]:
	"""
	Forward pass for Transformer.
	Args:
	idx: Tensor of token indices with shape (B, T)
	targets: Optional ground-truth target tokens with shape (B, T)
	Returns:
	logits: (B, T, N_vocab)
	loss: Cross-entropy loss + aux_loss (if targets provided)
	aux_loss: Load balancing auxiliary loss
	"""
	device = idx.device
	B, T = idx.shape
	assert T <= self.config.N_ctx, f"Sequence length {T} exceeds context window {self.config.N_ctx}"

	pos = torch.arange(0, T, dtype=torch.long, device=device).unsqueeze(0) # (1, T)

	# Token + positional embeddings
	tok_emb = self.wte(idx) # (B, T, D_model)
	pos_emb = self.wpe(pos) # (1, T, D_model)
	x = self.drop(tok_emb + pos_emb)

	# Pass through Transformer blocks
	for block in self.blocks:
	x = block(x)

	x = self.ln_f(x)
	logits = self.head(x) # (B, T, N_vocab)

	loss = None
	aux_loss = self.get_aux_loss()

	if targets is not None:
	# Shift tokens for next-token prediction
	if targets.shape[1] == logits.shape[1]:
	# Standard next-token prediction: position t predicts token t+1
	shift_logits = logits[:, :-1, :].contiguous().view(-1, self.config.N_vocab)
	shift_targets = targets[:, 1:].contiguous().view(-1)
	elif targets.shape[1] == logits.shape[1] - 1:
	# Targets are already pre-shifted by 1 token
	shift_logits = logits[:, :-1, :].contiguous().view(-1, self.config.N_vocab)
	shift_targets = targets.contiguous().view(-1)
	else:
	shift_logits = logits.contiguous().view(-1, self.config.N_vocab)
	shift_targets = targets.contiguous().view(-1)

	ce_loss = F.cross_entropy(shift_logits, shift_targets, ignore_index=-100)
	loss = ce_loss + aux_loss

	return logits, loss, aux_loss