hf_modeling.py

# hf_modeling.py: standalone transformers modeling file for released
# MetaDiffusion-600M artifacts (AutoModelForCausalLM with trust_remote_code).
# Copied into the export dir by export_hf.py.
#
# The generate() override runs LLaDA-style iterative denoising with
# left-to-right block commit: the leftmost masked positions are unmasked
# first, so an <|im_end|> cannot win at position 0 and produce empty output.
# Generation stops once a terminator is committed in the response region.

import math

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

from transformers import GenerationMixin, PretrainedConfig, PreTrainedModel


class MetaDiffusion600MConfig(PretrainedConfig):
    model_type = "metadiffusion"

    def __init__(
        self,
        hidden_size=1024,
        intermediate_size=3072,
        num_hidden_layers=28,
        num_attention_heads=16,
        num_key_value_heads=8,
        head_dim=128,
        vocab_size=151669,
        mask_vocab_size=151677,
        mask_token_id=151669,
        pad_token_id=151643,
        max_position_embeddings=32768,
        rope_theta=1000000.0,
        rms_norm_eps=1e-6,
        hidden_act="silu",
        qk_norm=True,
        timestep_emb_hidden=1024,
        tie_word_embeddings=False,
        eos_token_id=None,
        **kwargs,
    ):
        super().__init__(
            pad_token_id=pad_token_id,
            tie_word_embeddings=tie_word_embeddings,
            eos_token_id=eos_token_id,
            **kwargs,
        )
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim
        self.vocab_size = vocab_size
        self.mask_vocab_size = mask_vocab_size
        self.mask_token_id = mask_token_id
        self.max_position_embeddings = max_position_embeddings
        self.rope_theta = rope_theta
        self.rms_norm_eps = rms_norm_eps
        self.hidden_act = hidden_act
        self.qk_norm = qk_norm
        self.timestep_emb_hidden = timestep_emb_hidden


class RMSNorm(nn.Module):
    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):
        orig = x.dtype
        x = x.float()
        var = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(var + self.eps)
        return (self.weight.float() * x).to(orig)


class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=32768, base=1000000.0):
        super().__init__()
        self.dim = dim
        self.base = base

    def forward(self, x, position_ids):
        # computed fresh each call on purpose: a stored inv_freq buffer is
        # non-persistent, so it is NOT in the state dict and from_pretrained
        # leaves it UNINITIALIZED, producing garbage cos/sin and NaN logits
        # in the entire forward. Computing here is 28 tiny ops, immune to
        # whatever transformers does to buffers during loading.
        #
        # Numerics replicate training (model.py + train.py): init computes
        # inv_freq on CPU in fp32, then `model.to(device, dtype=bfloat16)`
        # rounds the buffer to bf16, and the forward upcasts it back to fp32.
        # Matching that here keeps the released file BITWISE-consistent with
        # the training implementation (plain fp32 inv_freq differs by ~1 bf16
        # ULP in the rotary, which drifts final logits by ~1.0 after 28
        # layers).
        inv_freq = 1.0 / (self.base ** (torch.arange(
            0, self.dim, 2).float() / self.dim))
        inv_freq = inv_freq.to(torch.bfloat16).to(torch.float32)
        inv = inv_freq[None, :, None].to(x.device).expand(position_ids.shape[0], -1, 1)
        pos = position_ids[:, None, :].float()
        freqs = (inv @ pos).transpose(1, 2)
        emb = torch.cat((freqs, freqs), dim=-1)
        return emb.cos().to(dtype=x.dtype), emb.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, sin = cos.unsqueeze(1), sin.unsqueeze(1)
    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)


class TimestepEmbedding(nn.Module):
    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):
        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)
        # cast to the MLP weight dtype: the model may be bf16 while t is fp32
        return self.mlp(emb.to(self.mlp[0].weight.dtype))


class TimestepModulation(nn.Module):
    """adaLN-style timestep conditioning: scale + shift the hidden state.

    Zero-init scale/shift so the model is identity at step 0. Gradient is
    dL/dscale = dL/dx * x (x nonzero), so the t-path trains: the old
    zero-init additive residual deadlocked (zero output x zero weight =
    zero gradient forever), leaving models noise-schedule-agnostic."""

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

    def forward(self, x, emb):
        scale, shift = self.proj(emb).chunk(2, dim=-1)
        scale, shift = scale[:, None, :], shift[:, None, :]
        return x * (1.0 + scale) + shift


class Attention(nn.Module):
    def __init__(self, config):
        super().__init__()
        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 * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) if config.qk_norm else nn.Identity()
        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) if config.qk_norm else nn.Identity()
        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)
        q, k = self.q_norm(q), self.k_norm(k)
        cos, sin = self.rotary_emb(x, position_ids)
        q, k = apply_rotary_pos_emb(q, k, cos, sin)
        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)
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask)
        return self.o_proj(out.transpose(1, 2).contiguous().view(batch, seq, -1))


