# 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