core/model.py

"""
Language Diffusion Transformer (DiT for text).

Public open-source version:
- pure PyTorch only
- state_dict key layout kept compatible with the internal model
"""
import math
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch._dynamo import disable as dynamo_disable


class FallbackRMSNorm(nn.Module):
    """Minimal RMSNorm implementation used by the public model."""

    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_dtype = x.dtype
        x_float = x.float()
        norm = torch.rsqrt(x_float.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        out = x_float * norm
        return (out.to(dtype=x_dtype) * self.weight).to(dtype=x_dtype)


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def _rope_cos_sin(
    seqlen: int,
    dim: int,
    theta: float,
    *,
    device: torch.device,
    dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
    inv_freq = 1.0 / (
        theta ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim)
    )
    t = torch.arange(seqlen, device=device, dtype=torch.float32)
    freqs = torch.outer(t, inv_freq)
    emb = torch.cat((freqs, freqs), dim=-1)
    cos = emb.cos().to(dtype=dtype)[None, :, None, :]
    sin = emb.sin().to(dtype=dtype)[None, :, None, :]
    return cos, sin


class TokenEmbedding(nn.Module):
    """Token embedding (untied from output)."""

    def __init__(self, vocab_size: int, dim: int):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.embed(x)


class TimestepEmbedding(nn.Module):
    """Sinusoidal timestep embedding -> conditioning vector."""

    def __init__(self, cond_dim: int, freq_dim: int = 256):
        super().__init__()
        self.freq_dim = freq_dim
        self.mlp = nn.Sequential(
            nn.Linear(freq_dim, cond_dim, bias=True),
            nn.SiLU(),
            nn.Linear(cond_dim, cond_dim, bias=True),
        )

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        if t.ndim == 2 and t.shape[1] == 1:
            t = t.squeeze(-1)
        half = self.freq_dim // 2
        freqs = torch.exp(
            -math.log(10000) * torch.arange(half, device=t.device, dtype=torch.float32) / half
        )
        args = t[:, None].to(dtype=torch.float32) * freqs[None]
        embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        embed = embed.to(dtype=self.mlp[0].weight.dtype)
        return self.mlp(embed)


class RotaryEmbedding(nn.Module):
    """Pure PyTorch RoPE implementation."""

    def __init__(self, dim: int, max_seq_len: int = 4096, theta: float = 10000.0):
        super().__init__()
        self.dim = int(dim)
        self.theta = float(theta)
        self.max_seq_len = max_seq_len

    def apply_bshd(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        cos, sin = _rope_cos_sin(
            q.shape[1],
            q.shape[-1],
            self.theta,
            device=q.device,
            dtype=q.dtype,
        )
        q = (q * cos) + (_rotate_half(q) * sin)
        k = (k * cos) + (_rotate_half(k) * sin)
        return q, k


class Attention(nn.Module):
    """
    Multi-head attention with expanded attention dimension.
    hidden_size -> attn_dim for Q,K,V -> hidden_size
    """

    def __init__(
        self,
        hidden_size: int,
        attn_dim: int,
        num_heads: int,
        head_dim: int = 128,
        attn_drop: float = 0.0,
        proj_drop: float = 0.0,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.attn_dim = attn_dim
        self.attn_drop = float(attn_drop)

        self.qkv = nn.Linear(hidden_size, attn_dim * 3, bias=False)
        self.proj = nn.Linear(attn_dim, hidden_size, bias=False)
        self.proj_drop = nn.Dropout(proj_drop)

    @dynamo_disable
    def forward(
        self,
        x: torch.Tensor,
        rope: Optional[RotaryEmbedding] = None,
        pack: Optional[object] = None,
    ) -> torch.Tensor:
        if pack is not None:
            raise RuntimeError("Packed attention is not included in the public torch-only model.")

        bsz, seqlen, _ = x.shape
        qkv = self.qkv(x).reshape(bsz, seqlen, 3, self.num_heads, self.head_dim)
        q, k, v = qkv.unbind(dim=2)

        if rope is not None:
            q, k = rope.apply_bshd(q, k)

        qh = q.permute(0, 2, 1, 3).contiguous()
        kh = k.permute(0, 2, 1, 3).contiguous()
        vh = v.permute(0, 2, 1, 3).contiguous()
        out = F.scaled_dot_product_attention(
            qh,
            kh,
            vh,
            dropout_p=self.attn_drop if self.training else 0.0,
            is_causal=False,
        )
        x = out.permute(0, 2, 1, 3).contiguous().reshape(bsz, seqlen, self.attn_dim)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class SwiGLU(nn.Module):
    """SwiGLU FFN with custom intermediate dimension."""

    def __init__(self, hidden_size: int, ffn_dim: int, dropout: float = 0.0):
        super().__init__()
        self.w12 = nn.Linear(hidden_size, 2 * ffn_dim, bias=False)
        self.w3 = nn.Linear(ffn_dim, hidden_size, bias=False)
        self.drop = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x1, x2 = self.w12(x).chunk(2, dim=-1)
        return self.w3(self.drop(F.silu(x1) * x2))


def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    """AdaLN modulation."""
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)


class DiTBlock(nn.Module):
    """
    DiT transformer block with AdaLN-Zero modulation.
    Pre-norm: RMSNorm -> modulate -> Attention/FFN -> gate
    """

    def __init__(
        self,
        hidden_size: int,
        attn_dim: int,
        ffn_dim: int,
        num_heads: int,
        head_dim: int = 128,
        cond_dim: int = 256,
        attn_drop: float = 0.0,
        drop: float = 0.0,
    ):
        super().__init__()
        self.norm1 = FallbackRMSNorm(hidden_size, eps=1e-6)
        self.attn = Attention(hidden_size, attn_dim, num_heads, head_dim, attn_drop, drop)
        self.norm2 = FallbackRMSNorm(hidden_size, eps=1e-6)
        self.mlp = SwiGLU(hidden_size, ffn_dim, dropout=drop)
        self.adaLN = nn.Sequential(
            nn.SiLU(),
            nn.Linear(cond_dim, 6 * hidden_size, bias=True),
        )

    def forward(
        self,
        x: torch.Tensor,
        c: torch.Tensor,
        rope: Optional[RotaryEmbedding] = None,
        pack: Optional[object] = None,
    ) -> torch.Tensor:
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN(c).chunk(6, dim=-1)
        x = x + gate_msa.unsqueeze(1) * self.attn(
            modulate(self.norm1(x), shift_msa, scale_msa),
            rope,
            pack=pack,
        )
        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
        return x


class FinalAdaLN(nn.Module):
    """Final AdaLN block that produces prelogits."""

    def __init__(self, hidden_size: int, cond_dim: int):
        super().__init__()
        self.norm = FallbackRMSNorm(hidden_size, eps=1e-6)
        self.adaLN = nn.Sequential(
            nn.SiLU(),
            nn.Linear(cond_dim, 2 * hidden_size, bias=True),
        )

    def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
        shift, scale = self.adaLN(c).chunk(2, dim=-1)
        return modulate(self.norm(x), shift, scale)


class LangDiT(nn.Module):
    """Language Diffusion Transformer."""

    def __init__(
        self,
        vocab_size: int = 64512,
        hidden_size: int = 2048,
        attn_dim: int = 3072,
        ffn_dim: int = 7168,
        depth: int = 48,
        num_heads: int = 24,
        head_dim: int = 128,
        max_seq_len: int = 4096,
        timestep_freq_dim: int = 256,
        rope_theta: float = 10000.0,
        cond_dim: int = 256,
        dropout: float = 0.0,
        attn_dropout: float = 0.0,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.vocab_size = vocab_size

        self.token_embed = TokenEmbedding(vocab_size, hidden_size)
        self.time_embed = TimestepEmbedding(cond_dim, freq_dim=timestep_freq_dim)
        self.rope = RotaryEmbedding(head_dim, max_seq_len, theta=rope_theta)
        self.blocks = nn.ModuleList(
            [
                DiTBlock(
                    hidden_size=hidden_size,
                    attn_dim=attn_dim,
                    ffn_dim=ffn_dim,
                    num_heads=num_heads,
                    head_dim=head_dim,
                    cond_dim=cond_dim,
                    attn_drop=attn_dropout,
                    drop=dropout,
                )
                for _ in range(depth)
            ]
        )
        self.final_ada = FinalAdaLN(hidden_size, cond_dim)
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

        self._init_weights()

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

        self.apply(init_fn)

        for block in self.blocks:
            nn.init.zeros_(block.adaLN[-1].weight)
            nn.init.zeros_(block.adaLN[-1].bias)
        nn.init.zeros_(self.final_ada.adaLN[-1].weight)
        nn.init.zeros_(self.final_ada.adaLN[-1].bias)

    def forward_hidden(
        self,
        input_ids: torch.Tensor,
        timesteps: torch.Tensor,
        *,
        pack: Optional[object] = None,
    ) -> torch.Tensor:
        if pack is not None:
            raise RuntimeError("Packed attention is not included in the public torch-only model.")

        x = self.token_embed(input_ids)
        c = self.time_embed(timesteps)

        for block in self.blocks:
            x = block(x, c, self.rope, pack=None)

        return self.final_ada(x, c)

    def logits_from_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.lm_head(hidden_states)

    def forward(
        self,
        input_ids: torch.Tensor,
        timesteps: torch.Tensor,
        *,
        pack: Optional[object] = None,
        return_hidden: bool = False,
    ) -> torch.Tensor:
        hidden = self.forward_hidden(input_ids, timesteps, pack=pack)
        if return_hidden:
            return hidden
        return self.logits_from_hidden(hidden)


def create_model(config: dict) -> LangDiT:
    """Create model from config dict."""
    model_cfg = config["model"]
    return LangDiT(
        vocab_size=model_cfg["vocab_size"],
        hidden_size=model_cfg["hidden_size"],
        attn_dim=model_cfg["attn_dim"],
        ffn_dim=model_cfg["ffn_dim"],
        depth=model_cfg["depth"],
        num_heads=model_cfg["num_heads"],
        head_dim=model_cfg["head_dim"],
        max_seq_len=model_cfg["max_seq_len"],
        timestep_freq_dim=model_cfg.get("timestep_freq_dim", 256),
        rope_theta=model_cfg.get("rope_theta", 10000.0),
        cond_dim=model_cfg["cond_dim"],
        dropout=model_cfg.get("dropout", 0.0),
        attn_dropout=model_cfg.get("attn_dropout", 0.0),
    )