class MLP(nn.Module):
    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 Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.self_attn = Attention(config)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp = MLP(config)
        self.timestep_modulation = TimestepModulation(config.hidden_size)

    def forward(self, x, t_emb, attention_mask=None, position_ids=None):
        residual = x
        x = self.input_layernorm(x)
        x = self.self_attn(x, attention_mask, position_ids)
        x = residual + x
        x = self.timestep_modulation(x, t_emb)
        residual = x
        x = self.post_attention_layernorm(x)
        x = self.mlp(x)
        x = residual + x
        x = self.timestep_modulation(x, t_emb)
        return x


class MetaDiffusion600MModel(PreTrainedModel):
    config_class = MetaDiffusion600MConfig

    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.embed_tokens = nn.Embedding(config.mask_vocab_size, config.hidden_size)
        self.timestep_emb = TimestepEmbedding(config.timestep_emb_hidden)
        self.layers = nn.ModuleList([Block(config) for _ in range(config.num_hidden_layers)])
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.mask_vocab_size, bias=False)
        self.post_init()

    def forward(self, input_ids, timesteps=None, attention_mask=None):
        batch, seq = input_ids.shape
        position_ids = torch.arange(seq, device=input_ids.device).unsqueeze(0).expand(batch, -1)
        if timesteps is None:
            timesteps = torch.full((batch,), 1.0, device=input_ids.device)
        x = self.embed_tokens(input_ids)
        t_emb = self.timestep_emb(timesteps)
        attn_mask = None
        if attention_mask is not None:
            attn_mask = ((1.0 - attention_mask[:, None, None, :].float()) * -1e9).to(x.dtype)
        for layer in self.layers:
            x = layer(x, t_emb, attn_mask, position_ids)
        x = self.norm(x)
        return F.linear(x, self.lm_head.weight)


