Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

chat_template.jinja

@@ -0,0 +1,125 @@
+{{- bos_token -}}
+{%- set preserve_thinking = preserve_thinking | default(false) -%}
+
+{%- macro format_arg_value(arg_value) -%}
+    {%- if arg_value is string -%}
+        {{- "'" + (arg_value | replace("\\", "\\\\") | replace("'", "\\'") | replace("\n", "\\n") | replace("\r", "\\r")) + "'" -}}
+    {%- elif arg_value is mapping or arg_value is iterable -%}
+        {{- arg_value | tojson -}}
+    {%- else -%}
+        {{- arg_value | string -}}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- macro parse_content(content) -%}
+    {%- if content is string -%}
+        {{- content -}}
+    {%- elif content is mapping -%}
+        {{- content | tojson -}}
+    {%- elif content is iterable -%}
+        {%- set _ns = namespace(result="") -%}
+        {%- for item in content -%}
+            {%- if item is string -%}
+                {%- set _ns.result = _ns.result + item -%}
+            {%- elif item is mapping and item.get("type") == "image" -%}
+                {%- set _ns.result = _ns.result + "<image>" -%}
+            {%- elif item is mapping and item.get("type") == "text" -%}
+                {%- set _ns.result = _ns.result + ((item.get("text") or "") | string) -%}
+            {%- else -%}
+                {%- set _ns.result = _ns.result + (item | tojson) -%}
+            {%- endif -%}
+        {%- endfor -%}
+        {{- _ns.result -}}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_tool_calls(tool_calls) -%}
+    {%- set tool_calls_ns = namespace(tool_calls=[]) -%}
+    {%- for tool_call in tool_calls -%}
+        {%- set func = tool_call["function"] if "function" in tool_call else tool_call -%}
+        {%- set func_name = func["name"] -%}
+        {%- set func_args = func.get("arguments") -%}
+        {%- set args_ns = namespace(arg_strings=[]) -%}
+        {%- if func_args is mapping -%}
+            {%- for arg_name, arg_value in func_args.items() -%}
+                {%- set args_ns.arg_strings = args_ns.arg_strings + [arg_name + "=" + format_arg_value(arg_value)] -%}
+            {%- endfor -%}
+        {%- elif func_args is string and (func_args | trim) not in ["", "{}", "null"] -%}
+            {{- raise_exception("Tool call arguments must be a mapping, got a JSON-encoded string: parse arguments with json.loads() before applying the chat template") -}}
+        {%- endif -%}
+        {%- set tool_calls_ns.tool_calls = tool_calls_ns.tool_calls + [func_name + "(" + (args_ns.arg_strings | join(", ")) + ")"] -%}
+    {%- endfor -%}
+    {{- "<|tool_call_start|>[" + (tool_calls_ns.tool_calls | join(", ")) + "]<|tool_call_end|>" -}}
+{%- endmacro -%}
+
+{%- set ns = namespace(system_prompt="", last_user_index=-1) -%}
+{%- if messages and messages[0]["role"] == "system" -%}
+    {%- if messages[0].get("content") -%}
+        {%- set ns.system_prompt = parse_content(messages[0]["content"]) -%}
+    {%- endif -%}
+    {%- set messages = messages[1:] -%}
+{%- endif -%}
+{%- if tools -%}
+    {%- set ns.system_prompt = ns.system_prompt + ("\n" if ns.system_prompt else "") + "List of tools: [" -%}
+    {%- for tool in tools -%}
+        {%- if tool is not string -%}
+            {%- set tool = tool | tojson -%}
+        {%- endif -%}
+        {%- set ns.system_prompt = ns.system_prompt + tool -%}
+        {%- if not loop.last -%}
+            {%- set ns.system_prompt = ns.system_prompt + ", " -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- set ns.system_prompt = ns.system_prompt + "]" -%}
+{%- endif -%}
+{%- if ns.system_prompt -%}
+    {{- "<|im_start|>system\n" + ns.system_prompt + "<|im_end|>\n" -}}
+{%- endif -%}
+{%- for message in messages -%}
+    {%- if message["role"] == "user" -%}
+        {%- set ns.last_user_index = loop.index0 -%}
+    {%- endif -%}
+{%- endfor -%}
+{%- for message in messages -%}
+    {{- "<|im_start|>" + message.role + "\n" -}}
+    {%- if message.role == "assistant" -%}
+        {%- generation -%}
+        {%- set keep_thinking = preserve_thinking or loop.index0 > ns.last_user_index -%}
+        {%- set thinking = message.thinking or message.reasoning or message.reasoning_content -%}
+        {%- set thinking = thinking if thinking is string else "" -%}
+        {%- if thinking and keep_thinking -%}
+            {{- "<think>" + thinking + "</think>" -}}
+        {%- endif -%}
+        {%- set _cfm_tag = "CONTINUE_FINAL_MESSAGE_TAG " -%}
+        {%- set _has_cfm = false -%}
+        {%- set content = "" -%}
+        {%- if message.get("content") -%}
+            {%- set content = parse_content(message.content) -%}
+        {%- endif -%}
+        {%- if not keep_thinking and "</think>" in content -%}
+            {%- set content = content.split("</think>")[-1] | trim -%}
+        {%- endif -%}
+        {%- if content.endswith(_cfm_tag) -%}
+            {%- set _has_cfm = true -%}
+            {%- set _trunc_len = (content | length) - (_cfm_tag | length) -%}
+            {%- set content = content[:_trunc_len] -%}
+        {%- endif -%}
+        {{- content -}}
+        {%- if message.tool_calls -%}
+            {{- render_tool_calls(message.tool_calls) -}}
+        {%- endif -%}
+        {%- if _has_cfm -%}
+            {{- _cfm_tag -}}
+        {%- endif -%}
+        {{- "<|im_end|>\n" -}}
+        {%- endgeneration -%}
+    {%- else %}
+        {%- if message.get("content") -%}
+            {{- parse_content(message["content"]) -}}
+        {%- endif -%}
+        {{- "<|im_end|>\n" -}}
+    {%- endif %}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{- "<|im_start|>assistant\n<think>" -}}
+{%- endif -%}

