File size: 5,999 Bytes
2c1a8e3 | 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 | 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",
)
|