From 420a2f1ce127596f86d82c0cca9af9282c988a8b Mon Sep 17 00:00:00 2001 From: Naomi Date: Thu, 14 May 2026 12:15:07 -0400 Subject: [PATCH 1/3] Add mergekit-check-compat pre-merge compatibility checker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a new CLI tool that loads model configs and tokenizers (no weights, no GPU required) and reports structured errors/warnings before a merge begins, so incompatibilities are caught early rather than mid-merge or in corrupt outputs. Checks performed: - Tensor shape params (hidden_size, num_layers, num_heads, intermediate_size) → ERROR if any differ; merge will fail at weight level - model_type mismatch → WARNING; shapes may match but semantics may not - RoPE theta and rope_scaling divergence → WARNING - Vocabulary size mismatch with FIM token drop advice → WARNING - FIM token presence in one model but not others → WARNING - Chat template / instruct-model mixing → INFO Handles multimodal wrapper configs (e.g. Mistral3, LLaVA) by transparently unwrapping text_config when top-level arch params are absent. Also handles rope_theta embedded inside rope_scaling dicts (Ministral3 style). --- mergekit/scripts/check_compat.py | 353 +++++++++++++++++++++++++++++++ pyproject.toml | 1 + 2 files changed, 354 insertions(+) create mode 100644 mergekit/scripts/check_compat.py diff --git a/mergekit/scripts/check_compat.py b/mergekit/scripts/check_compat.py new file mode 100644 index 00000000..caa17219 --- /dev/null +++ b/mergekit/scripts/check_compat.py @@ -0,0 +1,353 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only + +import logging +import sys +import textwrap +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple + +import click +import yaml +from transformers import AutoTokenizer + +from mergekit.common import ModelReference +from mergekit.config import MergeConfiguration +from mergekit.options import MergeOptions, PrettyPrintHelp, add_merge_options + +LOG = logging.getLogger(__name__) + +FIM_TOKENS = ["
", "", "", ""]
+
+
+@dataclass
+class Issue:
+    severity: str  # ERROR, WARNING, INFO
+    message: str
+
+
+def _short_name(ref: ModelReference) -> str:
+    return ref.model.path.rstrip("/").split("/")[-1]
+
+
+def _resolve_config(cfg: Any) -> Any:
+    """For multimodal configs, unwrap to text_config to access architecture params."""
+    if getattr(cfg, "hidden_size", None) is None and hasattr(cfg, "text_config"):
+        return cfg.text_config
+    return cfg
+
+
+def _safe_get(cfg: Any, attr: str, default: Any = None) -> Any:
+    return getattr(cfg, attr, default)
+
+
+def _get_rope_theta(cfg: Any) -> Optional[float]:
+    if hasattr(cfg, "rope_theta"):
+        return cfg.rope_theta
+    if hasattr(cfg, "rope_parameters") and isinstance(cfg.rope_parameters, dict):
+        return cfg.rope_parameters.get("rope_theta")
+    # Some configs (e.g. Ministral3) embed rope_theta inside rope_scaling
+    scaling = getattr(cfg, "rope_scaling", None)
+    if isinstance(scaling, dict) and "rope_theta" in scaling:
+        return scaling["rope_theta"]
+    return None
+
+
+def check_architecture(
+    configs: Dict[str, Any],
+) -> Tuple[Dict[str, Dict], List[Issue]]:
+    shape_params = [
+        "hidden_size",
+        "num_hidden_layers",
+        "num_attention_heads",
+        "num_key_value_heads",
+        "intermediate_size",
+    ]
+    table: Dict[str, Dict] = {}
+    issues: List[Issue] = []
+    resolved = {name: _resolve_config(cfg) for name, cfg in configs.items()}
+
+    for param in shape_params:
+        vals = {name: _safe_get(cfg, param) for name, cfg in resolved.items()}
+        table[param] = vals
+        unique = {v for v in vals.values() if v is not None}
+        if len(unique) > 1:
+            detail = ", ".join(f"{n}={v}" for n, v in vals.items())
+            issues.append(
+                Issue(
+                    "ERROR",
+                    f"Tensor shape mismatch in `{param}` ({detail}). "
+                    "Merge will fail or produce garbage output.",
+                )
+            )
+
+    model_types = {name: _safe_get(cfg, "model_type") for name, cfg in resolved.items()}
+    table["model_type"] = model_types
+    unique_types = {v for v in model_types.values() if v is not None}
+    if len(unique_types) > 1:
+        detail = ", ".join(f"{n}={v}" for n, v in model_types.items())
+        issues.append(
+            Issue(
+                "WARNING",
+                f"Model type mismatch ({detail}). Models may differ in activation "
+                "functions, attention variants, or normalization — merged weights may "
+                "be semantically incoherent even if tensor shapes match.",
+            )
+        )
+
+    return table, issues
+
+
+def check_rope(configs: Dict[str, Any]) -> Tuple[Dict[str, Dict], List[Issue]]:
+    table: Dict[str, Dict] = {}
+    issues: List[Issue] = []
+    resolved = {name: _resolve_config(cfg) for name, cfg in configs.items()}
+
+    thetas = {name: _get_rope_theta(cfg) for name, cfg in resolved.items()}
+    table["rope_theta"] = thetas
+    unique_thetas = {v for v in thetas.values() if v is not None}
+    if len(unique_thetas) > 1:
+        detail = ", ".join(f"{n}={v}" for n, v in thetas.items())
+        issues.append(
+            Issue(
+                "WARNING",
+                f"RoPE theta mismatch ({detail}). Positional encodings are "
+                "incompatible; merged model will likely have degraded coherence, "
+                "especially for longer outputs.",
+            )
+        )
+
+    scalings = {name: _safe_get(cfg, "rope_scaling") for name, cfg in resolved.items()}
+    table["rope_scaling"] = {
+        n: (s.get("type", "?") if isinstance(s, dict) else str(s) if s else "none")
+        for n, s in scalings.items()
+    }
+    scaling_types = {
+        n: (s.get("type") if isinstance(s, dict) else None) for n, s in scalings.items()
+    }
+    if any(v is not None for v in scaling_types.values()):
+        shown = {n: (v or "none") for n, v in scaling_types.items()}
+        if len(set(shown.values())) > 1:
+            issues.append(
+                Issue(
+                    "WARNING",
+                    f"RoPE scaling mismatch: {shown}. The output config inherits the "
+                    "base model's RoPE settings. If the base model lacks rope_scaling, "
+                    "the merged model may produce garbled output at longer sequence "
+                    "lengths. Manually add rope_scaling to the output config.json if needed.",
+                )
+            )
+
+    return table, issues
+
+
+def check_vocab(tokenizers: Dict[str, Any]) -> Tuple[Dict[str, Dict], List[Issue]]:
+    table: Dict[str, Dict] = {}
+    issues: List[Issue] = []
+
+    vocab_sizes = {name: tok.vocab_size for name, tok in tokenizers.items()}
+    table["vocab_size"] = vocab_sizes
+    unique_sizes = set(vocab_sizes.values())
+    if len(unique_sizes) > 1:
+        detail = ", ".join(f"{n}={v}" for n, v in vocab_sizes.items())
+        larger_has_fim = any(
+            vocab_sizes[n] > min(unique_sizes)
+            and any(tok in tokenizers[n].get_vocab() for tok in FIM_TOKENS)
+            for n in tokenizers
+        )
+        msg = (
+            f"Vocabulary size mismatch ({detail}). MergeKit will truncate to the "
+            "base model's vocab. "
+        )
+        if larger_has_fim:
+            msg += (
+                "FIM tokens (
, , , ) in the larger-vocab model "
+                "will be dropped. Consider: resize_tok_vocab.py or set "
+                "`tokenizer_source: union` in config."
+            )
+        issues.append(Issue("WARNING", msg))
+
+    return table, issues
+
+
+def check_fim_tokens(
+    tokenizers: Dict[str, Any],
+) -> Tuple[Dict[str, Dict], List[Issue]]:
+    table: Dict[str, Dict] = {}
+    issues: List[Issue] = []
+
+    fim_presence = {
+        name: any(tok in tokenizer.get_vocab() for tok in FIM_TOKENS)
+        for name, tokenizer in tokenizers.items()
+    }
+    table["FIM tokens"] = {n: "Yes" if v else "No" for n, v in fim_presence.items()}
+
+    if len(set(fim_presence.values())) > 1:
+        fim_models = [n for n, v in fim_presence.items() if v]
+        issues.append(
+            Issue(
+                "WARNING",
+                f"{', '.join(fim_models)} contains FIM tokens (
, , , "
+                ") but other models do not. FIM tokens may appear in natural "
+                "language completions after merging. Consider adding them to a "
+                "bad_words list at inference time.",
+            )
+        )
+
+    return table, issues
+
+
+def check_chat_template(
+    tokenizers: Dict[str, Any],
+) -> Tuple[Dict[str, Dict], List[Issue]]:
+    table: Dict[str, Dict] = {}
+    issues: List[Issue] = []
+
+    has_template = {
+        name: tok.chat_template is not None for name, tok in tokenizers.items()
+    }
+    table["chat_template"] = {n: "Yes" if v else "No" for n, v in has_template.items()}
+
+    if all(has_template.values()):
+        issues.append(
+            Issue(
+                "INFO",
+                "All models have chat templates. Merging two instruct models may "
+                "cause [INST]/[/INST] token bleed in outputs.",
+            )
+        )
+    elif any(has_template.values()):
+        instruct_models = [n for n, v in has_template.items() if v]
+        issues.append(
+            Issue(
+                "INFO",
+                f"{', '.join(instruct_models)} has a chat template but other models "
+                "do not. Output format may be unpredictable; consider using the base "
+                "model's tokenizer.",
+            )
+        )
+
+    return table, issues
+
+
+_STATUS = {"match": "✓", "warn": "⚠", "error": "✗"}
+
+
+def _format_table(rows: Dict[str, Dict], model_names: List[str]) -> str:
+    col_w = max(max(len(n) for n in model_names) + 2, 18)
+    label_w = max(max(len(k) for k in rows) + 2, 20)
+
+    lines = [" " * label_w + "".join(n.ljust(col_w) for n in model_names)]
+    for param, vals in rows.items():
+        row = param.ljust(label_w)
+        unique = {str(v) for v in vals.values() if v is not None}
+        for name in model_names:
+            v = vals.get(name)
+            cell = "N/A" if v is None else str(v)
+            row += cell.ljust(col_w)
+        row += _STATUS["match"] if len(unique) <= 1 else _STATUS["warn"]
+        lines.append(row)
+    return "\n".join(lines)
+
+
+def _format_issues(issues: List[Issue]) -> str:
+    lines = []
+    for issue in issues:
+        prefix = f"[{issue.severity}] "
+        wrapped = textwrap.fill(
+            issue.message,
+            width=78,
+            initial_indent=prefix,
+            subsequent_indent=" " * len(prefix),
+        )
+        lines.append(wrapped)
+        lines.append("")
+    return "\n".join(lines)
+
+
+def _verdict(issues: List[Issue]) -> Tuple[str, int]:
+    severities = {i.severity for i in issues}
+    if "ERROR" in severities:
+        return "MERGE WILL FAIL", 1
+    if "WARNING" in severities:
+        return "MERGE POSSIBLE WITH WARNINGS", 0
+    return "MERGE LOOKS COMPATIBLE", 0
+
+
+@click.command("mergekit-check-compat", cls=PrettyPrintHelp)
+@click.argument("config_file")
+@add_merge_options
+def main(config_file: str, merge_options: MergeOptions):
+    """Check model compatibility before merging.
+
+    Loads configs and tokenizers (no weights) for all models in CONFIG_FILE and
+    reports structural mismatches, RoPE incompatibilities, vocabulary differences,
+    and FIM/chat-template issues before any merge work begins.
+    """
+    merge_options.apply_global_options()
+
+    with open(config_file, "r", encoding="utf-8") as f:
+        merge_config = MergeConfiguration.model_validate(yaml.safe_load(f))
+
+    models = merge_config.referenced_models()
+    if not models:
+        print("No models found in config.")
+        sys.exit(1)
+
+    names = [_short_name(ref) for ref in models]
+    seen: Dict[str, int] = {}
+    for i, n in enumerate(names):
+        count = seen.get(n, 0) + 1
+        seen[n] = count
+        if count > 1:
+            names[i] = f"{n}_{count}"
+
+    trc = merge_options.trust_remote_code
+    print("Loading model configs and tokenizers (no weights)...")
+    arch_configs: Dict[str, Any] = {}
+    tokenizers: Dict[str, Any] = {}
+    for ref, name in zip(models, names):
+        try:
+            arch_configs[name] = ref.config(trust_remote_code=trc)
+        except Exception as exc:
+            print(f"ERROR: Could not load config for {name}: {exc}")
+            sys.exit(1)
+        try:
+            tokenizers[name] = AutoTokenizer.from_pretrained(
+                ref.model.path,
+                revision=ref.model.revision,
+                trust_remote_code=trc,
+            )
+        except Exception as exc:
+            print(f"ERROR: Could not load tokenizer for {name}: {exc}")
+            sys.exit(1)
+
+    all_rows: Dict[str, Dict] = {}
+    all_issues: List[Issue] = []
+
+    for fn in (check_architecture, check_rope):
+        rows, issues = fn(arch_configs)
+        all_rows.update(rows)
+        all_issues.extend(issues)
+
+    for fn in (check_vocab, check_fim_tokens, check_chat_template):
+        rows, issues = fn(tokenizers)
+        all_rows.update(rows)
+        all_issues.extend(issues)
+
+    print("\nModel Compatibility Report")
+    print("=" * 60)
+    print(_format_table(all_rows, names))
+
+    if all_issues:
+        print("\nIssues")
+        print("-" * 40)
+        print(_format_issues(all_issues))
+
+    verdict, exit_code = _verdict(all_issues)
+    print(f"Verdict: {verdict}")
+    sys.exit(exit_code)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/pyproject.toml b/pyproject.toml
index d5710eb7..2ab446fd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -50,6 +50,7 @@ mergekit-extract-lora = "mergekit.scripts.extract_lora:main"
 mergekit-evolve = "mergekit.scripts.evolve:main"
 mergekit-pytorch = "mergekit.scripts.merge_raw_pytorch:main"
 mergekit-multi = "mergekit.scripts.multimerge:main"