config.json

@@ -0,0 +1,1104 @@
+{
+  "architectures": [
+    "Fuse3ForCausalLM"
+  ],
+  "block_auto_adjust_ff_dim": false,
+  "block_dim": 2048,
+  "block_ffn_dim_multiplier": 1.0,
+  "block_mlp_init_scale": 1.0,
+  "block_multiple_of": 256,
+  "block_norm_eps": 1e-05,
+  "block_out_init_scale": 1.0,
+  "block_use_swiglu": true,
+  "block_use_xavier_init": true,
+  "bos_token_id": 124894,
+  "coding_enabled": true,
+  "conv_L_cache": 3,
+  "conv_bias": false,
+  "conv_dim": 2048,
+  "conv_use_xavier_init": true,
+  "dtype": "bfloat16",
+  "eos_token_id": 124900,
+  "expert_intermediate_size": 512,
+  "expert_scale_init": -5.0,
+  "experts_per_layer": {
+    "0": [
+      62,
+      75,
+      174,
+      231,
+      211,
+      36,
+      134,
+      159,
+      92,
+      70,
+      81,
+      188,
+      144,
+      138,
+      232,
+      241,
+      123,
+      254,
+      191,
+      128,
+      78,
+      37,
+      223,
+      187,
+      11,
+      124,
+      142,
+      48,
+      117,
+      226,
+      77,
+      116
+    ],
+    "1": [
+      193,
+      254,
+      51,
+      33,
+      80,
+      235,
+      107,
+      66,
+      169,
+      141,
+      13,
+      120,
+      209,
+      75,
+      133,
+      25,
+      40,
+      96,
+      48,
+      153,
+      222,
+      167,
+      162,
+      217,
+      214,
+      55,
+      23,
+      101,
+      74,
+      17,
+      121,
+      185
+    ],
+    "2": [
+      217,
+      39,
+      115,
+      92,
+      20,
+      216,
+      169,
+      164,
+      223,
+      111,
+      252,
+      182,
+      207,
+      191,
+      193,
+      90,
+      153,
+      157,
+      36,
+      10,
+      150,
+      149,
+      34,
+      42,
+      167,
+      103,
+      145,
+      22,
+      128,
+      187,
+      176,
+      41
+    ],
+    "3": [
+      196,
+      6,
+      69,
+      224,
+      201,
+      199,
+      149,
+      163,
+      251,
+      17,
+      36,
+      43,
+      212,
+      154,
+      92,
+      166,
+      119,
+      56,
+      197,
+      75,
+      126,
+      221,
+      109,
+      81,
+      113,
+      93,
+      183,
+      22,
+      31,
+      173,
+      255,
+      100
+    ],
+    "4": [
+      175,
+      253,
+      152,
+      90,
+      28,
+      38,
+      143,
+      36,
+      8,
+      221,
+      189,
+      228,
+      148,
+      160,
+      134,
+      121,
+      157,
+      6,
+      195,
+      235,
+      44,
+      187,
+      140,
+      206,
+      65,
+      46,
+      132,
+      70,
+      156,
+      161,
+      216,
+      94
+    ],
+    "5": [
+      153,
+      58,
+      251,
+      50,
+      70,
+      22,
+      177,
+      60,
+      218,
+      0,
+      77,
+      222,
+      159,
+      210,
+      163,
+      19,
+      146,
+      8,
+      53,
+      42,
+      74,
+      111,
+      172,
+      137,
+      61,
+      64,
+      7,
+      108,
+      136,
+      164,
+      185,
+      122
+    ],
+    "6": [
+      26,
+      12,
+      102,
+      159,
+      210,
+      125,
+      30,
+      212,
+      238,
+      64,
+      168,
+      34,
+      195,
+      249,
+      196,
+      79,
+      90,
+      92,
+      169,
+      186,
+      209,
+      237,
+      43,
+      72,
+      120,
+      60,
+      126,
+      218,
+      105,
+      69,
+      113,
+      202
+    ],
+    "7": [
+      185,
+      78,
+      245,
+      175,
+      227,
+      80,
+      92,
+      49,
+      233,
+      120,
+      174,
+      244,
+      3,
+      142,
+      5,
+      12,
+      104,
+      226,
+      246,
+      132,
+      184,
+      183,
+      240,
+      81,
+      39,
+      160,
+      131,
+      108,
+      75,
+      26,
+      140,
+      187
+    ],
+    "8": [
+      133,
+      72,
+      148,
+      23,
+      53,
+      16,
+      145,
+      224,
+      188,
+      233,
+      85,
+      76,
+      181,
+      195,
+      180,
+      95,
+      121,
+      71,
+      164,
+      249,
+      98,
+      41,
+      22,
+      193,
+      162,
+      1,
+      171,
+      73,
+      99,
+      192,
+      240,
+      151
+    ],
+    "9": [
+      99,
+      207,
+      1,
+      186,
+      161,
+      255,
+      130,
+      81,
+      229,
+      188,
+      68,
+      172,
+      254,
+      32,
+      28,
+      114,
+      137,
+      54,
+      179,
+      199,
+      225,
+      120,
+      95,
+      61,
+      251,
+      178,
+      176,
+      233,
+      140,
+      164,
+      20,
+      168
+    ],
+    "10": [
+      116,
+      132,
+      217,
+      156,
+      12,
+      207,
+      233,
+      146,
+      165,
+      134,
+      234,
+      242,
+      190,
+      148,
+      147,
+      250,
+      149,
+      126,
+      205,
+      176,
+      172,
+      184,
+      23,
+      124,
+      248,
+      137,
+      61,
+      38,
+      108,
+      54,
+      135,
+      14
+    ],
+    "11": [
+      60,
+      44,
+      141,
+      97,
+      32,
+      204,
+      223,
+      130,
+      107,
+      222,
+      37,
+      63,
+      142,
+      15,
+      218,
+      220,
+      230,
+      134,
+      153,
+      170,
+      10,
+      213,
+      98,
+      180,
+      139,
+      33,
+      20,
+      71,
+      215,
+      214,
+      162,
+      143
+    ],
+    "12": [
+      116,
+      250,
+      176,
+      161,
+      163,
+      143,
+      157,
+      10,
+      21,
+      38,
+      244,
+      196,
+      1,
+      15,
+      50,
+      137,
+      189,
+      2,
+      39,
+      186,
+      128,
+      14,
+      166,
+      67,
+      63,
+      211,
+      240,
+      17,
+      13,
+      145,
+      55,
+      115
+    ],
+    "13": [
+      93,
+      212,
+      45,
+      145,
+      219,
+      165,
+      14,
+      83,
+      208,
+      180,
+      183,
+      252,
+      42,
+      120,
+      0,
+      248,
+      151,
+      232,
+      170,
+      166,
+      58,
+      202,
+      66,
+      100,
+      194,
+      34,
+      92,
+      98,
+      158,
+      19,
+      174,
+      144
+    ],
+    "14": [
+      204,
+      44,
+      132,
+      253,
+      73,
+      194,
+      51,
+      164,
+      202,
+      90,
+      248,
+      214,
+      96,
+      166,
+      189,
+      54,
+      21,
+      64,
+      165,
+      17,
+      225,
+      117,
+      187,
+      113,
+      198,
+      67,
+      179,
+      115,
+      63,
+      163,
+      88,
+      167
+    ],
+    "15": [
+      147,
+      12,
+      139,
+      14,
+      169,
+      235,
+      206,
+      238,
+      18,
+      1,
+      98,
+      170,
+      114,
+      145,
+      228,
+      195,
+      50,
+      165,
+      242,
+      227,
+      148,
+      74,
+      134,
+      248,
+      3,
+      215,
+      28,
+      90,
+      21,
+      240,
+      81,
+      0
+    ],
+    "16": [
+      72,
+      172,
+      117,
+      191,
+      183,
+      252,
+      56,
+      112,
+      130,
+      28,
+      59,
+      104,
+      189,
+      155,
+      230,
+      247,
+      160,
+      234,
+      171,
+      195,
+      53,
+      228,
+      25,
+      35,
+      74,
+      116,
+      47,
+      147,
+      209,
+      37,
+      190,
+      5
+    ],
+    "17": [
+      89,
+      69,
+      38,
+      1,
+      185,
+      152,
+      233,
+      20,
+      19,
+      98,
+      209,
+      100,
+      8,
+      58,
+      44,
+      21,
+      62,
+      222,
+      140,
+      219,
+      125,
+      204,
+      117,
+      217,
+      135,
+      223,
+      189,
+      9,
+      54,
+      4,
+      126,
+      169
+    ],
+    "18": [
+      77,
+      91,
+      176,
+      138,
+      43,
+      139,
+      175,
+      126,
+      116,
+      225,
+      69,
+      212,
+      144,
+      245,
+      93,
+      152,
+      236,
+      204,
+      66,
+      171,
+      209,
+      146,
+      92,
+      64,
+      108,
+      13,
+      224,
+      241,
+      54,
+      123,
+      163,
+      223
+    ],
+    "19": [
+      212,
+      59,
+      122,
+      75,
+      115,
+      35,
+      101,
+      88,
+      164,
+      143,
+      203,
+      226,
+      61,
+      95,
+      60,
+      137,
+      99,
+      33,
+      219,
+      183,
+      22,
+      114,
+      116,
+      47,
+      85,
+      248,
+      67,
+      133,
+      231,
+      221,
+      11,
+      25
+    ],
+    "20": [
+      133,
+      72,
+      23,
+      148,
+      53,
+      195,
+      181,
+      224,
+      71,
+      171,
+      145,
+      95,
+      233,
+      76,
+      85,
+      41,
+      98,
+      188,
+      249,
+      99,
+      1,
+      172,
+      192,
+      197,
+      240,
+      116,
+      173,
+      206,
+      193,
+      151,
+      234,
+      121
+    ],
+    "21": [
+      207,
+      99,
+      186,
+      1,
+      188,
+      81,
+      251,
+      130,
+      229,
+      95,
+      172,
+      255,
+      68,
+      254,
+      137,
+      32,
+      199,
+      114,
+      179,
+      225,
+      140,
+      20,
+      28,
+      233,
+      54,
+      72,
+      243,
+      193,
+      178,
+      164,
+      183,
+      211
+    ],
+    "22": [
+      116,
+      156,
+      124,
+      132,
+      12,
+      40,
+      148,
+      217,
+      147,
+      233,
+      134,
+      149,
+      207,
+      165,
+      109,
+      172,
+      205,
+      126,
+      137,
+      250,
+      190,
+      242,
+      234,
+      65,
+      77,
+      182,
+      184,
+      163,
+      135,
+      146,
+      160,
+      108
+    ],
+    "23": [
+      60,
+      204,
+      141,
+      32,
+      44,
+      223,
+      222,
+      107,
+      37,
+      63,
+      213,
+      134,
+      193,
+      142,
+      220,
+      218,
+      15,
+      10,
+      130,
+      61,
+      153,
+      230,
+      139,
+      143,
+      121,
+      59,
+      20,
+      64,
+      162,
+      80,
+      117,
+      170
+    ],
+    "24": [
+      1,
+      116,
+      163,
+      176,
+      250,
+      157,
+      55,
+      143,
+      244,
+      39,
+      50,
+      166,
+      68,
+      10,
+      137,
+      196,
+      161,
+      2,
+      211,
+      38,
+      199,
+      167,
+      14,
+      67,
+      21,
+      158,
+      186,
+      93,
+      238,
+      217,
+      240,
+      99
+    ],
+    "25": [
+      93,
+      145,
+      165,
+      45,
+      208,
+      42,
+      219,
+      212,
+      14,
+      0,
+      47,
+      101,
+      183,
+      202,
+      83,
+      252,
+      166,
+      120,
+      100,
+      232,
+      151,
+      180,
+      155,
+      203,
+      248,
+      58,
+      251,
+      34,
+      170,
+      66,
+      243,
+      233
+    ],
+    "26": [
+      166,
+      204,
+      214,
+      132,
+      194,
+      90,
+      73,
+      202,
+      44,
+      248,
+      54,
+      96,
+      164,
+      67,
+      137,
+      17,
+      64,
+      165,
+      21,
+      117,
+      253,
+      189,
+      78,
+      88,
+      163,
+      113,
+      225,
+      198,
+      38,
+      115,
+      179,
+      94
+    ],
+    "27": [
+      147,
+      139,
+      98,
+      206,
+      12,
+      18,
+      114,
+      169,
+      14,
+      85,
+      1,
+      28,
+      50,
+      242,
+      170,
+      182,
+      240,
+      228,
+      235,
+      195,
+      145,
+      215,
+      134,
+      3,
+      99,
+      70,
+      238,
+      74,
+      183,
+      150,
+      190,
+      103
+    ],
+    "28": [
+      72,
+      152,
+      247,
+      183,
+      191,
+      56,
+      130,
+      59,
+      117,
+      234,
+      172,
+      28,
+      155,
+      112,
+      228,
+      104,
+      74,
+      165,
+      35,
+      252,
+      116,
+      171,
+      230,
+      53,
+      166,
+      37,
+      179,
+      54,
+      210,
+      80,
+      212,
+      238
+    ],
+    "29": [
+      89,
+      1,
+      38,
+      69,
+      185,
+      19,
+      125,
+      248,
+      98,
+      152,
+      100,
+      222,
+      117,
+      9,
+      4,
+      54,
+      204,
+      91,
+      209,
+      140,
+      44,
+      233,
+      167,
+      155,
+      238,
+      219,
+      79,
+      192,
+      21,
+      58,
+      150,
+      236
+    ]
+  },
+  "full_attn_idxs": null,
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 10752,
+  "layer_types": [
+    "conv",
+    "conv",
+    "full_attention",
+    "conv",
+    "conv",
+    "full_attention",
+    "conv",
+    "conv",
+    "conv",
+    "full_attention",
+    "conv",
+    "conv",
+    "conv",
+    "full_attention",
+    "conv",
+    "conv",
+    "conv",
+    "full_attention",
+    "conv",
+    "conv",
+    "conv",
+    "full_attention",
+    "conv",
+    "conv",
+    "full_attention",
+    "conv",
+    "conv",
+    "full_attention",
+    "conv",
+    "conv"
+  ],
+  "load_balance_coef": 0.001,
+  "max_position_embeddings": 131072,
+  "model_type": "fuse3",
+  "norm_eps": 1e-05,
+  "num_attention_heads": 32,
+  "num_augmented_layers": 30,
+  "num_heads": 32,
+  "num_hidden_layers": 30,
+  "num_key_value_heads": 8,
+  "pad_token_id": 124893,
+  "rope_parameters": {
+    "rope_theta": 10000000.0,
+    "rope_type": "default"
+  },
+  "router_init_scale": -2.0,
+  "swiglu_limit": 10.0,
+  "tie_word_embeddings": true,
+  "top_k_experts": 8,
+  "transformers_version": "5.14.1",
+  "use_cache": false,
+  "use_pos_enc": true,
+  "vocab_size": 128000
+}

