""" yk_diffusion: a from-scratch hybrid language model. A single Transformer runs in two modes, selected by a learned mode embedding: - AR mode (mode=0): causal attention -> standard next-token autoregressive LM ("normal") - DIFF mode (mode=1): bidirectional attention -> masked discrete diffusion denoising A time embedding conditions the diffusion mask ratio (MDLM / LLaDA-style absorbing-state training). The same weights serve both behaviours; a signal picks which one at inference time. """ import math import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.checkpoint import checkpoint # ---------------------------------------------------------------------------- # Building blocks # ---------------------------------------------------------------------------- class RMSNorm(nn.Module): def __init__(self, d, eps=1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(d)) def forward(self, x): var = x.pow(2).mean(-1, keepdim=True) x = x * torch.rsqrt(var + self.eps) return x * self.weight class RotaryEmbedding(nn.Module): def __init__(self, head_dim, max_len=8192, base=10000.0, rope_scale=1.0): super().__init__() # NTK-aware scaling: stretch the base so the trained (short) context # extends to longer inference windows without retraining. With # rope_scale = context_ratio (e.g. 128k/8k = 16) the model trained at # 8k still positions tokens correctly at 128k. base = base * (rope_scale ** (head_dim / (head_dim - 2))) inv = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) self.register_buffer("inv_freq", inv) def forward(self, seq_len, device): t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) freqs = torch.outer(t, self.inv_freq) # [L, head_dim/2] emb = torch.cat([freqs, freqs], dim=-1) # [L, head_dim] return emb.cos(), emb.sin() def rotate_half(x): x1, x2 = x.chunk(2, dim=-1) return torch.cat((-x2, x1), dim=-1) def apply_rope(x, cos, sin): # x: [B, h, L, hd]; cos/sin: [1, 1, L, hd] return x * cos + rotate_half(x) * sin class Attention(nn.Module): def __init__(self, d, n_heads): super().__init__() assert d % n_heads == 0 self.d = d self.n_heads = n_heads self.hd = d // n_heads self.scale = self.hd ** -0.5 self.qkv = nn.Linear(d, 3 * d, bias=False) self.proj = nn.Linear(d, d, bias=False) def forward(self, x, cos, sin, attn_mask=None): B, L, _ = x.shape qkv = self.qkv(x).reshape(B, L, 3, self.n_heads, self.hd) qkv = qkv.permute(2, 0, 3, 1, 4) # [3, B, h, L, hd] q, k, v = qkv[0], qkv[1], qkv[2] q = apply_rope(q, cos, sin) k = apply_rope(k, cos, sin) scores = (q @ k.transpose(-2, -1)) * self.scale if attn_mask is not None: scores = scores + attn_mask attn = scores.softmax(dim=-1) out = (attn @ v).transpose(1, 2).reshape(B, L, self.d) return self.proj(out) class MLP(nn.Module): def __init__(self, d, d_ff): super().__init__() self.w1 = nn.Linear(d, d_ff, bias=False) self.w3 = nn.Linear(d, d_ff, bias=False) self.w2 = nn.Linear(d_ff, d, bias=False) def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x)) class Block(nn.Module): def __init__(self, d, n_heads, d_ff): super().__init__() self.ln1 = RMSNorm(d) self.attn = Attention(d, n_heads) self.ln2 = RMSNorm(d) self.mlp = MLP(d, d_ff) def forward(self, x, cos, sin, attn_mask): x = x + self.attn(self.ln1(x), cos, sin, attn_mask) x = x + self.mlp(self.ln2(x)) return x # ---------------------------------------------------------------------------- # The model # ---------------------------------------------------------------------------- class YKDiff(nn.Module): def __init__(self, cfg): super().__init__() self.cfg = cfg d = cfg["d_model"] self.vocab = cfg["vocab_size"] self.max_len = cfg["max_len"] self.tok_emb = nn.Embedding(self.vocab, d) self.mode_emb = nn.Embedding(2, d) # 0 = AR, 1 = DIFF self.time_emb = nn.Linear(1, d, bias=False) # diffusion mask ratio conditioning self.rope = RotaryEmbedding(d // cfg["n_heads"], max_len=self.max_len, rope_scale=cfg.get("rope_scale", 1.0)) self.blocks = nn.ModuleList([ Block(d, cfg["n_heads"], cfg["d_ff"]) for _ in range(cfg["n_layers"]) ]) self.norm = RMSNorm(d) self.lm_head = nn.Linear(d, self.vocab, bias=False) with torch.no_grad(): self.lm_head.weight.copy_(self.tok_emb.weight) # weight tying self._causal = None @property def device(self): return next(self.parameters()).device def _causal_mask(self, L, device): if self._causal is None or self._causal.shape[-1] < L: m = torch.full((L, L), float("-inf"), device=device) m = torch.triu(m, diagonal=1) self._causal = m return self._causal[:L, :L] def forward(self, idx, mode, t=None, attn_mask=None): """ idx: [B, L] long token ids mode: [B] long (0 AR, 1 DIFF) t: [B] float|None diffusion mask ratio (None -> 0) attn_mask: [L, L]|None explicit mask; if None, AR uses causal, DIFF uses none """ B, L = idx.shape x = self.tok_emb(idx) x = x + self.mode_emb(mode).unsqueeze(1) if t is None: t = torch.zeros(B, device=idx.device) x = x + self.time_emb(t.unsqueeze(-1)).unsqueeze(1) cos, sin = self.rope(L, idx.device) cos = cos.unsqueeze(0).unsqueeze(0) # [1,1,L,hd] sin = sin.unsqueeze(0).unsqueeze(0) if attn_mask is None: # AR default causal; DIFF default bidirectional (None) attn_mask = self._causal_mask(L, idx.device) if mode[0].item() == 0 else None else: attn_mask = attn_mask.to(idx.device) # gradient checkpointing: trade ~20% compute for ~4x less # activation memory, so a big batch fits on 16 GB. training = self.training and torch.is_grad_enabled() for blk in self.blocks: if training: x = checkpoint(blk, x, cos, sin, attn_mask, use_reentrant=False) else: x = blk(x, cos, sin, attn_mask) x = self.norm(x) return self.lm_head(x) # [B, L, vocab]