File size: 9,123 Bytes
42072a7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | 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",
)
|