fuse3_model.py

@@ -0,0 +1,471 @@
+"""Fuse-3 model: LFM2-2.6B host + Qwen3.6-35B-A3B coding experts.
+
+Architecture: per-layer FFN augmentation with Qwen3.6 MoE experts.
+At each augmented host layer, coding experts from Qwen3.6 are added
+alongside the host's native SwiGLU FFN. A learned router decides which
+experts fire.
+
+Key design differences from Fuse-2:
+  - NO bridges needed: LFM2 hidden_size (2048) == Qwen3.6 expert input (2048).
+    Experts are 2048->512->2048, directly operating on host hidden states.
+  - LFM2 is NOT a Qwen model. Lfm2ForCausalLM has a hybrid conv+attention
+    architecture with 22 short-conv layers + 8 GQA layers.
+  - The augmented layer splits LFM2's forward into attention/conv + FFN,
+    then adds expert output as a parallel path to the dense FFN.
+  - expert_scale zero-init -> model starts as exact LFM2-2.6B
+  - Router initialized to low activation -> coding path fires rarely at first
+  - Frozen host, frozen experts -> only router + scale train
+  - use_cache = False initially (correctness first)
+  - Optional SwiGLU clamping (limit=10.0) for Qwen3.6 expert stability
+"""
+from __future__ import annotations
+
+import math
+from typing import Iterator
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import Lfm2Config, Lfm2ForCausalLM
+
+
+class Fuse3Config(Lfm2Config):
+    """LFM2 config extended with Fuse-3 MoE coding expert parameters."""
+
+    model_type = "fuse3"
+
+    def __init__(
+        self,
+        # Expert configuration
+        expert_intermediate_size: int = 512,   # Qwen3.6 expert intermediate
+        experts_per_layer: dict | None = None,  # layer_idx -> list of expert IDs
+        num_augmented_layers: int = 0,
+        top_k_experts: int = 8,                 # Qwen3.6 uses 8 routed
+        # Router configuration
+        router_init_scale: float = -2.0,        # low initial activation
+        load_balance_coef: float = 0.001,
+        # Expert stability
+        swiglu_limit: float = 10.0,             # clamp expert activations (from fuse-2 lessons)
+        coding_enabled: bool = True,
+        # Scale
+        expert_scale_init: float = -5.0,         # softplus(-5)≈0.007, small but active
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.expert_intermediate_size = expert_intermediate_size
+        self.experts_per_layer = experts_per_layer or {}
+        self.num_augmented_layers = num_augmented_layers
+        self.top_k_experts = top_k_experts
+        self.router_init_scale = router_init_scale
+        self.load_balance_coef = load_balance_coef
+        self.swiglu_limit = swiglu_limit
+        self.coding_enabled = coding_enabled
+        self.expert_scale_init = expert_scale_init
+        self.use_cache = False
+        self.architectures = ["Fuse3ForCausalLM"]
+
+
+class SwiGLUExpert(nn.Module):
+    """A single Qwen3.6 MoE expert (SwiGLU FFN).
+
+    gate_proj: (intermediate, hidden)  — w1
+    up_proj:   (intermediate, hidden)  — w3
+    down_proj: (hidden, intermediate)  — w2
+
+    Input:  (..., hidden_size)
+    Output: (..., hidden_size)
+    """
+
+    def __init__(self, hidden_size: int, intermediate_size: int, swiglu_limit: float = 0.0):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+        self.swiglu_limit = swiglu_limit
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate = self.gate_proj(x)
+        up = self.up_proj(x)
+        act = F.silu(gate) * up
+        if self.swiglu_limit > 0:
+            act = act.clamp(-self.swiglu_limit, self.swiglu_limit)
+        return self.down_proj(act)
+
+
+class Fuse3Router(nn.Module):
+    """Per-layer router for coding experts.
+
+    Uses sqrtsoftplus scoring (matching Qwen3.5 MoE's approach) with
+    top-k selection and optional load balancing.
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        num_experts: int,
+        top_k: int = 8,
+        init_scale: float = -2.0,
+    ):
+        super().__init__()
+        self.num_experts = num_experts
+        self.top_k = min(top_k, num_experts)
+        self.gate = nn.Linear(input_dim, num_experts, bias=False)
+        nn.init.normal_(self.gate.weight, mean=0.0, std=0.01)
+        self.init_scale = init_scale
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Route tokens to experts.
+
+        Args:
+            hidden_states: (batch*seq, hidden_size)
+
+        Returns:
+            router_weights: (batch*seq, top_k) — normalized weights for selected experts
+            expert_indices: (batch*seq, top_k) — which experts were selected
+            router_logits: (batch*seq, num_experts) — raw logits for load balancing
+        """
+        logits = self.gate(hidden_states)
+        scores = F.softplus(logits).sqrt()
+        topk_weights, topk_indices = scores.topk(self.top_k, dim=-1)
+        topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-8)
+        return topk_weights, topk_indices, logits
+
+
+class Fuse3AugmentedLayer(nn.Module):
+    """One LFM2 layer augmented with Qwen3.6 coding experts.
+
+    Forward flow (mirrors Lfm2DecoderLayer but with MoE added to FFN):
+      1. Attention or ShortConv (frozen host, same as LFM2)
+      2. Dense FFN (frozen host SwiGLU)
+      3. Router: select top-k coding experts from FFN input
+      4. Experts: parallel SwiGLU computation (frozen, from Qwen3.6)
+      5. expert_scale * expert_output added to residual
+      6. No bridge needed — hidden_size matches on both sides
+
+    The model starts as exact LFM2 (expert_scale=0) and learns to
+    incorporate coding experts through router + scale training.
+    """
+
+    def __init__(
+        self,
+        host_layer: nn.Module,
+        hidden_size: int,
+        expert_intermediate: int,
+        num_experts: int,
+        top_k: int = 8,
+        swiglu_limit: float = 10.0,
+        router_init_scale: float = -2.0,
+        coding_enabled: bool = True,
+        expert_scale_init: float = 0.0,
+    ):
+        super().__init__()
+        self.host_layer = host_layer
+        self.coding_enabled = coding_enabled
+        self.num_experts = num_experts
+        self.top_k = min(top_k, num_experts)
+        self.hidden_size = hidden_size
+
+        # Expose host layer attributes needed by the LFM2 model forward pass
+        self.is_attention_layer = getattr(host_layer, "is_attention_layer", False)
+
+        # Router (operates on hidden_size directly — no bridge)
+        self.router = Fuse3Router(
+            hidden_size, num_experts, top_k, router_init_scale
+        )
+
+        # Experts (frozen, loaded from Qwen3.6)
+        self.experts = nn.ModuleList([
+            SwiGLUExpert(hidden_size, expert_intermediate, swiglu_limit)
+            for _ in range(num_experts)
+        ])
+
+        # Scale factor for expert output (zero-init = exact LFM2)
+        self.expert_scale = nn.Parameter(torch.tensor(expert_scale_init))
+
+        # Store last router logits for load balancing loss
+        self._last_router_logits: torch.Tensor | None = None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values=None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        # ── 1. Run the original host layer (attention/conv + FFN) ──
+        # This delegates to the original Lfm2DecoderLayer.forward, ensuring
+        # perfect compatibility with the host model's conv/attention impl.
+        layer_output = self.host_layer(
+            hidden_states,
+            position_embeddings=position_embeddings,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        # If coding disabled or no experts, return original output unchanged
+        if not self.coding_enabled or self.num_experts == 0:
+            return layer_output
+
+        # ── 2. Always run router (for load balancing gradients) ──
+        original_shape = layer_output.shape
+        h_flat = layer_output.reshape(-1, original_shape[-1])
+
+        topk_weights, expert_indices, router_logits = self.router(h_flat)
+        self._last_router_logits = router_logits
+
+        # Skip expert computation if scale is effectively zero
+        # (router still ran, so LB loss gradients flow)
+        scale = F.softplus(self.expert_scale)
+        if scale.item() < 1e-4:
+            return layer_output
+
+        # ── 3. Compute expert outputs (sparse) ──
+        # Detach expert inputs/outputs — experts are frozen, so no gradients
+        # need to flow through the expert weights. Gradients only flow through
+        # topk_weights (router) and scale (expert_scale).
+        # Use index_add (out-of-place) to avoid in-place op autograd issues.
+        expert_output = torch.zeros_like(h_flat)
+
+        for k in range(self.top_k):
+            indices = expert_indices[:, k]
+            weights = topk_weights[:, k]
+
+            for eid in range(self.num_experts):
+                token_mask = indices == eid
+                if not token_mask.any():
+                    continue
+                expert_in = h_flat[token_mask].detach()
+                expert_out = self.experts[eid](expert_in).detach()
+                weighted_out = weights[token_mask].unsqueeze(-1) * expert_out
+                token_indices = torch.where(token_mask)[0]
+                expert_output = expert_output.index_add(
+                    0, token_indices, weighted_out
+                )
+
+        # ── 4. Normalize expert output to match host activation scale ──
+        # The experts come from Qwen3.6 which has different activation
+        # distributions. Rescale to match host std, but clamp the ratio
+        # to prevent amplification of sparse outputs.
+        host_std = layer_output.std().detach() + 1e-6
+        expert_std = expert_output.std().detach() + 1e-6
+        ratio = torch.clamp(host_std / expert_std, max=2.0)
+        expert_output = expert_output * ratio
+
+        # ── 5. Scale and add to residual ──
+        # Clamp scale to prevent runaway
+        scale = torch.clamp(scale, max=0.1)
+        expert_delta = scale * expert_output
+        expert_delta = expert_delta.reshape(original_shape)
+
+        return layer_output + expert_delta
+
+    def get_router_logits(self) -> torch.Tensor | None:
+        return self._last_router_logits
+
+
+class Fuse3ForCausalLM(Lfm2ForCausalLM):
+    """LFM2-2.6B host + Qwen3.6-35B-A3B coding experts.
+
+    The model starts as an exact LFM2-2.6B (expert_scale=0) and learns
+    to incorporate coding experts through router and scale training.
+    """
+
+    config_class = Fuse3Config
+    _no_split_modules = ["Lfm2DecoderLayer", "Fuse3AugmentedLayer"]
+
+    def __init__(self, config: Fuse3Config):
+        super().__init__(config)
+
+        # Replace specified layers with augmented versions
+        experts_per_layer = config.experts_per_layer or {}
+        augmented_count = 0
+
+        for layer_idx_str, expert_ids in experts_per_layer.items():
+            layer_idx = int(layer_idx_str)
+            if layer_idx >= len(self.model.layers):
+                raise ValueError(
+                    f"Layer {layer_idx} out of range "
+                    f"(model has {len(self.model.layers)} layers)"
+                )
+
+            num_experts = len(expert_ids)
+            if num_experts == 0:
+                continue
+
+            original_layer = self.model.layers[layer_idx]
+            self.model.layers[layer_idx] = Fuse3AugmentedLayer(
+                host_layer=original_layer,
+                hidden_size=config.hidden_size,
+                expert_intermediate=config.expert_intermediate_size,
+                num_experts=num_experts,
+                top_k=min(config.top_k_experts, num_experts),
+                swiglu_limit=config.swiglu_limit,
+                router_init_scale=config.router_init_scale,
+                coding_enabled=config.coding_enabled,
+                expert_scale_init=config.expert_scale_init,
+            )
+            augmented_count += 1
+
+        config.num_augmented_layers = augmented_count
+
+    def set_coding_enabled(self, enabled: bool) -> None:
+        for layer in self.model.layers:
+            if isinstance(layer, Fuse3AugmentedLayer):
+                layer.coding_enabled = enabled
+
+    def get_augmented_layers(self) -> list[tuple[int, Fuse3AugmentedLayer]]:
+        return [
+            (i, layer)
+            for i, layer in enumerate(self.model.layers)
+            if isinstance(layer, Fuse3AugmentedLayer)
+        ]
+
+    def get_trainable_params(self) -> dict[str, nn.Parameter]:
+        trainable = {}
+        for name, param in self.named_parameters():
+            if any(key in name for key in ("router", "expert_scale")):
+                trainable[name] = param
+        return trainable
+
+    def freeze_host_and_experts(self) -> None:
+        for name, param in self.named_parameters():
+            if any(key in name for key in ("router", "expert_scale")):
+                param.requires_grad = True
+            else:
+                param.requires_grad = False
+
+    def get_all_router_logits(self) -> list[torch.Tensor]:
+        """Collect router logits from all augmented layers after forward pass."""
+        logits = []
+        for layer in self.model.layers:
+            if hasattr(layer, '_last_router_logits') and layer._last_router_logits is not None:
+                logits.append(layer._last_router_logits)
+        return logits
+
+    def get_expert_scales(self) -> list[float]:
+        """Get current effective expert scale (softplus) from all augmented layers."""
+        import torch.nn.functional as F
+        scales = []
+        for layer in self.model.layers:
+            if hasattr(layer, 'expert_scale'):
+                scales.append(F.softplus(layer.expert_scale).item())
+        return scales
+
+    def reset_router_logits(self) -> None:
+        """Clear stored router logits (call before each forward pass during training)."""
+        for layer in self.model.layers:
+            if hasattr(layer, '_last_router_logits'):
+                layer._last_router_logits = None
+
+    def count_parameters(self) -> dict[str, int]:
+        counts = {
+            "host": 0, "experts": 0, "routers": 0, "scale": 0,
+            "total": 0, "trainable": 0,
+        }
+        for name, param in self.named_parameters():
+            n = param.numel()
+            counts["total"] += n
+            if param.requires_grad:
+                counts["trainable"] += n
+            if "router" in name:
+                counts["routers"] += n
+            elif "expert_scale" in name:
+                counts["scale"] += n
+            elif "experts" in name:
+                counts["experts"] += n
+            else:
+                counts["host"] += n
+        return counts
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values=None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs,
+    ):
+        # Delegate to Lfm2ForCausalLM.forward with original args.
+        # The augmented layers handle expert computation internally;
+        # KV cache works because Fuse3AugmentedLayer delegates to the
+        # original host layer's forward for attention/conv.
+        return super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+
+def load_expert_weights(
+    model: Fuse3ForCausalLM,
+    expert_dir: str,
+    expert_mapping: dict[int, list[int]],
+) -> dict:
+    """Load extracted Qwen3.6 expert weights into the Fuse3 model.
+
+    Args:
+        model: Fuse3 model with augmented layers
+        expert_dir: directory containing expert safetensors
+        expert_mapping: layer_idx -> list of expert IDs (matching selection order)
+
+    Returns:
+        Manifest of loaded tensors with hash verification
+    """
+    from safetensors.torch import load_file
+    import glob
+
+    shard_files = sorted(glob.glob(f"{expert_dir}/experts-*.safetensors"))
+    if not shard_files:
+        raise FileNotFoundError(f"No expert shards found in {expert_dir}")
+
+    all_tensors = {}
+    for shard in shard_files:
+        all_tensors.update(load_file(shard))
+
+    loaded = {}
+    for layer_idx, expert_ids in expert_mapping.items():
+        if layer_idx >= len(model.model.layers):
+            continue  # skip out-of-range layers
+        augmented = model.model.layers[layer_idx]
+        if not isinstance(augmented, Fuse3AugmentedLayer):
+            continue  # skip non-augmented layers
+
+        for local_idx, global_eid in enumerate(expert_ids):
+            prefix = f"layer{layer_idx:02d}_expert{global_eid:03d}"
+
+            for pname in ("gate_proj.weight", "up_proj.weight", "down_proj.weight"):
+                key = f"{prefix}.{pname}"
+                if key not in all_tensors:
+                    raise KeyError(f"Missing expert tensor: {key}")
+
+                tensor = all_tensors[key]
+                parts = pname.split(".")
+                module = augmented.experts[local_idx]
+                for part in parts[:-1]:
+                    module = getattr(module, part)
+                param = getattr(module, parts[-1])
+                param.data.copy_(tensor.to(param.dtype))
+
+                loaded[key] = {
+                    "shape": list(tensor.shape),
+                    "destination": f"layers.{layer_idx}.experts.{local_idx}.{pname}",
+                }
+
+    return loaded

generation_config.json

@@ -0,0 +1,16 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 124894,
+  "do_sample": true,
+  "eos_token_id": [
+    124900
+  ],
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 124893,
+  "repetition_penalty": 1.1,
+  "temperature": 0.1,
+  "top_k": 50,
+  "transformers_version": "5.14.1",
+  "use_cache": true
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:233662cf2cbf7a93b6fdbf60068de0d8f122c51f0f36147a2319a597f2e99d56
+size 11438511844

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b124e06b0d81002864892c3ebb866effe115aa1d5fbe8e6e9e23254655354a8
+size 17905696

tokenizer_config.json

@@ -0,0 +1,13 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "is_local": false,
+  "legacy": false,
+  "local_files_only": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "use_default_system_prompt": false
+}