File size: 7,706 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 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 | """
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
|