sensenova-loop / .meta /create_model.py
yujiepan's picture
Upload folder using huggingface_hub
2c1a8e3 verified
Raw History Blame Contribute Delete
6 kB
def create_model(source_model_id: str, source_revision: str, save_folder: str) -> str:
import copy
import json
from pathlib import Path
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import save_file
from transformers import AutoConfig, AutoModelForImageTextToText, set_seed
folder = Path(save_folder)
folder.mkdir(parents=True, exist_ok=True)
# Transformers 5.18 loads this custom Loop/mHC wrapper around the native Qwen3.5 backbone; ship it pinned.
for name in (
"modeling_sensenova_loop.py",
"chat_template.jinja",
"processor_config.json",
"preprocessor_config.json",
"video_preprocessor_config.json",
"tokenizer.json",
"tokenizer_config.json",
):
(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())
# Pre-loop 0-3, looped block 4-7 executed twice, post-loop 8-11: twelve physical layers / sixteen executions,
# four complete linear/linear/linear/full cycles. Residual width 8, head dim 32 with eight Q / four KV heads.
text_config = config["text_config"]
text_config.update(
head_dim=32,
hidden_size=8,
intermediate_size=32,
num_hidden_layers=12,
num_attention_heads=8,
num_key_value_heads=4,
linear_key_head_dim=32,
linear_num_key_heads=4,
linear_num_value_heads=8,
linear_value_head_dim=32,
)
text_config["layer_types"] = ["linear_attention", "linear_attention", "linear_attention", "full_attention"] * 3
text_config["rope_parameters"]["mrope_section"] = [1, 1, 2]
# Two-stream mHC spans exactly the looped block [4, 8), matching the source [4, 28) cadence.
config["loop_settings"] = {"counts": [2], "end": 8, "eval_count": 2, "start": 4}
config["mhc"] = {
"end": 8,
"fixed_identity": False,
"iterations": 20,
"pre_scale": 1.0,
"start": 4,
"streams": 2,
}
config["hidden_size"] = 8
# Four vision blocks keep the native Qwen3.5 tower and the fixed patch/temporal/merge and position contracts.
config["vision_config"].update(
hidden_size=64,
intermediate_size=128,
num_heads=2,
out_hidden_size=8,
depth=4,
)
(folder / "config.json").write_text(json.dumps(config, indent=2) + "\n")
set_seed(42)
model = AutoModelForImageTextToText.from_config(
AutoConfig.from_pretrained(folder, trust_remote_code=True),
trust_remote_code=True,
dtype=torch.bfloat16,
)
# The source keeps the hybrid linear-attention SSM parameters FP32 even under a BF16 load.
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5GatedDeltaNet
for layer in model.model.language_model.layers:
linear = getattr(layer, "linear_attn", None)
if isinstance(linear, Qwen3_5GatedDeltaNet):
with torch.no_grad():
linear.A_log = torch.nn.Parameter(linear.A_log.float())
linear.dt_bias = torch.nn.Parameter(linear.dt_bias.float())
linear.norm.weight = torch.nn.Parameter(linear.norm.weight.float())
linear.conv1d.weight = torch.nn.Parameter(linear.conv1d.weight.float())
# One BF16 MTP draft layer reusing a full-attention backbone layer outside the mHC range.
full_layer = model.model.language_model.layers[11]
model.mtp = torch.nn.ModuleDict(
{
"pre_fc_norm_embedding": torch.nn.RMSNorm(8),
"fc": torch.nn.Linear(16, 8, bias=False),
"layers": torch.nn.ModuleList([copy.deepcopy(full_layer)]),
"norm": torch.nn.RMSNorm(8),
"pre_fc_norm_hidden": torch.nn.RMSNorm(8),
}
)
model = model.cpu()
representation = str(model)
state = model.state_dict()
backbone, draft = {}, {}
for name, tensor in state.items():
if name.startswith("mtp."):
draft[name] = tensor.to(torch.bfloat16).contiguous()
else:
# The shipped checkpoint stores the complete backbone FP32, mHC routers included.
backbone[name] = tensor.to(torch.float32).contiguous()
save_file(backbone, str(folder / "model.safetensors"), metadata={"format": "pt"})
save_file(draft, str(folder / "mtp.safetensors"), metadata={"format": "pt"})
# Rewrite the edited config (save was manual) and pin the source generation config bytes.
(folder / "config.json").write_text(json.dumps(config, indent=2) + "\n")
(folder / "generation_config.json").write_bytes(
Path(hf_hub_download(source_model_id, "generation_config.json", revision=source_revision)).read_bytes()
)
expected_draft = {
"mtp.fc.weight",
"mtp.pre_fc_norm_embedding.weight",
"mtp.pre_fc_norm_hidden.weight",
"mtp.norm.weight",
"mtp.layers.0.input_layernorm.weight",
"mtp.layers.0.post_attention_layernorm.weight",
"mtp.layers.0.self_attn.q_proj.weight",
"mtp.layers.0.self_attn.k_proj.weight",
"mtp.layers.0.self_attn.v_proj.weight",
"mtp.layers.0.self_attn.o_proj.weight",
"mtp.layers.0.self_attn.q_norm.weight",
"mtp.layers.0.self_attn.k_norm.weight",
"mtp.layers.0.mlp.gate_proj.weight",
"mtp.layers.0.mlp.up_proj.weight",
"mtp.layers.0.mlp.down_proj.weight",
}
assert set(draft) == expected_draft, sorted(set(draft) ^ expected_draft)
assert all(
name.count("mhc_attn") + name.count("mhc_ffn") == 0 or backbone[name].dtype == torch.float32
for name in backbone
)
return representation
create_model(
source_model_id="sensenova/SenseNova-Loop-9B-Preview",
source_revision="9937e91c7301b7dcbb5bf17e3e4e95c499a11e99",
save_folder="/tmp/sensenova-loop",
)