""" 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