Upload 12 files

config.json

@@ -0,0 +1,24 @@
+{
+  "hidden_size": 768,
+  "intermediate_size": 2112,
+  "num_hidden_layers": 16,
+  "num_attention_heads": 12,
+  "num_key_value_heads": 6,
+  "head_dim": 64,
+  "vocab_size": 32010,
+  "mask_vocab_size": 32010,
+  "max_position_embeddings": 5120,
+  "rope_theta": 10000.0,
+  "rms_norm_eps": 1e-06,
+  "hidden_act": "silu",
+  "timestep_emb_hidden": 768,
+  "mask_token_id": 32000,
+  "pad_token_id": 1,
+  "mask_ratio_min": 0.0,
+  "mask_ratio_max": 1.0,
+  "tie_word_embeddings": false,
+  "model_type": "metadiffusion",
+  "architectures": [
+    "MetaDiffusionForCausalLM"
+  ]
+}

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "bos_token_id": 0,
+  "eos_token_id": 2,
+  "pad_token_id": 1,
+  "mask_token_id": 32000,
+  "temperature": 0.7,
+  "repetition_penalty": 1.5,
+  "re_mask": 0.1,
+  "num_steps": 128,
+  "max_new_tokens": 96,
+  "use_cache": false,
+  "transformers_version": "4.40.0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d7d0ce3745a3a99ff6863849dff20f2b536f0e2d2092fe263072056715c01d5
+size 339086208

scripts/chat.py

