convert.py

#!/usr/bin/env python3
"""
convert.py: convert Qwen3-0.6B (post-trained instruct) into a
MetaDiffusion-600M initialization checkpoint.

Usage:
    python convert.py \
        --source Qwen/Qwen3-0.6B \
        --output init/metadiffusion-600M-instruct.pt \
        --tokenizer-out data/tokenizer \
        --device cpu
"""

import argparse
import json
import sys
from pathlib import Path

import torch
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

sys.path.insert(0, str(Path(__file__).resolve().parent))
from model import MetaDiffusionConfig, MetaDiffusionLM  # noqa: E402

RAINBOW_TOKENS = [f"<|r{i}|>" for i in range(1, 8)]
NUM_NEW_TOKENS = 1 + len(RAINBOW_TOKENS)  # [MASK] + rainbow

PLAIN_CHAT_TEMPLATE = (
    "{% for message in messages %}"
    "{{ '<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>\\n' }}"
    "{% endfor %}"
)


def build_config(src_config, base_vocab: int) -> MetaDiffusionConfig:
    """base_vocab = real tokenizer length (config.vocab_size may be TP-padded,
    e.g. Qwen3: tokenizer 151,669 vs config 151,936)."""
    c = src_config
    config_vocab = int(getattr(c, "vocab_size", 0))
    if base_vocab > config_vocab:
        raise RuntimeError(
            f"tokenizer vocab {base_vocab} exceeds model vocab {config_vocab}")
    head_dim = int(getattr(c, "head_dim", 0) or (c.hidden_size // c.num_attention_heads))
    cfg = MetaDiffusionConfig(
        hidden_size=int(c.hidden_size),
        intermediate_size=int(c.intermediate_size),
        num_hidden_layers=int(c.num_hidden_layers),
        num_attention_heads=int(c.num_attention_heads),
        num_key_value_heads=int(c.num_key_value_heads),
        head_dim=head_dim,
        vocab_size=base_vocab,
        mask_vocab_size=base_vocab + NUM_NEW_TOKENS,
        mask_token_id=base_vocab,
        pad_token_id=int(getattr(c, "pad_token_id", None) or 151643),
        max_position_embeddings=int(c.max_position_embeddings),
        rope_theta=float(getattr(c, "rope_theta", None) or 1000000.0),
        rms_norm_eps=float(getattr(c, "rms_norm_eps", 1e-6)),
        hidden_act=str(getattr(c, "hidden_act", "silu")),
        qk_norm=bool(getattr(c, "qk_norm", True)),
        timestep_emb_hidden=int(c.hidden_size),
        tie_word_embeddings=False,
    )
    return cfg


def convert(source: str, output: str, tokenizer_out: str, device: str = "cpu"):
    logger = print
    logger(f"[*] Loading AR instruct model: {source}")

    src_config = AutoConfig.from_pretrained(source)
    tokenizer = AutoTokenizer.from_pretrained(source)
    base_vocab = len(tokenizer)
    if base_vocab != int(getattr(src_config, "vocab_size", 0)):
        logger(f"[*] Config vocab is TP-padded ({src_config.vocab_size}); "
               f"using real tokenizer vocab {base_vocab}")
    cfg = build_config(src_config, base_vocab)
    h = cfg.hidden_size
    logger(f"[*] Source: {cfg.num_hidden_layers}L x {h}W, vocab={cfg.vocab_size}, "
           f"heads={cfg.num_attention_heads}/{cfg.num_key_value_heads}, "
           f"head_dim={cfg.head_dim}, qk_norm={cfg.qk_norm}")

    ar_model = AutoModelForCausalLM.from_pretrained(source, torch_dtype=torch.float32, device_map=device)
    state = ar_model.state_dict()
    del ar_model
    torch.cuda.empty_cache() if torch.cuda.is_available() else None

    # --- Copy weights, strip the "model." prefix ---
    new_state = {}
    for key, value in state.items():
        new_key = key.replace("model.", "", 1) if key.startswith("model.") else key
        new_state[new_key] = value.clone()

    embed = new_state["embed_tokens.weight"]  # (V, H)

    # --- Trim untrained TP-padding rows, then append [MASK] + rainbow ---
    if embed.shape[0] > base_vocab:
        logger(f"[*] Trimming {embed.shape[0] - base_vocab} untrained padding rows "
               f"from embeddings / lm_head")
        embed = embed[:base_vocab]
    elif embed.shape[0] < base_vocab:
        raise RuntimeError(f"embedding rows {embed.shape[0]} < tokenizer vocab {base_vocab}")

    mean_row = embed.mean(dim=0, keepdim=True)  # (1, H)
    new_rows = mean_row.expand(NUM_NEW_TOKENS, -1).clone()  # (8, H)
    new_state["embed_tokens.weight"] = torch.cat([embed, new_rows], dim=0)

    # --- Untied lm_head: trim + extend identically ---
    if "lm_head.weight" in new_state:
        head = new_state["lm_head.weight"]
        head = head[:base_vocab] if head.shape[0] > base_vocab else head
        if head.shape[0] != base_vocab:
            raise RuntimeError(f"lm_head rows {head.shape[0]} != tokenizer vocab {base_vocab}")
        new_state["lm_head.weight"] = torch.cat([head, new_rows], dim=0)
    else:
        new_state["lm_head.weight"] = torch.cat([embed, new_rows], dim=0)
    logger(f"[*] embed_tokens / lm_head = {cfg.mask_vocab_size} rows "
           f"(mask={cfg.mask_token_id}, rainbow={cfg.mask_token_id + 1}..{cfg.mask_vocab_size - 1})")

    # --- New diffusion modules: adaLN modulation zero-init (identity at
    # step 0). The timestep embedding MLP keeps DEFAULT random init: with
    # both zero, zero output through the zero gate zeroes every t-path
    # gradient (deadlock: the model never learns noise conditioning). ---
    for i in range(cfg.num_hidden_layers):
        new_state[f"layers.{i}.timestep_modulation.proj.weight"] = torch.zeros(2 * h, h)
        new_state[f"layers.{i}.timestep_modulation.proj.bias"] = torch.zeros(2 * h)

    # --- Verify load. timestep_emb MLP stays at default (random) init so
    # the t-path is not deadlocked (zero MLP through a zero gate). Extra AR
    # keys (rotary inv_freq) are dropped. ---
    model = MetaDiffusionLM(cfg)
    missing, unexpected = model.load_state_dict(new_state, strict=False)
    missing = [k for k in missing if not k.startswith("timestep_emb.")]
    if missing:
        raise RuntimeError(f"unexpected missing keys after convert: {missing}")
    if unexpected:
        logger(f"[*] Dropping {len(unexpected)} unexpected AR keys "
               f"(e.g. {unexpected[0]})")
    n_params = sum(p.numel() for p in model.parameters())
    logger(f"[*] Verified load into MetaDiffusionLM: {n_params/1e6:.1f}M params "
           f"(AR transfer + {NUM_NEW_TOKENS} new token rows + timestep modules)")

    # --- Save checkpoint (full state, including default-init t-MLP) ---
    out = Path(output)
    out.parent.mkdir(parents=True, exist_ok=True)
    ckpt = {
        "config": cfg.__dict__,
        "model_state_dict": model.state_dict(),
        "metadata": {
            "source_model": source,
            "conversion_script": "convert.py",
            "mask_token_id": cfg.mask_token_id,
            "rainbow_token_ids": list(range(cfg.mask_token_id + 1, cfg.mask_vocab_size)),
            "num_new_tokens": NUM_NEW_TOKENS,
            "transferred": True,
        },
    }
    torch.save(ckpt, str(out))
    logger(f"[*] Saved checkpoint: {out}")

    sidecar = out.with_suffix(".json")
    with open(sidecar, "w") as f:
        json.dump(cfg.__dict__, f, indent=2)
    logger(f"[*] Saved config sidecar: {sidecar}")

    # --- Tokenizer: add [MASK] + rainbow at the expected ids ---
    tokenizer.add_special_tokens({"additional_special_tokens": ["[MASK]"] + RAINBOW_TOKENS})
    tokenizer.chat_template = PLAIN_CHAT_TEMPLATE
    mask_id = tokenizer.convert_tokens_to_ids("[MASK]")
    assert mask_id == cfg.mask_token_id, f"[MASK] landed at {mask_id}, expected {cfg.mask_token_id}"
    assert len(tokenizer) == cfg.mask_vocab_size, f"tokenizer vocab {len(tokenizer)} != {cfg.mask_vocab_size}"
    tok_dir = Path(tokenizer_out)
    tok_dir.mkdir(parents=True, exist_ok=True)
    tokenizer.save_pretrained(str(tok_dir))
    logger(f"[*] Saved tokenizer ({len(tokenizer)} tokens, [MASK]={mask_id}): {tok_dir}")

    logger("[*] Done. Next: python prepare_data.py && python train.py --init-checkpoint "
           f"{output}")


def main():
    parser = argparse.ArgumentParser(description="Convert Qwen3-0.6B (instruct) to MetaDiffusion-600M init")
    parser.add_argument("--source", default="Qwen/Qwen3-0.6B")
    parser.add_argument("--output", default="init/metadiffusion-600M-instruct.pt")
    parser.add_argument("--tokenizer-out", default="data/tokenizer")
    parser.add_argument("--device", default="cpu")
    args = parser.parse_args()
    convert(args.source, args.output, args.tokenizer_out, args.device)


if __name__ == "__main__":
    main()