model.py

"""
SmolGPT — 参考 HuggingFace SmolLM2-135M 的 "深而窄" 设计

设计理念（与同仓库 deep-gpt 对比）：
  - 深：n_layer 12 → 30（参考 SmolLM2 的 30 层）
  - 窄：n_embd 768 → 512，head_dim 保持 64
  - GQA：分组查询注意力 8 Q heads / 2 KV heads（4:1）
        类比 SmolLM2 的 9 Q / 3 KV，减小 KV cache 占用
  - ReLU² 前馈（Primer 论文），hidden_dim=1728；
    相比 SwiGLU 少 1 个矩阵，同参数预算下更宽
  - RMSNorm + QK-norm，RoPE theta=100000（参考 SmolLM2）
  - 残差出口零初始化（c_proj.weight = 0，起步即恒等映射）、tied embeddings

总参数：约 98.5M（满足 ≤100M 约束）
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


def precompute_freqs(head_dim: int, seq_len: int, theta: float = 100000.0):
    """预计算 RoPE 旋转频率（cos/sin），形状 (seq_len, head_dim)。"""
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
    t = torch.arange(seq_len, dtype=torch.float32)
    freqs = torch.outer(t, inv_freq)  # (seq_len, head_dim//2)
    cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1)
    sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1)
    return cos, sin


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def apply_rope(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
    cos = cos.unsqueeze(0).unsqueeze(0)  # (1, 1, T, head_dim)
    sin = sin.unsqueeze(0).unsqueeze(0)
    q = q * cos + rotate_half(q) * sin
    k = k * cos + rotate_half(k) * sin
    return q, k


class RMSNorm(nn.Module):
    """RMS 层归一化：相比 LayerNorm 无均值偏移，更快且数值稳定。"""

    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 * self.weight).type_as(x)


class GroupedQueryAttention(nn.Module):
    """分组查询注意力（GQA）+ QK-norm + RoPE + Flash SDPA。

    Q heads 数量 > KV heads 数量，K/V 在 head 维度上被多个 Q 共享。
    显著降低注意力的参数量与推理 KV cache 占用。
    """

    def __init__(self, n_embd: int, n_head: int, n_kv_head: int):
        super().__init__()
        assert n_embd % n_head == 0, f"n_embd={n_embd} 必须能被 n_head={n_head} 整除"
        assert n_head % n_kv_head == 0, f"n_head={n_head} 必须能被 n_kv_head={n_kv_head} 整除"

        self.n_head = n_head
        self.n_kv_head = n_kv_head
        self.head_dim = n_embd // n_head
        self.n_rep = n_head // n_kv_head  # 每个 KV head 服务多少个 Q head

        # 分开三组投影：Q 全宽，K/V 收窄到 n_kv_head * head_dim
        self.q_proj = nn.Linear(n_embd, n_head * self.head_dim, bias=False)
        self.k_proj = nn.Linear(n_embd, n_kv_head * self.head_dim, bias=False)
        self.v_proj = nn.Linear(n_embd, n_kv_head * self.head_dim, bias=False)
        self.o_proj = nn.Linear(n_embd, n_embd, bias=False)

        # QK-norm：每个 head 独立归一化，稳定大 batch / 长训练时的 attention logits
        self.q_norm = RMSNorm(self.head_dim)
        self.k_norm = RMSNorm(self.head_dim)

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape

        q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv_head, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv_head, self.head_dim).transpose(1, 2)

        # QK-norm（在 RoPE 之前）
        q = self.q_norm(q)
        k = self.k_norm(k)

        q, k = apply_rope(q, k, cos[:T], sin[:T])

        # GQA：把 K/V 在 head 维度上重复 n_rep 次以匹配 Q
        # 比 SDPA 的 enable_gqa 参数更通用（兼容旧版 PyTorch）
        if self.n_rep > 1:
            k = k.repeat_interleave(self.n_rep, dim=1)
            v = v.repeat_interleave(self.n_rep, dim=1)

        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.o_proj(y.transpose(1, 2).contiguous().view(B, T, C))


class ReLU2MLP(nn.Module):
    """ReLU² 前馈：c_proj(ReLU(c_fc(x))²)。

    来自 Primer 论文 (https://arxiv.org/abs/2109.08668v2)。相比 SwiGLU：
      - 只需 2 个矩阵（vs 3），同参数预算下 hidden_dim 可放大 1.5×
      - 没有门控分支，前向/反向都更快
      - LM 任务上质量与 SwiGLU 持平或略优（modded-nanogpt 验证）
    """

    def __init__(self, n_embd: int, hidden_dim: int):
        super().__init__()
        self.c_fc = nn.Linear(n_embd, hidden_dim, bias=False)
        self.c_proj = nn.Linear(hidden_dim, n_embd, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.c_proj(F.relu(self.c_fc(x)).square())


class Block(nn.Module):
    """Pre-norm 块：RMSNorm → GQA → 残差 → RMSNorm → ReLU² → 残差。"""

    def __init__(self, n_embd: int, n_head: int, n_kv_head: int, hidden_dim: int):
        super().__init__()
        self.norm1 = RMSNorm(n_embd)
        self.attn = GroupedQueryAttention(n_embd, n_head, n_kv_head)
        self.norm2 = RMSNorm(n_embd)
        self.ffn = ReLU2MLP(n_embd, hidden_dim)

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.norm1(x), cos, sin)
        x = x + self.ffn(self.norm2(x))
        return x


class SmolGPT(nn.Module):
    """
    深而窄的 GPT 模型，约 98.5M 参数。

    默认超参数：
      n_embd=512, n_layer=30, n_head=8, n_kv_head=2 (head_dim=64), hidden_dim=1728
    （hidden_dim 从原 SwiGLU 的 1152 调到 ReLU² 的 1728，保持 MLP 总参数
     不变：SwiGLU 用 3×512×1152 = ReLU² 用 2×512×1728）
    """

    cos: torch.Tensor
    sin: torch.Tensor

    def __init__(
        self,
        vocab_size: int = 50257,
        n_embd: int = 512,
        n_head: int = 8,
        n_kv_head: int = 2,
        n_layer: int = 30,
        block_size: int = 1024,
        hidden_dim: int = 1728,
        rope_theta: float = 100000.0,
    ):
        super().__init__()
        self.block_size = block_size

        self.token_emb = nn.Embedding(vocab_size, n_embd)

        head_dim = n_embd // n_head
        cos, sin = precompute_freqs(head_dim, block_size, theta=rope_theta)
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)

        self.blocks = nn.ModuleList(
            [Block(n_embd, n_head, n_kv_head, hidden_dim) for _ in range(n_layer)]
        )
        self.norm_f = RMSNorm(n_embd)

        # 权重绑定：lm_head 与 token_emb 共享，节省 ~26M 参数
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        self.lm_head.weight = self.token_emb.weight

        self._init_weights()

    def _init_weights(self):
        """GPT-2 风格初始化 + 残差出口零初始化。

        残差出口（GQA 的 o_proj、ReLU2MLP 的 c_proj）置零意味着训练起步时
        每个 Block 都是恒等映射 `x ← x + 0`，模型等价于 Embed → Norm → LM_Head
        的浅层模型。残差贡献从第一步开始由训练过程精确学出，loss 曲线起步
        更平滑（modded-nanogpt / Primer 等多个实验验证）。
        """
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
        # 残差流出口：attention 的 o_proj、MLP 的 c_proj
        for name, p in self.named_parameters():
            if name.endswith(("o_proj.weight", "c_proj.weight")):
                nn.init.zeros_(p)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        """
        Args:
            input_ids: LongTensor (B, T)
        Returns:
            logits: FloatTensor (B, T, vocab_size)
        """
        B, T = input_ids.shape
        assert T <= self.block_size, f"序列长度 {T} 超过 block_size {self.block_size}"

        x = self.token_emb(input_ids)
        cos = self.cos[:T]
        sin = self.sin[:T]

        for block in self.blocks:
            x = block(x, cos, sin)

        x = self.norm_f(x)
        return self.lm_head(x)


# --- 必须实现的接口 ---

def load_model(checkpoint_path: str, device: str = "cuda") -> torch.nn.Module:
    """从 checkpoint 加载 SmolGPT。支持纯 state_dict 或带 config 的字典格式。"""
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)

    if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
        state_dict = checkpoint["model_state_dict"]
        config = checkpoint.get("config", {})
        model = SmolGPT(**config)
    else:
        state_dict = checkpoint
        model = SmolGPT()

    # 去除 torch.compile 添加的 '_orig_mod.' 前缀
    unwanted_prefix = "_orig_mod."
    for k in list(state_dict.keys()):
        if k.startswith(unwanted_prefix):
            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)

    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()
    return model


# --- 快速完整性检查 ---

if __name__ == "__main__":
    model = SmolGPT()

    unique_params = sum(p.numel() for p in set(model.parameters()))
    all_params = sum(p.numel() for p in model.parameters())
    print(f"唯一参数量:  {unique_params:,} ({unique_params/1e6:.2f}M)")
    print(f"全部参数量:  {all_params:,} ({all_params/1e6:.2f}M)  ← model.parameters() 统计")

    assert unique_params <= 100_000_000, f"超过 100M 限制！实际 {unique_params/1e6:.2f}M"

    dummy = torch.randint(0, 50257, (2, 1024))
    logits = model(dummy)
    print(f"输入 shape:  {dummy.shape}")
    print(f"输出 shape:  {logits.shape}")
    assert logits.shape == (2, 1024, 50257), "输出形状错误！"
    print("接口检查通过。")