#!/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()