+mergekit-check-compat = "mergekit.scripts.check_compat:main"
 
 [tool.setuptools]
 packages = [

From 27b8d04ef834035e83696ba971af18d07a525b74 Mon Sep 17 00:00:00 2001
From: Naomi 
Date: Tue, 7 Jul 2026 10:11:54 -0400
Subject: [PATCH 2/3] Fix bugbot findings: _get_rope_theta dead fallback and
 wrong table status symbol
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- _get_rope_theta: remove unreachable-code structure; check rope_scaling
  dict first (Ministral3 style) then fall back via getattr to avoid
  returning a class-level default as if it were a real config value
- _format_table: use ✗ for shape-param mismatches (ERROR level) and ⚠
  for all others; introduce _SHAPE_PARAMS frozenset shared between
  check_architecture and _format_table so the sets can't drift
---
 mergekit/scripts/check_compat.py | 32 ++++++++++++++++++--------------
 1 file changed, 18 insertions(+), 14 deletions(-)

diff --git a/mergekit/scripts/check_compat.py b/mergekit/scripts/check_compat.py
index caa17219..a18b012f 100644
--- a/mergekit/scripts/check_compat.py
+++ b/mergekit/scripts/check_compat.py
@@ -19,6 +19,14 @@
 
 FIM_TOKENS = ["
", "", "", ""]
 
+_SHAPE_PARAMS = frozenset([
+    "hidden_size",
+    "num_hidden_layers",
+    "num_attention_heads",
+    "num_key_value_heads",
+    "intermediate_size",
+])
+
 
 @dataclass
 class Issue:
@@ -42,27 +50,18 @@ def _safe_get(cfg: Any, attr: str, default: Any = None) -> Any:
 
 
 def _get_rope_theta(cfg: Any) -> Optional[float]:
-    if hasattr(cfg, "rope_theta"):
-        return cfg.rope_theta
-    if hasattr(cfg, "rope_parameters") and isinstance(cfg.rope_parameters, dict):
-        return cfg.rope_parameters.get("rope_theta")
-    # Some configs (e.g. Ministral3) embed rope_theta inside rope_scaling
+    # Some configs (e.g. Ministral3) embed rope_theta inside rope_scaling dict;
+    # check there first so we don't return a class-level default instead.
     scaling = getattr(cfg, "rope_scaling", None)
     if isinstance(scaling, dict) and "rope_theta" in scaling:
         return scaling["rope_theta"]
-    return None
+    return getattr(cfg, "rope_theta", None)
 
 
 def check_architecture(
     configs: Dict[str, Any],
 ) -> Tuple[Dict[str, Dict], List[Issue]]:
-    shape_params = [
-        "hidden_size",
-        "num_hidden_layers",
-        "num_attention_heads",
-        "num_key_value_heads",
-        "intermediate_size",
-    ]
+    shape_params = list(_SHAPE_PARAMS)
     table: Dict[str, Dict] = {}
     issues: List[Issue] = []
     resolved = {name: _resolve_config(cfg) for name, cfg in configs.items()}
@@ -245,7 +244,12 @@ def _format_table(rows: Dict[str, Dict], model_names: List[str]) -> str:
             v = vals.get(name)
             cell = "N/A" if v is None else str(v)
             row += cell.ljust(col_w)
-        row += _STATUS["match"] if len(unique) <= 1 else _STATUS["warn"]
+        if len(unique) <= 1:
+            row += _STATUS["match"]
+        elif param in _SHAPE_PARAMS:
+            row += _STATUS["error"]
+        else:
+            row += _STATUS["warn"]
         lines.append(row)
     return "\n".join(lines)
 

From aae614d68ff2a6314b80ed6e4d3de59287caa581 Mon Sep 17 00:00:00 2001
From: Naomi 
Date: Tue, 7 Jul 2026 17:28:15 -0400
Subject: [PATCH 3/3] Fix second round of bugbot findings: None-masking and
 union tokenizer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- None values no longer silently pass mismatch checks: shape param and
  RoPE theta comparisons now fire when one model has a value and another
  has None (missing field), rather than treating None as a match.
  _format_table also includes N/A as a distinct value for the status
  symbol, so a mixed None/non-None row shows ⚠/✗ instead of ✓.
- check_vocab now accepts tokenizer_source and suppresses the truncation
  WARNING (emits INFO instead) when tokenizer_source is "union", since
  no truncation occurs in that mode. main() derives the effective source
  from either the legacy tokenizer_source field or the newer
  tokenizer.source field and passes it through.
---
 mergekit/scripts/check_compat.py | 73 ++++++++++++++++++++++----------
 1 file changed, 51 insertions(+), 22 deletions(-)

diff --git a/mergekit/scripts/check_compat.py b/mergekit/scripts/check_compat.py
index a18b012f..b81d3661 100644
--- a/mergekit/scripts/check_compat.py
+++ b/mergekit/scripts/check_compat.py
@@ -69,8 +69,9 @@ def check_architecture(
     for param in shape_params:
         vals = {name: _safe_get(cfg, param) for name, cfg in resolved.items()}
         table[param] = vals
-        unique = {v for v in vals.values() if v is not None}
-        if len(unique) > 1:
+        non_none = {v for v in vals.values() if v is not None}
+        has_none = any(v is None for v in vals.values())
+        if len(non_none) > 1 or (non_none and has_none):
             detail = ", ".join(f"{n}={v}" for n, v in vals.items())
             issues.append(
                 Issue(
@@ -104,8 +105,9 @@ def check_rope(configs: Dict[str, Any]) -> Tuple[Dict[str, Dict], List[Issue]]:
 
     thetas = {name: _get_rope_theta(cfg) for name, cfg in resolved.items()}
     table["rope_theta"] = thetas
-    unique_thetas = {v for v in thetas.values() if v is not None}
-    if len(unique_thetas) > 1:
+    non_none_thetas = {v for v in thetas.values() if v is not None}
+    has_none_theta = any(v is None for v in thetas.values())
+    if len(non_none_thetas) > 1 or (non_none_thetas and has_none_theta):
         detail = ", ".join(f"{n}={v}" for n, v in thetas.items())
         issues.append(
             Issue(
@@ -140,7 +142,9 @@ def check_rope(configs: Dict[str, Any]) -> Tuple[Dict[str, Dict], List[Issue]]:
     return table, issues
 
 
-def check_vocab(tokenizers: Dict[str, Any]) -> Tuple[Dict[str, Dict], List[Issue]]:
+def check_vocab(
+    tokenizers: Dict[str, Any], tokenizer_source: Optional[str] = None
+) -> Tuple[Dict[str, Dict], List[Issue]]:
     table: Dict[str, Dict] = {}
     issues: List[Issue] = []
 
@@ -149,22 +153,31 @@ def check_vocab(tokenizers: Dict[str, Any]) -> Tuple[Dict[str, Dict], List[Issue
     unique_sizes = set(vocab_sizes.values())
     if len(unique_sizes) > 1:
         detail = ", ".join(f"{n}={v}" for n, v in vocab_sizes.items())
-        larger_has_fim = any(
-            vocab_sizes[n] > min(unique_sizes)
-            and any(tok in tokenizers[n].get_vocab() for tok in FIM_TOKENS)
-            for n in tokenizers
-        )
-        msg = (
-            f"Vocabulary size mismatch ({detail}). MergeKit will truncate to the "
-            "base model's vocab. "
-        )
-        if larger_has_fim:
-            msg += (
-                "FIM tokens (
, , , ) in the larger-vocab model "
-                "will be dropped. Consider: resize_tok_vocab.py or set "
-                "`tokenizer_source: union` in config."
+        if tokenizer_source == "union":
+            issues.append(
+                Issue(
+                    "INFO",
+                    f"Vocabulary size mismatch ({detail}). Tokenizer source is "
+                    "'union' — all tokens from both models will be kept.",
+                )
             )
-        issues.append(Issue("WARNING", msg))
+        else:
+            larger_has_fim = any(
+                vocab_sizes[n] > min(unique_sizes)
+                and any(tok in tokenizers[n].get_vocab() for tok in FIM_TOKENS)
+                for n in tokenizers
+            )
+            msg = (
+                f"Vocabulary size mismatch ({detail}). MergeKit will truncate to the "
+                "base model's vocab. "
+            )
+            if larger_has_fim:
+                msg += (
+                    "FIM tokens (
, , , ) in the larger-vocab model "
+                    "will be dropped. Consider: resize_tok_vocab.py or set "
+                    "`tokenizer_source: union` in config."
+                )
+            issues.append(Issue("WARNING", msg))
 
     return table, issues
 
@@ -239,7 +252,7 @@ def _format_table(rows: Dict[str, Dict], model_names: List[str]) -> str:
     lines = [" " * label_w + "".join(n.ljust(col_w) for n in model_names)]
     for param, vals in rows.items():
         row = param.ljust(label_w)
-        unique = {str(v) for v in vals.values() if v is not None}
+        unique = {str(v) if v is not None else "N/A" for v in vals.values()}
         for name in model_names:
             v = vals.get(name)
             cell = "N/A" if v is None else str(v)
@@ -334,7 +347,23 @@ def main(config_file: str, merge_options: MergeOptions):
         all_rows.update(rows)
         all_issues.extend(issues)
 
-    for fn in (check_vocab, check_fim_tokens, check_chat_template):
+    if merge_config.tokenizer_source is not None:
+        tok_source = (
+            merge_config.tokenizer_source
+            if isinstance(merge_config.tokenizer_source, str)
+            else str(merge_config.tokenizer_source)
+        )
+    elif merge_config.tokenizer is not None:
+        src = merge_config.tokenizer.source
+        tok_source = src if isinstance(src, str) else str(src)
+    else:
+        tok_source = None
+
+    rows, issues = check_vocab(tokenizers, tokenizer_source=tok_source)
+    all_rows.update(rows)
+    all_issues.extend(issues)
+
+    for fn in (check_fim_tokens, check_chat_template):
         rows, issues = fn(tokenizers)
         all_rows.update(rows)
         all_issues.extend(issues)