@@ -0,0 +1,284 @@
+#!/usr/bin/env python3
+"""
+chat.py: Multi-turn ChatML chat with MetaDiffusion chat-SFT models.
+
+Works with:
+  - an exported dir (config.json + model.safetensors + tokenizer/)
+  - a training checkpoint (step_*.pt) with --tokenizer
+
+Usage:
+    Interactive:
+        python3 chat.py --model-path MetaDiffusion-150M-ChatBase/
+        python3 chat.py --model-path checkpoints_chat/best.pt \
+                        --tokenizer data/tokenizer
+
+    One-shot:
+        python3 chat.py --model-path MetaDiffusion-150M-ChatBase/ \
+                        --prompt "What is 2+2?" --max-new-tokens 128
+"""
+
+import argparse
+import json
+import math
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from safetensors.torch import load_file
+from transformers import AutoTokenizer
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from model import MetaDiffusionLM, MetaDiffusionConfig  # noqa: E402
+
+MASK_TOKEN_ID = 32000
+CHAT_TOKENS = ["<|im_start|>", "<|im_end|>"] + [f"<|r{i}|>" for i in range(1, 8)]
+IM_START, IM_END = "<|im_start|>", "<|im_end|>"
+
+
+def build_config(config_dict):
+    valid = {k: v for k, v in config_dict.items()
+             if k in MetaDiffusionConfig.__dataclass_fields__}
+    config = MetaDiffusionConfig(**valid)
+    config.tie_word_embeddings = False
+    return config
+
+
+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)
+        sd = load_file(path / "model.safetensors")
+        sd = {k[len("model."):] if k.startswith("model.") else k: v
+              for k, v in sd.items()}
+        missing, unexpected = model.load_state_dict(sd, strict=False)
+        if missing or unexpected:
+            print(f"  Warning: missing={missing[:3]} unexpected={unexpected[:3]}")
+    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)
+    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_chat_tokens(tokenizer):
+    """Add ChatML + rainbow tokens if missing (base tokenizer case).
+
+    Id 32000 is the diffusion [MASK] id, so a reserved filler token takes it
+    first; chat tokens must land at 32001..32009.
+    """
+    if tokenizer.convert_tokens_to_ids(IM_START) == tokenizer.unk_token_id:
+        if len(tokenizer) == 32000:
+            tokenizer.add_special_tokens({"additional_special_tokens": ["<|reserved|>"]})
+        tokenizer.add_special_tokens({"additional_special_tokens": CHAT_TOKENS})
+        assert tokenizer.convert_tokens_to_ids(IM_END) == 32002, \
+            "chat token ids wrong (collide with mask id 32000)"
+    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):
+    """Denoise a block of [MASK] tokens after the prompt (inference.py schedule).
+
+    With stop_on_end=True, denoising halts as soon as <|im_end|> or </s> is
+    committed in the response region: committed tokens never change, so the
+    output is identical, but we skip filling garbage after the terminator.
+    """
+    model.eval()
+    total_len = prompt_ids.shape[1] + gen_len
+    x = torch.full((1, total_len), MASK_TOKEN_ID, device=device, dtype=torch.long)
+    x[0, : prompt_ids.shape[1]] = prompt_ids
+    mask_id = MASK_TOKEN_ID
+    im_end_id = tokenizer.convert_tokens_to_ids(IM_END)
+    eos_id = tokenizer.eos_token_id
+    prompt_len = prompt_ids.shape[1]
+
+    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()
+        n_total = int((frac_next - frac_now) * gen_len + 0.5)
+        if i == num_steps - 1:
+            n_unmask = n_masked
+        else:
+            n_unmask = max(n_total, 1) if n_masked > 0 else 0
+
+        t = 1.0 - frac_now
+        logits = model(x, torch.full((1,), t, device=device))
+
+        # Never predict the mask token
+        logits[:, :, mask_id] = -1e9
+
+        # Repetition penalty over everything already on the sequence
+        if repetition_penalty != 1.0:
+            for tok in x[0].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
+        mask_logits = logits[mask_positions]
+        probs = F.softmax(mask_logits / temperature, dim=-1)
+        sampled = torch.multinomial(probs, 1).squeeze(-1)
+        mask_flat = mask_positions.nonzero(as_tuple=False)
+
+        if n_unmask < mask_positions.sum():
+            # Left-to-right commit: fill the leftmost masked positions first
+            # (semi-autoregressive block generation, as in LLaDA-MoE eval).
+            # Confidence-based commit lets <|im_end|> win the race at ANY
+            # position and commits mid-block tokens before position 0, which
+            # produced empty responses and fragment-style output on this model.
+            fill_positions = mask_flat[:n_unmask]
+            for idx, tok in zip(fill_positions, sampled[:n_unmask]):
+                x[idx[0], idx[1]] = tok
+        else:
+            x[mask_positions] = sampled
+
+        if watch:
+            remaining = (x == mask_id).sum().item()
+            live_toks = [t for t in x[0, prompt_len:].tolist() if t != mask_id]
+            partial = tokenizer.decode(
+                cut_response(live_toks, tokenizer), skip_special_tokens=True
+            ).strip()
+            if len(partial) > 70:
+                partial = partial[:70] + "..."
+            line = (f"step {i+1:3d}/{num_steps} | t={t:.3f} "
+                    f"| masks={remaining:3d} | {partial}")
+            if sys.stdout.isatty():
+                # live in-place update
+                sys.stdout.write("\r" + line[:99].ljust(99))
+                sys.stdout.flush()
+            elif i % max(1, num_steps // 8) == 0:
+                print(line)
+
+        # Stop early once the model commits a terminator in the response
+        if stop_on_end and ((x[0, prompt_len:] == im_end_id).any() or
+                            (x[0, prompt_len:] == eos_id).any()):
+            break
+
+    if watch and sys.stdout.isatty():
+        sys.stdout.write("\n")
+    return x
+
+
+def cut_response(tokens, tokenizer):
+    """Cut generated token list at <|im_end|> or </s>; drop rainbow/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(f"<|r{i}|>") for i in range(1, 8)}
+    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 run_turn(model, tokenizer, messages, args, device):
+    prompt = format_messages(messages) + f"\n{IM_START}assistant\n"
+    prompt_ids = torch.tensor([tokenizer.encode(prompt, add_special_tokens=False)],
+                              device=device)
+    # Retry on empty responses: small models occasionally commit <|im_end|>
+    # as the first token. Bump temperature per attempt
+    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)
+        response_tokens = x[0, prompt_ids.shape[1]:].tolist()
+        response_tokens = cut_response(response_tokens, tokenizer)
+        text = tokenizer.decode(response_tokens, skip_special_tokens=True).strip()
+        if text:
+            return text
+    return "(empty response)"
+
+
+def main():
+    parser = argparse.ArgumentParser(description="MetaDiffusion chat (ChatML)")
+    parser.add_argument("--model-path", required=True,
+                        help="Exported dir (config+safetensors+tokenizer) or step_*.pt")
+    parser.add_argument("--tokenizer", default=None,
+                        help="Tokenizer dir (needed when --model-path is a step_*.pt)")
+    parser.add_argument("--prompt", default=None, help="One-shot prompt (else REPL)")
+    parser.add_argument("--system", default="You are a helpful assistant.",
+                        help="System prompt for the REPL")
+    parser.add_argument("--max-new-tokens", type=int, default=96,
+                        help="Response block size)")
+    parser.add_argument("--num-steps", type=int, default=128, help="Denoising steps")
+    parser.add_argument("--temperature", type=float, default=0.7)
+    parser.add_argument("--repetition-penalty", type=float, default=1.5,
+                        help="Small diffusion models loop without a strong penalty")
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--watch", action="store_true", help="Show denoising progress")
+    args = parser.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)
+
+    model_path = Path(args.model_path)
+    tok_path = args.tokenizer
+    if tok_path is None:
+        if model_path.is_dir():
+            cand = model_path / "tokenizer"
+            if not cand.exists() and (model_path / "tokenizer.json").exists():
+                cand = model_path  # exported dirs keep the tokenizer at root
+            tok_path = str(cand)
+    if tok_path is None or not Path(tok_path).exists():
+        raise Exception("No tokenizer")
+    tokenizer = AutoTokenizer.from_pretrained(str(tok_path))
+    tokenizer = ensure_chat_tokens(tokenizer)
+    print(f"[*] Tokenizer: {tok_path} (vocab {len(tokenizer)})")
+
+    if args.prompt:
+        messages = [{"role": "user", "content": args.prompt}]
+        text = run_turn(model, tokenizer, messages, args, device)
+        print(f"\nUser: {args.prompt}\nAssistant: {text}\n")
+        return
+
+    print("\nMetaDiffusion chat, type 'exit', 'quit' or Ctrl-D 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()

scripts/convert_data.py

@@ -0,0 +1,208 @@
+#!/usr/bin/env python3
+"""
+convert_data.py: turn local ChatML datasets into tokenized train.pt/val.pt
+for train_chat.py.
+
+Accepts .jsonl, .json, and .parquet files (auto-detected by extension), and
+these record shapes (one conversation per record):
+
+    {"messages": [{"role": "user", "content": ...}, ...]}   # HF ChatML style
+    [{"role": ..., "content": ...}, ...]                     # bare message list
+    {"conversation": [...]} or {"chat": [...]}               # aliases
+    {"instruction": ..., "input": ..., "output": ...}        # alpaca style (converted)
+    {"data": [...]} / {"conversations": [...]}               # JSON containers of the above
+
+Parquet rows may store the messages column as a list of dicts or as a JSON
+string (both work).
+
+Usage:
+    python3 convert_data.py --model-path . --input my_data.jsonl \
+        --output data/my_data
+    python3 convert_data.py --model-path . --input a.jsonl b.jsonl \
+        --output data/mixed
+    python3 convert_data.py --model-path . --input ./folder \
+        --output data/folder
+
+Then train:
+    python3 train_chat.py --model-path . --data-dir data/my_data \
+        --output-dir my_checkpoints --lr 7e-5 --epochs 3
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+from transformers import AutoTokenizer
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+from prepare_data import (  # noqa: E402
+    add_chat_tokens,
+    build_conv_segments,
+    save_dataset,
+    tokenize_and_split,
+)
+
+SUPPORTED_EXTS = {".jsonl", ".json", ".parquet"}
+CONTAINER_KEYS = ("conversations", "data", "rows", "examples")
+MESSAGE_KEYS = ("messages", "conversation", "chat")
+
+
+# ---------------------------------------------------------------------------
+# File reading
+# ---------------------------------------------------------------------------
+
+def collect_files(paths):
+    """Expand --input args (files and/or dirs) into a sorted file list."""
+    files = []
+    for p in paths:
+        p = Path(p)
+        if p.is_dir():
+            files.extend(f for f in sorted(p.iterdir())
+                         if f.suffix.lower() in SUPPORTED_EXTS)
+        elif p.suffix.lower() in SUPPORTED_EXTS:
+            files.append(p)
+        else:
+            print(f"[!] Skipping unsupported file: {p} (want .jsonl/.json/.parquet)")
+    return files
+
+
+def iter_records(path):
+    """Yield one raw record (dict or list) per conversation from a file."""
+    ext = path.suffix.lower()
+    if ext == ".jsonl":
+        with open(path) as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                yield json.loads(line)
+    elif ext == ".json":
+        with open(path) as f:
+            obj = json.load(f)
+        if isinstance(obj, list):
+            yield from obj
+        elif isinstance(obj, dict):
+            for key in CONTAINER_KEYS:
+                if isinstance(obj.get(key), list):
+                    yield from obj[key]
+                    return
+            yield obj  # single-conversation file
+    elif ext == ".parquet":
+        import pandas as pd  # lazy: only needed for parquet
+        df = pd.read_parquet(path)
+        for _, row in df.iterrows():
+            yield dict(row)
+    else:
+        raise ValueError(f"Unsupported extension: {path}")
+
+
+def normalize_record(rec):
+    """Turn one record into a list of {role, content} messages, or None."""
+    if isinstance(rec, list):
+        msgs = [m for m in rec
+                if isinstance(m, dict) and m.get("content")]
+        return msgs or None
+
+    if not isinstance(rec, dict):
+        return None
+
+    # ChatML-style keys (value may be a list of dicts, a numpy array of dicts
+    # from parquet round-trips, or a JSON string)
+    for key in MESSAGE_KEYS:
+        v = rec.get(key)
+        if isinstance(v, str):
+            try:
+                v = json.loads(v)
+            except json.JSONDecodeError:
+                continue
+        if isinstance(v, (list, np.ndarray)):
+            msgs = [m for m in v
+                    if isinstance(m, dict) and m.get("content")]
+            if msgs:
+                return msgs
+
+    # Alpaca-style record: instruction / input / output
+    if rec.get("instruction") and rec.get("output"):
+        user = rec["instruction"]
+        if rec.get("input"):
+            user += f"\n\n{rec['input']}"
+        return [{"role": "user", "content": user},
+                {"role": "assistant", "content": rec["output"]}]
+
+    return None
+
+
+def load_convs_from_files(files):
+    """Build (roles, segs) conversations from all files."""
+    convs = []
+    for path in files:
+        n_before = len(convs)
+        for rec in iter_records(path):
+            msgs = normalize_record(rec)
+            if msgs is None:
+                continue
+            segs, roles = build_conv_segments(msgs)
+            if not segs:
+                continue
+            convs.append((roles, segs))
+        print(f"[*] {path.name}: {len(convs) - n_before} conversations")
+    return convs
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert local ChatML data (jsonl/json/parquet) to tokenized .pt")
+    parser.add_argument("--model-path", default=".",
+                        help="Dir with tokenizer.json (the release dir works)")
+    parser.add_argument("--input", nargs="+", required=True,
+                        help="File(s) or dir(s): .jsonl, .json, .parquet")
+    parser.add_argument("--output", default="data/converted",
+                        help="Output dir (train.pt, val.pt, tokenizer/, stats.json)")
+    parser.add_argument("--val-size", type=int, default=500, help="Held-out examples")
+    parser.add_argument("--max-len", type=int, default=1024, help="Max tokens per example")
+    parser.add_argument("--max-resp-tokens", type=int, default=256,
+                        help="Cap on target response tokens (keeps <|im_end|>)")
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+
+    out = Path(args.output)
+    out.mkdir(parents=True, exist_ok=True)
+
+    print(f"[*] Loading tokenizer from {args.model_path}")
+    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
+    add_chat_tokens(tokenizer)
+    tokenizer.save_pretrained(out / "tokenizer")
+
+    files = collect_files(args.input)
+    if not files:
+        print("[!] No .jsonl/.json/.parquet files found in the inputs.")
+        sys.exit(1)
+    print(f"[*] Files: {', '.join(f.name for f in files)}")
+
+    convs = load_convs_from_files(files)
+    if not convs:
+        print("[!] No conversations parsed (check the record shapes in the docstring).")
+        sys.exit(1)
+    print(f"[*] Total: {len(convs)} conversations")
+
+    # Auto-scale the val split: never take everything for small datasets
+    val_size = min(args.val_size, max(1, len(convs) // 10))
+    if val_size != args.val_size:
+        print(f"[*] Small dataset: using val_size={val_size}")
+
+    train_examples, val_examples, skipped, dupes = tokenize_and_split(
+        convs, set(), tokenizer, val_size, args.max_len,
+        args.max_resp_tokens, args.seed)
+
+    save_dataset(out, train_examples, val_examples, tokenizer, skipped, dupes)
+
+
+if __name__ == "__main__":
+    main()

scripts/export_hf.py

@@ -0,0 +1,181 @@
+#!/usr/bin/env python3
+"""
+export_hf.py: Export a chat-SFT checkpoint to a HuggingFace-style dir.
+
+Output dir contains:
+    model.safetensors       (weights, "model."-prefixed keys)
+    config.json             (architecture, mask_vocab_size = 32010)
+    generation_config.json  (sampling defaults for chat)
+    tokenizer.json / tokenizer_config.json   (with ChatML + rainbow tokens)
+    README.md               (minimal)
+
+Usage:
+    python3 export_hf.py \
+        --checkpoint checkpoints_chat/step_3000.pt \
+        --tokenizer data/no_robots_chatml/tokenizer \
+        --output MetaDiffusion-150M-Chat/
+"""
+
+import argparse
+import json
+import shutil
+import sys
+from dataclasses import asdict
+from pathlib import Path
+
+import torch
+from safetensors.torch import save_file
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from model import MetaDiffusionConfig  # noqa: E402
+
+KEY_MAP = {
+    "embed_tokens.weight": "model.embed_tokens.weight",
+    "norm.weight": "model.norm.weight",
+    "lm_head.weight": "model.lm_head.weight",
+}
+
+LAYER_KEY_MAP = {
+    "input_layernorm.weight": "input_layernorm.weight",
+    "self_attn.q_proj.weight": "self_attn.q_proj.weight",
+    "self_attn.k_proj.weight": "self_attn.k_proj.weight",
+    "self_attn.v_proj.weight": "self_attn.v_proj.weight",
+    "self_attn.o_proj.weight": "self_attn.o_proj.weight",
+    "post_attention_layernorm.weight": "post_attention_layernorm.weight",
+    "mlp.gate_proj.weight": "mlp.gate_proj.weight",
+    "mlp.up_proj.weight": "mlp.up_proj.weight",
+    "mlp.down_proj.weight": "mlp.down_proj.weight",
+    "timestep_residual.proj.weight": "timestep_residual.proj.weight",
+    "timestep_residual.proj.bias": "timestep_residual.proj.bias",
+}
+
+TIMESTEP_KEY_MAP = {
+    "timestep_emb.mlp.0.weight": "model.timestep_emb.mlp.0.weight",
+    "timestep_emb.mlp.0.bias": "model.timestep_emb.mlp.0.bias",
+    "timestep_emb.mlp.2.weight": "model.timestep_emb.mlp.2.weight",
+    "timestep_emb.mlp.2.bias": "model.timestep_emb.mlp.2.bias",
+}
+
+GENERATION_CONFIG = {
+    "bos_token_id": 0,
+    "eos_token_id": 2,
+    "pad_token_id": 1,
+    "mask_token_id": 32000,
+    "temperature": 0.7,
+    "repetition_penalty": 1.5,
+    "re_mask": 0.1,
+    "num_steps": 128,
+    "max_new_tokens": 96,
+    "use_cache": False,
+    "transformers_version": "4.40.0",
+}
+
+
+def remap_state_dict(state_dict, half=False):
+    # Strip torch.compile's _orig_mod. prefix (old checkpoints saved from a
+    # compiled model have it; training now saves clean keys)
+    state_dict = {k.replace("_orig_mod.", "", 1) if k.startswith("_orig_mod.") else k: v
+                  for k, v in state_dict.items()}
+    new_dict = {}
+    for key, tensor in state_dict.items():
+        if key in KEY_MAP:
+            new_key = KEY_MAP[key]
+        elif key in TIMESTEP_KEY_MAP:
+            new_key = TIMESTEP_KEY_MAP[key]
+        elif key.startswith("layers."):
+            parts = key.split(".")
+            # layers.N.<component>.<rest>
+            layer_idx, component = parts[1], parts[2]
+            rest = ".".join(parts[3:])
+            comp_key = f"{component}.{rest}" if rest else component
+            new_key = f"model.layers.{layer_idx}.{comp_key}"
+        else:
+            new_key = key
+        new_dict[new_key] = tensor.half() if half else tensor.float()
+    return new_dict
+
+
+def package_scripts(out):
+    """Copy the runnable pipeline into out/scripts/ so the release is
+    self-contained: chat, finetune, and re-export work from the artifact."""
+    src = Path(__file__).resolve().parent
+    root = src.parent  # model.py lives in the project root
+    scripts_dir = out / "scripts"
+    scripts_dir.mkdir(parents=True, exist_ok=True)
+    for name in ["model.py", "chat.py", "prepare_data.py", "train_chat.py",
+                 "export_hf.py", "convert_data.py"]:
+        cand = src / name if (src / name).exists() else root / 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.40\nsafetensors>=0.4\n"
+                       "datasets>=2.18\nnumpy>=1.26\n"
+                       "pandas>=2.0\npyarrow>=14.0\n")
+    print(f"[*] Packaged scripts -> {scripts_dir}")
+
+
+def export(checkpoint_path, tokenizer_dir, output_dir, half=False):
+    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__}
+    )
+    config.tie_word_embeddings = False
+
+    print("[*] Remapping state dict...")
+    state_dict = remap_state_dict(ckpt["model_state_dict"], half=half)
+    save_file(state_dict, out / "model.safetensors")
+    print(f"[*] Saved {len(state_dict)} tensors -> {out / 'model.safetensors'} "
+          f"({'fp16' if half else 'fp32'})")
+
+    # Keep vocab fields consistent with the actual weights: the chat pipeline
+    # resizes embed_tokens/lm_head to 32010 (Supra 32000 + chat tokens), and
+    # the ckpt config still carries vocab_size=32000 from the base. Any loader
+    # using vocab_size to size embeddings would fail with a size mismatch.
+    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)")
+
+    config_dict = asdict(config)
+    config_dict.pop("dtype", None)  # transformers chokes on "torch.float32" strings
+    config_dict["model_type"] = "metadiffusion"
+    config_dict["architectures"] = ["MetaDiffusionForCausalLM"]
+    with open(out / "config.json", "w") as f:
+        json.dump(config_dict, f, indent=2)
+
+    with open(out / "generation_config.json", "w") as f:
+        json.dump(GENERATION_CONFIG, f, indent=2)
+
+    # Copy tokenizer (has ChatML + rainbow tokens)
+    tok_dir = Path(tokenizer_dir)
+    for name in ["tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"]:
+        src = tok_dir / name
+        if src.exists():
+            shutil.copy2(src, out / name)
+
+    package_scripts(out)
+    print(f"[*] Exported to {out}")
+    print("    Vocab:", len(state_dict.get("model.lm_head.weight", [])),
+          "| step:", ckpt.get("step"))
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Export chat-SFT checkpoint to HF-style dir")
+    parser.add_argument("--checkpoint", required=True, help="step_*.pt or best.pt")
+    parser.add_argument("--tokenizer", default="data/no_robots_chatml/tokenizer")
+    parser.add_argument("--output", required=True, help="Output dir")
+    parser.add_argument("--fp16", action="store_true",
+                        help="Save weights as fp16 (half the size; matches the "
+                             "base release format)")
+    args = parser.parse_args()
+    export(args.checkpoint, args.tokenizer, args.output, half=args.fp16)
+
+
+if __name__ == "__main__":
+    main()

