Upload 9 files

chat.py

@@ -0,0 +1,548 @@
+#!/usr/bin/env python3
+"""chat.py: ChatML chat with MetaDiffusion-600M checkpoints.
+
+Left-to-right block commit (semi-autoregressive): the leftmost masked
+positions are filled first, so <|im_end|> cannot win the race at position 0.
+
+Usage:
+    Interactive:      python chat.py --model-path checkpoints/step_30000.pt
+    One-shot:         python chat.py --model-path checkpoints/step_30000.pt \
+                          --prompt "What is 2+2?" --watch
+"""
+
+import argparse
+import json
+import math
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from model import MetaDiffusionConfig, MetaDiffusionLM  # noqa: E402
+
+IM_START, IM_END = "<|im_start|>", "<|im_end|>"
+RAINBOW_TOKENS = [f"<|r{i}|>" for i in range(1, 8)]
+
+
+def build_config(config_dict):
+    valid = {k: v for k, v in config_dict.items()
+             if k in MetaDiffusionConfig.__dataclass_fields__}
+    return MetaDiffusionConfig(**valid)
+
+
+def load_model(model_path, device):
+    path = Path(model_path)
+    if path.is_dir():
+        with open(path / "config.json") as f:
+            config = build_config(json.load(f))
+        model = MetaDiffusionLM(config).to(device)
+        from safetensors.torch import load_file
+        sd = load_file(path / "model.safetensors")
+        sd = {k[len("model."):] if k.startswith("model.") else k: v for k, v in sd.items()}
+        model.load_state_dict(sd, strict=True)
+    else:
+        ckpt = torch.load(model_path, map_location=device, weights_only=False)
+        config = build_config(ckpt["config"])
+        model = MetaDiffusionLM(config).to(device)
+        sd = {k.replace("_orig_mod.", "", 1) if k.startswith("_orig_mod.") else k: v
+              for k, v in ckpt["model_state_dict"].items()}
+        model.load_state_dict(sd, strict=True)
+    model.eval()
+    print(f"  Loaded {sum(p.numel() for p in model.parameters())/1e6:.1f}M params, "
+          f"vocab={config.mask_vocab_size}")
+    return model
+
+
+def ensure_special_tokens(tokenizer):
+    """Add [MASK] + rainbow if missing"""
+    added = []
+    if tokenizer.convert_tokens_to_ids("[MASK]") == tokenizer.unk_token_id:
+        added.append("[MASK]")
+    missing = [t for t in RAINBOW_TOKENS
+               if tokenizer.convert_tokens_to_ids(t) == tokenizer.unk_token_id]
+    if missing:
+        added.extend(missing)
+    if added:
+        tokenizer.add_special_tokens({"additional_special_tokens": added})
+    return tokenizer
+
+
+def format_messages(messages):
+    parts = []
+    for m in messages:
+        parts.append(f"{IM_START}{m['role']}\n{m['content']}{IM_END}")
+    return "\n".join(parts)
+
+
+def cumulative_unmask_frac(i, N):
+    return 0.5 * (1 - math.cos(math.pi * i / N))
+
+
+@torch.no_grad()
+def generate_response(model, tokenizer, prompt_ids, gen_len, num_steps,
+                      temperature, repetition_penalty, device, watch=False,
+                      stop_on_end=True, cfg_scale=0.0, ban_ids=None,
+                      top_p=0.0, min_p=0.0, refine=False, refine_frac=0.3,
+                      refine_steps=16, im_end_bias=0.0, im_end_bias_t=0.3,
+                      smart_remask=False, smart_remask_thresh=0.5,
+                      smart_remask_iters=2):
+    mask_id = model.config.mask_token_id
+    im_end_id = tokenizer.convert_tokens_to_ids(IM_END)
+    eos_id = tokenizer.eos_token_id
+    rainbow_ids = [tokenizer.convert_tokens_to_ids(t) for t in RAINBOW_TOKENS]
+    prompt_len = prompt_ids.shape[1]
+    total_len = prompt_len + gen_len
+
+    x = torch.full((1, total_len), mask_id, device=device, dtype=torch.long)
+    x[0, :prompt_len] = prompt_ids
+    # commit-confidence map for --smart-remask (top-1 prob at commit time);
+    # 1.0 for prompt/uncommitted so only real commits can fall below the bar
+    conf = (torch.ones((1, total_len), dtype=torch.float32, device=device)
+            if smart_remask else None)
+
+    terminated = False
+    for i in range(num_steps):
+        frac_now = cumulative_unmask_frac(i, num_steps)
+        frac_next = 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_total = int((frac_next - frac_now) * gen_len + 0.5)
+            n_unmask = max(n_total, 1) if n_masked > 0 else 0
+        if n_unmask == 0:
+            break
+
+        t = 1.0 - frac_now
+        t_val = torch.full((1,), t, device=device)
+        logits = model(x, t_val).float()  # fp32 sampling path: stable softmax
+        # Mid-run curriculum probes extrapolate t beyond what the model has
+        # seen (e.g. t=0.99 at step 8.5K when the ramp max is ~0.48). The
+        # timestep embedding can then blow up to NaN/inf inside the bf16
+        # forward. Sanitize once here so CFG, softmax and multinomial never
+        # see a poisoned distribution.
+        logits = torch.nan_to_num(logits, nan=0.0, posinf=50.0, neginf=-50.0)
+        if cfg_scale > 0:
+            # classifier-free guidance: unconditional branch sees the prompt
+            # region masked too; logits = cond + s*(cond - uncond).
+            # The all-mask input is off the training manifold, so its logits
+            # can be extreme; extrapolating them in bf16 overflows to inf and
+            # poisons softmax/multinomial. Compute in fp32 and clamp.
+            uncond_x = torch.full_like(x, mask_id)
+            uncond_logits = model(uncond_x, t_val).float()
+            uncond_logits = torch.nan_to_num(uncond_logits, nan=0.0,
+                                             posinf=50.0, neginf=-50.0)
+            logits = (logits + cfg_scale * (logits - uncond_logits)).clamp(-50.0, 50.0)
+        # padding placeholders (rainbow) and [MASK] are never legitimate output
+        logits[:, :, mask_id] = -1e9
+        logits[:, :, rainbow_ids] = -1e9
+        if ban_ids:
+            # partial-byte vocab entries that cannot decode to valid UTF-8:
+            # the literal "�" characters; never legitimate output either
+            logits[:, :, ban_ids] = -1e9
+        if im_end_bias != 0.0 and t < im_end_bias_t:
+            # pragmatic terminator nudge at the end of denoising: the model's
+            # continuation knowledge runs out before its terminator probability
+            # rises, so make <|im_end|> competitive in the frontier distribution
+            logits[:, :, im_end_id] = logits[:, :, im_end_id] + 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
+        sampled, probs = sample_masked(logits, mask_positions, temperature,
+                                       top_p, min_p)
+        p_max = probs.max(dim=-1).values if conf is not None else None
+        mask_flat = mask_positions.nonzero(as_tuple=False)
+
+        if n_unmask < mask_positions.sum():
+            # Left-to-right commit: fill the leftmost masked positions first
+            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 watch:
+            remaining = (x == mask_id).sum().item()
+            live = [t for t in x[0, prompt_len:].tolist() if t != mask_id]
+            partial = tokenizer.decode(cut_response(live, tokenizer),
+                                       skip_special_tokens=True).strip()[:70]
+            line = f"step {i+1:3d}/{num_steps} | t={t:.3f} | masks={remaining:3d} | {partial}"
+            if sys.stdout.isatty():
+                sys.stdout.write("\r" + line[:99].ljust(99))
+                sys.stdout.flush()
+            elif i % max(1, num_steps // 8) == 0:
+                print(line)
+
+        if stop_on_end and ((x[0, prompt_len:] == im_end_id).any() or
+                            (x[0, prompt_len:] == eos_id).any()):
+            terminated = True
+            break
+
+    if smart_remask and conf is not None:
+        # confidence-gated remasking: re-mask only the low-confidence commits
+        # (the junk-prone tokens) and re-denoise them with the head fixed
+        x = smart_remask_pass(model, x, prompt_len, gen_len, conf, im_end_id,
+                         eos_id, mask_id, rainbow_ids, ban_ids, refine_steps,
+                         smart_remask_thresh, smart_remask_iters, temperature,
+                         repetition_penalty, device, top_p, min_p,
+                         im_end_bias, im_end_bias_t)
+    elif refine and not terminated:
+        x = refine_tail(model, x, prompt_len, gen_len, im_end_id, eos_id,
+                        mask_id, rainbow_ids, ban_ids, refine_steps,
+                        refine_frac, temperature, repetition_penalty, device,
+                        top_p, min_p)
+
+    if watch and sys.stdout.isatty():
+        sys.stdout.write("\n")
+    return x
+
+
+def sample_masked(logits, mask_positions, temperature, top_p=0.0, min_p=0.0):
+    """Truncation-sampled tokens for the masked positions.
+
+    top-p nucleus (Holtzman 2020) or Min-P (Nguyen 2024) truncate the
+    unreliable tail of the distribution, which is exactly where junk and rare
+    tokens live when the model runs out of budget. Min-P scales the cutoff by
+    the top token's probability (pbase 0.05-0.1 recommended; use ONE of them).
+    Also guards degenerate rows so multinomial never sees inf/nan/negatives.
+    Returns (sampled, probs)."""
+    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)
+    # degenerate rows (all-zero after truncation/nan handling) fall back to
+    # uniform so multinomial never sees an invalid distribution
+    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)
+    return torch.multinomial(probs, 1).squeeze(-1), probs
+
+
+def refine_tail(model, x, prompt_len, gen_len, im_end_id, eos_id, mask_id,
+                rainbow_ids, ban_ids, steps, frac, temperature,
+                repetition_penalty, device, top_p=0.0, min_p=0.0,
+                im_end_bias=0.0, im_end_bias_t=0.3):
+    """PURE/TOLERATOR-style post-hoc refinement: when a response never
+    committed <|im_end|>, re-mask the tail (keeping the head fixed) and
+    re-denoise it with a short chain. Training-free; converts leftover
+    compute into coherence instead of letting committed junk stick."""
+    cut = prompt_len + int(gen_len * (1.0 - frac))
+    x[0, cut:] = mask_id
+    tail_len = gen_len - (cut - prompt_len)
+    for i in range(steps):
+        n_masked = (x[0, cut:] == mask_id).sum().item()
+        if n_masked == 0:
+            break
+        if i == steps - 1:
+            n_unmask = n_masked
+        else:
+            n_unmask = max(int((cumulative_unmask_frac(i + 1, steps)
+                                - cumulative_unmask_frac(i, steps)) * tail_len + 0.5), 1)
+        t_val = torch.full((1,), 1.0 - cumulative_unmask_frac(i, steps), device=device)
+        logits = model(x, t_val).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 ban_ids:
+            logits[:, :, ban_ids] = -1e9
+        t_now = 1.0 - cumulative_unmask_frac(i, steps)
+        if im_end_bias != 0.0 and t_now < im_end_bias_t:
+            logits[:, :, im_end_id] = logits[:, :, im_end_id] + 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
+        sampled, _ = sample_masked(logits, mask_positions, temperature, top_p, min_p)
+        mask_flat = mask_positions.nonzero(as_tuple=False)
+        if n_unmask < mask_positions.sum():
+            for idx, tok in zip(mask_flat[:n_unmask], sampled[:n_unmask]):
+                x[idx[0], idx[1]] = tok
+        else:
+            x[mask_positions] = sampled
+        if (x[0, prompt_len:] == im_end_id).any() or (x[0, prompt_len:] == eos_id).any():
+            break
+    return x
+
+
+def smart_remask_pass(model, x, prompt_len, gen_len, conf, im_end_id, eos_id,
+                      mask_id, rainbow_ids, ban_ids, steps, thresh, max_iters,
+                 temperature, repetition_penalty, device, top_p=0.0,
+                 min_p=0.0, im_end_bias=0.0, im_end_bias_t=0.3):
+    """Confidence-gated re-denoising (PURE-style smart remasking).
+
+    The blind --refine tail remask wastes budget on tokens the model already
+    committed with high confidence. Here, re-mask exactly the tokens whose
+    top-1 commit probability fell below `thresh` (the junk-prone ones, often
+    the budget-tail fillers) and re-denoise them with the head fixed. Runs
+    even when im_end committed: it also cleans low-confidence junk sitting
+    before the terminator. Repeats up to max_iters rounds and stops early
+    once the terminator commits or nothing is below the bar."""
+    lo = prompt_len
+    hi = prompt_len + gen_len
+    for _ in range(max_iters):
+        resp = x[0, lo:hi]
+        term = (resp == im_end_id) | (resp == eos_id)
+        if term.any():
+            # never touch the terminator or anything past it
+            hi = lo + term.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(steps):
+            n_masked = (x[0, lo:hi] == mask_id).sum().item()
+            if n_masked == 0:
+                break
+            if i == steps - 1:
+                n_unmask = n_masked
+            else:
+                n_unmask = max(int((cumulative_unmask_frac(i + 1, steps)
+                                    - cumulative_unmask_frac(i, steps))
+                                   * n_remask + 0.5), 1)
+            n_unmask = min(n_unmask, n_masked)
+            t_now = 1.0 - cumulative_unmask_frac(i, steps)
+            t_val = torch.full((1,), t_now, device=device)
+            logits = model(x, t_val).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 ban_ids:
+                logits[:, :, ban_ids] = -1e9
+            if im_end_bias != 0.0 and t_now < im_end_bias_t:
+                logits[:, :, im_end_id] = logits[:, :, im_end_id] + 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
+            sampled, probs = sample_masked(logits, mask_positions,
+                                           temperature, top_p, min_p)
+            p_max = probs.max(dim=-1).values
+            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 (x[0, lo:hi] == im_end_id).any() or \
+               (x[0, lo:hi] == eos_id).any():
+                break
+        if (x[0, lo:hi] == im_end_id).any() or (x[0, lo:hi] == eos_id).any():
+            break
+    return x
+
+
+def cut_response(tokens, tokenizer):
+    """Cut at <|im_end|> / eos; drop rainbow and pad tokens."""
+    im_end_id = tokenizer.convert_tokens_to_ids(IM_END)
+    eos_id = tokenizer.eos_token_id
+    rainbow_ids = {tokenizer.convert_tokens_to_ids(t) for t in RAINBOW_TOKENS}
+    out = []
+    for t in tokens:
+        if t == im_end_id or t == eos_id:
+            break
+        if t in rainbow_ids or t == tokenizer.pad_token_id:
+            continue
+        out.append(t)
+    return out
+
+
+def invalid_utf8_ids(tokenizer):
+    """Ids whose decode is *only* U+FFFD. Byte-fallback tokens that merely
+    contain a replacement char when decoded alone stay; those are how Qwen
+    builds rare unicode."""
+    ban = []
+    for i in range(len(tokenizer)):
+        s = tokenizer.decode([i], skip_special_tokens=True)
+        if s and all(c == "\uFFFD" for c in s):
+            ban.append(i)
+    return ban
+
+
+def trim_messages(messages, tokenizer, max_context, max_new_tokens):
+    """Drop oldest non-system turns until prompt + gen budget fits."""
+    budget = max(32, max_context - max_new_tokens)
+    kept = list(messages)
+    while kept:
+        prompt = format_messages(kept) + f"\n{IM_START}assistant\n"
+        n = len(tokenizer.encode(prompt, add_special_tokens=False))
+        if n <= budget:
+            return kept
+        drop_at = next((i for i, m in enumerate(kept) if m["role"] != "system"), None)
+        if drop_at is None:
+            return kept
+        del kept[drop_at]
+    return kept
+
+
+def run_turn(model, tokenizer, messages, args, device):
+    max_ctx = getattr(args, "max_context", 4096)
+    cap = min(getattr(model.config, "max_position_embeddings", 40960), max_ctx)
+    messages = trim_messages(messages, tokenizer, cap, args.max_new_tokens)
+    prompt = format_messages(messages) + f"\n{IM_START}assistant\n"
+    prompt_ids = torch.tensor([tokenizer.encode(prompt, add_special_tokens=False)],
+                              device=device)
+    for attempt in range(3):
+        x = generate_response(model, tokenizer, prompt_ids, args.max_new_tokens,
+                              args.num_steps,
+                              args.temperature * (1 + 0.15 * attempt),
+                              args.repetition_penalty, device, watch=args.watch,
+                              cfg_scale=args.cfg_scale,
+                              ban_ids=getattr(args, "bad_token_ids", None),
+                              top_p=args.top_p, min_p=args.min_p,
+                              refine=args.refine, refine_frac=args.refine_frac,
+                              refine_steps=args.refine_steps,
+                              im_end_bias=args.im_end_bias,
+                              im_end_bias_t=args.im_end_bias_t,
+                              smart_remask=args.smart_remask,
+                              smart_remask_thresh=args.smart_remask_thresh,
+                              smart_remask_iters=args.smart_remask_iters)
+        text = tokenizer.decode(cut_response(x[0, prompt_ids.shape[1]:].tolist(),
+                                             tokenizer),
+                                skip_special_tokens=True).strip()
+        if text:
+            return text
+    return "(empty response)"
+
+
+def main():
+    p = argparse.ArgumentParser(description="MetaDiffusion-600M chat")
+    p.add_argument("--model-path", required=True)
+    p.add_argument("--tokenizer", default=None, help="Tokenizer dir (needed for .pt checkpoints)")
+    p.add_argument("--prompt", default=None)
+    p.add_argument("--system", default="You are a helpful assistant.")
+    p.add_argument("--max-new-tokens", type=int, default=96)
+    p.add_argument("--num-steps", type=int, default=128)
+    p.add_argument("--temperature", type=float, default=0.7)
+    p.add_argument("--repetition-penalty", type=float, default=1.5)
+    p.add_argument("--cfg-scale", type=float, default=0.0,
+                   help="Classifier-free guidance scale (0 = off; try 0.5-1.2). "
+                        "Unconditional branch masks the prompt too.")
+    p.add_argument("--top-p", type=float, default=0.0,
+                   help="Nucleus sampling: keep tokens covering this mass (0=off; "
+                        "use EITHER --top-p or --min-p, not both)")
+    p.add_argument("--min-p", type=float, default=0.1,
+                   help="Min-P truncation: keep tokens >= min_p x top-token prob "
+                        "(0=off; 0.05-0.1 recommended). Truncates the junk tail.")
+    p.add_argument("--refine", action="store_true",
+                   help="Post-hoc tail refinement: if im_end never commits, re-mask "
+                        "the tail and re-denoise it (PURE/TOLERATOR-style)")
+    p.add_argument("--refine-frac", type=float, default=0.3,
+                   help="Fraction of the response tail to re-denoise (--refine)")
+    p.add_argument("--refine-steps", type=int, default=16,
+                   help="Denoising steps for the refinement pass")
+    p.add_argument("--smart-remask", action="store_true",
+                   help="Confidence-gated remasking (PURE-style): re-mask only "
+                        "the tokens committed with low top-1 probability and "
+                        "re-denoise them with the head fixed. Runs even when "
+                        "im_end committed (cleans pre-terminator junk); takes "
+                        "precedence over --refine.")
+    p.add_argument("--smart-remask-thresh", type=float, default=0.5,
+                   help="Commit-confidence bar (top-1 token prob at commit "
+                        "time); tokens below it are re-masked (--smart-remask)")
+    p.add_argument("--smart-remask-iters", type=int, default=2,
+                   help="Max refinement rounds; stops early when the "
+                        "terminator commits or nothing is below the bar")
+    p.add_argument("--im-end-bias", type=float, default=0.0,
+                   help="Logit bonus on <|im_end|> when t < --im-end-bias-t "
+                        "(pragmatic terminator nudge; try 1.5-3.0)")
+    p.add_argument("--im-end-bias-t", type=float, default=0.3,
+                   help="t threshold below which --im-end-bias applies")
+    p.add_argument("--device", default="cuda")
+    p.add_argument("--max-context", type=int, default=4096,
+                   help="Trim multi-turn history so prompt+gen fits this many tokens")
+    p.add_argument("--watch", action="store_true")
+    args = p.parse_args()
+
+    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
+    print(f"[*] Loading model from {args.model_path}")
+    model = load_model(args.model_path, device)
+
+    tok_path = args.tokenizer
+    if tok_path is None:
+        model_path = Path(args.model_path)
+        if model_path.is_dir():
+            cand = model_path / "tokenizer"
+            if not cand.exists() and (model_path / "tokenizer.json").exists():
+                cand = model_path
+            tok_path = str(cand)
+    if not tok_path or not Path(tok_path).exists():
+        raise SystemExit("No tokenizer found; pass --tokenizer (data/tokenizer)")
+    tokenizer = ensure_special_tokens(AutoTokenizer.from_pretrained(str(tok_path)))
+    print(f"[*] Tokenizer: {tok_path} (vocab {len(tokenizer)})")
+    # hard-ban vocab entries that cannot decode to valid UTF-8 (partial-byte
+    # tokens): they surface as "�" garbage and are never legitimate output
+    args.bad_token_ids = invalid_utf8_ids(tokenizer)
+    if args.bad_token_ids:
+        print(f"[*] Banning {len(args.bad_token_ids)} standalone-U+FFFD tokens")
+
+    if args.prompt:
+        text = run_turn(model, tokenizer, [{"role": "user", "content": args.prompt}],
+                        args, device)
+        print(f"\nUser: {args.prompt}\nAssistant: {text}\n")
+        return
+
+    print("\nMetaDiffusion-600M chat. Type 'exit' to leave.\n")
+    messages = [{"role": "system", "content": args.system}]
+    while True:
+        try:
+            user_input = input("You: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print()
+            break
+        if user_input.lower() in ("exit", "quit"):
+            break
+        if not user_input:
+            continue
+        messages.append({"role": "user", "content": user_input})
+        text = run_turn(model, tokenizer, messages, args, device)
+        print(f"Assistant: {text}\n")
+        messages.append({"role": "assistant", "content": text})
+
+
+if __name__ == "__main__":
+    main()

convert.py

@@ -0,0 +1,191 @@
+#!/usr/bin/env python3
+"""
+convert.py: convert Qwen3-0.6B (post-trained instruct) into a
+MetaDiffusion-600M initialization checkpoint.
+
+Usage:
+    python convert.py \
+        --source Qwen/Qwen3-0.6B \
+        --output init/metadiffusion-600M-instruct.pt \
+        --tokenizer-out data/tokenizer \
+        --device cpu
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from model import MetaDiffusionConfig, MetaDiffusionLM  # noqa: E402
+
+RAINBOW_TOKENS = [f"<|r{i}|>" for i in range(1, 8)]
+NUM_NEW_TOKENS = 1 + len(RAINBOW_TOKENS)  # [MASK] + rainbow
+
+PLAIN_CHAT_TEMPLATE = (
+    "{% for message in messages %}"
+    "{{ '<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>\\n' }}"
+    "{% endfor %}"
+)
+
+
+def build_config(src_config, base_vocab: int) -> MetaDiffusionConfig:
+    """base_vocab = real tokenizer length (config.vocab_size may be TP-padded,
+    e.g. Qwen3: tokenizer 151,669 vs config 151,936)."""
+    c = src_config
+    config_vocab = int(getattr(c, "vocab_size", 0))
+    if base_vocab > config_vocab:
+        raise RuntimeError(
+            f"tokenizer vocab {base_vocab} exceeds model vocab {config_vocab}")
+    head_dim = int(getattr(c, "head_dim", 0) or (c.hidden_size // c.num_attention_heads))
+    cfg = MetaDiffusionConfig(
+        hidden_size=int(c.hidden_size),
+        intermediate_size=int(c.intermediate_size),
+        num_hidden_layers=int(c.num_hidden_layers),
+        num_attention_heads=int(c.num_attention_heads),
+        num_key_value_heads=int(c.num_key_value_heads),
+        head_dim=head_dim,
+        vocab_size=base_vocab,
+        mask_vocab_size=base_vocab + NUM_NEW_TOKENS,
+        mask_token_id=base_vocab,
+        pad_token_id=int(getattr(c, "pad_token_id", None) or 151643),
+        max_position_embeddings=int(c.max_position_embeddings),
+        rope_theta=float(getattr(c, "rope_theta", None) or 1000000.0),
+        rms_norm_eps=float(getattr(c, "rms_norm_eps", 1e-6)),
+        hidden_act=str(getattr(c, "hidden_act", "silu")),
+        qk_norm=bool(getattr(c, "qk_norm", True)),
+        timestep_emb_hidden=int(c.hidden_size),
+        tie_word_embeddings=False,
+    )
+    return cfg
+
+
+def convert(source: str, output: str, tokenizer_out: str, device: str = "cpu"):
+    logger = print
+    logger(f"[*] Loading AR instruct model: {source}")
+
+    src_config = AutoConfig.from_pretrained(source)
+    tokenizer = AutoTokenizer.from_pretrained(source)
+    base_vocab = len(tokenizer)
+    if base_vocab != int(getattr(src_config, "vocab_size", 0)):
+        logger(f"[*] Config vocab is TP-padded ({src_config.vocab_size}); "
+               f"using real tokenizer vocab {base_vocab}")
+    cfg = build_config(src_config, base_vocab)
+    h = cfg.hidden_size
+    logger(f"[*] Source: {cfg.num_hidden_layers}L x {h}W, vocab={cfg.vocab_size}, "
+           f"heads={cfg.num_attention_heads}/{cfg.num_key_value_heads}, "
+           f"head_dim={cfg.head_dim}, qk_norm={cfg.qk_norm}")
+
+    ar_model = AutoModelForCausalLM.from_pretrained(source, torch_dtype=torch.float32, device_map=device)
+    state = ar_model.state_dict()
+    del ar_model
+    torch.cuda.empty_cache() if torch.cuda.is_available() else None
+
+    # --- Copy weights, strip the "model." prefix ---
+    new_state = {}
+    for key, value in state.items():
+        new_key = key.replace("model.", "", 1) if key.startswith("model.") else key
+        new_state[new_key] = value.clone()
+
+    embed = new_state["embed_tokens.weight"]  # (V, H)
+
+    # --- Trim untrained TP-padding rows, then append [MASK] + rainbow ---
+    if embed.shape[0] > base_vocab:
+        logger(f"[*] Trimming {embed.shape[0] - base_vocab} untrained padding rows "
+               f"from embeddings / lm_head")
+        embed = embed[:base_vocab]
+    elif embed.shape[0] < base_vocab:
+        raise RuntimeError(f"embedding rows {embed.shape[0]} < tokenizer vocab {base_vocab}")
+
+    mean_row = embed.mean(dim=0, keepdim=True)  # (1, H)
+    new_rows = mean_row.expand(NUM_NEW_TOKENS, -1).clone()  # (8, H)
+    new_state["embed_tokens.weight"] = torch.cat([embed, new_rows], dim=0)
+
+    # --- Untied lm_head: trim + extend identically ---
+    if "lm_head.weight" in new_state:
+        head = new_state["lm_head.weight"]
+        head = head[:base_vocab] if head.shape[0] > base_vocab else head
+        if head.shape[0] != base_vocab:
+            raise RuntimeError(f"lm_head rows {head.shape[0]} != tokenizer vocab {base_vocab}")
+        new_state["lm_head.weight"] = torch.cat([head, new_rows], dim=0)
+    else:
+        new_state["lm_head.weight"] = torch.cat([embed, new_rows], dim=0)
+    logger(f"[*] embed_tokens / lm_head = {cfg.mask_vocab_size} rows "
+           f"(mask={cfg.mask_token_id}, rainbow={cfg.mask_token_id + 1}..{cfg.mask_vocab_size - 1})")
+
+    # --- New diffusion modules: adaLN modulation zero-init (identity at
+    # step 0). The timestep embedding MLP keeps DEFAULT random init: with
+    # both zero, zero output through the zero gate zeroes every t-path
+    # gradient (deadlock: the model never learns noise conditioning). ---
+    for i in range(cfg.num_hidden_layers):
+        new_state[f"layers.{i}.timestep_modulation.proj.weight"] = torch.zeros(2 * h, h)
+        new_state[f"layers.{i}.timestep_modulation.proj.bias"] = torch.zeros(2 * h)
+
+    # --- Verify load. timestep_emb MLP stays at default (random) init so
+    # the t-path is not deadlocked (zero MLP through a zero gate). Extra AR
+    # keys (rotary inv_freq) are dropped. ---
+    model = MetaDiffusionLM(cfg)
+    missing, unexpected = model.load_state_dict(new_state, strict=False)
+    missing = [k for k in missing if not k.startswith("timestep_emb.")]
+    if missing:
+        raise RuntimeError(f"unexpected missing keys after convert: {missing}")
+    if unexpected:
+        logger(f"[*] Dropping {len(unexpected)} unexpected AR keys "
+               f"(e.g. {unexpected[0]})")
+    n_params = sum(p.numel() for p in model.parameters())
+    logger(f"[*] Verified load into MetaDiffusionLM: {n_params/1e6:.1f}M params "
+           f"(AR transfer + {NUM_NEW_TOKENS} new token rows + timestep modules)")
+
+    # --- Save checkpoint (full state, including default-init t-MLP) ---
+    out = Path(output)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    ckpt = {
+        "config": cfg.__dict__,
+        "model_state_dict": model.state_dict(),
+        "metadata": {
+            "source_model": source,
+            "conversion_script": "convert.py",
+            "mask_token_id": cfg.mask_token_id,
+            "rainbow_token_ids": list(range(cfg.mask_token_id + 1, cfg.mask_vocab_size)),
+            "num_new_tokens": NUM_NEW_TOKENS,
+            "transferred": True,
+        },
+    }
+    torch.save(ckpt, str(out))
+    logger(f"[*] Saved checkpoint: {out}")
+
+    sidecar = out.with_suffix(".json")
+    with open(sidecar, "w") as f:
+        json.dump(cfg.__dict__, f, indent=2)
+    logger(f"[*] Saved config sidecar: {sidecar}")
+
+    # --- Tokenizer: add [MASK] + rainbow at the expected ids ---
+    tokenizer.add_special_tokens({"additional_special_tokens": ["[MASK]"] + RAINBOW_TOKENS})
+    tokenizer.chat_template = PLAIN_CHAT_TEMPLATE
+    mask_id = tokenizer.convert_tokens_to_ids("[MASK]")
+    assert mask_id == cfg.mask_token_id, f"[MASK] landed at {mask_id}, expected {cfg.mask_token_id}"
+    assert len(tokenizer) == cfg.mask_vocab_size, f"tokenizer vocab {len(tokenizer)} != {cfg.mask_vocab_size}"
+    tok_dir = Path(tokenizer_out)
+    tok_dir.mkdir(parents=True, exist_ok=True)
+    tokenizer.save_pretrained(str(tok_dir))
+    logger(f"[*] Saved tokenizer ({len(tokenizer)} tokens, [MASK]={mask_id}): {tok_dir}")
+
+    logger("[*] Done. Next: python prepare_data.py && python train.py --init-checkpoint "
+           f"{output}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert Qwen3-0.6B (instruct) to MetaDiffusion-600M init")
+    parser.add_argument("--source", default="Qwen/Qwen3-0.6B")
+    parser.add_argument("--output", default="init/metadiffusion-600M-instruct.pt")
+    parser.add_argument("--tokenizer-out", default="data/tokenizer")
+    parser.add_argument("--device", default="cpu")
+    args = parser.parse_args()
+    convert(args.source, args.output, args.tokenizer_out, args.device)
+
+
+if __name__ == "__main__":
+    main()

eval.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python3
+"""eval.py: lm-evaluation-harness wrapper for MetaDiffusion-600M.
+
+Scoring: single-step diffusion (mask the continuation, forward once at
+t=1.0, log-prob of the true tokens at masked positions).
+
+Usage (needs lm-eval in the environment):
+    python eval.py --checkpoint checkpoints/step_30000.pt \
+        --tasks hellaswag,arc_easy,arc_challenge,piqa \
+        --tokenizer data/tokenizer
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer
+
+from lm_eval.api.model import LM
+from lm_eval.api.registry import register_model
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from model import MetaDiffusionLM, MetaDiffusionConfig  # noqa: E402
+
+import logging
+logger = logging.getLogger(__name__)
+
+
+@register_model("metadiffusion_600m")
+class MetaDiffusion600MWrapper(LM):
+    def __init__(self, checkpoint: str, dtype: str = "float32",
+                 device: str = "cuda", tokenizer_name: str = "Qwen/Qwen3-0.6B",
+                 max_length: int = 1024, batch_size: int = 4, **kwargs):
+        super().__init__()
+        self._device = torch.device(device)
+        self._max_length = max_length
+        self._batch_size = batch_size
+        dtype_map = {"float32": torch.float32, "float16": torch.float16,
+                     "bfloat16": torch.bfloat16}
+        self._dtype = dtype_map.get(dtype, torch.float32)
+
+        logger.info(f"Loading checkpoint: {checkpoint}")
+        ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
+        config = MetaDiffusionConfig(
+            **{k: v for k, v in ckpt.get("config", ckpt).items()
+               if k in MetaDiffusionConfig.__dataclass_fields__})
+        self.model = MetaDiffusionLM(config)
+        sd = ckpt.get("model_state_dict", ckpt)
+        sd = {k.replace("_orig_mod.", "", 1) if isinstance(k, str) and k.startswith("_orig_mod.") else k: v
+              for k, v in sd.items()}
+        self.model.load_state_dict(sd, strict=True)
+        self.model = self.model.to(device=self._device, dtype=self._dtype).eval()
+
+        self._mask_token_id = config.mask_token_id
+        self._pad_token_id = config.pad_token_id
+        self._tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+        if self._tokenizer.pad_token_id is None:
+            self._tokenizer.pad_token_id = self._tokenizer.eos_token_id
+        self._eos_token_id = self._tokenizer.eos_token_id
+        logger.info(f"Loaded {config.num_hidden_layers}L x {config.hidden_size}W, "
+                    f"vocab={config.mask_vocab_size}, {self._dtype}")
+
+    def _score_pair(self, context_tokens, continuation_tokens):
+        full_ids = context_tokens + continuation_tokens
+        if len(full_ids) > self._max_length:
+            excess = len(full_ids) - self._max_length
+            context_tokens = context_tokens[excess:] if len(context_tokens) > excess else []
+            full_ids = full_ids[excess:]
+        ctx_len = len(context_tokens)
+        seq_len = len(full_ids)
+
+        input_ids = torch.tensor([full_ids], device=self._device)
+        for i in range(ctx_len, seq_len):
+            input_ids[0, i] = self._mask_token_id
+        t = torch.tensor([1.0], device=self._device)
+
+        with torch.no_grad():
+            logits = self.model(input_ids, t)
+        log_probs = F.log_softmax(logits[0], dim=-1)
+
+        total = 0.0
+        is_greedy = True
+        for pos in range(ctx_len, seq_len):
+            true_token = full_ids[pos]
+            total += log_probs[pos, true_token].item()
+            if log_probs[pos].argmax().item() != true_token:
+                is_greedy = False
+        return total, is_greedy
+
+    def loglikelihood(self, requests, disable_tqdm=False):
+        results = []
+        for request in requests:
+            context, continuation = request.arguments
+            ctx = self._tokenizer.encode(context, add_special_tokens=False)
+            cont = self._tokenizer.encode(continuation, add_special_tokens=False)
+            if not cont:
+                cont = [self._eos_token_id]
+            results.append(self._score_pair(ctx, cont))
+        return results
+
+    def loglikelihood_rolling(self, requests, disable_tqdm=False):
+        results = []
+        for request in requests:
+            tokens = self._tokenizer.encode(request.arguments[0], add_special_tokens=False)
+            if len(tokens) <= 1:
+                results.append(0.0)
+                continue
+            lp, _ = self._score_pair(tokens[:1], tokens[1:])
+            results.append(lp)
+        return results
+
+    def generate_until(self, requests, disable_tqdm=False):
+        from chat import generate_response
+        results = []
+        for request in requests:
+            prompt = request.arguments[0]
+            prompt_ids = torch.tensor(
+                [self._tokenizer.encode(prompt, add_special_tokens=False)],
+                device=self._device)
+            x = generate_response(self.model, self._tokenizer, prompt_ids,
+                                  gen_len=128, num_steps=32, temperature=0.2,
+                                  repetition_penalty=1.2,
+                                  device=self._device, stop_on_end=True,
+                                  min_p=0.1)
+            out = x[0, prompt_ids.shape[1]:].cpu().tolist()
+            results.append(self._tokenizer.decode(out, skip_special_tokens=True))
+        return results
+
+
+def main():
+    p = argparse.ArgumentParser(description="Evaluate MetaDiffusion-600M with lm-eval")
+    p.add_argument("--checkpoint", required=True)
+    p.add_argument("--tasks", default="hellaswag,arc_easy,arc_challenge,piqa")
+    p.add_argument("--tokenizer", default="data/tokenizer")
+    p.add_argument("--device", default="cuda:0")
+    p.add_argument("--dtype", default="bfloat16")
+    p.add_argument("--limit", default=None, help="Sample limit (smoke test)")
+    p.add_argument("--output", default=None, help="Result JSON path")
+    args = p.parse_args()
+
+    from lm_eval import simple_evaluate
+    results = simple_evaluate(
+        model="metadiffusion_600m",
+        model_args=f"checkpoint={args.checkpoint},dtype={args.dtype},"
+                   f"device={args.device},tokenizer_name={args.tokenizer}",
+        tasks=args.tasks.split(","),
+        limit=float(args.limit) if args.limit is not None else None,
+    )
+    for task, res in results["results"].items():
+        acc = res.get("acc_norm,none") or res.get("acc,none")
+        print(f"{task}: {acc:.4f}" if acc is not None else f"{task}: {res}")
+    if args.output:
+        import json
+        with open(args.output, "w") as f:
+            json.dump(results["results"], f, indent=2)
+        print(f"Wrote {args.output}")
+
+
+if __name__ == "__main__":
+    main()

export_hf.py

@@ -0,0 +1,188 @@
+#!/usr/bin/env python3
+"""export_hf.py: export a MetaDiffusion-600M checkpoint to a release dir.
+
+Output:
+    model.safetensors          fp16, "model."-prefixed keys
+    config.json                no dtype key; vocab fields = weight rows;
+                               auto_map -> hf_modeling.py custom classes
+    generation_config.json     denoising defaults
+    tokenizer/                 Qwen3 tokenizer + [MASK] + rainbow tokens
+    hf_modeling.py             standalone modeling (AutoModelForCausalLM,
+                               GenerationMixin with iterative denoising)
+    scripts/                   self-contained pipeline copy
+
+Usage:
+    python export_hf.py --checkpoint checkpoints/step_30000.pt \
+        --tokenizer data/tokenizer --output MetaDiffusion-600M-Instruct-v1
+
+Then load with:
+    AutoModelForCausalLM.from_pretrained(dir, trust_remote_code=True)
+"""
+
+import argparse
+import json
+import shutil
+import sys
+from pathlib import Path
+
+import torch
+from safetensors.torch import save_file
+from transformers import AutoTokenizer
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from model import MetaDiffusionConfig  # noqa: E402
+
+GENERATION_CONFIG = {
+    "temperature": 0.7,
+    "repetition_penalty": 1.5,
+    "num_steps": 128,
+    "max_new_tokens": 96,
+    "top_p": 0.0,          # truncation sampling: min-p is the release default
+    "min_p": 0.1,
+    "im_end_bias": 2.0,
+    "im_end_bias_t": 0.3,
+    "do_sample": True,
+    "transformers_version": "4.49.0",
+}
+
+PLAIN_CHAT_TEMPLATE = (
+    "{% for message in messages %}{{ '<|im_start|>' + message['role'] }}\n"
+    "{{ message['content'] }}<|im_end|>\n"
+    "{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n"
+    "{% endif %}"
+)
+
+
+def remap_state_dict(state_dict, dtype="bf16"):
+    cast = {"bf16": torch.bfloat16, "fp16": torch.float16,
+            "fp32": torch.float32}[dtype]
+    new_dict = {}
+    for key, tensor in state_dict.items():
+        key = key.replace("_orig_mod.", "", 1) if key.startswith("_orig_mod.") else key
+        new_dict["model." + key] = tensor.to(cast)
+    return new_dict
+
+
+def package_scripts(out):
+    src = Path(__file__).resolve().parent
+    scripts_dir = out / "scripts"
+    scripts_dir.mkdir(parents=True, exist_ok=True)
+    for name in ["model.py", "convert.py", "prepare_data.py", "train.py",
+                 "chat.py", "eval.py", "export_hf.py", "hf_modeling.py"]:
+        cand = src / name
+        if cand.exists():
+            shutil.copy2(cand, scripts_dir / name)
+    req = scripts_dir / "requirements.txt"
+    if not req.exists():
+        req.write_text("torch>=2.2\ntransformers>=4.49\nsafetensors>=0.4\n"
+                       "datasets>=2.18\nnumpy>=1.26\n")
+    print(f"[*] Packaged scripts -> {scripts_dir}")
+
+
+def export(checkpoint_path, tokenizer_dir, output_dir, dtype="bf16"):
+    out = Path(output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    print(f"[*] Loading checkpoint {checkpoint_path}")
+    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    config = MetaDiffusionConfig(
+        **{k: v for k, v in ckpt["config"].items()
+           if k in MetaDiffusionConfig.__dataclass_fields__})
+
+    print("[*] Remapping state dict...")
+    state_dict = remap_state_dict(ckpt["model_state_dict"], dtype=dtype)
+    save_file(state_dict, out / "model.safetensors")
+    print(f"[*] Saved {len(state_dict)} tensors -> {out / 'model.safetensors'} "
+          f"({dtype})")
+
+    # Vocab fields must match the actual weight rows
+    n_vocab = state_dict["model.lm_head.weight"].shape[0]
+    config.vocab_size = n_vocab
+    config.mask_vocab_size = n_vocab
+    print(f"[*] Vocab in config: {n_vocab} (matches weights)")
+
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
+    eos_ids = [tokenizer.eos_token_id] if tokenizer.eos_token_id is not None else []
+    im_end = tokenizer.convert_tokens_to_ids("<|im_end|>")
+    if im_end != tokenizer.unk_token_id and im_end not in eos_ids:
+        eos_ids.append(im_end)
+    eos_ids = [e for e in eos_ids if e is not None]
+    print(f"[*] eos ids: {eos_ids}")
+
+    # partial-byte vocab entries that cannot decode to valid UTF-8 (the
+    # literal replacement-char garbage): hard-banned at generation
+    bad_ids = []
+    for i in range(len(tokenizer)):
+        s = tokenizer.decode([i], skip_special_tokens=True)
+        if s and all(c == "\uFFFD" for c in s):
+            bad_ids.append(i)
+
+    config_dict = config.__dict__.copy()
+    config_dict.pop("dtype", None)  # transformers chokes on "torch.float32" strings
+    config_dict["model_type"] = "metadiffusion"
+    config_dict["architectures"] = ["MetaDiffusion600MForCausalLM"]
+    config_dict["auto_map"] = {
+        "AutoConfig": "hf_modeling.MetaDiffusion600MConfig",
+        "AutoModelForCausalLM": "hf_modeling.MetaDiffusion600MForCausalLM",
+    }
+    config_dict["eos_token_id"] = eos_ids
+    config_dict["rainbow_token_ids"] = list(range(config.mask_token_id + 1,
+                                                  config.mask_token_id + 8))
+    config_dict["invalid_utf8_token_ids"] = bad_ids
+    with open(out / "config.json", "w") as f:
+        json.dump(config_dict, f, indent=2)
+    print(f"[*] Saved config.json (mask_token_id={config.mask_token_id}, "
+          f"{len(bad_ids)} banned invalid-UTF8 tokens)")
+
+    gen_config = dict(GENERATION_CONFIG)
+    gen_config["eos_token_id"] = eos_ids
+    gen_config["pad_token_id"] = config.pad_token_id
+    gen_config["mask_token_id"] = config.mask_token_id
+    with open(out / "generation_config.json", "w") as f:
+        json.dump(gen_config, f, indent=2)
+
+    shutil.copytree(tokenizer_dir, out / "tokenizer", dirs_exist_ok=True)
+    print(f"[*] Copied tokenizer -> {out / 'tokenizer'}")
+    # pin the plain chat template in the exported tokenizer (the saved one is
+    # empty, which silently falls back to the think-injecting Qwen3 default)
+    tok_cfg_path = out / "tokenizer" / "tokenizer_config.json"
+    if tok_cfg_path.exists():
+        tc = json.loads(tok_cfg_path.read_text())
+        tc["chat_template"] = PLAIN_CHAT_TEMPLATE
+        tok_cfg_path.write_text(json.dumps(tc, indent=2, ensure_ascii=False))
+        print("[*] Pinned plain chat_template in exported tokenizer")
+
+    # chat_template.jinja takes precedence over the config string since
+    # transformers 5.x; overwrite it so both sources carry the plain
+    # template (the Qwen3 default jinja injects <think> blocks).
+    (out / "tokenizer" / "chat_template.jinja").write_text(
+        PLAIN_CHAT_TEMPLATE)
+    print("[*] Pinned plain chat_template.jinja")
+
+
+    src = Path(__file__).resolve().parent
+    shutil.copy2(src / "hf_modeling.py", out / "hf_modeling.py")
+    print("[*] Copied hf_modeling.py (trust_remote_code)")
+
+    package_scripts(out)
+    print(f"[*] Done: {out}")
+    print("    Load with: AutoModelForCausalLM.from_pretrained("
+          f"'{out}', trust_remote_code=True)")
+    print("    (Write the README yourself; export never touches it.)")
+
+
+def main():
+    p = argparse.ArgumentParser(description="Export MetaDiffusion-600M release dir")
+    p.add_argument("--checkpoint", required=True)
+    p.add_argument("--tokenizer", default="data/tokenizer")
+    p.add_argument("--output", required=True)
+    p.add_argument("--dtype", default="bf16", choices=["bf16", "fp16", "fp32"],
+                   help="Weight dtype (default bf16: matches training dtype and "
+                        "survives timestep extrapolation; fp16 overflows to NaN "
+                        "on curriculum-trained checkpoints)")
+    args = p.parse_args()
+    export(args.checkpoint, args.tokenizer, args.output, dtype=args.dtype)
+
+
+if __name__ == "__main__":
+    main()

model.py

@@ -0,0 +1,296 @@
+"""
+MetaDiffusionLM: masked-diffusion LM converted from Qwen3-0.6B-Instruct.
+
+Architecture: 28L x 1024W, GQA (16Q / 8KV, head_dim 128), QK-norm, RoPE,
+timestep conditioning (sinusoidal MLP + zero-init per-block residual).
+Bidirectional attention (no causal mask) is the key difference from the AR source.
+
+Parameter init:
+- Copied from AR checkpoint: token embeddings, all transformer blocks, norms,
+  QK-norm, RoPE buffers.
+- New, zero-init (identity at step 0): timestep embedding MLP, per-block
+  timestep residual. The model starts as exactly the AR model; diffusion
+  behavior is learned on top.
+- New, mean-init: [MASK] token row and the 7 rainbow padding token rows
+  (appended to both embed_tokens and the untied lm_head).
+
+Training loss (train.py):
+- CE on masked positions (the diffusion objective).
+"""
+
+import math
+from dataclasses import asdict, dataclass, field
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class MetaDiffusionConfig:
+    hidden_size: int = 1024
+    intermediate_size: int = 3072
+    num_hidden_layers: int = 28
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 8
+    head_dim: int = 128
+    vocab_size: int = 151669          # Qwen3 real tokenizer vocab (config 151936 is TP-padded)
+    mask_vocab_size: int = 151677     # + [MASK] + 7 rainbow tokens
+    mask_token_id: int = 151669
+    pad_token_id: int = 151643        # <|endoftext|> in Qwen3
+    max_position_embeddings: int = 32768
+    rope_theta: float = 1000000.0
+    rms_norm_eps: float = 1e-6
+    hidden_act: str = "silu"
+    qk_norm: bool = True
+    timestep_emb_hidden: int = 1024
+    tie_word_embeddings: bool = False
+    mask_ratio_min: float = 0.0
+    mask_ratio_max: float = 1.0
+    dtype: str = "float32"
+
+
+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
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.max_position_embeddings = max_position_embeddings
+
+    def forward(self, x, position_ids):
+        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 @ position_ids_expanded).transpose(1, 2)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos = emb.cos().to(dtype=x.dtype)
+        sin = emb.sin().to(dtype=x.dtype)
+        return cos, sin
+
+
+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)
+    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 MLP 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):
+        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 a pure copy of the AR model at
+    step 0. Unlike the old zero-init ADDITIVE residual (TimestepResidual),
+    the gradient here is dL/dscale = dL/dx * x with x nonzero, so the
+    t-path trains: the additive version deadlocked (zero output through a
+    zero weight = zero outer-product gradient forever), leaving the model
+    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 SelfAttention(nn.Module):
+    """GQA attention, bidirectional (no causal mask), optional QK-norm."""
+
+    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 * 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 = self.q_norm(q)
+        k = 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)
+        out = out.transpose(1, 2).contiguous().view(batch, seq, -1)
+        return self.o_proj(out)
+
+
+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 TransformerBlock(nn.Module):
+    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_modulation = TimestepModulation(config.hidden_size)
+
+    def forward(self, x, timestep_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, timestep_emb)
+
+        residual = x
+        x = self.post_attention_layernorm(x)
+        x = self.mlp(x)
+        x = residual + x
+        x = self.timestep_modulation(x, timestep_emb)
+        return x
+
+
+class MetaDiffusionLM(nn.Module):
+    def __init__(self, config: MetaDiffusionConfig):
+        super().__init__()
+        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([TransformerBlock(config) for _ in range(config.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        if config.tie_word_embeddings:
+            self.lm_head = None
+        else:
+            self.lm_head = nn.Linear(config.hidden_size, config.mask_vocab_size, bias=False)
+
+    def forward(self, input_ids, timesteps, attention_mask=None):
+        batch, seq = input_ids.shape
+        position_ids = torch.arange(seq, device=input_ids.device).unsqueeze(0).expand(batch, -1)
+
+        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)
+        if self.lm_head is not None:
+            logits = self.lm_head(x)
+        else:
+            logits = F.linear(x, self.embed_tokens.weight)
+        return logits
+
+    def compute_loss(self, logits, labels, mask_positions, pad_token_id=None,
+                     eos_token_id=None, eos_weight=1.0):
+        """CE on masked positions only (the masked-diffusion objective).
+
+        eos_token_id/eos_weight: boost the loss on the terminator token when
+        it is a masked target, so the model learns to emit it (EOS-weighting,
+        arXiv 2506.05017)."""
+        logits_masked = logits[mask_positions]
+        labels_masked = labels[mask_positions]
+        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
+        ce = F.cross_entropy(logits_masked, labels_masked, reduction="none")
+        if eos_token_id is not None and eos_weight != 1.0:
+            w = torch.where(labels_masked == eos_token_id, eos_weight, 1.0)
+            ce = ce * w
+        return ce.mean(), labels_masked.numel()
+
+    @classmethod
+    def from_checkpoint(cls, checkpoint_path, device="cpu"):
+        ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+        config_dict = ckpt.get("config", ckpt)
+        config = MetaDiffusionConfig(
+            **{k: v for k, v in config_dict.items() if k in MetaDiffusionConfig.__dataclass_fields__}
+        )
+        model = cls(config)
+        sd = {k.replace("_orig_mod.", "", 1) if k.startswith("_orig_mod.") else k: v
+              for k, v in ckpt.get("model_state_dict", ckpt).items()}
+        model.load_state_dict(sd, strict=True)
+        return model, ckpt

prepare_data.py

@@ -0,0 +1,388 @@
+#!/usr/bin/env python3
+"""
+prepare_data.py: build the chat SFT corpus for MetaDiffusion-600M.
+
+Datasets:
+- HuggingFaceTB/smol-smoltalk      : 460K general instruction conversations
+- OpenCoder-LLM/opc-sft-stage1[lang:python] + stage2 : code instruction data
+- OpenMathInstruct-2               : math (CoT solutions + answers)
+
+Pipeline per conversation:
+1. Normalize row -> messages [{role, content}]
+2. Format with the Qwen chat template (<|im_start|>...<|im_end|>)
+3. Tokenize; mark assistant-content tokens + rainbow padding as "response"
+4. Truncate to seq_len (drop samples whose prompt alone overflows)
+5. Rainbow-pad (cyclic <|r1|>..<|r7|>) so the model never sees repeated <eos>
+
+Output: data/ids.bin (uint32 token ids; vocab is 151677, does NOT fit uint16),
+data/resp.bin (uint8: 1 = assistant content / rainbow region), data/meta.json.
+
+train.py masks ONLY the resp==1 region by default (prompt stays clean), which
+matches the proven ChatDataset convention.
+
+Usage:
+    python prepare_data.py --out data --seq-len 512
+    python prepare_data.py --datasets smol --max-samples 50000 --out data/smol-small
+"""
+
+import argparse
+import itertools
+import json
+import logging
+import random
+import re
+from pathlib import Path
+
+import numpy as np
+from transformers import AutoTokenizer
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+RAINBOW_TOKENS = [f"<|r{i}|>" for i in range(1, 8)]
+NUM_NEW_TOKENS = 1 + len(RAINBOW_TOKENS)  # [MASK] + rainbow
+
+# Per-dataset default sample caps (None = everything)
+DEFAULT_CAPS = {"smol": None, "opc": 150_000, "math": 100_000, "no_robots": None}
+
+
+def ensure_special_tokens(tokenizer):
+    """Add [MASK] + rainbow tokens if missing.
+
+    Works for both a fresh tokenizer (adds all 8 at the end) and the
+    already-extended one saved by convert.py: [MASK] is always the first of
+    the 8 appended tokens, so mask_id == len(tokenizer) - 8."""
+    added = []
+    if tokenizer.convert_tokens_to_ids("[MASK]") == tokenizer.unk_token_id:
+        added.append("[MASK]")
+    missing_rainbow = [t for t in RAINBOW_TOKENS
+                       if tokenizer.convert_tokens_to_ids(t) == tokenizer.unk_token_id]
+    if missing_rainbow:
+        added.extend(missing_rainbow)
+    if added:
+        tokenizer.add_special_tokens({"additional_special_tokens": added})
+    mask_id = tokenizer.convert_tokens_to_ids("[MASK]")
+    expected = len(tokenizer) - NUM_NEW_TOKENS
+    assert mask_id == expected, f"[MASK] at {mask_id}, expected {expected}"
+    return tokenizer
+
+
+def get_messages(row) -> list | None:
+    """Normalize a dataset row into [{role, content}] or None."""
+    if isinstance(row, dict):
+        for key in ("messages", "conversations", "conversation"):
+            val = row.get(key)
+            if isinstance(val, list) and val:
+                msgs = []
+                for m in val:
+                    role = str(m.get("role", "")).lower()
+                    if role in ("human", "prompt", "user"):
+                        role = "user"
+                    elif role in ("gpt", "assistant", "response", "bot", "output"):
+                        role = "assistant"
+                    if role not in ("user", "assistant", "system"):
+                        continue
+                    content = m.get("content", m.get("value", ""))
+                    if isinstance(content, list):
+                        content = " ".join(str(c.get("text", c)) for c in content)
+                    if content:
+                        msgs.append({"role": role, "content": str(content).strip()})
+                if msgs and any(m["role"] == "assistant" for m in msgs):
+                    return msgs
+        # Instruction/output shapes
+        inst = row.get("instruction") or row.get("prompt") or row.get("question") or row.get("problem")
+        out = row.get("output") or row.get("response") or row.get("answer") or row.get("solution")
+        if inst and out:
+            return [{"role": "user", "content": str(inst).strip()},
+                    {"role": "assistant", "content": str(out).strip()}]
+    return None
+
+
+def math_messages(row) -> list | None:
+    """OpenMathInstruct-2 shape: problem + generated_solution + expected_answer."""
+    if not isinstance(row, dict):
+        return None
+    problem = row.get("problem") or row.get("question")
+    solution = row.get("generated_solution") or row.get("solution")
+    if not problem or not solution:
+        return None
+    answer = row.get("expected_answer")
+    if answer and str(answer).strip():
+        solution = f"{solution}\n\nFinal answer: {answer}"
+    return [{"role": "user", "content": str(problem).strip()},
+            {"role": "assistant", "content": str(solution).strip()}]
+
+
+_THINK_RE = re.compile(r"<think>.*?</think>", re.S)
+
+PLAIN_CHAT_TEMPLATE = (
+    "{% for message in messages %}"
+    "{{ '<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>\\n' }}"
+    "{% endfor %}"
+)
+
+
+def is_junk_content(content):
+    if "\uFFFD" in content:
+        return True
+    if not content:
+        return False
+    n_ascii = sum(1 for c in content if ord(c) < 128)
+    return (len(content) - n_ascii) / len(content) > 0.25
+
+
+def strip_think(messages):
+    out = []
+    for m in messages:
+        content = m["content"]
+        if m["role"] == "assistant" and isinstance(content, str):
+            content = _THINK_RE.sub("", content)
+            content = re.sub(r"\n{3,}", "\n\n", content).strip()
+        out.append({"role": m["role"], "content": content})
+    return out
+
+
+def format_conversation(tokenizer, messages, seq_len, min_response_tokens):
+    """Tokenize a conversation. Returns (ids, resp_flags) truncated/padded to
+    seq_len, or None if the prompt alone cannot fit.
+
+    The assistant span INCLUDES the trailing <|im_end|> token: a masked
+    diffusion model can only learn to emit the terminator if it is a masked
+    training target (VoidPadding format: [prompt][response][im_end][pad]*).
+    """
+    messages = strip_think(messages)
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
+    enc = tokenizer(text, return_offsets_mapping=True, add_special_tokens=False)
+    ids = enc["input_ids"]
+    offsets = enc["offset_mapping"]
+
+    resp_flags = [0] * len(ids)
+    pos = 0
+    for m in messages:
+        if m["role"] == "assistant":
+            marker = "<|im_start|>assistant\n"
+            idx = text.find(marker, pos)
+            if idx == -1:
+                return None
+            content_start = idx + len(marker)
+            content_end = content_start + len(m["content"])
+            # extend the span through the <|im_end|> token (the terminator is
+            # a first-class masked target; this is the im_end-hardening fix)
+            im_end_at = text.find("<|im_end|>", content_end)
+            if im_end_at != -1:
+                content_end = im_end_at + len("<|im_end|>")
+            for t_i, (cs, ce) in enumerate(offsets):
+                if cs >= content_start and ce <= content_end and t_i < len(resp_flags):
+                    resp_flags[t_i] = 1
+            pos = content_end
+    if not any(resp_flags):
+        return None
+
+    # Find the first response token; prompt must fit with room for a response
+    resp_start = resp_flags.index(1)
+    if resp_start > seq_len - min_response_tokens:
+        return None
+
+    ids = ids[:seq_len]
+    resp_flags = resp_flags[:seq_len]
+
+    # Rainbow pad the tail (all maskable, teaches the model to stop-and-pad)
+    rainbow_ids = [tokenizer.convert_tokens_to_ids(t) for t in RAINBOW_TOKENS]
+    n_pad = seq_len - len(ids)
+    if n_pad > 0:
+        ids = ids + [rainbow_ids[i % 7] for i in range(n_pad)]
+        resp_flags = resp_flags + [1] * n_pad
+    return ids, resp_flags
+
+
+def load_rows(dataset_name, max_samples):
+    from datasets import load_dataset
+
+    if dataset_name == "smol":
+        ds = load_dataset("HuggingFaceTB/smol-smoltalk", split="train")
+    elif dataset_name == "opc":
+        rows = []
+        for repo in ("OpenCoder-LLM/opc-sft-stage1", "OpenCoder-LLM/opc-sft-stage2"):
+            loaded = None
+            for config in ("lang:python", "lang:generic", None):
+                try:
+                    loaded = load_dataset(repo, config, split="train") if config else \
+                        load_dataset(repo, split="train")
+                    logger.info(f"  loaded {repo} config={config}: {len(loaded)} rows")
+                    break
+                except Exception:
+                    continue
+            if loaded is not None:
+                rows.append(loaded)
+        if not rows:
+            raise RuntimeError("Could not load any opc-sft config")
+        ds = rows[0] if len(rows) == 1 else None
+    elif dataset_name == "math":
+        ds = load_dataset("nvidia/OpenMathInstruct-2", split="train",
+                          streaming=True)
+    elif dataset_name == "no_robots":
+        ds = load_dataset("HuggingFaceH4/no_robots", split="train")
+    else:
+        raise ValueError(f"unknown dataset: {dataset_name}")
+
+    if ds is None:
+        # opc multi-repo path: chain them
+        def gen():
+            for r in rows:
+                yield from r
+        return gen()
+    return ds
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Build MetaDiffusion-600M chat SFT corpus")
+    parser.add_argument("--out", default="data", help="Output dir (ids.bin, resp.bin, meta.json)")
+    parser.add_argument("--tokenizer", default="data/tokenizer",
+                        help="Tokenizer dir (from convert.py) or HF id")
+    parser.add_argument("--seq-len", type=int, default=512)
+    parser.add_argument("--datasets", default="smol,opc,math",
+                        help="Comma list of: smol, opc, math, no_robots")
+    parser.add_argument("--jsonl", default=None,
+                        help="Local JSONL of {\"messages\": [...]} rows")
+    parser.add_argument("--max-samples", type=int, default=0,
+                        help="Per-dataset cap (0 = dataset default)")
+    parser.add_argument("--math-repeat", type=int, default=1,
+                        help="Process the math dataset N times")
+    parser.add_argument("--filter-junk", action=argparse.BooleanOptionalAction, default=True,
+                        help="Drop samples whose assistant content is >25%% non-ASCII. On by default; --no-filter-junk to keep them.")
+    parser.add_argument("--length-balance", action="store_true",
+                        help="Duplicate samples whose response-token count is in "
+                             "[--length-balance-min, --length-balance-max] by "
+                             "--length-balance-mult (targets the 60-160 gen budget).")
+    parser.add_argument("--length-balance-min", type=int, default=60)
+    parser.add_argument("--length-balance-max", type=int, default=160)
+    parser.add_argument("--length-balance-mult", type=int, default=3)
+    parser.add_argument("--min-response-tokens", type=int, default=8)
+    parser.add_argument("--val-fraction", type=float, default=0.05,
+                        help="Hold out this fraction as ids_val.bin/resp_val.bin for early stopping")
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
+    tokenizer = ensure_special_tokens(tokenizer)
+    tokenizer.chat_template = PLAIN_CHAT_TEMPLATE
+    logger.info(f"Tokenizer: {len(tokenizer)} tokens, [MASK]={tokenizer.convert_tokens_to_ids('[MASK]')}")
+
+    rng = random.Random(args.seed)
+    all_ids = []
+    all_resp = []
+    n_samples = 0
+
+    def process_rows(ds_name, rows, cap):
+        nonlocal n_samples
+        ds_n = 0
+        ds_skipped = 0
+        for row in rows:
+            if cap is not None and ds_n >= cap:
+                break
+            msgs = math_messages(row) if ds_name == "math" else get_messages(row)
+            if msgs is None:
+                ds_skipped += 1
+                continue
+            if args.filter_junk:
+                ac = msgs[-1]["content"] if msgs and msgs[-1]["role"] == "assistant" else ""
+                if isinstance(ac, str) and is_junk_content(ac):
+                    ds_skipped += 1
+                    continue
+            result = format_conversation(tokenizer, msgs, args.seq_len, args.min_response_tokens)
+            if result is None:
+                ds_skipped += 1
+                continue
+            ids, resp = result
+            copies = 1
+            if args.length_balance:
+                n_resp = int(sum(resp))
+                if args.length_balance_min <= n_resp <= args.length_balance_max:
+                    copies = args.length_balance_mult
+            for _ in range(copies):
+                all_ids.append(np.array(ids, dtype=np.uint32))
+                all_resp.append(np.array(resp, dtype=np.uint8))
+                ds_n += 1
+            if ds_n % 20000 == 0:
+                logger.info(f"  [{ds_name}] {ds_n:,} samples")
+        logger.info(f"[{ds_name}] kept {ds_n:,}, skipped {ds_skipped:,}")
+        n_samples += ds_n
+
+    for ds_name in [d.strip() for d in args.datasets.split(",") if d.strip()]:
+        cap = args.max_samples if args.max_samples > 0 else DEFAULT_CAPS.get(ds_name)
+        logger.info(f"[{ds_name}] loading (cap={cap})...")
+        rows = load_rows(ds_name, cap)
+        repeats = args.math_repeat if ds_name == "math" else 1
+        if repeats > 1:
+            cache_path = out / f".math_cache{'_' + str(cap) if cap else ''}.jsonl"
+            if cache_path.exists():
+                logger.info(f"[math] using local cache {cache_path}")
+                with open(cache_path) as f:
+                    rows = [json.loads(line) for line in f if line.strip()]
+            else:
+                rows = list(itertools.islice(rows, cap) if cap else rows)
+                with open(cache_path, "w") as f:
+                    for r in rows:
+                        f.write(json.dumps(r) + "\n")
+                logger.info(f"[math] cached {len(rows):,} rows to {cache_path}")
+        for rep in range(repeats):
+            if repeats > 1:
+                logger.info(f"[{ds_name}] pass {rep + 1}/{repeats}")
+            process_rows(ds_name, rows, cap)
+
+    if args.jsonl:
+        with open(args.jsonl) as f:
+            rows = [json.loads(line) for line in f if line.strip()]
+        logger.info(f"[jsonl] {len(rows)} rows from {args.jsonl}")
+        cap = args.max_samples if args.max_samples > 0 else None
+        process_rows("jsonl", rows, cap)
+
+    if n_samples == 0:
+        raise SystemExit("no samples kept; check --datasets / --jsonl / filters")
+    logger.info(f"Shuffling {n_samples:,} samples (seed {args.seed})...")
+    order = list(range(n_samples))
+    rng.shuffle(order)
+    n_val = int(n_samples * args.val_fraction)
+    train_order = order[n_val:]
+    val_order = order[:n_val]
+    if not train_order:
+        raise SystemExit(f"val-fraction {args.val_fraction} left zero train samples")
+    ids_flat = np.concatenate([all_ids[i] for i in train_order])
+    resp_flat = np.concatenate([all_resp[i] for i in train_order])
+
+    ids_flat.tofile(out / "ids.bin")
+    resp_flat.tofile(out / "resp.bin")
+
+    if n_val > 0:
+        ids_val = np.concatenate([all_ids[i] for i in val_order])
+        resp_val = np.concatenate([all_resp[i] for i in val_order])
+        ids_val.tofile(out / "ids_val.bin")
+        resp_val.tofile(out / "resp_val.bin")
+        logger.info(f"Val held out: {n_val:,} samples -> {out / 'ids_val.bin'} / {out / 'resp_val.bin'}")
+
+    meta = {
+        "seq_len": args.seq_len,
+        "n_samples": n_samples - n_val,
+        "n_val_samples": n_val,
+        "val_held_out": True,
+        "n_tokens": int(ids_flat.shape[0]),
+        "mask_token_id": tokenizer.convert_tokens_to_ids("[MASK]"),
+        "rainbow_token_ids": [tokenizer.convert_tokens_to_ids(t) for t in RAINBOW_TOKENS],
+        "vocab_size": len(tokenizer),
+        "datasets": args.datasets + (f",jsonl:{args.jsonl}" if args.jsonl else ""),
+        "tokenizer_dir": str(args.tokenizer),
+        "seed": args.seed,
+        "filter_junk": args.filter_junk,
+        "length_balance": args.length_balance,
+    }
+    with open(out / "meta.json", "w") as f:
+        json.dump(meta, f, indent=2)
+    logger.info(f"Wrote {out/'ids.bin'} ({ids_flat.nbytes/1e9:.2f} GB, {meta['n_tokens']:,} tokens), "
+                f"{out/'resp.bin'}, {out/'meta.json'}")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,5 @@
+torch>=2.2
+transformers>=4.49
+safetensors>=0.4
+datasets>=2.18
+numpy>=1.26

train.py

@@ -0,0 +1,702 @@
+#!/usr/bin/env python3
+"""
+train.py: masked-diffusion chat SFT for MetaDiffusion-600M.
+
+Usage:
+    python train.py --init-checkpoint init/metadiffusion-600M-instruct.pt \
+        --data-dir data --output-dir checkpoints --max-steps 30000
+    python train.py ... --keep-knowledge --keep-mult 0.05
+"""
+
+import argparse
+import datetime
+import glob
+import heapq
+import json
+import logging
+import math
+import os
+import re
+import signal
+import sys
+import time
+from dataclasses import asdict
+
+import numpy as np
+import torch
+from torch.optim.lr_scheduler import LambdaLR
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def format_bytes(b):
+    for unit in ["B", "KB", "MB", "GB", "TB"]:
+        if b < 1024:
+            return f"{b:.1f} {unit}"
+        b /= 1024
+    return f"{b:.1f} PB"
+
+
+def format_eta(seconds):
+    seconds = max(0, int(seconds))
+    h, rem = divmod(seconds, 3600)
+    m, s = divmod(rem, 60)
+    if h > 0:
+        return f"{h}h{m:02d}m"
+    if m > 0:
+        return f"{m}m{s:02d}s"
+    return f"{s}s"
+
+
+def detect_max_batch_size(model, seq_len, device, keep_free_fraction=0.1,
+                          vocab=151677, reserve_bytes=0):
+    """Forward+backward probe to find the largest batch that fits.
+
+    reserve_bytes: memory the optimizer states need during real training
+    (2 x n_params x 4 for fp32 Adam, x 1 for 8-bit Adam). Without this the
+    probe over-predicts and the first real step OOMs."""
+    if not torch.cuda.is_available():
+        return 8
+    # mem_get_info() returns (free, total). Base the budget on ACTUAL free
+    # memory: other processes may hold part of the card.
+    free_mem, total_mem = torch.cuda.mem_get_info()
+    logger.info(f"GPU: {format_bytes(total_mem)} total, {format_bytes(free_mem)} free")
+    mem_limit = free_mem - int(total_mem * keep_free_fraction) - reserve_bytes
+    logger.info(f"VRAM budget: {format_bytes(mem_limit)} "
+                f"(based on {format_bytes(free_mem)} free, "
+                f"reserving {format_bytes(reserve_bytes)} for optimizer + accum grads)")
+    model = model.to(device)
+    model.train()
+
+    last_working = 1
+    for bs in [1, 2, 4, 8, 16, 24, 32, 48]:
+        torch.cuda.empty_cache()
+        try:
+            ids = torch.randint(0, vocab, (bs, seq_len), device=device)
+            labels = ids.clone()
+            mask = torch.rand(bs, seq_len, device=device) < 0.5
+            t = torch.rand(bs, device=device)
+            logits = model(ids, t)
+            loss, n = model.compute_loss(logits, labels, mask)
+            if n > 0:
+                (loss / 4).backward()
+            torch.cuda.synchronize()
+            peak = torch.cuda.max_memory_allocated()
+            logger.info(f"  batch={bs:>2d}: peak {format_bytes(peak)} "
+                        f"({peak/total_mem*100:.0f}%)")
+            if peak >= mem_limit:
+                break
+            last_working = bs
+        except RuntimeError as e:
+            if "out of memory" in str(e).lower():
+                logger.info(f"  batch={bs:>2d}: OOM")
+                break
+            raise
+        finally:
+            model.zero_grad(set_to_none=True)
+            torch.cuda.empty_cache()
+    logger.info(f"Detected max batch_size={last_working}")
+    return last_working
+
+
+def get_step_from_filename(filename):
+    basename = os.path.basename(filename)
+    m = re.match(r"step_(\d+)(?:_\w+)?\.pt$", basename)
+    return int(m.group(1)) if m else None
+
+def load_best_val_steps(stats_path, max_n):
+    """Return the steps with the N lowest val losses from stats.jsonl.
+
+    Falls back to train loss when val never ran (patience=0)."""
+    if not os.path.exists(stats_path):
+        return set()
+    val_entries, loss_entries = [], []
+    with open(stats_path) as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                e = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            if "step" not in e:
+                continue
+            if e.get("val_loss") is not None:
+                val_entries.append((float(e["val_loss"]), int(e["step"])))
+            elif "loss" in e:
+                loss_entries.append((float(e["loss"]), int(e["step"])))
+    pool = val_entries if val_entries else loss_entries
+    return {s for _, s in heapq.nsmallest(max_n, pool)}
+
+
+def cleanup_checkpoints(output_dir, keep_last_n, keep_best_n, stats_path):
+    """Keep the N latest regular checkpoints + the N best-by-val-loss
+    checkpoints; delete everything else to save storage."""
+    # sort by step NUMBER, not filename: lexicographic order breaks once
+    # steps hit 5 digits (step_30000.pt < step_4000.pt alphabetically)
+    all_ckpts = sorted(glob.glob(os.path.join(output_dir, "step_*.pt")),
+                       key=lambda c: (get_step_from_filename(c) or -1, c))
+    if len(all_ckpts) <= keep_last_n + keep_best_n:
+        return
+
+    last_steps = {get_step_from_filename(c) for c in all_ckpts[-keep_last_n:]}
+    best_steps = load_best_val_steps(stats_path, keep_best_n)
+
+    kept = set()
+    # latest N: regular files only (no _valbest suffix)
+    for s in last_steps:
+        for c in all_ckpts:
+            if "valbest" not in c and get_step_from_filename(c) == s:
+                kept.add(c)
+    # best N by val: prefer the regular file at that step, else its valbest
+    for s in best_steps:
+        cands = [c for c in all_ckpts if get_step_from_filename(c) == s]
+        if not cands:
+            continue
+        regular = [c for c in cands if "valbest" not in c]
+        kept.add(regular[0] if regular else cands[0])
+
+    deleted = 0
+    for c in all_ckpts:
+        if c not in kept:
+            try:
+                os.remove(c)
+                deleted += 1
+            except OSError:
+                pass
+    if deleted:
+        logger.info(f"Cleaned {deleted} checkpoints "
+                    f"(kept {len(kept)}: {keep_last_n} latest + {keep_best_n} best-val)")
+
+
+def cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps, min_lr_ratio=0.1):
+    def lr_lambda(step):
+        if step < warmup_steps:
+            return float(step) / float(max(1, warmup_steps))
+        progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps))
+        return max(min_lr_ratio, 0.5 * (1.0 + math.cos(math.pi * progress)))
+    return LambdaLR(optimizer, lr_lambda)
+
+
+def last_eos_mask(clean, eos_token_id, resp_bool):
+    """True only at the rightmost <|im_end|> in the response region per row.
+
+    Historical multi-turn terminators stay unforced; the final terminator
+    still gets a gradient every step under --eos-mask-always."""
+    is_eos = (clean == eos_token_id) & resp_bool
+    rev = torch.flip(is_eos, [1]).int()
+    last = torch.flip(rev.cumsum(1) == 1, [1]) & is_eos
+    return last
+
+
+def no_decay_name(name, param):
+    return param.ndim < 2 or "norm" in name or name.endswith(".bias")
+
+
+def main():
+    p = argparse.ArgumentParser(description="MetaDiffusion-600M chat SFT")
+    p.add_argument("--init-checkpoint", required=True)
+    p.add_argument("--data-dir", default="data")
+    p.add_argument("--output-dir", default="checkpoints")
+    p.add_argument("--max-steps", type=int, default=30000)
+    p.add_argument("--seq-len", type=int, default=512)
+    p.add_argument("--batch-size", type=int, default=0, help="0 = auto-detect")
+    p.add_argument("--grad-accum-steps", type=int, default=2)
+    p.add_argument("--lr", type=float, default=5e-5)
+    p.add_argument("--min-lr-ratio", type=float, default=0.1)
+    p.add_argument("--warmup-steps", type=int, default=200)
+    p.add_argument("--weight-decay", type=float, default=0.01)
+    p.add_argument("--transferred-lr-mult", type=float, default=0.33)
+    p.add_argument("--max-grad-norm", type=float, default=1.0)
+    p.add_argument("--optimizer", default="auto",
+                   choices=["auto", "adamw", "adamw8bit"],
+                   help="auto = torchao/bitsandbytes 8-bit Adam if installed, "
+                        "else AdamW. adamw8bit forces 8-bit (warns if missing).")
+    p.add_argument("--dtype", default="bfloat16", choices=["float32", "bfloat16", "float16"])
+    p.add_argument("--mask-all", action="store_true", help="mask the whole sequence, not just responses")
+    p.add_argument("--mask-ratio-min", type=float, default=0.0)
+    p.add_argument("--mask-ratio-max", type=float, default=1.0)
+    p.add_argument("--curriculum", action="store_true",
+                   help="Ramp the mask ratio from --mask-ratio-min to 1.0 over the "
+                        "run. Val stays at fixed --val-t.")
+    p.add_argument("--curriculum-early-stop-gate",
+                   action=argparse.BooleanOptionalAction, default=True,
+                   help="While the curriculum ramp has not yet covered --val-t, "
+                        "val is out-of-domain (extrapolation): compute + log it, "
+                        "but do NOT update best-val / patience / valbest. Without "
+                        "this, early stopping can fire mid-ramp on OOD noise and "
+                        "kill the run before the model learns high-mask denoising. "
+                        "--no-curriculum-early-stop-gate restores the old behavior.")
+    p.add_argument("--eos-weight", type=float, default=1.0,
+                   help="Loss multiplier on the <|im_end|> token when it is a "
+                        "masked target (teach termination; try 10, hammer 25)")
+    p.add_argument("--eos-token-id", type=int, default=151645,
+                   help="<|im_end|> token id (Qwen3: 151645)")
+    p.add_argument("--eos-mask-always", action="store_true",
+                   help="Force-mask the last <|im_end|> terminator in every window, "
+                        "regardless of t: the terminator gets a gradient on every "
+                        "step and the curriculum can never starve it")
+    p.add_argument("--keep-knowledge", action="store_true", help="RND1-style LR split")
+    p.add_argument("--keep-mult", type=float, default=0.05, help="LR mult for MLP/norm/embed when --keep-knowledge")
+    p.add_argument("--save-every", type=int, default=2000)
+    p.add_argument("--log-every", type=int, default=50)
+    p.add_argument("--val-every", type=int, default=500,
+                   help="Val loss check interval (needs ids_val.bin/resp_val.bin)")
+    p.add_argument("--patience", type=int, default=0,
+                   help="Early stop after N val checks without improvement (0 = off)")
+    p.add_argument("--min-delta", type=float, default=1e-4,
+                   help="Min val loss improvement to count as improvement")
+    p.add_argument("--val-batches", type=int, default=50,
+                   help="Batches averaged per val check")
+    p.add_argument("--val-t", type=float, default=0.5,
+                   help="Fixed mask ratio for val checks (stable early stopping; "
+                        "t~U(0,1) makes val swing +-0.5 and patience fires on noise)")
+    p.add_argument("--keep-last-n", type=int, default=3,
+                   help="Keep the N latest regular checkpoints")
+    p.add_argument("--keep-best-n", type=int, default=3,
+                   help="Keep the N lowest-val-loss checkpoints (train-loss fallback)")
+    p.add_argument("--resume-from", default=None)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--device", default="cuda:0")
+    args = p.parse_args()
+
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
+    dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[args.dtype]
+
+    # --- Model ---
+    from model import MetaDiffusionConfig, MetaDiffusionLM
+    logger.info(f"Loading init checkpoint: {args.init_checkpoint}")
+    ckpt = torch.load(args.init_checkpoint, map_location="cpu", weights_only=False)
+    config = MetaDiffusionConfig(
+        **{k: v for k, v in ckpt["config"].items() if k in MetaDiffusionConfig.__dataclass_fields__}
+    )
+    model = MetaDiffusionLM(config)
+    model.load_state_dict(ckpt["model_state_dict"], strict=True)
+    model = model.to(device=device, dtype=dtype)
+    logger.info(f"Model: {config.num_hidden_layers}L x {config.hidden_size}W, "
+                f"vocab={config.mask_vocab_size}, params={sum(p.numel() for p in model.parameters())/1e6:.1f}M, {args.dtype}")
+
+    # --- Data ---
+    data_dir = args.data_dir
+    ids_arr = np.memmap(os.path.join(data_dir, "ids.bin"), dtype=np.uint32, mode="r")
+    resp_arr = np.memmap(os.path.join(data_dir, "resp.bin"), dtype=np.uint8, mode="r")
+    assert ids_arr.shape == resp_arr.shape, "ids.bin / resp.bin length mismatch"
+    n_tokens = len(ids_arr)
+    logger.info(f"Data: {n_tokens:,} tokens, seq_len={args.seq_len}")
+
+    meta_path = os.path.join(data_dir, "meta.json")
+    data_meta = {}
+    if os.path.exists(meta_path):
+        with open(meta_path) as f:
+            data_meta = json.load(f)
+    if not data_meta.get("val_held_out") and data_meta.get("n_val_samples", 0) > 0:
+        skip = int(data_meta["n_val_samples"]) * int(data_meta.get("seq_len", args.seq_len))
+        if skip > 0 and skip < n_tokens:
+            ids_arr = ids_arr[skip:]
+            resp_arr = resp_arr[skip:]
+            n_tokens = len(ids_arr)
+            logger.info(f"Legacy leaked val: skipped first {skip:,} tokens of "
+                        f"ids.bin ({data_meta['n_val_samples']} samples). "
+                        f"Train tokens now {n_tokens:,}")
+
+    # --- Validation split (for early stopping) ---
+    val_ids_arr = val_resp_arr = None
+    val_n = 0
+    if args.patience > 0:
+        v_ids = os.path.join(data_dir, "ids_val.bin")
+        v_resp = os.path.join(data_dir, "resp_val.bin")
+        if os.path.exists(v_ids) and os.path.exists(v_resp):
+            val_ids_arr = np.memmap(v_ids, dtype=np.uint32, mode="r")
+            val_resp_arr = np.memmap(v_resp, dtype=np.uint8, mode="r")
+            val_n = len(val_ids_arr)
+            logger.info(f"Val data: {val_n:,} tokens (early stopping active, "
+                        f"patience={args.patience}, val_every={args.val_every})")
+        else:
+            logger.warning("--patience set but no ids_val.bin/resp_val.bin found; "
+                           "re-run prepare_data.py (writes a val split by default). "
+                           "Early stopping disabled.")
+            args.patience = 0
+
+    def sample_windows(batch_size, ids_src, resp_src, n_src, rng=None):
+        n_windows = n_src // args.seq_len
+        if n_windows < 1:
+            raise RuntimeError(
+                f"Need at least {args.seq_len} tokens, have {n_src}")
+        if rng is None:
+            idx = np.random.randint(0, n_windows, size=batch_size)
+        else:
+            idx = rng.randint(0, n_windows, size=batch_size)
+        starts = idx * args.seq_len
+        ids = np.stack([ids_src[s:s + args.seq_len] for s in starts])
+        resp = np.stack([resp_src[s:s + args.seq_len] for s in starts])
+        return torch.from_numpy(ids).long(), torch.from_numpy(resp).bool()
+
+    def sample_batch(batch_size, ids_arr_local, resp_arr_local, n_tokens_local,
+                     max_ratio=None, rng=None):
+        """Aligned windows; mask within the response region."""
+        clean, resp_bool = sample_windows(
+            batch_size, ids_arr_local, resp_arr_local, n_tokens_local, rng=rng)
+        max_r = args.mask_ratio_max if max_ratio is None else max_ratio
+        if rng is None:
+            t = torch.rand(batch_size) * (max_r - args.mask_ratio_min) + args.mask_ratio_min
+            mask_u = torch.rand(batch_size, args.seq_len)
+        else:
+            t = torch.from_numpy(
+                rng.rand(batch_size) * (max_r - args.mask_ratio_min) + args.mask_ratio_min
+            ).float()
+            mask_u = torch.from_numpy(rng.rand(batch_size, args.seq_len)).float()
+        mask_positions = mask_u < t[:, None]
+        if not args.mask_all:
+            mask_positions = mask_positions & resp_bool
+        if args.eos_mask_always:
+            mask_positions = mask_positions | last_eos_mask(
+                clean, args.eos_token_id, resp_bool)
+        input_ids = clean.clone()
+        input_ids[mask_positions] = config.mask_token_id
+        return input_ids, clean, mask_positions, t
+
+    # --- Optimizer selection (before batch detection: it reserves memory) ---
+    OptimizerCls = None
+    if args.optimizer in ("auto", "adamw8bit"):
+        for mod_name, cls_name, label in (
+                ("torchao.optim", "AdamW8bit", "torchao"),
+                ("bitsandbytes.optim", "AdamW8bit", "bitsandbytes")):
+            try:
+                mod = __import__(mod_name, fromlist=[cls_name])
+                OptimizerCls = getattr(mod, cls_name)
+                logger.info(f"Optimizer: 8-bit Adam ({label})")
+                break
+            except ImportError:
+                continue
+    if OptimizerCls is None:
+        from torch.optim import AdamW as OptimizerCls
+        if args.optimizer == "adamw8bit":
+            logger.warning("adamw8bit requested but neither torchao nor "
+                           "bitsandbytes is installed; using AdamW "
+                           "(pip install torchao)")
+
+    # --- Auto batch ---
+    if args.batch_size <= 0 and torch.cuda.is_available():
+        state_bytes = 1 if OptimizerCls.__name__ == "AdamW8bit" else 4
+        n_params = sum(p.numel() for p in model.parameters())
+        reserve_bytes = 2 * n_params * state_bytes  # optimizer states
+        # with grad accumulation, the (accum-1) earlier micro-batch grads are
+        # still alive when the last micro-batch's backward peaks
+        reserve_bytes += (args.grad_accum_steps - 1) * n_params * 2  # bf16 grads
+        logger.info(f"Optimizer/grad reserve: {format_bytes(reserve_bytes)} "
+                    f"(states {state_bytes} B/param + "
+                    f"{args.grad_accum_steps - 1} extra bf16 grads)")
+        logger.info("Auto-detecting max batch size...")
+        args.batch_size = detect_max_batch_size(
+            model, args.seq_len, device,
+            keep_free_fraction=0.1, vocab=config.mask_vocab_size,
+            reserve_bytes=reserve_bytes)
+    if args.batch_size <= 0:
+        args.batch_size = 8
+    logger.info(f"batch_size={args.batch_size} x grad_accum={args.grad_accum_steps} "
+                f"= effective {args.batch_size * args.grad_accum_steps}")
+
+    # --- Parameter groups ---
+    # Only the diffusion-new modules get full LR. embed_tokens / lm_head are
+    # almost entirely transferred AR rows (MASK + rainbow are 8 of 151677).
+    new_keys = ("timestep_emb", "timestep_modulation")
+    attn_key = "self_attn"
+    buckets = {("new", True): [], ("new", False): [],
+               ("attn", True): [], ("attn", False): [],
+               ("non_attn", True): [], ("non_attn", False): []}
+    for name, param in model.named_parameters():
+        if any(k in name for k in new_keys):
+            kind = "new"
+        elif attn_key in name:
+            kind = "attn"
+        else:
+            kind = "non_attn"
+        buckets[(kind, not no_decay_name(name, param))].append(param)
+
+    def make_group(kind, decay, lr, label):
+        params = buckets[(kind, decay)]
+        if not params:
+            return None
+        return {"params": params, "lr": lr,
+                "weight_decay": args.weight_decay if decay else 0.0,
+                "name": label}
+
+    if args.keep_knowledge:
+        raw = [
+            make_group("new", True, args.lr, "new"),
+            make_group("new", False, args.lr, "new_nodecay"),
+            make_group("attn", True, args.lr, "attention"),
+            make_group("attn", False, args.lr, "attention_nodecay"),
+            make_group("non_attn", True, args.lr * args.keep_mult, "mlp_norm_embed"),
+            make_group("non_attn", False, args.lr * args.keep_mult, "mlp_norm_embed_nodecay"),
+        ]
+    else:
+        tr_lr = args.lr * args.transferred_lr_mult
+        raw = [
+            make_group("new", True, args.lr, "new"),
+            make_group("new", False, args.lr, "new_nodecay"),
+            make_group("attn", True, tr_lr, "transferred"),
+            make_group("attn", False, tr_lr, "transferred_nodecay"),
+            make_group("non_attn", True, tr_lr, "transferred"),
+            make_group("non_attn", False, tr_lr, "transferred_nodecay"),
+        ]
+    param_groups = [g for g in raw if g is not None]
+    for g in param_groups:
+        n = sum(p.numel() for p in g["params"])
+        logger.info(f"  group {g['name']}: {n:,} params, lr={g['lr']:.2e}, "
+                    f"wd={g['weight_decay']}")
+
+    optimizer = OptimizerCls(param_groups)
+    scheduler = cosine_schedule_with_warmup(optimizer, args.warmup_steps, args.max_steps, args.min_lr_ratio)
+
+    start_step = 0
+    if args.resume_from:
+        ckpt_r = torch.load(args.resume_from, map_location="cpu", weights_only=False)
+        model.load_state_dict(ckpt_r["model_state_dict"])
+        start_step = int(ckpt_r.get("step", get_step_from_filename(args.resume_from) or 0))
+        if "optimizer_state" in ckpt_r:
+            try:
+                optimizer.load_state_dict(ckpt_r["optimizer_state"])
+                scheduler.load_state_dict(ckpt_r.get("scheduler_state", {}))
+                logger.info("Restored optimizer + scheduler state (seamless continuation)")
+            except Exception as e:
+                logger.warning(f"Could not restore optimizer state ({e}); "
+                               "rebuilding fresh. (Changed hyperparams / param group "
+                               "layout between segments?)")
+                optimizer = OptimizerCls(param_groups)
+                scheduler = cosine_schedule_with_warmup(
+                    optimizer, args.warmup_steps, args.max_steps, args.min_lr_ratio)
+        else:
+            logger.warning("Checkpoint has no optimizer state (older run); "
+                           "rebuilding fresh.")
+        logger.info(f"Resumed from {args.resume_from} (continuing from step {start_step})")
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    stats_path = os.path.join(args.output_dir, "stats.jsonl")
+    stats_file = open(stats_path, "a")
+
+    # --- Train loop ---
+    @torch.no_grad()
+    def compute_val_loss():
+        """Masked CE over --val-batches of the held-out split.
+
+        Fixed mask ratio (--val-t) + deterministic aligned windows AND
+        seeded masks: torch.rand made the estimate depend on the training
+        RNG and patience fired on noise."""
+        model.eval()
+        total_ce, total_n = 0.0, 0
+        rng = np.random.RandomState(args.seed)
+        for _ in range(args.val_batches):
+            clean, resp_bool = sample_windows(
+                args.batch_size, val_ids_arr, val_resp_arr, val_n, rng=rng)
+            t = torch.full((args.batch_size,), args.val_t)
+            mask_u = torch.from_numpy(rng.rand(args.batch_size, args.seq_len)).float()
+            mask_positions = mask_u < t[:, None]
+            if not args.mask_all:
+                mask_positions = mask_positions & resp_bool
+            if args.eos_mask_always:
+                mask_positions = mask_positions | last_eos_mask(
+                    clean, args.eos_token_id, resp_bool)
+            input_ids = clean.clone()
+            input_ids[mask_positions] = config.mask_token_id
+            input_ids = input_ids.to(device)
+            clean = clean.to(device)
+            mask_positions = mask_positions.to(device)
+            t = t.to(device)
+            logits = model(input_ids, t)
+            ce, n = model.compute_loss(logits, clean, mask_positions,
+                                       pad_token_id=config.pad_token_id,
+                                       eos_token_id=args.eos_token_id,
+                                       eos_weight=args.eos_weight)
+            total_ce += ce.item() * n
+            total_n += n
+        model.train()
+        return total_ce / max(1, total_n)
+
+    def save_ckpt(step, path):
+        torch.save({
+            "config": asdict(config),
+            "model_state_dict": model.state_dict(),
+            "optimizer_state": optimizer.state_dict(),
+            "scheduler_state": scheduler.state_dict(),
+            "step": step,
+            "metadata": {
+                "init_checkpoint": args.init_checkpoint,
+                "keep_knowledge": args.keep_knowledge,
+                "dtype": args.dtype,
+            },
+        }, path)
+        logger.info(f"Saved {path}")
+
+    model.train()
+    try:
+        scaler = torch.amp.GradScaler(
+            "cuda", enabled=(args.dtype == "float16" and torch.cuda.is_available()))
+    except AttributeError:  # torch < 2.3
+        scaler = torch.cuda.amp.GradScaler(
+            enabled=(args.dtype == "float16" and torch.cuda.is_available()))
+    eff_bs = args.batch_size * args.grad_accum_steps
+    best_val, best_step, no_improve = float("inf"), 0, 0
+
+    # Ctrl+C handling: first press finishes the current step and saves it,
+    # second press force-quits without saving.
+    interrupted = False
+
+    def _handle_sigint(sig, frame):
+        nonlocal interrupted
+        if interrupted:
+            print("\nSecond Ctrl+C: force quitting without save.", flush=True)
+            os._exit(130)
+        interrupted = True
+        print("\nCtrl+C received: finishing current step, then saving...",
+              flush=True)
+
+    signal.signal(signal.SIGINT, _handle_sigint)
+
+    if interrupted:
+        logger.info("Ctrl+C before training started; nothing to save.")
+        stats_file.close()
+        sys.exit(130)
+
+    if args.curriculum:
+        ramp_floor = max(args.mask_ratio_min, 0.1)
+        logger.info(f"Curriculum: mask ratio ramps {ramp_floor:.2f} -> 1.00 "
+                    f"over the run (val stays at fixed t={args.val_t})")
+        if args.curriculum_early_stop_gate and args.patience > 0:
+            logger.info(f"Curriculum early-stop gate ON: val at t={args.val_t} is "
+                        f"OOD until the ramp covers it; early stopping / valbest "
+                        f"inactive until then (--no-curriculum-early-stop-gate to "
+                        f"disable)")
+    if args.eos_weight != 1.0:
+        logger.info(f"EOS weighting: <|im_end|> (id {args.eos_token_id}) "
+                    f"loss x{args.eos_weight}")
+    if args.eos_mask_always:
+        logger.info("EOS mask-always: last <|im_end|> in each window is a "
+                    "masked target (curriculum cannot starve the terminator)")
+    logger.info(f"Training {args.max_steps:,} steps... (Ctrl+C saves current step)")
+    t0 = time.time()
+    step = start_step
+    for step in range(start_step + 1, args.max_steps + 1):
+        optimizer.zero_grad(set_to_none=True)
+        total_ce = 0.0
+        total_n = 0
+        val_loss = None
+        cur_max_ratio = None
+        if args.curriculum:
+            # easy-to-hard: mask ratio ramps from a 0.1 floor to 1.0 across the
+            # run. Starting at 0.0 leaves whole micro-batches unmasked (no grad
+            # path); the floor keeps t ~ U(0, 0.1) at step 1 (~100 masked
+            # tokens per 2048-token batch). Segment-relative: any resume
+            # segment covers the full ramp, so a shortened segment still trains
+            # the high-mask regime (where im_end lives) end to end.
+            seg_steps = max(1, args.max_steps - start_step)
+            progress = (step - start_step) / seg_steps
+            ramp_floor = max(args.mask_ratio_min, 0.1)
+            cur_max_ratio = ramp_floor + (1.0 - ramp_floor) * progress
+        did_backward = False
+        for _ in range(args.grad_accum_steps):
+            input_ids, clean, mask_positions, t = sample_batch(
+                args.batch_size, ids_arr, resp_arr, n_tokens,
+                max_ratio=cur_max_ratio)
+            input_ids = input_ids.to(device)
+            clean = clean.to(device)
+            mask_positions = mask_positions.to(device)
+            t = t.to(device)
+
+            logits = model(input_ids, t)
+            ce, n = model.compute_loss(logits, clean, mask_positions, pad_token_id=config.pad_token_id,
+                                       eos_token_id=args.eos_token_id,
+                                       eos_weight=args.eos_weight)
+            if n == 0:
+                continue  # no masked tokens this micro-batch (early curriculum);
+                          # compute_loss returns a constant, not a grad path
+            loss = ce / args.grad_accum_steps
+            total_ce += ce.item() * n  # weighted by masked-token count: true per-token CE
+            total_n += n
+            scaler.scale(loss).backward()
+            did_backward = True
+
+        if did_backward:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+            scaler.step(optimizer)
+            scaler.update()
+            scheduler.step()
+
+        if interrupted:
+            save_ckpt(step, os.path.join(args.output_dir, f"step_{step}.pt"))
+            logger.info(f"Ctrl+C: saved current step -> step_{step}.pt")
+            break
+
+        if step % args.log_every == 0:
+            elapsed = time.time() - t0
+            # segment-relative step count: cumulative `step` over segment
+            # elapsed would print a fake 100K+ tok/s after --resume-from
+            tok_per_s = eff_bs * args.seq_len * (step - start_step) / max(1.0, elapsed)
+            vram = torch.cuda.max_memory_allocated() / 1e9 if torch.cuda.is_available() else 0
+            steps_done = max(1, step - start_step)
+            eta_s = (args.max_steps - step) * (elapsed / steps_done)
+            finish = datetime.datetime.now() + datetime.timedelta(seconds=eta_s)
+            eta_str = (f"ETA {format_eta(eta_s)} (el {format_eta(elapsed)}, "
+                       f"done {finish.strftime('%I:%M %p').lstrip('0')})")
+            val_loss = None
+            if args.patience > 0 and step % args.val_every == 0:
+                val_loss = compute_val_loss()
+                logger.info(f"step {step:>6d}/{args.max_steps} | loss={total_ce/max(1,total_n):.4f}"
+                            f" | val={val_loss:.4f} | lr={scheduler.get_last_lr()[0]:.2e} | "
+                            f"{tok_per_s:.0f} tok/s | vram={vram:.1f}GB | {eta_str}")
+            else:
+                logger.info(f"step {step:>6d}/{args.max_steps} | loss={total_ce/max(1,total_n):.4f}"
+                            f" | lr={scheduler.get_last_lr()[0]:.2e} | "
+                            f"{tok_per_s:.0f} tok/s | vram={vram:.1f}GB | {eta_str}")
+            stats_file.write(json.dumps({
+                "step": step, "loss": total_ce / max(1, total_n),
+                "lr": float(scheduler.get_last_lr()[0]), "tok_per_s": tok_per_s,
+                "vram_gb": vram, "val_loss": val_loss,
+            }) + "\n")
+            stats_file.flush()
+
+        if step % args.save_every == 0:
+            save_ckpt(step, os.path.join(args.output_dir, f"step_{step}.pt"))
+            cleanup_checkpoints(args.output_dir, args.keep_last_n,
+                                args.keep_best_n, stats_path)
+
+        if args.patience > 0 and step % args.val_every == 0:
+            if val_loss is None:
+                val_loss = compute_val_loss()
+            # Curriculum gate: before the ramp covers --val-t the val estimate
+            # is extrapolation, not generalization. Keep the number for the
+            # curve but do not let it move best_val / patience / valbest.
+            gated = (args.curriculum_early_stop_gate and args.curriculum
+                     and cur_max_ratio is not None and cur_max_ratio < args.val_t)
+            if gated:
+                logger.info(f"  [val] {val_loss:.4f} gated (t_max {cur_max_ratio:.3f} "
+                            f"< val-t {args.val_t}); early stopping inactive")
+            elif val_loss < best_val - args.min_delta:
+                best_val, best_step, no_improve = val_loss, step, 0
+                save_ckpt(step, os.path.join(args.output_dir, f"step_{step}_valbest.pt"))
+                logger.info(f"  [val] new best {val_loss:.4f} -> step_{step}_valbest.pt")
+            else:
+                no_improve += 1
+                logger.info(f"  [val] no improvement ({no_improve}/{args.patience}), "
+                            f"best={best_val:.4f} @ step {best_step}")
+                if no_improve >= args.patience:
+                    logger.info(f"Early stop at step {step}; best val {best_val:.4f} "
+                                f"at step {best_step} (step_{best_step}_valbest.pt)")
+                    save_ckpt(step, os.path.join(args.output_dir, f"step_{step}.pt"))
+                    break
+
+    stats_file.close()
+    if interrupted:
+        logger.info(f"Stopped early (Ctrl+C) at step {step}. "
+                    f"Resume with --resume-from {args.output_dir}/step_{step}.pt")
+    else:
+        logger.info(f"Done. Final checkpoint: {args.output_dir}/step_{step}.pt")
+
+
+if __name__ == "__main__":
+    main()