class MetaDiffusion600MForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = MetaDiffusion600MConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = MetaDiffusion600MModel(config)
        self.post_init()

    def forward(self, input_ids, timesteps=None, attention_mask=None, **kwargs):
        logits = self.model(input_ids, timesteps, attention_mask)
        return type("MDOutput", (), {"logits": logits})()

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids}

    def _cumulative_unmask_frac(self, i, n):
        return 0.5 * (1 - math.cos(math.pi * i / n))

    def generate(self, input_ids, max_new_tokens=None, num_steps=None,
                 temperature=None, repetition_penalty=None, top_p=None,
                 min_p=None, im_end_bias=None, im_end_bias_t=None,
                 smart_remask=None, smart_remask_thresh=None,
                 smart_remask_iters=None, **kwargs):
        """LLaDA-style iterative denoising with left-to-right commit.

        All sampling params fall back to generation_config.json values when
        not passed explicitly (release defaults ship in the config):
        top_p/min_p: truncation sampling that cuts the unreliable tail of the
        distribution (the junk-token source); use at most one (min_p 0.05-0.1
        recommended, Nguyen 2024; top-p 0.9 Holtzman 2020).
        im_end_bias: pragmatic logit nudge on the stop tokens when t is low
        (release guardrail for terminator reliability)."""
        device = input_ids.device
        config = self.config
        gc = self.generation_config
        max_new_tokens = max_new_tokens if max_new_tokens is not None else getattr(gc, "max_new_tokens", 96)
        num_steps = num_steps if num_steps is not None else getattr(gc, "num_steps", 128)
        temperature = temperature if temperature is not None else getattr(gc, "temperature", 0.7)
        repetition_penalty = repetition_penalty if repetition_penalty is not None else getattr(gc, "repetition_penalty", 1.5)
        top_p = top_p if top_p is not None else getattr(gc, "top_p", 0.0)
        min_p = min_p if min_p is not None else getattr(gc, "min_p", 0.1)
        im_end_bias = im_end_bias if im_end_bias is not None else getattr(gc, "im_end_bias", 0.0)
        im_end_bias_t = im_end_bias_t if im_end_bias_t is not None else getattr(gc, "im_end_bias_t", 0.3)
        smart_remask = smart_remask if smart_remask is not None else getattr(gc, "smart_remask", False)
        smart_remask_thresh = smart_remask_thresh if smart_remask_thresh is not None else getattr(gc, "smart_remask_thresh", 0.5)
        smart_remask_iters = smart_remask_iters if smart_remask_iters is not None else getattr(gc, "smart_remask_iters", 2)
        refine_steps = getattr(gc, "refine_steps", 16)
        mask_id = config.mask_token_id
        eos_ids = self.generation_config.eos_token_id
        if not isinstance(eos_ids, (list, tuple)):
            eos_ids = [eos_ids] if eos_ids is not None else []
        eos_ids = [int(e) for e in eos_ids if e is not None]

        prompt_len = input_ids.shape[1]
        x = torch.full((1, prompt_len + max_new_tokens), mask_id, device=device, dtype=torch.long)
        x[0, :prompt_len] = input_ids[0]
        # commit-confidence map for smart remasking (top-1 prob at commit time)
        conf = (torch.ones((1, x.shape[1]), dtype=torch.float32, device=device)
                if smart_remask else None)

        self.eval()
        with torch.no_grad():
            for i in range(num_steps):
                frac_now = self._cumulative_unmask_frac(i, num_steps)
                frac_next = self._cumulative_unmask_frac(i + 1, num_steps)
                n_masked = (x == mask_id).sum().item()
                if i == num_steps - 1:
                    n_unmask = n_masked
                else:
                    n_unmask = max(int((frac_next - frac_now) * max_new_tokens + 0.5), 1) if n_masked > 0 else 0
                if n_unmask == 0:
                    break
                t = torch.full((1,), 1.0 - frac_now, device=device)
                logits = self.model(x, t)
                logits = logits.logits if hasattr(logits, "logits") else logits
                # fp32 sampling path + sanitize (chat.py parity): models
                # trained with a mask-ratio curriculum have never seen t near
                # 1.0, so the timestep embedding can emit NaN/inf in bf16 when
                # generating; softmax/multinomial must never see them
                logits = logits.float()
                logits = torch.nan_to_num(logits, nan=0.0, posinf=50.0, neginf=-50.0)
                logits[:, :, mask_id] = -1e9
                rainbow_ids = getattr(config, "rainbow_token_ids", None) or \
                    list(range(mask_id + 1, mask_id + 8))
                logits[:, :, rainbow_ids] = -1e9
                bad_ids = getattr(config, "invalid_utf8_token_ids", None)
                if bad_ids:
                    logits[:, :, bad_ids] = -1e9
                if repetition_penalty != 1.0:
                    committed = x[0, prompt_len:]
                    committed = committed[committed != mask_id]
                    if committed.numel() > 0:
                        for tok in committed.unique():
                            ti = tok.item()
                            logits[0, :, ti] = torch.where(
                                logits[0, :, ti] < 0,
                                logits[0, :, ti] * repetition_penalty,
                                logits[0, :, ti] / repetition_penalty)
                if im_end_bias != 0.0 and 1.0 - frac_now < im_end_bias_t:
                    # pragmatic terminator nudge: BEFORE softmax so it actually
                    # shapes the sampled distribution
                    for eid in eos_ids:
                        logits[0, :, eid] = logits[0, :, eid] + im_end_bias
                mask_positions = x == mask_id
                probs = F.softmax(logits[mask_positions] / max(temperature, 1e-8), dim=-1)
                probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0)
                if top_p > 0.0:
                    sorted_probs, indices = probs.sort(dim=-1, descending=True)
                    drop = (sorted_probs.cumsum(dim=-1) - sorted_probs) > top_p
                    sorted_probs = sorted_probs.masked_fill(drop, 0.0)
                    sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True).clamp(min=1e-12)
                    probs = torch.zeros_like(probs).scatter_(-1, indices, sorted_probs)
                elif min_p > 0.0:
                    threshold = min_p * probs.max(dim=-1, keepdim=True).values
                    probs = probs.masked_fill(probs < threshold, 0.0)
                    probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-12)
                zero_rows = probs.sum(dim=-1, keepdim=True) <= 0
                if zero_rows.any():
                    probs = probs + zero_rows.to(probs.dtype)
                probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-12)
                p_max = probs.max(dim=-1).values
                sampled = torch.multinomial(probs, 1).squeeze(-1)
                mask_flat = mask_positions.nonzero(as_tuple=False)
                if n_unmask < mask_positions.sum():
                    fill_positions = mask_flat[:n_unmask]
                    for idx, tok in zip(fill_positions, sampled[:n_unmask]):
                        x[idx[0], idx[1]] = tok
                    if conf is not None:
                        conf[fill_positions[:, 0], fill_positions[:, 1]] = p_max[:n_unmask]
                else:
                    x[mask_positions] = sampled
                    if conf is not None:
                        conf[mask_positions] = p_max
                if any((x[0, prompt_len:] == e).any().item() for e in eos_ids):
                    break
        if smart_remask and conf is not None:
            rainbow_ids = getattr(config, "rainbow_token_ids", None) or \
                list(range(mask_id + 1, mask_id + 8))
            bad_ids = getattr(config, "invalid_utf8_token_ids", None)
            x = _smart_remask(self.model, x, prompt_len, max_new_tokens, conf,
                              eos_ids, mask_id, rainbow_ids, bad_ids,
                              smart_remask_thresh, smart_remask_iters,
                              refine_steps, temperature, repetition_penalty,
                              top_p, min_p, im_end_bias, im_end_bias_t,
                              self._cumulative_unmask_frac)
        return x