scripts/model.py

@@ -0,0 +1,526 @@
+"""
+MetaDiffusion: Convert AR LLMs to Masked Diffusion LLMs
+Based on Supra-1.5-50M-Base-exp architecture
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass, field
+from typing import Optional, Tuple
+
+
+@dataclass
+class MetaDiffusionConfig:
+    # Architecture (matching Supra-1.5-50M)
+    hidden_size: int = 512
+    intermediate_size: int = 1408
+    num_hidden_layers: int = 12
+    num_attention_heads: int = 8
+    num_key_value_heads: int = 4
+    head_dim: int = 64
+    vocab_size: int = 32000          # original vocab
+    mask_vocab_size: int = 32001     # vocab + [MASK] token
+    max_position_embeddings: int = 5120
+    rope_theta: float = 10000.0
+    rms_norm_eps: float = 1e-6
+    hidden_act: str = "silu"
+
+    # Diffusion-specific
+    timestep_emb_hidden: int = 512
+    mask_token_id: int = 32000       # index of [MASK] in embedding table
+    pad_token_id: int = 1            # Supra pad token
+
+    # Masking strategy
+    mask_ratio_min: float = 0.0
+    mask_ratio_max: float = 1.0
+
+    # Training
+    dtype: torch.dtype = torch.float32
+    tie_word_embeddings: bool = True
+
+class RotaryEmbedding(nn.Module):
+    """RoPE - position embeddings for the attention layers."""
+
+    def __init__(self, dim, max_position_embeddings=5120, base=10000.0, device=None):
+        super().__init__()
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device).float() / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        # x: (batch, seq, hidden) - used for dtype/device only
+        # position_ids: (batch, seq)
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(
+            position_ids.shape[0], -1, 1
+        )
+        position_ids_expanded = position_ids[:, None, :].float()
+        freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos = emb.cos()
+        sin = emb.sin()
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin):
+    cos = cos.unsqueeze(1)  # (batch, 1, seq, dim)
+    sin = sin.unsqueeze(1)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class TimestepEmbedding(nn.Module):
+    """Sinusoidal timestep embedding with learned projection."""
+
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.mlp = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size * 4),
+            nn.SiLU(),
+            nn.Linear(hidden_size * 4, hidden_size),
+        )
+
+    def forward(self, t):
+        # t: (batch,) timesteps in [0, 1]
+        half_dim = self.hidden_size // 2
+        emb = math.log(10000.0) / (half_dim - 1)
+        emb = torch.exp(
+            torch.arange(half_dim, device=t.device, dtype=torch.float32) * -emb
+        )
+        emb = t[:, None].float() * emb[None, :]
+        emb = torch.cat([emb.sin(), emb.cos()], dim=-1)  # (batch, hidden_size)
+        return self.mlp(emb).to(t.dtype)  # (batch, hidden_size)
+
+
+class TimestepResidual(nn.Module):
+    """Add timestep embedding to hidden state at each block.
+    Initialized to zero so it starts as identity (no disruption to pretrained weights)."""
+
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.proj = nn.Linear(hidden_size, hidden_size)
+        nn.init.zeros_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+
+    def forward(self, x, emb):
+        # x: (batch, seq, hidden)
+        # emb: (batch, hidden)
+        return x + self.proj(emb)[:, None, :]
+
+
+class RMSNorm(nn.Module):
+    """Llama-style RMSNorm."""
+
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x):
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps)
+        return self.weight * x
+
+
+class SelfAttention(nn.Module):
+    """Multi-head attention with GQA and RoPE.
+    BIDIRECTIONAL (no causal mask) - this is the key difference from AR."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, self.num_heads * config.head_dim, bias=False
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, self.num_kv_heads * config.head_dim, bias=False
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, self.num_kv_heads * config.head_dim, bias=False
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * config.head_dim, config.hidden_size, bias=False
+        )
+        self.rotary_emb = RotaryEmbedding(
+            config.head_dim,
+            max_position_embeddings=config.max_position_embeddings,
+            base=config.rope_theta,
+        )
+
+    def forward(self, x, attention_mask=None, position_ids=None):
+        batch, seq, _ = x.shape
+
+        q = self.q_proj(x).view(batch, seq, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(batch, seq, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(batch, seq, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        cos, sin = self.rotary_emb(x, position_ids)
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        # GQA: repeat KV heads to match Q heads
+        if self.num_kv_groups > 1:
+            k = k.repeat_interleave(self.num_kv_groups, dim=1)
+            v = v.repeat_interleave(self.num_kv_groups, dim=1)
+
+        # Bidirectional attention - no causal mask!
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask)
+        out = out.transpose(1, 2).contiguous().view(batch, seq, -1)
+        return self.o_proj(out)
+
+
+class MLP(nn.Module):
+    """Llama-style gated FFN (SwiGLU)."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.gate_proj = nn.Linear(
+            config.hidden_size, config.intermediate_size, bias=False
+        )
+        self.up_proj = nn.Linear(
+            config.hidden_size, config.intermediate_size, bias=False
+        )
+        self.down_proj = nn.Linear(
+            config.intermediate_size, config.hidden_size, bias=False
+        )
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TransformerBlock(nn.Module):
+    """Llama transformer block adapted for diffusion.
+    - Pre-norm architecture
+    - Bidirectional attention
+    - Timestep conditioning via residual addition
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.self_attn = SelfAttention(config)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.mlp = MLP(config)
+        self.timestep_residual = TimestepResidual(config.hidden_size)
+
+    def forward(self, x, timestep_emb, attention_mask=None, position_ids=None):
+        # Pre-norm + attention + residual + timestep
+        residual = x
+        x = self.input_layernorm(x)
+        x = self.self_attn(x, attention_mask, position_ids)
+        x = residual + x
+        x = self.timestep_residual(x, timestep_emb)
+
+        # Pre-norm + FFN + residual + timestep
+        residual = x
+        x = self.post_attention_layernorm(x)
+        x = self.mlp(x)
+        x = residual + x
+        x = self.timestep_residual(x, timestep_emb)
+
+        return x
+
+class MetaDiffusionLM(nn.Module):
+    """Masked Diffusion Language Model.
+
+    Converts an AR Llama-style model to a masked-diffusion LM.
+    Transfers: embeddings, all transformer blocks, RoPE, norms.
+    New: timestep embedding, [MASK] token.
+    Output head is tied with embeddings by default (set tie_word_embeddings=False to untie).
+    """
+
+    def __init__(self, config: MetaDiffusionConfig):
+        super().__init__()
+        self.config = config
+
+        # Embeddings (vocab + 1 for [MASK])
+        self.embed_tokens = nn.Embedding(
+            config.mask_vocab_size, config.hidden_size, padding_idx=config.pad_token_id
+        )
+
+        # Timestep conditioning
+        self.timestep_emb = TimestepEmbedding(config.timestep_emb_hidden)
+
+        # Transformer stack
+        self.layers = nn.ModuleList(
+            [TransformerBlock(config) for _ in range(config.num_hidden_layers)]
+        )
+
+        # Final norm
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # Output projection (can be tied with embeddings for parameter efficiency)
+        if config.tie_word_embeddings:
+            self.lm_head = None  # Will use embed_tokens.weight in forward
+        else:
+            self.lm_head = nn.Linear(
+                config.hidden_size, config.mask_vocab_size, bias=False
+            )
+
+        self.post_init()
+
+    def post_init(self):
+        if self.lm_head is not None:
+            nn.init.normal_(self.lm_head.weight, std=0.02)
+
+    def forward(self, input_ids, timesteps, attention_mask=None):
+        """
+        Forward pass for training.
+
+        Args:
+            input_ids: (batch, seq) - tokens with masked positions replaced by mask_token_id
+            timesteps: (batch,) - diffusion timestep in [0, 1]
+            attention_mask: optional (batch, seq) - 1 for real tokens, 0 for padding
+
+        Returns:
+            logits: (batch, seq, mask_vocab_size)
+        """
+        batch, seq = input_ids.shape
+
+        # Position IDs (0, 1, 2, ...)
+        position_ids = (
+            torch.arange(seq, device=input_ids.device).unsqueeze(0).expand(batch, -1)
+        )
+
+        # Embed tokens
+        x = self.embed_tokens(input_ids)
+
+        # Get timestep embedding
+        t_emb = self.timestep_emb(timesteps)
+
+        # Convert attention mask for SDPA (0 -> keep, -inf -> mask out)
+        attn_mask = None
+        if attention_mask is not None:
+            attn_mask = ((1.0 - attention_mask[:, None, None, :].float()) * -1e9).to(
+                x.dtype
+            )
+
+        # Pass through transformer blocks
+        for layer in self.layers:
+            x = layer(x, t_emb, attn_mask, position_ids)
+
+        # Final norm + project to vocab
+        x = self.norm(x)
+        if self.lm_head is not None:
+            logits = self.lm_head(x)
+        else:
+            # Tied embeddings: use embed_tokens.weight transposed
+            logits = F.linear(x, self.embed_tokens.weight)
+
+        return logits
+
+    def compute_loss(self, logits, labels, mask_positions, pad_token_id=None):
+        """
+        Compute cross-entropy loss on masked positions only.
+
+        Args:
+            logits: (batch, seq, vocab_size)
+            labels: (batch, seq) - original token IDs (before masking)
+            mask_positions: (batch, seq) - bool tensor, True where token was masked
+            pad_token_id: ignore these positions in loss
+
+        Returns:
+            loss: scalar
+            num_masked: number of positions in loss
+        """
+        logits_masked = logits[mask_positions]
+        labels_masked = labels[mask_positions]
+
+        # Filter out padding tokens
+        if pad_token_id is not None:
+            valid = labels_masked != pad_token_id
+            logits_masked = logits_masked[valid]
+            labels_masked = labels_masked[valid]
+
+        if labels_masked.numel() == 0:
+            return torch.tensor(0.0, device=logits.device), 0
+
+        loss = F.cross_entropy(logits_masked, labels_masked)
+        return loss, labels_masked.numel()
+
+    @torch.no_grad()
+    def generate(self, batch_size, seq_len, num_steps=256, device="cuda"):
+        """
+        Iterative denoising generation (LLaDA-style).
+
+        Starts from all-mask tokens and progressively unmaskes the most confident predictions.
+
+        Args:
+            batch_size: number of sequences to generate
+            seq_len: length of each sequence
+            num_steps: number of denoising iterations
+
+        Returns:
+            tokens: (batch, seq) - generated token IDs
+        """
+        mask_token_id = self.config.mask_token_id
+        x = torch.full(
+            (batch_size, seq_len), mask_token_id, device=device, dtype=torch.long
+        )
+
+        # Linear schedule from t=1 (all mask) to t=0 (no mask)
+        timesteps = torch.linspace(1.0, 0.0, num_steps + 1, device=device)
+
+        for i in range(num_steps):
+            t = timesteps[i]
+            t_next = timesteps[i + 1]
+            t_batch = torch.full((batch_size,), t, device=device)
+
+            # Get predictions
+            logits = self.forward(x, t_batch)
+            pred_tokens = logits.argmax(dim=-1)
+
+            # Confidence of predicted tokens
+            probs = F.softmax(logits, dim=-1)
+            confidence = probs.gather(-1, pred_tokens.unsqueeze(-1)).squeeze(-1)
+
+            # Number of tokens to unmask this step
+            num_unmask = max(1, int(seq_len * (t - t_next)))
+
+            # Only consider currently-masked positions
+            is_mask = x == mask_token_id
+            confidence_masked = confidence.clone()
+            confidence_masked[~is_mask] = -1.0
+
+            # Unmask the most confident predictions
+            _, top_indices = confidence_masked.topk(num_unmask, dim=-1)
+            batch_idx = (
+                torch.arange(batch_size, device=device).unsqueeze(-1).expand_as(top_indices)
+            )
+            x[batch_idx, top_indices] = pred_tokens[batch_idx, top_indices]
+
+        return x
+
+    @classmethod
+    def from_pretrained_ar(
+        cls,
+        model_name_or_path: str,
+        **kwargs,
+    ):
+        """
+        Initialize a MetaDiffusionLM from a pretrained AR Llama model.
+
+        Transfers all AR weights and initializes new diffusion components.
+        """
+        from transformers import LlamaForCausalLM
+
+        print(f"Loading AR model: {model_name_or_path}")
+        ar_model = LlamaForCausalLM.from_pretrained(model_name_or_path)
+        ar_config = ar_model.config
+
+        # Build diffusion config from AR config
+        head_dim = getattr(
+            ar_config,
+            "head_dim",
+            ar_config.hidden_size // ar_config.num_attention_heads,
+        )
+        config = MetaDiffusionConfig(
+            hidden_size=ar_config.hidden_size,
+            intermediate_size=ar_config.intermediate_size,
+            num_hidden_layers=ar_config.num_hidden_layers,
+            num_attention_heads=ar_config.num_attention_heads,
+            num_key_value_heads=ar_config.num_key_value_heads,
+            head_dim=head_dim,
+            vocab_size=ar_config.vocab_size,
+            mask_vocab_size=ar_config.vocab_size + 1,
+            max_position_embeddings=ar_config.max_position_embeddings,
+            rope_theta=getattr(ar_config, "rope_theta", 10000.0),
+            rms_norm_eps=ar_config.rms_norm_eps,
+            mask_token_id=ar_config.vocab_size,
+            pad_token_id=getattr(ar_config, "pad_token_id", 1),
+            **kwargs,
+        )
+
+        model = cls(config)
+
+        # Build state dict mapping
+        state_dict = ar_model.state_dict()
+        new_state_dict = {}
+
+        # --- Embeddings ---
+        # Copy original vocab embeddings
+        ar_embeds = state_dict["model.embed_tokens.weight"]
+        mask_embed = ar_embeds.mean(dim=0, keepdim=True)  # [MASK] = mean of all embeds
+        new_state_dict["embed_tokens.weight"] = torch.cat(
+            [ar_embeds, mask_embed], dim=0
+        )
+
+        # --- Transformer blocks ---
+        for i in range(config.num_hidden_layers):
+            ar_prefix = f"model.layers.{i}"
+            new_prefix = f"layers.{i}"
+
+            # Attention
+            new_state_dict[f"{new_prefix}.self_attn.q_proj.weight"] = state_dict[
+                f"{ar_prefix}.self_attn.q_proj.weight"
+            ]
+            new_state_dict[f"{new_prefix}.self_attn.k_proj.weight"] = state_dict[
+                f"{ar_prefix}.self_attn.k_proj.weight"
+            ]
+            new_state_dict[f"{new_prefix}.self_attn.v_proj.weight"] = state_dict[
+                f"{ar_prefix}.self_attn.v_proj.weight"
+            ]
+            new_state_dict[f"{new_prefix}.self_attn.o_proj.weight"] = state_dict[
+                f"{ar_prefix}.self_attn.o_proj.weight"
+            ]
+
+            # MLP
+            new_state_dict[f"{new_prefix}.mlp.gate_proj.weight"] = state_dict[
+                f"{ar_prefix}.mlp.gate_proj.weight"
+            ]
+            new_state_dict[f"{new_prefix}.mlp.up_proj.weight"] = state_dict[
+                f"{ar_prefix}.mlp.up_proj.weight"
+            ]
+            new_state_dict[f"{new_prefix}.mlp.down_proj.weight"] = state_dict[
+                f"{ar_prefix}.mlp.down_proj.weight"
+            ]
+
+            # Layer norms
+            new_state_dict[f"{new_prefix}.input_layernorm.weight"] = state_dict[
+                f"{ar_prefix}.input_layernorm.weight"
+            ]
+            new_state_dict[
+                f"{new_prefix}.post_attention_layernorm.weight"
+            ] = state_dict[f"{ar_prefix}.post_attention_layernorm.weight"]
+
+        # --- Final norm ---
+        new_state_dict["norm.weight"] = state_dict["model.norm.weight"]
+
+        # --- Output head ---
+        # Initialize from AR embeddings (since AR used tied embeddings, E^T was the output)
+        if not config.tie_word_embeddings:
+            new_state_dict["lm_head.weight"] = torch.cat(
+                [ar_embeds, torch.zeros(1, config.hidden_size, device=ar_embeds.device)],
+                dim=0,
+            ).clone()
+
+        # Load with strict=False for new diffusion params
+        missing, unexpected = model.load_state_dict(new_state_dict, strict=False)
+
+        print(f"Weights transferred from AR model")
+        print(f"  Missing (new diffusion params): {len(missing)}")
+        print(f"  Unexpected: {len(unexpected)}")
+        if missing:
+            for k in missing:
+                print(f"    NEW: {k}")
+
+        return model

scripts/prepare_data.py

@@ -0,0 +1,411 @@
+#!/usr/bin/env python3
+"""
+prepare_data.py: Download HuggingFaceH4/no_robots and pre-tokenize it as ChatML
+for MetaDiffusion chat SFT.
+
+Output (in --data-dir):
+    train.pt              list of {"input_ids": LongTensor, "assistant_start": int,
+                                   "assistant_end": int}
+    val.pt                same, held-out
+    tokenizer/            Supra tokenizer with ChatML + rainbow tokens added
+    stats.json            counts and length stats
+
+The assistant region (content + trailing <|im_end|>) is the only part that
+will be masked during training; everything before it is visible prompt.
+
+Usage:
+    python3 prepare_data.py --model-path ../hf_release --data-dir data/no_robots_chatml
+"""
+
+import argparse
+import hashlib
+import json
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+IM_START = "<|im_start|>"
+IM_END = "<|im_end|>"
+RESERVED = "<|reserved|>"  # filler so id 32000 stays free for [MASK]
+CHAT_TOKENS = [IM_START, IM_END] + [f"<|r{i}|>" for i in range(1, 8)]  # rainbow pads
+BASE_VOCAB = 32000  # Supra tokenizer entry count (ids 0..31999)
+
+
+def add_chat_tokens(tokenizer):
+    """Add chat + rainbow tokens at ids 32001..32009.
+
+    The base tokenizer has 32000 entries, so the first added token would take
+    id 32000, which is the diffusion [MASK] id. A reserved filler token takes
+    that slot first; mask id 32000 must never exist in the tokenizer.
+    """
+    if len(tokenizer) == BASE_VOCAB:
+        tokenizer.add_special_tokens({"additional_special_tokens": [RESERVED]})
+    n = tokenizer.add_special_tokens(
+        {"additional_special_tokens": CHAT_TOKENS}
+    )
+    im_start = tokenizer.convert_tokens_to_ids(IM_START)
+    im_end = tokenizer.convert_tokens_to_ids(IM_END)
+    assert im_start == 32001, f"im_start id {im_start} != 32001 (collides with [MASK])"
+    assert im_end == 32002, f"im_end id {im_end} != 32002"
+    print(f"[*] Added {n} special tokens, vocab now {len(tokenizer)}")
+    print(f"[*] im_start={im_start} im_end={im_end} "
+          f"rainbow={[tokenizer.convert_tokens_to_ids(f'<|r{i}|>') for i in range(1, 8)]}")
+    return n
+
+
+def format_segment(role: str, content: str) -> str:
+    return f"{IM_START}{role}\n{content}{IM_END}"
+
+
+def load_messages(dataset_name: str):
+    """Return a list of message lists for one dataset name."""
+    if dataset_name == "no_robots":
+        ds = load_dataset("HuggingFaceH4/no_robots", split="train")
+        return [list(row["messages"]) for row in ds]
+    if dataset_name == "alpaca":
+        ds = load_dataset("yahma/alpaca-cleaned", split="train")
+        out = []
+        for row in ds:
+            user = row["instruction"]
+            if row.get("input"):
+                user += f"\n\n{row['input']}"
+            out.append([{"role": "user", "content": user},
+                        {"role": "assistant", "content": row["output"]}])
+        return out
+    if dataset_name == "dolly":
+        ds = load_dataset("databricks/databricks-dolly-15k", split="train")
+        out = []
+        for row in ds:
+            user = row["instruction"]
+            if row.get("context"):
+                user += f"\n\n{row['context']}"
+            out.append([{"role": "user", "content": user},
+                        {"role": "assistant", "content": row["response"]}])
+        return out
+    if dataset_name == "smol-smoltalk":
+        # The actual SFT set for SmolLM2-135M-Instruct: 484K short, high-quality
+        # conversations designed for <1B models (no function calling, no
+        # advanced math). Cap with --max-examples (e.g. 60000) for a 150M model.
+        ds = load_dataset("HuggingFaceTB/smol-smoltalk", split="train")
+        return [list(row["messages"]) for row in ds]
+    if dataset_name == "math":
+        return load_math_data()
+    raise ValueError(f"Unknown dataset: {dataset_name} "
+                     f"(choose from: no_robots, alpaca, dolly, smol-smoltalk, math)")
+
+MATH_QUESTION_TEMPLATES = {
+    "add": ["What is {a} + {b}?", "What is {a} plus {b}?", "Add {a} and {b}.",
+            "What does {a} + {b} equal?"],
+    "sub": ["What is {a} - {b}?", "What is {a} minus {b}?", "Subtract {b} from {a}.",
+            "What does {a} - {b} equal?"],
+    "mul": ["What is {a} × {b}?", "What is {a} times {b}?", "Multiply {a} by {b}.",
+            "What does {a} × {b} equal?"],
+    "div": ["What is {a} ÷ {b}?", "What is {a} divided by {b}?", "Divide {a} by {b}.",
+            "What does {a} ÷ {b} equal?"],
+}
+MATH_ANSWER_TEMPLATES = ["The answer is {r}.", "It is {r}.", "{r}"]
+
+
+def load_math_data(seed: int = 42):
+    """Exhaustive basic-arithmetic QA pairs (add/sub/mul/div), ChatML messages.
+
+    A 150M model learns arithmetic by memorization, so cover EVERY pair in a
+    small range rather than sampling: add a<=b in 1..99, sub b<a in 1..99,
+    times tables 1..12, exact divisions 1..12.
+    """
+    rng = random.Random(seed)
+    pairs = []  # (op, a, b, result)
+    for a in range(1, 100):
+        for b in range(a, 100):
+            pairs.append(("add", a, b, a + b))
+    for a in range(2, 100):
+        for b in range(1, a):
+            pairs.append(("sub", a, b, a - b))
+    for a in range(1, 13):
+        for b in range(1, 13):
+            pairs.append(("mul", a, b, a * b))
+    for b in range(1, 13):
+        for q in range(1, 13):
+            pairs.append(("div", b * q, b, q))
+    rng.shuffle(pairs)
+
+    out = []
+    for op, a, b, r in pairs:
+        question = rng.choice(MATH_QUESTION_TEMPLATES[op]).format(a=a, b=b)
+        answer = rng.choice(MATH_ANSWER_TEMPLATES).format(r=r)
+        out.append([{"role": "user", "content": question},
+                    {"role": "assistant", "content": answer}])
+    print(f"    (synthetic arithmetic: {len(out)} pairs, "
+          f"add {sum(1 for p in pairs if p[0]=='add')}, "
+          f"sub {sum(1 for p in pairs if p[0]=='sub')}, "
+          f"mul {sum(1 for p in pairs if p[0]=='mul')}, "
+          f"div {sum(1 for p in pairs if p[0]=='div')})")
+    return out
+
+
+def example_key(input_ids_tensor):
+    """Compact dedup key: SHA-256 of token ids.
+
+    Tuple keys for 484K x 600-token conversations cost ~5 GB of set memory;
+    digests cost ~50 MB. (Hashing was never the bottleneck; per-segment
+    tokenizer.encode() calls were.)
+    """
+    return hashlib.sha256(
+        input_ids_tensor.numpy().astype(np.uint32).tobytes()
+    ).digest()
+
+
+def build_conv_segments(messages):
+    """Split a conversation into ChatML segment strings + roles."""
+    segs, roles = [], []
+    for m in messages:
+        role = m.get("role", "user")
+        content = m.get("content", "")
+        if not content.strip():
+            continue
+        segs.append(format_segment(role, content))
+        roles.append(role)
+    return segs, roles
+
+
+def reconstruct_example(seg_ids, roles, tokenizer, max_len, max_resp):
+    """Rebuild a conversation from pre-encoded segments (tokenize_example
+    logic, but with tokenization already done in batch)."""
+    ids = []
+    assistant_ranges = []
+    for role, seg in zip(roles, seg_ids):
+        start = len(ids)
+        ids.extend(seg)
+        if role == "assistant":
+            assistant_ranges.append((start, len(ids)))
+
+    if not ids or not assistant_ranges:
+        return None
+    a0, a1 = assistant_ranges[-1]
+    if a1 <= a0:
+        return None
+
+    # Cap the response length, always keeping the trailing <|im_end|>
+    if a1 - a0 > max_resp:
+        im_end_id = tokenizer.convert_tokens_to_ids(IM_END)
+        ids = ids[:a0] + ids[a0:a0 + max_resp - 1] + [im_end_id]
+        a1 = a0 + max_resp
+
+    # Truncate from the front (history), keep the target response intact
+    resp = ids[a0:a1]
+    if len(resp) > max_len:
+        resp = resp[:max_len]
+        a1 = a0 + len(resp)
+    hist = ids[:a0]
+    room = max_len - len(resp)
+    if len(hist) > room:
+        hist = hist[len(hist) - room:] if room > 0 else []
+    ids = hist + resp
+    return {
+        "input_ids": torch.tensor(ids, dtype=torch.long),
+        "assistant_start": len(hist),
+        "assistant_end": len(ids),
+    }
+
+
+def tokenize_and_split(convs, math_indices, tokenizer, val_size, max_len,
+                       max_resp, seed=42):
+    """Batched tokenize + reconstruct + dedup + val split.
+
+    convs: list of (roles, segs) from build_conv_segments.
+    math_indices: conversation indices exempt from dedup (math repeats are
+    intentional). Returns (train_examples, val_examples, skipped, dupes).
+    """
+    all_seg_strs = [s for _, segs in convs for s in segs]
+    print(f"[*] Encoding {len(all_seg_strs)} segments (batched)...")
+    encoded = []
+    CHUNK = 200_000
+    for i in range(0, len(all_seg_strs), CHUNK):
+        chunk = all_seg_strs[i:i + CHUNK]
+        encoded.extend(tokenizer(chunk, add_special_tokens=False)["input_ids"])
+
+    rng = random.Random(seed)
+    idxs = list(range(len(convs)))
+    rng.shuffle(idxs)
+    val_idxs = set(idxs[:val_size])
+
+    train_examples, val_examples = [], []
+    skipped = dupes = 0
+    seen = set()
+    ptr = 0
+    for i, (roles, segs) in enumerate(convs):
+        seg_ids = encoded[ptr:ptr + len(segs)]
+        ptr += len(segs)
+        ex = reconstruct_example(seg_ids, roles, tokenizer, max_len, max_resp)
+        if ex is None:
+            skipped += 1
+            continue
+        if i not in math_indices:
+            key = example_key(ex["input_ids"])
+            if key in seen:
+                dupes += 1
+                continue
+            seen.add(key)
+        (val_examples if i in val_idxs else train_examples).append(ex)
+    return train_examples, val_examples, skipped, dupes
+
+
+def save_dataset(out_dir, train_examples, val_examples, tokenizer,
+                 skipped, dupes):
+    """Save train.pt / val.pt / stats.json and print the summary."""
+    out = Path(out_dir)
+    torch.save({"examples": train_examples}, out / "train.pt")
+    torch.save({"examples": val_examples}, out / "val.pt")
+
+    lens = [e["input_ids"].numel() for e in train_examples]
+    stats = {
+        "train": len(train_examples),
+        "val": len(val_examples),
+        "skipped": skipped,
+        "dupes": dupes,
+        "avg_tokens": sum(lens) / len(lens) if lens else 0,
+        "max_tokens": max(lens) if lens else 0,
+        "vocab_size": len(tokenizer),
+        "im_start_id": tokenizer.convert_tokens_to_ids(IM_START),
+        "im_end_id": tokenizer.convert_tokens_to_ids(IM_END),
+        "rainbow_ids": [tokenizer.convert_tokens_to_ids(f"<|r{i}|>") for i in range(1, 8)],
+        "mask_token_id": 32000,
+    }
+    with open(out / "stats.json", "w") as f:
+        json.dump(stats, f, indent=2)
+
+    print(f"[*] Done: {stats['train']} train / {stats['val']} val "
+          f"(skipped {skipped}, dupes {dupes})")
+    print(f"[*] Avg tokens per example: {stats['avg_tokens']:.0f} (max {stats['max_tokens']})")
+    print(f"[*] im_start={stats['im_start_id']} im_end={stats['im_end_id']} "
+          f"rainbow={stats['rainbow_ids']}")
+    return stats
+
+
+def tokenize_example(tokenizer, messages, max_len: int, max_resp: int = 256):
+    """Tokenize a conversation segment-wise; return ids + assistant bounds.
+
+    The target (last assistant response) is capped at `max_resp` tokens
+    including its trailing <|im_end|>: long web-text targets teach rambling,
+    and the terminator must stay in the target so the model learns to emit it.
+    """
+    ids = []
+    assistant_ranges = []  # (start, end) per message, in token space
+    for m in messages:
+        role = m.get("role", "user")
+        content = m.get("content", "")
+        if not content.strip():
+            continue
+        seg = format_segment(role, content)
+        seg_ids = tokenizer.encode(seg, add_special_tokens=False)
+        start = len(ids)
+        ids.extend(seg_ids)
+        if role == "assistant":
+            assistant_ranges.append((start, len(ids)))
+
+    if not ids or not assistant_ranges:
+        return None
+
+    # Target turn = the LAST assistant response (LLaDA-MoE style)
+    a0, a1 = assistant_ranges[-1]
+    if a1 <= a0:
+        return None
+
+    # Cap the response length, always keeping the trailing <|im_end|>
+    if a1 - a0 > max_resp:
+        im_end_id = tokenizer.convert_tokens_to_ids(IM_END)
+        ids = ids[:a0] + ids[a0:a0 + max_resp - 1] + [im_end_id]
+        a1 = a0 + max_resp
+
+    # Truncate from the front (history), keep the target response intact
+    resp = ids[a0:a1]
+    if len(resp) > max_len:
+        resp = resp[:max_len]
+        a1 = a0 + len(resp)
+    hist = ids[:a0]
+    room = max_len - len(resp)
+    if len(hist) > room:
+        hist = hist[len(hist) - room:] if room > 0 else []
+    ids = hist + resp
+    return {
+        "input_ids": torch.tensor(ids, dtype=torch.long),
+        "assistant_start": len(hist),
+        "assistant_end": len(ids),
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Prepare no_robots as ChatML for MetaDiffusion SFT")
+    parser.add_argument("--model-path", default="../hf_release", help="Dir with tokenizer.json (base model)")
+    parser.add_argument("--data-dir", default="data/no_robots_chatml", help="Output dir")
+    parser.add_argument("--datasets", default="no_robots",
+                        help="Comma list: no_robots, alpaca, dolly, smol-smoltalk, "
+                             "math (e.g. no_robots,alpaca,dolly,smol-smoltalk,math)")
+    parser.add_argument("--max-examples", type=int, default=0,
+                        help="Cap per downloaded dataset (0 = no cap). smol-smoltalk "
+                             "is 484K; use ~60000 for a 150M model. Does not apply "
+                             "to synthetic math (exhaustive coverage is the point).")
+    parser.add_argument("--val-size", type=int, default=500, help="Held-out examples")
+    parser.add_argument("--max-len", type=int, default=1024, help="Max tokens per example")
+    parser.add_argument("--max-resp-tokens", type=int, default=256,
+                        help="Cap on target response tokens (keeps <|im_end|>)")
+    parser.add_argument("--math-repeat", type=int, default=1,
+                        help="Upsample synthetic math: N passes over all pairs "
+                             "(different phrasings per pass; 3-4 helps a 150M "
+                             "model memorize the mapping)")
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+
+    out = Path(args.data_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    print(f"[*] Loading tokenizer from {args.model_path}")
+    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
+    add_chat_tokens(tokenizer)
+    tokenizer.save_pretrained(out / "tokenizer")
+
+    convs = []  # (roles, segs) per conversation
+    math_indices = set()
+    for name in args.datasets.split(","):
+        name = name.strip()
+        print(f"[*] Loading dataset: {name}")
+        if name == "math":
+            # Upsample: several passes over the same pairs, each pass re-rolling
+            # phrasings (fresh seed). Math repeats are NOT deduped: identical
+            # copies are intentional extra gradient steps (memorization needs
+            # repetition, not variety).
+            start = len(convs)
+            for rep in range(args.math_repeat):
+                msgs = load_math_data(seed=args.seed + rep)
+                for messages in msgs:
+                    segs, roles = build_conv_segments(messages)
+                    convs.append((roles, segs))
+                print(f"    pass {rep + 1}/{args.math_repeat}: {len(msgs)}")
+            math_indices.update(range(start, len(convs)))
+        else:
+            msgs = load_messages(name)
+            if args.max_examples > 0 and len(msgs) > args.max_examples:
+                rng_cap = random.Random(args.seed)
+                rng_cap.shuffle(msgs)
+                msgs = msgs[: args.max_examples]
+                print(f"    capped to {len(msgs)}")
+            for messages in msgs:
+                segs, roles = build_conv_segments(messages)
+                convs.append((roles, segs))
+            print(f"    {len(msgs)} conversations")
+    print(f"[*] Total: {len(convs)} conversations")
+
+    train_examples, val_examples, skipped, dupes = tokenize_and_split(
+        convs, math_indices, tokenizer, args.val_size, args.max_len,
+        args.max_resp_tokens, args.seed)
+
+    save_dataset(out, train_examples, val_examples, tokenizer, skipped, dupes)
+
+
+if __name__ == "__main__":
+    main()

scripts/requirements.txt

@@ -0,0 +1,7 @@
+torch>=2.2
+transformers>=4.40
+safetensors>=0.4
+datasets>=2.18
+numpy>=1.26
+pandas>=2.0
+pyarrow>=14.0

scripts/train_chat.py

@@ -0,0 +1,698 @@
+#!/usr/bin/env python3
+"""
+train_chat.py: Chat SFT for MetaDiffusion-150M-exp (LLaDA Algorithm 2 style).
+
+Checkpoints use the same format as train.py:
+    {step, model_state_dict, optimizer_state_dict, scheduler_state_dict, config}
+
+Usage:
+    python3 train_chat.py \
+        --model-path ../hf_release \
+        --data-dir data/no_robots_chatml \
+        --output-dir checkpoints_chat \
+        --epochs 8
+"""
+
+import argparse
+import glob
+import heapq
+import json
+import logging
+import math
+import os
+import re
+import shutil
+import sys
+import time
+from dataclasses import asdict
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from safetensors.torch import load_file
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import LambdaLR
+from torch.utils.data import DataLoader, Dataset
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from model import MetaDiffusionLM, MetaDiffusionConfig  # noqa: E402
+
+MASK_TOKEN_ID = 32000
+BASE_VOCAB = 32000          # Supra tokenizer entry count
+RESERVED = "<|reserved|>"   # filler so id 32000 stays free for [MASK]
+BASE_VOCAB_WITH_RESERVED = BASE_VOCAB + 1
+CHAT_TOKENS = ["<|im_start|>", "<|im_end|>"] + [f"<|r{i}|>" for i in range(1, 8)]
+CHAT_VOCAB = BASE_VOCAB + 1 + len(CHAT_TOKENS)  # 32010
+
+
+def ensure_chat_tokens(tokenizer):
+    """Make sure chat tokens live at ids 32001..32009
+
+    Handles both a fresh base tokenizer (adds <|reserved|> at 32000 first) and
+    an already-prepared one (no-op).
+    """
+    if tokenizer.convert_tokens_to_ids("<|im_start|>") == tokenizer.unk_token_id:
+        if len(tokenizer) == BASE_VOCAB:
+            tokenizer.add_special_tokens({"additional_special_tokens": [RESERVED]})
+        tokenizer.add_special_tokens({"additional_special_tokens": CHAT_TOKENS})
+    return tokenizer
+
+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_duration(seconds):
+    seconds = max(0, int(seconds))
+    hours, rem = divmod(seconds, 3600)
+    minutes, secs = divmod(rem, 60)
+    if hours > 0:
+        return f"{hours}h{minutes:02d}m{secs:02d}s"
+    if minutes > 0:
+        return f"{minutes}m{secs:02d}s"
+    return f"{secs}s"
+
+
+def get_gpu_memory_info():
+    if not torch.cuda.is_available():
+        return None, None
+    torch.cuda.synchronize()
+    free, total = torch.cuda.mem_get_info()
+    return total, free
+
+
+def detect_max_batch_size(model, seq_len, device, keep_free_fraction=0.1,
+                          amp_dtype=None):
+    """Largest batch size that fits in VRAM with headroom (from train.py)."""
+    if not torch.cuda.is_available():
+        return 8
+    total_mem, free_mem = get_gpu_memory_info()
+    if total_mem is None:
+        return 8
+    logger.info(f"GPU memory: {format_bytes(total_mem)} total, {format_bytes(free_mem)} free")
+    model = model.to(device).train()
+    mem_limit = total_mem - int(total_mem * keep_free_fraction)
+    last_working, first_oom = 1, None
+    for bs in [1, 2, 4, 8, 16, 32, 64, 128, 256]:
+        torch.cuda.synchronize()
+        if torch.cuda.memory_allocated() >= mem_limit:
+            first_oom = bs
+            break
+        try:
+            input_ids = torch.randint(0, 32000, (bs, seq_len), device=device)
+            labels = torch.randint(0, 32000, (bs, seq_len), device=device)
+            mask_positions = torch.rand(bs, seq_len, device=device) < 0.5
+            timesteps = torch.rand(bs, device=device)
+            with torch.autocast("cuda", dtype=amp_dtype or torch.float16):
+                logits = model(input_ids, timesteps)
+                loss, num_masked = model.compute_loss(logits, labels, mask_positions)
+            if num_masked > 0:
+                (loss / 4).backward()
+            torch.cuda.synchronize()
+            peak_mem = torch.cuda.max_memory_allocated()
+            logger.info(f"  batch_size={bs:>3d}: peak VRAM={format_bytes(peak_mem)} "
+                        f"(limit={format_bytes(mem_limit)})")
+            if peak_mem >= mem_limit:
+                first_oom = bs
+                break
+            last_working = bs
+        except RuntimeError as e:
+            if "out of memory" in str(e).lower():
+                first_oom = bs
+                break
+            raise
+        finally:
+            model.zero_grad(set_to_none=True)
+            torch.cuda.empty_cache()
+    if first_oom is not None and last_working < first_oom - 1:
+        lo, hi = last_working, first_oom
+        while lo + 1 < hi:
+            mid = (lo + hi) // 2
+            try:
+                input_ids = torch.randint(0, 32000, (mid, seq_len), device=device)
+                labels = torch.randint(0, 32000, (mid, seq_len), device=device)
+                mask_positions = torch.rand(mid, seq_len, device=device) < 0.5
+                timesteps = torch.rand(mid, device=device)
+                with torch.autocast("cuda", dtype=amp_dtype or torch.float16):
+                    logits = model(input_ids, timesteps)
+                    loss, num_masked = model.compute_loss(logits, labels, mask_positions)
+                if num_masked > 0:
+                    (loss / 4).backward()
+                torch.cuda.synchronize()
+                if torch.cuda.max_memory_allocated() < mem_limit:
+                    lo = mid
+                else:
+                    hi = mid
+            except RuntimeError as e:
+                if "out of memory" in str(e).lower():
+                    hi = mid
+                else:
+                    raise
+            finally:
+                model.zero_grad(set_to_none=True)
+                torch.cuda.empty_cache()
+        last_working = lo
+    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_steps(stats_path, max_n):
+    if not os.path.exists(stats_path):
+        return set()
+    entries = []
+    with open(stats_path) as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                entry = json.loads(line)
+                if "step" in entry and "loss" in entry:
+                    entries.append((entry["loss"], entry["step"]))
+            except json.JSONDecodeError:
+                continue
+    return {step for _, step in heapq.nsmallest(max_n, entries)}
+
+
+def cleanup_checkpoints(output_dir, keep_first_n, keep_last_n, keep_best_n, stats_path):
+    all_ckpts = sorted(glob.glob(os.path.join(output_dir, "step_*.pt")))
+    if len(all_ckpts) <= keep_first_n + keep_last_n + keep_best_n:
+        return
+    first_steps = {get_step_from_filename(c) for c in all_ckpts[:keep_first_n]}
+    last_steps = {get_step_from_filename(c) for c in all_ckpts[-keep_last_n:]}
+    best_steps = load_best_steps(stats_path, keep_best_n)
+    keep_steps = (first_steps | last_steps | best_steps) - {None}
+    for ckpt in all_ckpts:
+        s = get_step_from_filename(ckpt)
+        if s is not None and s not in keep_steps:
+            try:
+                os.remove(ckpt)
+            except OSError:
+                pass
+    logger.info(f"Cleaned old checkpoints (kept {len(keep_steps)}: "
+                f"{len(first_steps)} first, {len(last_steps)} last, {len(best_steps)} best)")
+
+
+def get_free_disk_space(path):
+    return shutil.disk_usage(path).free
+
+def build_config(config_dict):
+    """Build MetaDiffusionConfig, ignoring non-dataclass keys (model_type, ...)."""
+    valid = {k: v for k, v in config_dict.items() if k in MetaDiffusionConfig.__dataclass_fields__}
+    config = MetaDiffusionConfig(**valid)
+    config.tie_word_embeddings = False  # released checkpoint has an untied lm_head
+    return config
+
+
+def load_model(model_path, device):
+    """Load MetaDiffusionLM from a dir (config.json + model.safetensors) or a step_*.pt."""
+    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)
+        sd = load_file(path / "model.safetensors")
+        sd = {k[len("model."):] if k.startswith("model.") else k: v for k, v in sd.items()}
+        missing, unexpected = model.load_state_dict(sd, strict=False)
+        if missing or unexpected:
+            logger.warning(f"missing={missing[:5]} unexpected={unexpected[:5]}")
+    else:
+        ckpt = torch.load(path, map_location=device, weights_only=False)
+        config = build_config(ckpt["config"])
+        model = MetaDiffusionLM(config).to(device)
+        model.load_state_dict(clean_state_dict(ckpt["model_state_dict"]))
+    return model, config
+
+
+def clean_state_dict(state_dict):
+    """Strip torch.compile's _orig_mod. prefix from checkpoint keys."""
+    return {k.replace("_orig_mod.", "", 1) if k.startswith("_orig_mod.") else k: v
+            for k, v in state_dict.items()}
+
+
+def expand_embeddings(model, new_vocab):
+    """Mean-init new rows (ChatML + rainbow tokens) in embed_tokens and lm_head.
+
+    New modules are created on the model's device/dtype: nn.Embedding/nn.Linear
+    default to CPU, which would crash the first forward ("Tensor device
+    mismatch") unless something else moves the model afterwards.
+    """
+    old_vocab = model.config.mask_vocab_size
+    if new_vocab <= old_vocab:
+        return
+    device = model.embed_tokens.weight.device
+    dtype = model.embed_tokens.weight.dtype
+    mean_emb = model.embed_tokens.weight.data.mean(dim=0, keepdim=True)
+    n_new = new_vocab - old_vocab
+
+    emb = torch.cat([model.embed_tokens.weight.data, mean_emb.expand(n_new, -1)], dim=0)
+    model.embed_tokens = nn.Embedding(new_vocab, model.config.hidden_size,
+                                      padding_idx=model.config.pad_token_id).to(device, dtype)
+    model.embed_tokens.weight.data.copy_(emb)
+
+    head = torch.cat([model.lm_head.weight.data, mean_emb.expand(n_new, -1)], dim=0)
+    model.lm_head = nn.Linear(model.config.hidden_size, new_vocab, bias=False).to(device, dtype)
+    model.lm_head.weight.data.copy_(head)
+
+    model.config.mask_vocab_size = new_vocab
+    logger.info(f"Expanded embeddings {old_vocab} -> {new_vocab} (mean init)")
+
+
+class ChatDataset(Dataset):
+    """no_robots ChatML examples; masks ONLY the last assistant response.
+
+    Examples are stored as plain int lists (NOT tensors): with forkserver
+    workers (torch's default once CUDA is initialized), every tensor in the
+    dataset is transferred through shared memory at worker spawn, and 9000
+    tensors blows the open-file limit. Lists pickle as bytes.
+    """
+
+    def __init__(self, data_path, tokenizer, seq_len, seed=42):
+        raw = torch.load(data_path, weights_only=True)["examples"]
+        self.examples = [
+            {
+                "ids": ex["input_ids"].tolist(),
+                "a0": int(ex["assistant_start"]),
+                "a1": int(ex["assistant_end"]),
+            }
+            for ex in raw
+        ]
+        self.seq_len = seq_len
+        self.mask_id = MASK_TOKEN_ID
+        self.rainbow_ids = [
+            tokenizer.convert_tokens_to_ids(f"<|r{i}|>") for i in range(1, 8)
+        ]
+        self.seed = seed
+        logger.info(f"Loaded {len(self.examples)} examples from {data_path}")
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, idx):
+        ex = self.examples[idx]
+        ids = ex["ids"]
+        a0, a1 = ex["a0"], ex["a1"]
+
+        # Guard truncation
+        if len(ids) > self.seq_len:
+            resp = ids[a0:a1]
+            if len(resp) > self.seq_len:
+                resp = resp[: self.seq_len]
+            room = self.seq_len - len(resp)
+            hist = ids[:a0]
+            hist = hist[len(hist) - room:] if room > 0 else []
+            ids = hist + resp
+            a0, a1 = len(hist), len(ids)
+
+        # Rainbow padding (cyclic, never masked, never in loss)
+        n = len(ids)
+        pad = self.seq_len - n
+        full = ids + [self.rainbow_ids[j % 7] for j in range(pad)]
+
+        can_mask = torch.zeros(self.seq_len, dtype=torch.bool)
+        can_mask[a0:a1] = True
+
+        t = torch.rand(1).item()
+        rand = torch.rand(self.seq_len)
+        mask_pos = (rand < t) & can_mask
+
+        input_ids = torch.tensor(full, dtype=torch.long)
+        input_ids[mask_pos] = self.mask_id
+        attention = torch.ones(self.seq_len, dtype=torch.long)
+        attention[n:] = 0
+
+        return {
+            "input_ids": input_ids,
+            "labels": torch.tensor(full, dtype=torch.long),
+            "mask_positions": mask_pos,
+            "timesteps": torch.tensor(t, dtype=torch.float32),
+            "attention_mask": attention,
+            "resp_len": torch.tensor(max(a1 - a0, 1), dtype=torch.float32),
+        }
+
+
+def collate_fn(batch):
+    return {
+        k: torch.stack([b[k] for b in batch]) for k in batch[0]
+    }
+
+
+def worker_init_fn(worker_id):
+    torch.manual_seed(42 + worker_id)
+
+def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps,
+                                    min_lr_ratio=0.1):
+    def lr_lambda(current_step):
+        if current_step < num_warmup_steps:
+            return float(current_step) / float(max(1, num_warmup_steps))
+        progress = float(current_step - num_warmup_steps) / float(
+            max(1, num_training_steps - num_warmup_steps)
+        )
+        return max(min_lr_ratio, 0.5 * (1.0 + math.cos(math.pi * progress)))
+    return LambdaLR(optimizer, lr_lambda)
+
+
+@torch.no_grad()
+def evaluate(model, val_dataset, batch_size, device, dtype, n_max=100):
+    model.eval()
+    losses, n_seen = [], 0
+    for start in range(0, min(len(val_dataset), n_max), batch_size):
+        idxs = list(range(start, min(start + batch_size, n_max)))
+        batch = collate_fn([val_dataset[i] for i in idxs])
+        input_ids = batch["input_ids"].to(device)
+        labels = batch["labels"].to(device)
+        mask_positions = batch["mask_positions"].to(device)
+        timesteps = batch["timesteps"].to(device)
+        attention = batch["attention_mask"].to(device)
+        resp_len = batch["resp_len"].to(device)
+        with torch.autocast("cuda", dtype=dtype):
+            logits = model(input_ids, timesteps, attention_mask=attention)
+        if mask_positions.any():
+            ce = F.cross_entropy(logits.float()[mask_positions],
+                                 labels[mask_positions], reduction="none")
+            w = (1.0 / (timesteps * resp_len)).unsqueeze(1).expand_as(labels)
+            loss = (ce * w[mask_positions]).sum() / input_ids.shape[0]
+            losses.append(loss.item())
+            n_seen += 1
+    model.train()
+    return sum(losses) / len(losses) if losses else float("nan"), n_seen
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Chat SFT for MetaDiffusion (LLaDA Algorithm 2)")
+    parser.add_argument("--model-path", default="../hf_release",
+                        help="Dir with config.json + model.safetensors, or a step_*.pt")
+    parser.add_argument("--data-dir", default="data/no_robots_chatml")
+    parser.add_argument("--output-dir", default="checkpoints_chat")
+    parser.add_argument("--seq-len", type=int, default=512)
+    parser.add_argument("--batch-size", type=int, default=0, help="0 = auto-detect")
+    parser.add_argument("--grad-accum-steps", type=int, default=4)
+    parser.add_argument("--num-workers", type=int, default=4,
+                        help="DataLoader workers (0 if forkserver shm issues)")
+    parser.add_argument("--lr", type=float, default=3e-5)
+    parser.add_argument("--min-lr-ratio", type=float, default=0.1)
+    parser.add_argument("--warmup-steps", type=int, default=100)
+    parser.add_argument("--weight-decay", type=float, default=0.1)
+    parser.add_argument("--epochs", type=int, default=8)
+    parser.add_argument("--max-steps", type=int, default=0, help="0 = epochs only")
+    parser.add_argument("--transferred-lr-mult", type=float, default=0.33)
+    parser.add_argument("--new-lr-mult", type=float, default=1.0)
+    parser.add_argument("--max-grad-norm", type=float, default=1.0)
+    parser.add_argument("--save-every", type=int, default=500)
+    parser.add_argument("--log-every", type=int, default=50)
+    parser.add_argument("--val-every", type=int, default=200)
+    parser.add_argument("--patience", type=int, default=3,
+                        help="Early stop after N val checks without improvement (0 = off)")
+    parser.add_argument("--min-delta", type=float, default=0.001,
+                        help="Relative val-loss improvement required to count as progress")
+    parser.add_argument("--resume-from", type=str, default=None)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--device", default="cuda", help="cuda, cuda:1, cpu")
+    parser.add_argument("--bf16", action="store_true", help="Use bf16 instead of fp16")
+    parser.add_argument("--no-compile", action="store_true")
+    parser.add_argument("--keep-free-vram", type=float, default=0.1)
+    parser.add_argument("--keep-first-n", type=int, default=2)
+    parser.add_argument("--keep-last-n", type=int, default=2)
+    parser.add_argument("--keep-best-n", type=int, default=2)
+    parser.add_argument("--disk-min-gb", type=float, default=5.0)
+    parser.add_argument("--export-dir", type=str, default=None,
+                        help="Export final dir (config+safetensors+tokenizer)")
+    args = parser.parse_args()
+
+    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
+    torch.manual_seed(args.seed)
+
+    model, config = load_model(args.model_path, device)
+    logger.info(f"Loaded: {config.num_hidden_layers}L x {config.hidden_size}W, "
+                f"vocab={config.mask_vocab_size}")
+
+    if args.resume_from is None:
+        expand_embeddings(model, CHAT_VOCAB)
+    else:
+        logger.info(f"Resuming: keeping expanded vocab {config.mask_vocab_size}")
+
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained(os.path.join(args.data_dir, "tokenizer"))
+    ensure_chat_tokens(tokenizer)
+    im_end = tokenizer.convert_tokens_to_ids("<|im_end|>")
+    assert im_end == 32002, (
+        f"Tokenizer has im_end={im_end}, expected 32002. "
+        f"Data dir is stale (pre-fix ids): re-run prepare_data.py first."
+    )
+    logger.info(f"Tokenizer vocab: {len(tokenizer)} | im_end={im_end}")
+
+    if args.bf16:
+        model = model.to(torch.bfloat16)
+        amp_dtype = torch.bfloat16
+    else:
+        # fp16 AMP: keep fp32 master weights, autocast does the fp16 compute.
+        # GradScaler requires fp32 gradients; fp16 weights would produce fp16
+        # grads and unscale_ raises "Attempting to unscale FP16 gradients".
+        amp_dtype = torch.float16
+
+    train_ds = ChatDataset(os.path.join(args.data_dir, "train.pt"), tokenizer,
+                           args.seq_len, seed=args.seed)
+    val_ds = ChatDataset(os.path.join(args.data_dir, "val.pt"), tokenizer,
+                         args.seq_len, seed=args.seed)
+
+    transferred_names, new_names = set(), set()
+    for name, p in model.named_parameters():
+        if any(k in name for k in ["timestep_emb", "timestep_residual", "lm_head",
+                                   "embed_tokens.weight"]):
+            new_names.add(name)
+        else:
+            transferred_names.add(name)
+    param_groups = [
+        {"params": [p for n, p in model.named_parameters() if n in transferred_names],
+         "lr": args.lr * args.transferred_lr_mult, "name": "transferred"},
+        {"params": [p for n, p in model.named_parameters() if n in new_names],
+         "lr": args.lr * args.new_lr_mult, "name": "new"},
+    ]
+    for pg in param_groups:
+        logger.info(f"  {pg['name']}: {sum(p.numel() for p in pg['params']):,} params, "
+                    f"lr={pg['lr']:.2e}")
+
+    optimizer = AdamW(param_groups, weight_decay=args.weight_decay)
+
+    batch_size = args.batch_size
+    if batch_size <= 0 and torch.cuda.is_available():
+        batch_size = detect_max_batch_size(model, args.seq_len, device,
+                                           args.keep_free_vram, amp_dtype)
+    if batch_size <= 0:
+        batch_size = 8
+    eff_batch = batch_size * args.grad_accum_steps
+    steps_per_epoch = max(1, math.ceil(len(train_ds) / eff_batch))
+    total_steps = args.max_steps if args.max_steps > 0 else steps_per_epoch * args.epochs
+    logger.info(f"batch={batch_size} accum={args.grad_accum_steps} "
+                f"eff={eff_batch} steps/epoch={steps_per_epoch} total={total_steps}")
+
+    # Compile AFTER batch detection + resume (avoids recompiles per probe
+    # batch size, and lets resume load clean keys into a plain nn.Module)
+    scheduler = get_cosine_schedule_with_warmup(
+        optimizer, args.warmup_steps, total_steps, args.min_lr_ratio
+    )
+
+    dataloader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
+                            num_workers=args.num_workers, pin_memory=True,
+                            drop_last=False, collate_fn=collate_fn,
+                            worker_init_fn=worker_init_fn)
+
+    global_step = 0
+    if args.resume_from:
+        ckpt = torch.load(args.resume_from, map_location=device, weights_only=False)
+        model.load_state_dict(clean_state_dict(ckpt["model_state_dict"]))
+        optim_state = ckpt.get("optimizer_state_dict", {})
+        if optim_state and "param_groups" in optim_state and "state" in optim_state:
+            try:
+                optimizer.load_state_dict(optim_state)
+            except (ValueError, KeyError) as e:
+                logger.warning(f"Optimizer state not loaded: {e}")
+        sched_state = ckpt.get("scheduler_state_dict", {})
+        if sched_state and sched_state.get("last_epoch", 0) == ckpt.get("step", 0):
+            try:
+                scheduler.load_state_dict(sched_state)
+            except (ValueError, KeyError) as e:
+                logger.warning(f"Scheduler state not loaded: {e}")
+        global_step = ckpt.get("step", 0)
+        logger.info(f"Resumed from step {global_step}")
+
+    if not args.no_compile:
+        logger.info("Compiling model...")
+        model = torch.compile(model)
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    stats_path = os.path.join(args.output_dir, "stats.jsonl")
+    stats_file = open(stats_path, "a")
+    with open(os.path.join(args.output_dir, "config.json"), "w") as f:
+        json.dump(asdict(model.config), f, indent=2, default=str)
+
+    scaler = torch.amp.GradScaler("cuda", enabled=not args.bf16)
+    free_disk = get_free_disk_space(args.output_dir)
+    if free_disk < args.disk_min_gb * 1e9:
+        stats_file.close()
+        raise RuntimeError(f"Insufficient disk space: {format_bytes(free_disk)}")
+
+    model.train()
+    optimizer.zero_grad()
+    loss_total, loss_count = 0.0, 0
+    start_time = time.time()
+    last_log_time = start_time
+    data_iter = iter(dataloader)
+    epoch = 0
+    best_val = float("inf")
+    no_improve = 0
+    early_stopped = False
+
+    while global_step < total_steps:
+        if global_step % steps_per_epoch == 0 and global_step > 0:
+            epoch += 1
+        try:
+            batch = next(data_iter)
+        except StopIteration:
+            epoch += 1
+            data_iter = iter(dataloader)
+            batch = next(data_iter)
+
+        input_ids = batch["input_ids"].to(device)
+        labels = batch["labels"].to(device)
+        mask_positions = batch["mask_positions"].to(device)
+        timesteps = batch["timesteps"].to(device)
+        attention = batch["attention_mask"].to(device)
+        resp_len = batch["resp_len"].to(device)
+
+        with torch.autocast("cuda", dtype=amp_dtype):
+            logits = model(input_ids, timesteps, attention_mask=attention)
+
+        num_masked = mask_positions.sum().item()
+        if num_masked > 0:
+            # LLaDA GUIDELINES loss: CE/(t * response_len) summed over masked
+            # response tokens, mean over batch. Expected value ~ per-token CE.
+            ce = F.cross_entropy(logits.float()[mask_positions],
+                                 labels[mask_positions], reduction="none")
+            w = (1.0 / (timesteps * resp_len)).unsqueeze(1).expand_as(labels)
+            loss = (ce * w[mask_positions]).sum() / input_ids.shape[0]
+            scaler.scale(loss / args.grad_accum_steps).backward()
+            loss_total += loss.item()
+        else:
+            loss = torch.tensor(0.0, device=device)
+
+        global_step += 1
+
+        if global_step % args.grad_accum_steps == 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+            skipped = scaler.step(optimizer)  # True when grads had inf/nan
+            scaler.update()
+            if not skipped:
+                scheduler.step()  # don't advance LR on skipped steps
+            optimizer.zero_grad()
+
+        loss_count += 1
+
+        if global_step % args.log_every == 0:
+            now = time.time()
+            elapsed = now - start_time
+            avg_loss = loss_total / max(1, loss_count)
+            ppl = math.exp(min(avg_loss, 20))
+            lr = scheduler.get_last_lr()[0]
+            steps_per_sec = args.log_every / max(now - last_log_time, 1e-6)
+            eta = format_duration((total_steps - global_step) / steps_per_sec)
+            logger.info(f"Step {global_step:>6d} | epoch {epoch:.1f} | loss={avg_loss:.4f} | "
+                        f"ppl={ppl:.1f} | lr={lr:.2e} | {steps_per_sec:.1f} steps/s | "
+                        f"elapsed={elapsed:.0f}s | eta={eta}")
+            loss_total, loss_count = 0.0, 0
+            last_log_time = now
+
+            stats_entry = {"step": global_step, "epoch": round(epoch, 2),
+                           "loss": round(avg_loss, 4), "ppl": round(ppl, 1), "lr": lr}
+            stats_file.write(json.dumps(stats_entry) + "\n")
+            stats_file.flush()
+
+        # Validation + early stopping (independent of log cadence)
+        if global_step % args.val_every == 0:
+            val_loss, _ = evaluate(model, val_ds, batch_size, device, amp_dtype)
+            logger.info(f"  val_loss={val_loss:.4f}")
+            stats_file.write(json.dumps({"step": global_step,
+                                         "val_loss": round(val_loss, 4)}) + "\n")
+            stats_file.flush()
+            if val_loss < best_val * (1.0 - args.min_delta):
+                best_val = val_loss
+                no_improve = 0
+                ckpt_path = os.path.join(args.output_dir, "best.pt")
+                torch.save({"step": global_step,
+                            "model_state_dict": clean_state_dict(model.state_dict()),
+                            "optimizer_state_dict": optimizer.state_dict(),
+                            "scheduler_state_dict": scheduler.state_dict(),
+                            "config": asdict(model.config)}, ckpt_path)
+                logger.info(f"  Best val loss, saved {ckpt_path}")
+            else:
+                no_improve += 1
+                logger.info(f"  No val improvement ({no_improve}/{args.patience} checks, "
+                            f"best={best_val:.4f})")
+                if args.patience > 0 and no_improve >= args.patience:
+                    logger.info(f"Early stopping at step {global_step}: no val loss "
+                                f"improvement for {args.patience} checks "
+                                f"(best={best_val:.4f})")
+                    ckpt_path = os.path.join(args.output_dir, f"step_{global_step}.pt")
+                    torch.save({"step": global_step,
+                                "model_state_dict": clean_state_dict(model.state_dict()),
+                                "optimizer_state_dict": optimizer.state_dict(),
+                                "scheduler_state_dict": scheduler.state_dict(),
+                                "config": asdict(model.config)}, ckpt_path)
+                    logger.info(f"Saved final checkpoint: {ckpt_path}")
+                    stats_file.write(json.dumps(
+                        {**stats_entry, "best_val": round(best_val, 4),
+                         "early_stopped": True}) + "\n")
+                    stats_file.flush()
+                    stats_file.close()
+                    early_stopped = True
+                    break
+
+        if global_step % args.save_every == 0:
+            ckpt_path = os.path.join(args.output_dir, f"step_{global_step}.pt")
+            torch.save({"step": global_step,
+                        "model_state_dict": clean_state_dict(model.state_dict()),
+                        "optimizer_state_dict": optimizer.state_dict(),
+                        "scheduler_state_dict": scheduler.state_dict(),
+                        "config": asdict(model.config)}, ckpt_path)
+            logger.info(f"Saved checkpoint: {ckpt_path}")
+            cleanup_checkpoints(args.output_dir, args.keep_first_n,
+                                args.keep_last_n, args.keep_best_n, stats_path)
+            if get_free_disk_space(args.output_dir) < args.disk_min_gb * 1e9:
+                logger.warning("Low disk after save; stopping")
+                stats_file.close()
+                return
+
+    stats_file.close()
+    if early_stopped:
+        logger.info(f"Early stopping triggered; best val loss {best_val:.4f} "
+                    f"saved as best.pt")
+    else:
+        logger.info(f"Training complete at step {global_step}")
+
+    if args.export_dir:
+        from export_hf import export
+        export(os.path.join(args.output_dir, f"step_{global_step}.pt")
+               if not os.path.exists(os.path.join(args.output_dir, "best.pt"))
+               else os.path.join(args.output_dir, "best.pt"),
+               os.path.join(args.data_dir, "tokenizer"),
+               args.export_dir)
+
+
+if __name__ == "__main__":
+    main()

tokenizer_config.json

@@ -0,0 +1,23 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|r1|>",
+    "<|r2|>",
+    "<|r3|>",
+    "<|r4|>",
+    "<|r5|>",
+    "<|r6|>",
+    "<|r7|>"
+  ],
+  "is_local": true,
+  "local_files_only": false,
+  "model_max_length": 5120,
+  "pad_token": "<pad>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<unk>"
+}