training_code/model/transformer.py

"""
1B Parameter Decoder-Only Transformer — built from scratch.

Techniques:
  - RoPE (Rotary Position Embeddings)
  - Grouped Query Attention (GQA)
  - SwiGLU Feed-Forward
  - RMSNorm (pre-norm architecture)
  - Flash Attention 2 (via PyTorch SDPA)
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from .config import ModelConfig


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * norm).type_as(x) * self.weight


def precompute_rope_freqs(dim: int, max_seq_len: int, theta: float = 10000.0) -> torch.Tensor:
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
    t = torch.arange(max_seq_len, dtype=torch.float32)
    freqs = torch.outer(t, freqs)
    return torch.polar(torch.ones_like(freqs), freqs)  # complex64


def apply_rope(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor):
    B, S, H, D = xq.shape
    xq_c = torch.view_as_complex(xq.float().reshape(B, S, H, D // 2, 2))
    xk_c = torch.view_as_complex(xk.float().reshape(B, S, xk.shape[2], D // 2, 2))
    freqs = freqs_cis[:S].clone().unsqueeze(0).unsqueeze(2)
    xq_out = torch.view_as_real(xq_c * freqs).flatten(3)
    xk_out = torch.view_as_real(xk_c * freqs).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)


class GroupedQueryAttention(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_kv_heads
        self.head_dim = config.head_dim
        self.num_groups = self.num_heads // self.num_kv_heads

        self.wq = nn.Linear(config.hidden_dim, self.num_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(config.hidden_dim, self.num_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(config.hidden_dim, self.num_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(self.num_heads * self.head_dim, config.hidden_dim, bias=False)

    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
        B, S, _ = x.shape

        q = self.wq(x).view(B, S, self.num_heads, self.head_dim)
        k = self.wk(x).view(B, S, self.num_kv_heads, self.head_dim)
        v = self.wv(x).view(B, S, self.num_kv_heads, self.head_dim)

        q, k = apply_rope(q, k, freqs_cis)

        # Expand KV heads for GQA
        if self.num_groups > 1:
            k = k.unsqueeze(3).expand(B, S, self.num_kv_heads, self.num_groups, self.head_dim)
            k = k.reshape(B, S, self.num_heads, self.head_dim)
            v = v.unsqueeze(3).expand(B, S, self.num_kv_heads, self.num_groups, self.head_dim)
            v = v.reshape(B, S, self.num_heads, self.head_dim)

        # (B, num_heads, S, head_dim) for SDPA
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = out.transpose(1, 2).contiguous().view(B, S, -1)
        return self.wo(out)


class SwiGLUFFN(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.w_gate = nn.Linear(config.hidden_dim, config.intermediate_dim, bias=False)
        self.w_up = nn.Linear(config.hidden_dim, config.intermediate_dim, bias=False)
        self.w_down = nn.Linear(config.intermediate_dim, config.hidden_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class TransformerBlock(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.attention_norm = RMSNorm(config.hidden_dim, eps=config.rms_norm_eps)
        self.attention = GroupedQueryAttention(config)
        self.ffn_norm = RMSNorm(config.hidden_dim, eps=config.rms_norm_eps)
        self.ffn = SwiGLUFFN(config)

    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
        x = x + self.attention(self.attention_norm(x), freqs_cis)
        x = x + self.ffn(self.ffn_norm(x))
        return x


class Transformer(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config

        self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_dim)
        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_layers)])
        self.norm = RMSNorm(config.hidden_dim, eps=config.rms_norm_eps)
        self.output = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)

        # Pre-compute RoPE frequencies
        self.register_buffer(
            "freqs_cis",
            precompute_rope_freqs(config.head_dim, config.max_seq_len * 2, config.rope_theta),
            persistent=False,
        )

        self._init_weights()

    def _init_weights(self):
        """Initialize with scaled normal, following GPT-NeoX / LLaMA conventions."""
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)

        # Scale residual projections by 1/sqrt(2*num_layers)
        scale = (2 * self.config.num_layers) ** -0.5
        for layer in self.layers:
            nn.init.normal_(layer.attention.wo.weight, mean=0.0, std=0.02 * scale)
            nn.init.normal_(layer.ffn.w_down.weight, mean=0.0, std=0.02 * scale)

    def forward(self, tokens: torch.Tensor, targets: torch.Tensor = None):
        B, S = tokens.shape
        h = self.tok_embeddings(tokens)

        freqs_cis = self.freqs_cis[:S]
        for layer in self.layers:
            h = layer(h, freqs_cis)
        h = self.norm(h)
        logits = self.output(h)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                ignore_index=-100,
            )
        return logits, loss