def create_model(source_model_id: str, source_revision: str, save_folder: str) -> str: import json import os from pathlib import Path import torch from huggingface_hub import hf_hub_download from safetensors.torch import save_file folder = Path(save_folder) folder.mkdir(parents=True, exist_ok=True) for name in ( "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "generation_config.json", "LICENSE", ): (folder / name).write_bytes(Path(hf_hub_download(source_model_id, name, revision=source_revision)).read_bytes()) config = json.loads(Path(hf_hub_download(source_model_id, "config.json", revision=source_revision)).read_text()) # Internal-only geometry; task vocabulary, token ids, attention method and all LSA/MTP schedules are fixed. # LONGCAT_TINY_NUM_HEADS lets the same generator emit a TP=1 evidence variant (full 32-head attention # fan-out; indexer, MoE and embedding tensors are TP-neutral and stay byte-identical). # LONGCAT_TINY_NUM_LAYERS sets the logical layer count (the TP=1 indexer-topk evidence model needs # 8 logical layers so its 16 physical attention sublayers match the TP=4 checkpoint's 8). tiny_heads = int(os.environ.get("LONGCAT_TINY_NUM_HEADS", "8")) tiny_layers = int(os.environ.get("LONGCAT_TINY_NUM_LAYERS", "4")) config.update( num_layers=tiny_layers, hidden_size=8, ffn_hidden_size=64, expert_ffn_hidden_size=32, num_attention_heads=tiny_heads, kv_lora_rank=256, q_lora_rank=32, qk_nope_head_dim=64, qk_rope_head_dim=64, v_head_dim=128, n_routed_experts=32, zero_expert_num=128, moe_topk=12, oe_vocab_size_ratio=4, oe_neighbor_num=2, oe_split_num=2, ) # LongCat has two attention sublayers per logical layer (physical ids 0..2*num_layers-1). # The DSA indexer-topk capturer defaults to num_hidden_layers slots; expose the real # physical attention-sublayer count so the return_indexer_topk buffer is sized correctly. config["num_indexer_layers"] = 2 * config["num_layers"] # The published checkpoint ships one draft layer regardless of mtp_num_layers; keep its runtime schedule field. (folder / "config.json").write_text(json.dumps(config, indent=2) + "\n") hidden = config["hidden_size"] ffn = config["ffn_hidden_size"] expert_ffn = config["expert_ffn_hidden_size"] num_layers = config["num_layers"] num_heads = config["num_attention_heads"] q_lora = config["q_lora_rank"] kv_lora = config["kv_lora_rank"] nope = config["qk_nope_head_dim"] rope = config["qk_rope_head_dim"] value = config["v_head_dim"] n_routed = config["n_routed_experts"] n_zero = config["zero_expert_num"] index_heads = config["index_n_heads"] index_head_dim = config["index_head_dim"] vocab = config["vocab_size"] qk = nope + rope torch.manual_seed(42) def weight(shape, dtype=torch.bfloat16): return (torch.randn(*shape) * 0.02).to(dtype).contiguous() tensors = {} tensors["model.embed_tokens.weight"] = weight((vocab, hidden)) tensors["lm_head.weight"] = weight((vocab, hidden)) tensors["model.norm.weight"] = weight((hidden,)) for layer in range(num_layers): prefix = f"model.layers.{layer}" for sublayer in (0, 1): attn = f"{prefix}.self_attn.{sublayer}" tensors[f"{attn}.q_a_proj.weight"] = weight((q_lora, hidden)) tensors[f"{attn}.q_a_layernorm.weight"] = weight((q_lora,)) tensors[f"{attn}.q_b_proj.weight"] = weight((num_heads * qk, q_lora)) tensors[f"{attn}.kv_a_proj_with_mqa.weight"] = weight((kv_lora + rope, hidden)) tensors[f"{attn}.kv_a_layernorm.weight"] = weight((kv_lora,)) tensors[f"{attn}.kv_b_proj.weight"] = weight((num_heads * (nope + value), kv_lora)) tensors[f"{attn}.o_proj.weight"] = weight((hidden, num_heads * value)) tensors[f"{prefix}.input_layernorm.{sublayer}.weight"] = weight((hidden,)) tensors[f"{prefix}.post_attention_layernorm.{sublayer}.weight"] = weight((hidden,)) tensors[f"{prefix}.mlps.{sublayer}.gate_proj.weight"] = weight((ffn, hidden)) tensors[f"{prefix}.mlps.{sublayer}.up_proj.weight"] = weight((ffn, hidden)) tensors[f"{prefix}.mlps.{sublayer}.down_proj.weight"] = weight((hidden, ffn)) # LSA indexer lives on the first attention sublayer of every logical layer (CLI factor 2). indexer = f"{prefix}.self_attn.0.indexer" tensors[f"{indexer}.wq_b.weight"] = weight((index_heads * index_head_dim, q_lora)) tensors[f"{indexer}.wk.weight"] = weight((index_head_dim, hidden)) tensors[f"{indexer}.weights_proj.weight"] = weight((index_heads, hidden)) tensors[f"{indexer}.k_norm.weight"] = weight((index_head_dim,)) # Routers are FP32 in the source checkpoint; identity zero experts carry no weights. tensors[f"{prefix}.mlp.router.classifier.weight"] = weight((n_routed + n_zero, hidden), torch.float32) tensors[f"{prefix}.mlp.router.e_score_correction_bias"] = torch.zeros((n_routed + n_zero,), dtype=torch.float32) for expert in range(n_routed): tensors[f"{prefix}.mlp.experts.{expert}.gate_proj.weight"] = weight((expert_ffn, hidden)) tensors[f"{prefix}.mlp.experts.{expert}.up_proj.weight"] = weight((expert_ffn, hidden)) tensors[f"{prefix}.mlp.experts.{expert}.down_proj.weight"] = weight((hidden, expert_ffn)) # Over-embeddings: two n-gram windows, modulus ratio 4 like tiny-random/longcat-flash-lite. modulus = config["oe_vocab_size_ratio"] * vocab over_width = hidden // (config["oe_split_num"] * (config["oe_neighbor_num"] - 1)) for index in range(config["oe_split_num"] * (config["oe_neighbor_num"] - 1)): rows = modulus + 2 * index + 1 tensors[f"model.ngram_embeddings.embedders.{index}.weight"] = weight((rows, over_width)) tensors[f"model.ngram_embeddings.post_projs.{index}.weight"] = weight((hidden, over_width)) # One MTP/NextN draft layer: tied embedding copy, EH projection, norms, dense MLA with its own indexer. tensors["model.mtp.embed_tokens.weight"] = tensors["model.embed_tokens.weight"].clone() tensors["model.mtp.layers.0.eh_proj.weight"] = weight((hidden, 2 * hidden)) tensors["model.mtp.layers.0.enorm.m.weight"] = weight((hidden,)) tensors["model.mtp.layers.0.hnorm.m.weight"] = weight((hidden,)) tensors["model.mtp.layers.0.input_layernorm.weight"] = weight((hidden,)) tensors["model.mtp.layers.0.post_attention_layernorm.weight"] = weight((hidden,)) draft = "model.mtp.layers.0.self_attn" tensors[f"{draft}.q_a_proj.weight"] = weight((q_lora, hidden)) tensors[f"{draft}.q_a_layernorm.weight"] = weight((q_lora,)) tensors[f"{draft}.q_b_proj.weight"] = weight((num_heads * qk, q_lora)) tensors[f"{draft}.kv_a_proj_with_mqa.weight"] = weight((kv_lora + rope, hidden)) tensors[f"{draft}.kv_a_layernorm.weight"] = weight((kv_lora,)) tensors[f"{draft}.kv_b_proj.weight"] = weight((num_heads * (nope + value), kv_lora)) tensors[f"{draft}.o_proj.weight"] = weight((hidden, num_heads * value)) tensors[f"{draft}.indexer.wq_b.weight"] = weight((index_heads * index_head_dim, q_lora)) tensors[f"{draft}.indexer.wk.weight"] = weight((index_head_dim, hidden)) tensors[f"{draft}.indexer.weights_proj.weight"] = weight((index_heads, hidden)) tensors[f"{draft}.indexer.k_norm.weight"] = weight((index_head_dim,)) tensors["model.mtp.layers.0.transformer_layer.mlp.gate_proj.weight"] = weight((ffn, hidden)) tensors["model.mtp.layers.0.transformer_layer.mlp.up_proj.weight"] = weight((ffn, hidden)) tensors["model.mtp.layers.0.transformer_layer.mlp.down_proj.weight"] = weight((hidden, ffn)) tensors["model.mtp.norm.weight"] = weight((hidden,)) save_file(tensors, str(folder / "model.safetensors"), metadata={"format": "pt"}) expected_prefixes = ( "model.embed_tokens.", "model.norm.", "lm_head.", "model.layers.", "model.ngram_embeddings.", "model.mtp.", ) assert all(name.startswith(expected_prefixes) for name in tensors) assert len(tensors["model.mtp.embed_tokens.weight"].shape) == 2 return ( "LongcatCausalLM (direct SGLang-schema tiny checkpoint): 4 dual-sublayer LongCat Flash decoder " "layers with shortcut MoE (32 routed + 128 identity zero experts, top-k 12), compressed MLA " "(q_lora_rank=32, kv_lora_rank=32, 32 heads, nope 32 / rope 64 / value 32), LSA indexers " "(16 heads, head dim 128, CLI factor 2), two over-embedding windows, one MTP/NextN draft layer, " "131,072-row untied token tables." ) create_model( source_model_id="meituan-longcat/LongCat-Flash-Lite-Sparse", source_revision="35de83a18b5838ec1657edb3f885a59bbdd7888f", save_folder="/tmp/longcat-flash-lite-sparse", )