scripts/model.py

"""
MetaDiffusion: Convert AR LLMs to Masked Diffusion LLMs
Based on Supra-1.5-50M-Base-exp architecture
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass, field
from typing import Optional, Tuple


@dataclass
class MetaDiffusionConfig:
    # Architecture (matching Supra-1.5-50M)
    hidden_size: int = 512
    intermediate_size: int = 1408
    num_hidden_layers: int = 12
    num_attention_heads: int = 8
    num_key_value_heads: int = 4
    head_dim: int = 64
    vocab_size: int = 32000          # original vocab
    mask_vocab_size: int = 32001     # vocab + [MASK] token
    max_position_embeddings: int = 5120
    rope_theta: float = 10000.0
    rms_norm_eps: float = 1e-6
    hidden_act: str = "silu"

    # Diffusion-specific
    timestep_emb_hidden: int = 512
    mask_token_id: int = 32000       # index of [MASK] in embedding table
    pad_token_id: int = 1            # Supra pad token

    # Masking strategy
    mask_ratio_min: float = 0.0
    mask_ratio_max: float = 1.0

    # Training
    dtype: torch.dtype = torch.float32
    tie_word_embeddings: bool = True

class RotaryEmbedding(nn.Module):
    """RoPE - position embeddings for the attention layers."""

    def __init__(self, dim, max_position_embeddings=5120, base=10000.0, device=None):
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base

        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    @torch.no_grad()
    def forward(self, x, position_ids):
        # x: (batch, seq, hidden) - used for dtype/device only
        # position_ids: (batch, seq)
        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(
            position_ids.shape[0], -1, 1
        )
        position_ids_expanded = position_ids[:, None, :].float()
        freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos()
        sin = emb.sin()
        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)


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


def apply_rotary_pos_emb(q, k, cos, sin):
    cos = cos.unsqueeze(1)  # (batch, 1, seq, dim)
    sin = sin.unsqueeze(1)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


class TimestepEmbedding(nn.Module):
    """Sinusoidal timestep embedding with learned projection."""

    def __init__(self, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 4),
            nn.SiLU(),
            nn.Linear(hidden_size * 4, hidden_size),
        )

    def forward(self, t):
        # t: (batch,) timesteps in [0, 1]
        half_dim = self.hidden_size // 2
        emb = math.log(10000.0) / (half_dim - 1)
        emb = torch.exp(
            torch.arange(half_dim, device=t.device, dtype=torch.float32) * -emb
        )
        emb = t[:, None].float() * emb[None, :]
        emb = torch.cat([emb.sin(), emb.cos()], dim=-1)  # (batch, hidden_size)
        return self.mlp(emb).to(t.dtype)  # (batch, hidden_size)


class TimestepResidual(nn.Module):
    """Add timestep embedding to hidden state at each block.
    Initialized to zero so it starts as identity (no disruption to pretrained weights)."""

    def __init__(self, hidden_size):
        super().__init__()
        self.proj = nn.Linear(hidden_size, hidden_size)
        nn.init.zeros_(self.proj.weight)
        nn.init.zeros_(self.proj.bias)

    def forward(self, x, emb):
        # x: (batch, seq, hidden)
        # emb: (batch, hidden)
        return x + self.proj(emb)[:, None, :]


class RMSNorm(nn.Module):
    """Llama-style RMSNorm."""

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

    def forward(self, x):
        var = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(var + self.eps)
        return self.weight * x


class SelfAttention(nn.Module):
    """Multi-head attention with GQA and RoPE.
    BIDIRECTIONAL (no causal mask) - this is the key difference from AR."""

    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.num_kv_groups = self.num_heads // self.num_kv_heads

        self.q_proj = nn.Linear(
            config.hidden_size, self.num_heads * config.head_dim, bias=False
        )
        self.k_proj = nn.Linear(
            config.hidden_size, self.num_kv_heads * config.head_dim, bias=False
        )
        self.v_proj = nn.Linear(
            config.hidden_size, self.num_kv_heads * config.head_dim, bias=False
        )
        self.o_proj = nn.Linear(
            self.num_heads * config.head_dim, config.hidden_size, bias=False
        )
        self.rotary_emb = RotaryEmbedding(
            config.head_dim,
            max_position_embeddings=config.max_position_embeddings,
            base=config.rope_theta,
        )

    def forward(self, x, attention_mask=None, position_ids=None):
        batch, seq, _ = x.shape

        q = self.q_proj(x).view(batch, seq, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch, seq, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch, seq, self.num_kv_heads, self.head_dim).transpose(1, 2)

        cos, sin = self.rotary_emb(x, position_ids)
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # GQA: repeat KV heads to match Q heads
        if self.num_kv_groups > 1:
            k = k.repeat_interleave(self.num_kv_groups, dim=1)
            v = v.repeat_interleave(self.num_kv_groups, dim=1)

        # Bidirectional attention - no causal mask!
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask)
        out = out.transpose(1, 2).contiguous().view(batch, seq, -1)
        return self.o_proj(out)


class MLP(nn.Module):
    """Llama-style gated FFN (SwiGLU)."""

    def __init__(self, config):
        super().__init__()
        self.gate_proj = nn.Linear(
            config.hidden_size, config.intermediate_size, bias=False
        )
        self.up_proj = nn.Linear(
            config.hidden_size, config.intermediate_size, bias=False
        )
        self.down_proj = nn.Linear(
            config.intermediate_size, config.hidden_size, bias=False
        )

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class TransformerBlock(nn.Module):
    """Llama transformer block adapted for diffusion.
    - Pre-norm architecture
    - Bidirectional attention
    - Timestep conditioning via residual addition
    """

    def __init__(self, config):
        super().__init__()
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.self_attn = SelfAttention(config)
        self.post_attention_layernorm = RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.mlp = MLP(config)
        self.timestep_residual = TimestepResidual(config.hidden_size)

    def forward(self, x, timestep_emb, attention_mask=None, position_ids=None):
        # Pre-norm + attention + residual + timestep
        residual = x
        x = self.input_layernorm(x)
        x = self.self_attn(x, attention_mask, position_ids)
        x = residual + x
        x = self.timestep_residual(x, timestep_emb)

        # Pre-norm + FFN + residual + timestep
        residual = x
        x = self.post_attention_layernorm(x)
        x = self.mlp(x)
        x = residual + x
        x = self.timestep_residual(x, timestep_emb)

        return x

class MetaDiffusionLM(nn.Module):
    """Masked Diffusion Language Model.

    Converts an AR Llama-style model to a masked-diffusion LM.
    Transfers: embeddings, all transformer blocks, RoPE, norms.
    New: timestep embedding, [MASK] token.
    Output head is tied with embeddings by default (set tie_word_embeddings=False to untie).
    """

    def __init__(self, config: MetaDiffusionConfig):
        super().__init__()
        self.config = config

        # Embeddings (vocab + 1 for [MASK])
        self.embed_tokens = nn.Embedding(
            config.mask_vocab_size, config.hidden_size, padding_idx=config.pad_token_id
        )

        # Timestep conditioning
        self.timestep_emb = TimestepEmbedding(config.timestep_emb_hidden)

        # Transformer stack
        self.layers = nn.ModuleList(
            [TransformerBlock(config) for _ in range(config.num_hidden_layers)]
        )

        # Final norm
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        # Output projection (can be tied with embeddings for parameter efficiency)
        if config.tie_word_embeddings:
            self.lm_head = None  # Will use embed_tokens.weight in forward
        else:
            self.lm_head = nn.Linear(
                config.hidden_size, config.mask_vocab_size, bias=False
            )

        self.post_init()

    def post_init(self):
        if self.lm_head is not None:
            nn.init.normal_(self.lm_head.weight, std=0.02)

    def forward(self, input_ids, timesteps, attention_mask=None):
        """
        Forward pass for training.

        Args:
            input_ids: (batch, seq) - tokens with masked positions replaced by mask_token_id
            timesteps: (batch,) - diffusion timestep in [0, 1]
            attention_mask: optional (batch, seq) - 1 for real tokens, 0 for padding

        Returns:
            logits: (batch, seq, mask_vocab_size)
        """
        batch, seq = input_ids.shape

        # Position IDs (0, 1, 2, ...)
        position_ids = (
            torch.arange(seq, device=input_ids.device).unsqueeze(0).expand(batch, -1)
        )

        # Embed tokens
        x = self.embed_tokens(input_ids)

        # Get timestep embedding
        t_emb = self.timestep_emb(timesteps)

        # Convert attention mask for SDPA (0 -> keep, -inf -> mask out)
        attn_mask = None
        if attention_mask is not None:
            attn_mask = ((1.0 - attention_mask[:, None, None, :].float()) * -1e9).to(
                x.dtype
            )

        # Pass through transformer blocks
        for layer in self.layers:
            x = layer(x, t_emb, attn_mask, position_ids)

        # Final norm + project to vocab
        x = self.norm(x)
        if self.lm_head is not None:
            logits = self.lm_head(x)
        else:
            # Tied embeddings: use embed_tokens.weight transposed
            logits = F.linear(x, self.embed_tokens.weight)

        return logits

    def compute_loss(self, logits, labels, mask_positions, pad_token_id=None):
        """
        Compute cross-entropy loss on masked positions only.

        Args:
            logits: (batch, seq, vocab_size)
            labels: (batch, seq) - original token IDs (before masking)
            mask_positions: (batch, seq) - bool tensor, True where token was masked
            pad_token_id: ignore these positions in loss

        Returns:
            loss: scalar
            num_masked: number of positions in loss
        """
        logits_masked = logits[mask_positions]
        labels_masked = labels[mask_positions]

        # Filter out padding tokens
        if pad_token_id is not None:
            valid = labels_masked != pad_token_id
            logits_masked = logits_masked[valid]
            labels_masked = labels_masked[valid]

        if labels_masked.numel() == 0:
            return torch.tensor(0.0, device=logits.device), 0

        loss = F.cross_entropy(logits_masked, labels_masked)
        return loss, labels_masked.numel()

    @torch.no_grad()
    def generate(self, batch_size, seq_len, num_steps=256, device="cuda"):
        """
        Iterative denoising generation (LLaDA-style).

        Starts from all-mask tokens and progressively unmaskes the most confident predictions.

        Args:
            batch_size: number of sequences to generate
            seq_len: length of each sequence
            num_steps: number of denoising iterations

        Returns:
            tokens: (batch, seq) - generated token IDs
        """
        mask_token_id = self.config.mask_token_id
        x = torch.full(
            (batch_size, seq_len), mask_token_id, device=device, dtype=torch.long
        )

        # Linear schedule from t=1 (all mask) to t=0 (no mask)
        timesteps = torch.linspace(1.0, 0.0, num_steps + 1, device=device)

        for i in range(num_steps):
            t = timesteps[i]
            t_next = timesteps[i + 1]
            t_batch = torch.full((batch_size,), t, device=device)

            # Get predictions
            logits = self.forward(x, t_batch)
            pred_tokens = logits.argmax(dim=-1)

            # Confidence of predicted tokens
            probs = F.softmax(logits, dim=-1)
            confidence = probs.gather(-1, pred_tokens.unsqueeze(-1)).squeeze(-1)

            # Number of tokens to unmask this step
            num_unmask = max(1, int(seq_len * (t - t_next)))

            # Only consider currently-masked positions
            is_mask = x == mask_token_id
            confidence_masked = confidence.clone()
            confidence_masked[~is_mask] = -1.0

            # Unmask the most confident predictions
            _, top_indices = confidence_masked.topk(num_unmask, dim=-1)
            batch_idx = (
                torch.arange(batch_size, device=device).unsqueeze(-1).expand_as(top_indices)
            )
            x[batch_idx, top_indices] = pred_tokens[batch_idx, top_indices]

        return x

    @classmethod
    def from_pretrained_ar(
        cls,
        model_name_or_path: str,
        **kwargs,
    ):
        """
        Initialize a MetaDiffusionLM from a pretrained AR Llama model.

        Transfers all AR weights and initializes new diffusion components.
        """
        from transformers import LlamaForCausalLM

        print(f"Loading AR model: {model_name_or_path}")
        ar_model = LlamaForCausalLM.from_pretrained(model_name_or_path)
        ar_config = ar_model.config

        # Build diffusion config from AR config
        head_dim = getattr(
            ar_config,
            "head_dim",
            ar_config.hidden_size // ar_config.num_attention_heads,
        )
        config = MetaDiffusionConfig(
            hidden_size=ar_config.hidden_size,
            intermediate_size=ar_config.intermediate_size,
            num_hidden_layers=ar_config.num_hidden_layers,
            num_attention_heads=ar_config.num_attention_heads,
            num_key_value_heads=ar_config.num_key_value_heads,
            head_dim=head_dim,
            vocab_size=ar_config.vocab_size,
            mask_vocab_size=ar_config.vocab_size + 1,
            max_position_embeddings=ar_config.max_position_embeddings,
            rope_theta=getattr(ar_config, "rope_theta", 10000.0),
            rms_norm_eps=ar_config.rms_norm_eps,
            mask_token_id=ar_config.vocab_size,
            pad_token_id=getattr(ar_config, "pad_token_id", 1),
            **kwargs,
        )

        model = cls(config)

        # Build state dict mapping
        state_dict = ar_model.state_dict()
        new_state_dict = {}

        # --- Embeddings ---
        # Copy original vocab embeddings
        ar_embeds = state_dict["model.embed_tokens.weight"]
        mask_embed = ar_embeds.mean(dim=0, keepdim=True)  # [MASK] = mean of all embeds
        new_state_dict["embed_tokens.weight"] = torch.cat(
            [ar_embeds, mask_embed], dim=0
        )

        # --- Transformer blocks ---
        for i in range(config.num_hidden_layers):
            ar_prefix = f"model.layers.{i}"
            new_prefix = f"layers.{i}"

            # Attention
            new_state_dict[f"{new_prefix}.self_attn.q_proj.weight"] = state_dict[
                f"{ar_prefix}.self_attn.q_proj.weight"
            ]
            new_state_dict[f"{new_prefix}.self_attn.k_proj.weight"] = state_dict[
                f"{ar_prefix}.self_attn.k_proj.weight"
            ]
            new_state_dict[f"{new_prefix}.self_attn.v_proj.weight"] = state_dict[
                f"{ar_prefix}.self_attn.v_proj.weight"
            ]
            new_state_dict[f"{new_prefix}.self_attn.o_proj.weight"] = state_dict[
                f"{ar_prefix}.self_attn.o_proj.weight"
            ]

            # MLP
            new_state_dict[f"{new_prefix}.mlp.gate_proj.weight"] = state_dict[
                f"{ar_prefix}.mlp.gate_proj.weight"
            ]
            new_state_dict[f"{new_prefix}.mlp.up_proj.weight"] = state_dict[
                f"{ar_prefix}.mlp.up_proj.weight"
            ]
            new_state_dict[f"{new_prefix}.mlp.down_proj.weight"] = state_dict[
                f"{ar_prefix}.mlp.down_proj.weight"
            ]

            # Layer norms
            new_state_dict[f"{new_prefix}.input_layernorm.weight"] = state_dict[
                f"{ar_prefix}.input_layernorm.weight"
            ]
            new_state_dict[
                f"{new_prefix}.post_attention_layernorm.weight"
            ] = state_dict[f"{ar_prefix}.post_attention_layernorm.weight"]

        # --- Final norm ---
        new_state_dict["norm.weight"] = state_dict["model.norm.weight"]

        # --- Output head ---
        # Initialize from AR embeddings (since AR used tied embeddings, E^T was the output)
        if not config.tie_word_embeddings:
            new_state_dict["lm_head.weight"] = torch.cat(
                [ar_embeds, torch.zeros(1, config.hidden_size, device=ar_embeds.device)],
                dim=0,
            ).clone()

        # Load with strict=False for new diffusion params
        missing, unexpected = model.load_state_dict(new_state_dict, strict=False)

        print(f"Weights transferred from AR model")
        print(f"  Missing (new diffusion params): {len(missing)}")
        print(f"  Unexpected: {len(unexpected)}")
        if missing:
            for k in missing:
                print(f"    NEW: {k}")

        return model