From cd491068836b754dd8a0cbb96eea16f0eb59603c Mon Sep 17 00:00:00 2001 From: Ryan Gillespie Date: Mon, 21 Sep 2026 23:24:51 +0200 Subject: [PATCH 1/4] architecture: canonical parameter-name mapping (first increment of #708) A shared role-key space so parameters from different architecture families can be addressed uniformly before any merge method runs: model.layers.0.self_attn.q_proj.weight -> layer_0.attn_q.weight model.layers.11.mlp.up_proj.weight -> layer_11.ffn_up.weight Dependency-free name mapping only: values pass through by reference, no tensor ops, no new requirements. Ships the Llama-family detector (Llama, Qwen, Mistral, Gemma, TinyLlama layouts) with detection, canonicalization, and an exact decanonicalize inverse (round-trip tested). Remaining families (gpt2, bert, neox, opt, t5, phi) and the shape bridge follow in the same pattern per #708. The detector set is exercised end-to-end by the cross-family merge pipeline behind the published Optitransfer converge collectives (9 models across 4 architecture families in one set of weights). Signed-off-by: Ryan Gillespie --- mergekit/architecture/canonical.py | 144 +++++++++++++++++++++++++++++ tests/test_canonical.py | 88 ++++++++++++++++++ 2 files changed, 232 insertions(+) create mode 100644 mergekit/architecture/canonical.py create mode 100644 tests/test_canonical.py diff --git a/mergekit/architecture/canonical.py b/mergekit/architecture/canonical.py new file mode 100644 index 00000000..9e98e855 --- /dev/null +++ b/mergekit/architecture/canonical.py @@ -0,0 +1,144 @@ +# Copyright 2026 Ryan Gillespie (mgillr). Apache-2.0 (mergekit licence). +"""Canonical parameter-name mapping for cross-architecture merging. + +First increment of the cross-architecture proposal (#708): a shared +role-key space so that parameters from different architecture families +can be addressed uniformly before any merge method runs. + + model.layers.0.self_attn.q_proj.weight -> layer_0.attn_q.weight + model.layers.11.mlp.up_proj.weight -> layer_11.ffn_up.weight + model.embed_tokens.weight -> embed_tokens.weight + model.norm.weight -> final_norm.weight + +This module is deliberately dependency-free: it maps names to names and +passes values through by reference. The Llama-family layout (Llama, +Qwen, Mistral, Gemma, TinyLlama and friends) is the highest-traffic +case and ships first; the remaining families from the source system +(gpt2, bert, neox, opt, t5, phi) follow in the same pattern once the +module's home is settled. The shape bridge (pad/truncate + optional +per-key alignment) is a separate increment. + +The detector set is exercised end-to-end by the cross-family merge +pipeline that produced the published `Optitransfer/*converge*` +collectives (9 models across 4 architecture families, merged into one +set of weights). +""" +from __future__ import annotations + +import re +from typing import Any, Mapping + +# Canonical roles, one schema for every family: +DECODER_ROLES = ( + "pre_attn_norm", + "attn_q", + "attn_k", + "attn_v", + "attn_o", + "pre_ffn_norm", + "ffn_gate", + "ffn_up", + "ffn_down", +) + +GLOBAL_KEYS = ("embed_tokens.weight", "final_norm.weight", "head_lm.weight") + +_LAYER_RE = re.compile(r"^layer_(\d+)\.([a-z_]+)\.(weight|bias)$") + + +def layer_key(layer_idx: int, role: str, param: str = "weight") -> str: + """Canonical key for a per-layer role: ``layer_3.attn_q.weight``.""" + return f"layer_{layer_idx}.{role}.{param}" + + +def global_key(name: str) -> str: + """Canonical key for a non-layer parameter.""" + return name + + +def is_canonical(key: str) -> bool: + m = _LAYER_RE.match(key) + if m and m.group(2) in DECODER_ROLES: + return True + return key in GLOBAL_KEYS + + +class LlamaFamilyDetector: + """Llama-layout checkpoints: Llama, Qwen, Mistral, Gemma, TinyLlama. + + Detects the ``model.layers.N.self_attn.*`` layout (distinguished + from OPT's ``model.decoder.layers.``) and rewrites it into the + canonical schema. Pure name remapping; values pass through by + reference and unmapped keys (rotary buffers, position ids, optimizer + state) are dropped. + """ + + name = "llama_family" + family = "llama" + + _layer_re = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + _TAIL = { + "input_layernorm.weight": ("pre_attn_norm", "weight"), + "post_attention_layernorm.weight": ("pre_ffn_norm", "weight"), + "self_attn.q_proj.weight": ("attn_q", "weight"), + "self_attn.q_proj.bias": ("attn_q", "bias"), + "self_attn.k_proj.weight": ("attn_k", "weight"), + "self_attn.k_proj.bias": ("attn_k", "bias"), + "self_attn.v_proj.weight": ("attn_v", "weight"), + "self_attn.v_proj.bias": ("attn_v", "bias"), + "self_attn.o_proj.weight": ("attn_o", "weight"), + "self_attn.o_proj.bias": ("attn_o", "bias"), + "mlp.gate_proj.weight": ("ffn_gate", "weight"), + "mlp.up_proj.weight": ("ffn_up", "weight"), + "mlp.down_proj.weight": ("ffn_down", "weight"), + } + + _GLOBALS = { + "model.embed_tokens.weight": "embed_tokens.weight", + "model.norm.weight": "final_norm.weight", + "lm_head.weight": "head_lm.weight", + } + + def detect(self, state_dict: Mapping[str, Any]) -> bool: + if any(k.startswith("model.decoder.") for k in state_dict): + return False # OPT layout + return any(k.startswith("model.layers.") for k in state_dict) and any( + ".self_attn.q_proj." in k for k in state_dict + ) + + def canonicalize( + self, state_dict: Mapping[str, Any] + ) -> dict[str, Any]: + out: dict[str, Any] = {} + for key, value in state_dict.items(): + g = self._GLOBALS.get(key) + if g is not None: + out[g] = value + continue + m = self._layer_re.match(key) + if m is None: + continue + tail = self._TAIL.get(m.group(2)) + if tail is not None: + out[layer_key(int(m.group(1)), tail[0], tail[1])] = value + return out + + def decanonicalize( + self, canonical: Mapping[str, Any] + ) -> dict[str, Any]: + """Inverse of :meth:`canonicalize` (round-trippable).""" + rev_tail = {v: k for k, v in self._TAIL.items()} + rev_glob = {v: k for k, v in self._GLOBALS.items()} + out: dict[str, Any] = {} + for key, value in canonical.items(): + g = rev_glob.get(key) + if g is not None: + out[g] = value + continue + m = _LAYER_RE.match(key) + if m and (m.group(2), m.group(3)) in rev_tail: + out[ + f"model.layers.{m.group(1)}.{rev_tail[(m.group(2), m.group(3))]}" + ] = value + return out diff --git a/tests/test_canonical.py b/tests/test_canonical.py new file mode 100644 index 00000000..8bce1813 --- /dev/null +++ b/tests/test_canonical.py @@ -0,0 +1,88 @@ +import pytest + +try: + from mergekit.architecture.canonical import ( + LlamaFamilyDetector, + is_canonical, + layer_key, + ) +except ImportError: # torch not installed: load the module standalone + import importlib.util + from pathlib import Path + + _spec = importlib.util.spec_from_file_location( + "canonical", + Path(__file__).parents[1] / "mergekit" / "architecture" / "canonical.py", + ) + _mod = importlib.util.module_from_spec(_spec) + _spec.loader.exec_module(_mod) + LlamaFamilyDetector = _mod.LlamaFamilyDetector + is_canonical = _mod.is_canonical + layer_key = _mod.layer_key + + +@pytest.fixture +def llama_sd(): + return { + "model.embed_tokens.weight": object(), + "model.layers.0.self_attn.q_proj.weight": object(), + "model.layers.0.self_attn.k_proj.bias": object(), + "model.layers.0.input_layernorm.weight": object(), + "model.layers.0.mlp.up_proj.weight": object(), + "model.layers.11.self_attn.o_proj.weight": object(), + "model.layers.11.post_attention_layernorm.weight": object(), + "model.norm.weight": object(), + "lm_head.weight": object(), + "model.layers.0.self_attn.rotary_emb.inv_freq": object(), # dropped + } + + +def test_detect_accepts_llama_layout(llama_sd): + assert LlamaFamilyDetector().detect(llama_sd) + + +def test_detect_rejects_opt_layout(): + assert not LlamaFamilyDetector().detect( + {"model.decoder.layers.0.self_attn.q_proj.weight": object()} + ) + + +def test_detect_rejects_gpt2_layout(): + assert not LlamaFamilyDetector().detect( + {"transformer.h.0.attn.c_attn.weight": object()} + ) + + +def test_canonicalize_maps_roles(llama_sd): + out = LlamaFamilyDetector().canonicalize(llama_sd) + assert out["layer_0.attn_q.weight"] is llama_sd[ + "model.layers.0.self_attn.q_proj.weight" + ] + assert out["layer_11.attn_o.weight"] is llama_sd[ + "model.layers.11.self_attn.o_proj.weight" + ] + assert out["embed_tokens.weight"] is llama_sd["model.embed_tokens.weight"] + assert out["final_norm.weight"] is llama_sd["model.norm.weight"] + + +def test_canonicalize_drops_buffers(llama_sd): + out = LlamaFamilyDetector().canonicalize(llama_sd) + assert not any("inv_freq" in k for k in out) + assert all(is_canonical(k) for k in out) + + +def test_round_trip(llama_sd): + det = LlamaFamilyDetector() + canon = det.canonicalize(llama_sd) + back = det.decanonicalize(canon) + # exact key-set and value identity on the mappable subset; + # dropped buffers (inv_freq) are gone by contract + expected = {k: v for k, v in llama_sd.items() if "inv_freq" not in k} + assert back == expected + + +def test_layer_key_shape(): + assert layer_key(3, "ffn_down") == "layer_3.ffn_down.weight" + assert is_canonical("layer_3.ffn_down.weight") + assert not is_canonical("layer_3.not_a_role.weight") + assert not is_canonical("model.layers.3.mlp.down_proj.weight") From ae930973bb19c54b8043b673a0738d3f2a4f1b3a Mon Sep 17 00:00:00 2001 From: Ryan Gillespie Date: Mon, 21 Sep 2026 23:36:02 +0200 Subject: [PATCH 2/4] architecture: bert and opt detectors (second increment of #708) BertFamilyDetector covers BERT and RoBERTa encoder layouts (post-norm LayerNorms mapped onto the post_attn_norm / post_ffn_norm canonical slots); OPTDetector covers the OPT decoder layout. Same pattern as the Llama-family detector in #709: pure name mapping, values by reference, no new dependencies. Cross-detector exclusivity tested (a Llama-layout dict matches neither). Signed-off-by: Ryan Gillespie --- mergekit/architecture/canonical.py | 117 ++++++++++++++++++++++++++++- tests/test_canonical.py | 43 +++++++++++ 2 files changed, 159 insertions(+), 1 deletion(-) diff --git a/mergekit/architecture/canonical.py b/mergekit/architecture/canonical.py index 9e98e855..5dd55f3b 100644 --- a/mergekit/architecture/canonical.py +++ b/mergekit/architecture/canonical.py @@ -39,9 +39,18 @@ "ffn_gate", "ffn_up", "ffn_down", + "post_attn_norm", + "post_ffn_norm", ) -GLOBAL_KEYS = ("embed_tokens.weight", "final_norm.weight", "head_lm.weight") +GLOBAL_KEYS = ( + "embed_tokens.weight", + "embed_positions.weight", + "final_norm.weight", + "final_norm.bias", + "head_lm.weight", + "head_lm.bias", +) _LAYER_RE = re.compile(r"^layer_(\d+)\.([a-z_]+)\.(weight|bias)$") @@ -142,3 +151,109 @@ def decanonicalize( f"model.layers.{m.group(1)}.{rev_tail[(m.group(2), m.group(3))]}" ] = value return out + + +class BertFamilyDetector: + """BERT / RoBERTa encoder layout (post-norm; DistilBERT excluded). + + Post-norm LayerNorms map onto the post_attn_norm / post_ffn_norm + slots so comparison within a BERT set is consistent. + """ + + name = "bert" + family = "bert" + + _layer_re = re.compile(r"^(?:bert|roberta)\.encoder\.layer\.(\d+)\.(.+)$") + + _TAIL = { + "attention.self.query.weight": ("attn_q", "weight"), + "attention.self.query.bias": ("attn_q", "bias"), + "attention.self.key.weight": ("attn_k", "weight"), + "attention.self.key.bias": ("attn_k", "bias"), + "attention.self.value.weight": ("attn_v", "weight"), + "attention.self.value.bias": ("attn_v", "bias"), + "attention.output.dense.weight": ("attn_o", "weight"), + "attention.output.dense.bias": ("attn_o", "bias"), + "attention.output.LayerNorm.weight": ("post_attn_norm", "weight"), + "attention.output.LayerNorm.bias": ("post_attn_norm", "bias"), + "intermediate.dense.weight": ("ffn_up", "weight"), + "intermediate.dense.bias": ("ffn_up", "bias"), + "output.dense.weight": ("ffn_down", "weight"), + "output.dense.bias": ("ffn_down", "bias"), + "output.LayerNorm.weight": ("post_ffn_norm", "weight"), + "output.LayerNorm.bias": ("post_ffn_norm", "bias"), + } + + def detect(self, state_dict: Mapping[str, Any]) -> bool: + return any( + k.startswith(("bert.", "roberta.")) for k in state_dict + ) and any(".attention.self.query." in k for k in state_dict) + + def canonicalize(self, state_dict: Mapping[str, Any]) -> dict[str, Any]: + out: dict[str, Any] = {} + for key, value in state_dict.items(): + if key.endswith("embeddings.word_embeddings.weight"): + out["embed_tokens.weight"] = value + continue + m = self._layer_re.match(key) + if m is None: + continue + tail = self._TAIL.get(m.group(2)) + if tail is not None: + out[layer_key(int(m.group(1)), tail[0], tail[1])] = value + return out + + +class OPTDetector: + """OPT decoder layout: ``model.decoder.layers.N.*``.""" + + name = "opt" + family = "opt" + + _layer_re = re.compile(r"^model\.decoder\.layers\.(\d+)\.(.+)$") + + _TAIL = { + "self_attn_layer_norm.weight": ("pre_attn_norm", "weight"), + "self_attn_layer_norm.bias": ("pre_attn_norm", "bias"), + "self_attn.q_proj.weight": ("attn_q", "weight"), + "self_attn.q_proj.bias": ("attn_q", "bias"), + "self_attn.k_proj.weight": ("attn_k", "weight"), + "self_attn.k_proj.bias": ("attn_k", "bias"), + "self_attn.v_proj.weight": ("attn_v", "weight"), + "self_attn.v_proj.bias": ("attn_v", "bias"), + "self_attn.out_proj.weight": ("attn_o", "weight"), + "self_attn.out_proj.bias": ("attn_o", "bias"), + "final_layer_norm.weight": ("pre_ffn_norm", "weight"), # per-layer fc norm name in OPT + "fc1.weight": ("ffn_up", "weight"), + "fc1.bias": ("ffn_up", "bias"), + "fc2.weight": ("ffn_down", "weight"), + "fc2.bias": ("ffn_down", "bias"), + } + + _GLOBALS = { + "model.decoder.embed_tokens.weight": "embed_tokens.weight", + "model.decoder.embed_positions.weight": "embed_positions.weight", + "model.decoder.final_layer_norm.weight": "final_norm.weight", + "model.decoder.final_layer_norm.bias": "final_norm.bias", + "lm_head.weight": "head_lm.weight", + } + + def detect(self, state_dict: Mapping[str, Any]) -> bool: + return any( + k.startswith("model.decoder.layers.") for k in state_dict + ) and any(".self_attn.q_proj." in k for k in state_dict) + + def canonicalize(self, state_dict: Mapping[str, Any]) -> dict[str, Any]: + out: dict[str, Any] = {} + for key, value in state_dict.items(): + g = self._GLOBALS.get(key) + if g is not None: + out[g] = value + continue + m = self._layer_re.match(key) + if m is None: + continue + tail = self._TAIL.get(m.group(2)) + if tail is not None: + out[layer_key(int(m.group(1)), tail[0], tail[1])] = value + return out diff --git a/tests/test_canonical.py b/tests/test_canonical.py index 8bce1813..5cccea04 100644 --- a/tests/test_canonical.py +++ b/tests/test_canonical.py @@ -19,6 +19,7 @@ LlamaFamilyDetector = _mod.LlamaFamilyDetector is_canonical = _mod.is_canonical layer_key = _mod.layer_key + _canon = _mod @pytest.fixture @@ -86,3 +87,45 @@ def test_layer_key_shape(): assert is_canonical("layer_3.ffn_down.weight") assert not is_canonical("layer_3.not_a_role.weight") assert not is_canonical("model.layers.3.mlp.down_proj.weight") + + +def test_bert_detect_and_map(): + sd = { + "bert.embeddings.word_embeddings.weight": object(), + "bert.encoder.layer.0.attention.self.query.weight": object(), + "bert.encoder.layer.0.attention.output.LayerNorm.weight": object(), + "bert.encoder.layer.1.output.dense.bias": object(), + } + assert _canon.BertFamilyDetector().detect(sd) + out = _canon.BertFamilyDetector().canonicalize(sd) + assert out["layer_0.attn_q.weight"] is sd[ + "bert.encoder.layer.0.attention.self.query.weight" + ] + assert out["layer_0.post_attn_norm.weight"] is sd[ + "bert.encoder.layer.0.attention.output.LayerNorm.weight" + ] + assert out["layer_1.ffn_down.bias"] is sd["bert.encoder.layer.1.output.dense.bias"] + assert all(is_canonical(k) for k in out) + + +def test_opt_detect_and_map(): + sd = { + "model.decoder.embed_tokens.weight": object(), + "model.decoder.layers.0.self_attn.q_proj.weight": object(), + "model.decoder.layers.2.fc1.weight": object(), + "model.decoder.final_layer_norm.weight": object(), + } + assert _canon.OPTDetector().detect(sd) + out = _canon.OPTDetector().canonicalize(sd) + assert out["layer_0.attn_q.weight"] is sd[ + "model.decoder.layers.0.self_attn.q_proj.weight" + ] + assert out["layer_2.ffn_up.weight"] is sd["model.decoder.layers.2.fc1.weight"] + assert out["final_norm.weight"] is sd["model.decoder.final_layer_norm.weight"] + assert all(is_canonical(k) for k in out) + + +def test_cross_detector_exclusivity(): + llama = {"model.layers.0.self_attn.q_proj.weight": object()} + assert not _canon.OPTDetector().detect(llama) + assert not _canon.BertFamilyDetector().detect(llama) From b1dbd503bf35a97d3077ee2ea4db12daeb757611 Mon Sep 17 00:00:00 2001 From: Ryan Gillespie Date: Mon, 21 Sep 2026 23:43:33 +0200 Subject: [PATCH 3/4] fix(bugbot): map OPT per-layer FFN norm bias; bind _canon in both test import branches Signed-off-by: Ryan Gillespie --- mergekit/architecture/canonical.py | 1 + tests/test_canonical.py | 8 ++++++++ 2 files changed, 9 insertions(+) diff --git a/mergekit/architecture/canonical.py b/mergekit/architecture/canonical.py index 5dd55f3b..57d7a55c 100644 --- a/mergekit/architecture/canonical.py +++ b/mergekit/architecture/canonical.py @@ -224,6 +224,7 @@ class OPTDetector: "self_attn.out_proj.weight": ("attn_o", "weight"), "self_attn.out_proj.bias": ("attn_o", "bias"), "final_layer_norm.weight": ("pre_ffn_norm", "weight"), # per-layer fc norm name in OPT + "final_layer_norm.bias": ("pre_ffn_norm", "bias"), "fc1.weight": ("ffn_up", "weight"), "fc1.bias": ("ffn_up", "bias"), "fc2.weight": ("ffn_down", "weight"), diff --git a/tests/test_canonical.py b/tests/test_canonical.py index 5cccea04..30326126 100644 --- a/tests/test_canonical.py +++ b/tests/test_canonical.py @@ -122,6 +122,14 @@ def test_opt_detect_and_map(): ] assert out["layer_2.ffn_up.weight"] is sd["model.decoder.layers.2.fc1.weight"] assert out["final_norm.weight"] is sd["model.decoder.final_layer_norm.weight"] + +def test_opt_maps_ffn_norm_bias(): + # regression: Bugbot finding — per-layer FFN LayerNorm bias was dropped + sd = {"model.decoder.layers.0.final_layer_norm.bias": object()} + out = _canon.OPTDetector().canonicalize(sd) + assert out["layer_0.pre_ffn_norm.bias"] is sd[ + "model.decoder.layers.0.final_layer_norm.bias" + ] assert all(is_canonical(k) for k in out) From f3b02aa3d16c4b3d1054986aaee9c4e7e4ca4e6a Mon Sep 17 00:00:00 2001 From: Ryan Gillespie Date: Mon, 21 Sep 2026 23:43:57 +0200 Subject: [PATCH 4/4] fix(bugbot): bind _canon in the installed-import branch too (NameError in CI) Signed-off-by: Ryan Gillespie --- tests/test_canonical.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_canonical.py b/tests/test_canonical.py index 30326126..1c7e7b9d 100644 --- a/tests/test_canonical.py +++ b/tests/test_canonical.py @@ -6,6 +6,7 @@ is_canonical, layer_key, ) + import mergekit.architecture.canonical as _canon # noqa: F401 except ImportError: # torch not installed: load the module standalone import importlib.util from pathlib import Path