| """Pre-normalized decoder-only transformer with interchangeable FFNs.""" |
|
|
| import torch |
| from torch import Tensor, nn |
| import torch.nn.functional as F |
|
|
| from moe import TransformerConfig, MLP, MoE |
|
|
|
|
| class Block(nn.Module): |
| """Apply pre-normalized causal attention and the selected FFN.""" |
|
|
| def __init__(self, config): |
| """Create attention projections, normalization layers and the FFN.""" |
| super().__init__() |
| self.norm1 = nn.LayerNorm(config.d_model) |
| self.norm2 = nn.LayerNorm(config.d_model) |
| self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False) |
| self.proj = nn.Linear(config.d_model, config.d_model, bias=False) |
| self.ffn = MLP(config) if config.num_experts == 1 else MoE(config) |
| self.heads = config.n_heads |
| self.dropout = config.dropout |
|
|
| def forward(self, x, attention_mask, routing_mask=None): |
| """Apply causal attention and the FFN to (batch, tokens, hidden) states.""" |
| b, t, d = x.shape |
| q, k, v = self.qkv(self.norm1(x)).chunk(3, -1) |
| q, k, v = [ |
| z.reshape(b, t, self.heads, d // self.heads).transpose(1, 2) |
| for z in (q, k, v) |
| ] |
| mask = None |
| if attention_mask is not None and not bool(attention_mask.all()): |
| mask = ( |
| torch.ones(t, t, dtype=torch.bool, device=x.device).tril()[None, None] |
| & attention_mask[:, None, None, :].bool() |
| ) |
| attended = F.scaled_dot_product_attention( |
| q, |
| k, |
| v, |
| attn_mask=mask, |
| is_causal=mask is None, |
| dropout_p=self.dropout if self.training else 0.0, |
| ) |
| x = x + F.dropout( |
| self.proj(attended.transpose(1, 2).reshape(b, t, d)), |
| self.dropout, |
| self.training, |
| ) |
| if isinstance(self.ffn, MoE): |
| self.ffn.tracking_mask = ( |
| attention_mask if routing_mask is None else routing_mask |
| ) |
| return x + F.dropout(self.ffn(self.norm2(x)), self.dropout, self.training) |
|
|
|
|
| class CausalLM(nn.Module): |
| """Decode tokens using tied embeddings and dense or MoE blocks.""" |
|
|
| def __init__(self, config: TransformerConfig): |
| """Build the decoder, initialize its weights and tie the output embedding.""" |
| super().__init__() |
| if config.d_model % config.n_heads: |
| raise ValueError("d_model must be divisible by n_heads") |
| self.config = config |
| self.context_length = config.context_length |
| self.usage_mask = None |
| self.token_embedding = nn.Embedding(config.vocab_size, config.d_model) |
| self.position_embedding = nn.Embedding(config.context_length, config.d_model) |
| self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layers)]) |
| self.norm = nn.LayerNorm(config.d_model) |
| self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) |
| self.apply(self._initialize) |
| self.lm_head.weight = self.token_embedding.weight |
|
|
| @staticmethod |
| def _initialize(module): |
| """Initialize projection and embedding weights normally, with zero biases.""" |
| if isinstance(module, (nn.Linear, nn.Embedding)): |
| nn.init.normal_(module.weight, std=0.02) |
| if getattr(module, "bias", None) is not None: |
| nn.init.zeros_(module.bias) |
|
|
| @property |
| def aux_loss(self): |
| """Return the mean balancing loss, or zero for a dense model.""" |
| losses = [ |
| block.ffn.aux_loss for block in self.blocks if isinstance(block.ffn, MoE) |
| ] |
| return ( |
| torch.stack(losses).mean() |
| if losses |
| else self.token_embedding.weight.new_zeros(()) |
| ) |
|
|
| def forward( |
| self, input_ids: Tensor, attention_mask: Tensor | None = None |
| ) -> Tensor: |
| """Return logits of shape (batch, tokens, vocabulary) for a token sequence. |
| |
| Use learned absolute positions and causal attention. The optional mask |
| marks valid input positions; sequences must fit the configured context. |
| """ |
| if input_ids.ndim != 2 or input_ids.shape[1] > self.context_length: |
| raise ValueError("Expected (batch, sequence) within context_length") |
| positions = torch.arange(input_ids.shape[1], device=input_ids.device) |
| hidden = ( |
| self.token_embedding(input_ids) + self.position_embedding(positions)[None] |
| ) |
| for block in self.blocks: |
| hidden = block(hidden, attention_mask, self.usage_mask) |
| return self.lm_head(self.norm(hidden)) |
|
|
| def parameter_counts(self): |
| """Count stored and active parameters, counting tied weights once.""" |
| total = sum(p.numel() for p in self.parameters()) |
| inactive = sum( |
| sum(p.numel() for p in b.ffn.parameters()) - b.ffn.active_parameter_count() |
| for b in self.blocks |
| if isinstance(b.ffn, MoE) |
| ) |
| return {"total": total, "active": total - inactive} |
|
|
|
|
| class DenseCausalLM(CausalLM): |
| """Provide the dense decoder through the earlier constructor interface.""" |
|
|
| def __init__( |
| self, vocab_size, d_model, n_layers, n_heads, d_ff, context_length, dropout |
| ): |
| """Create a decoder configuration with one dense feedforward expert.""" |
| super().__init__( |
| TransformerConfig( |
| d_model, |
| d_ff, |
| vocab_size, |
| context_length, |
| n_layers=n_layers, |
| n_heads=n_heads, |
| dropout=dropout, |
| ) |
| ) |
|
|
|
|
| def build_model(config): |
| """Construct the configured dense or MoE causal language model.""" |
| m = config.model |
| f = m.ffn |
| return CausalLM( |
| TransformerConfig( |
| m.d_model, |
| m.d_ff, |
| m.vocab_size, |
| m.context_length, |
| f.num_experts, |
| f.top_k, |
| f.shared_expert, |
| f.match_active_parameters, |
| m.n_layers, |
| m.n_heads, |
| m.dropout, |
| ) |
| ) |
|
|