def _smart_remask(model, x, prompt_len, gen_len, conf, eos_ids, mask_id,
                  rainbow_ids, bad_ids, thresh, max_iters, refine_steps,
                  temperature, repetition_penalty, top_p, min_p, im_end_bias,
                  im_end_bias_t, cumfrac):
    """Confidence-gated re-denoising (PURE-style smart remasking): re-mask
    exactly the tokens whose top-1 commit probability fell below `thresh`
    and re-denoise them with the head fixed (chat.py parity). Runs even when
    a terminator committed, cleaning low-confidence junk before it. Stops
    early once a terminator commits or nothing is below the bar."""
    device = x.device
    lo = prompt_len
    hi = prompt_len + gen_len
    for _ in range(max_iters):
        if eos_ids:
            term_mask = (x[0, lo:hi] == eos_ids[0])
            for e in eos_ids[1:]:
                term_mask = term_mask | (x[0, lo:hi] == e)
            if term_mask.any():
                # never touch the terminator or anything past it
                hi = lo + term_mask.nonzero(as_tuple=True)[0][0].item()
                if hi <= lo:
                    break
        low = (conf[0, lo:hi] < thresh).nonzero(as_tuple=True)[0]
        if low.numel() == 0:
            break
        n_remask = low.numel()
        x[0, lo + low] = mask_id
        conf[0, lo + low] = 1.0  # re-commits below the bar get caught again
        for i in range(refine_steps):
            n_masked = (x[0, lo:hi] == mask_id).sum().item()
            if n_masked == 0:
                break
            if i == refine_steps - 1:
                n_unmask = n_masked
            else:
                n_unmask = max(int((cumfrac(i + 1, refine_steps)
                                    - cumfrac(i, refine_steps))
                                   * n_remask + 0.5), 1)
            n_unmask = min(n_unmask, n_masked)
            t_now = 1.0 - cumfrac(i, refine_steps)
            t_val = torch.full((1,), t_now, device=device)
            out = model(x, t_val)
            logits = out.logits if hasattr(out, "logits") else out
            logits = logits.float()
            logits = torch.nan_to_num(logits, nan=0.0, posinf=50.0, neginf=-50.0)
            logits[:, :, mask_id] = -1e9
            logits[:, :, rainbow_ids] = -1e9
            if bad_ids:
                logits[:, :, bad_ids] = -1e9
            if im_end_bias != 0.0 and t_now < im_end_bias_t:
                for e in eos_ids:
                    logits[0, :, e] = logits[0, :, e] + im_end_bias
            if repetition_penalty != 1.0:
                committed = x[0, prompt_len:]
                committed = committed[committed != mask_id]
                if committed.numel() > 0:
                    for tok in committed.unique():
                        ti = tok.item()
                        logits[0, :, ti] = torch.where(
                            logits[0, :, ti] < 0,
                            logits[0, :, ti] * repetition_penalty,
                            logits[0, :, ti] / repetition_penalty)
            mask_positions = x == mask_id
            probs = F.softmax(logits[mask_positions] / max(temperature, 1e-8), dim=-1)
            probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0)
            if top_p > 0.0:
                sorted_probs, indices = probs.sort(dim=-1, descending=True)
                drop = (sorted_probs.cumsum(dim=-1) - sorted_probs) > top_p
                sorted_probs = sorted_probs.masked_fill(drop, 0.0)
                sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True).clamp(min=1e-12)
                probs = torch.zeros_like(probs).scatter_(-1, indices, sorted_probs)
            elif min_p > 0.0:
                threshold = min_p * probs.max(dim=-1, keepdim=True).values
                probs = probs.masked_fill(probs < threshold, 0.0)
                probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-12)
            zero_rows = probs.sum(dim=-1, keepdim=True) <= 0
            if zero_rows.any():
                probs = probs + zero_rows.to(probs.dtype)
            probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-12)
            p_max = probs.max(dim=-1).values
            sampled = torch.multinomial(probs, 1).squeeze(-1)
            mask_flat = mask_positions.nonzero(as_tuple=False)
            n_fill = min(n_unmask, mask_flat.shape[0])
            if n_fill:
                idxs = mask_flat[:n_fill]
                x[idxs[:, 0], idxs[:, 1]] = sampled[:n_fill]
                conf[idxs[:, 0], idxs[:, 1]] = p_max[:n_fill]
            if eos_ids and any((x[0, lo:hi] == e).any().item() for e in eos_ids):
                break
        if eos_ids and any((x[0, lo:hi] == e).any().item() for e in eos_ids):
            break
    return x