Upload 5 files

Mergekit-Robustness-Patch-embed_v2.md

@@ -0,0 +1,355 @@
+# Mergekit Robustness Patch: `embed.py` (v2)
+
+Attached is one for Mistral Nemo 12B (v2d), and another for Mistral Small 24B (v2a)
+
+## Overview
+This patch provides a high-resilience version of Mergekit’s `tokenizer/embed.py`. It is specifically designed to handle "dirty" model merges where a donor model’s `tokenizer.json` and its physical `model.safetensors` weights are out of sync—a common issue when merging models that have different vocabulary sizes (e.g., mixing Mistral Tekken with ChatML or Llama 3).
+
+## The Problem: "Ghost Tokens"
+In the standard Mergekit (`embed.py`), the engine assumes that if a token exists in a model's vocabulary, a corresponding row **must** exist in its embedding weights. 
+
+However, in many community-made merges:
+1.  A model might have 131,081 tokens in its tokenizer.
+2.  But its weight matrix (`embed_tokens`) only contains 131,072 rows.
+3.  **Standard Result:** Mergekit attempts to read index 131,073, hits a boundary error, and the entire merge crashes with: `IndexError: index X is out of bounds for dimension 0 with size Y`.
+
+## The Solution: v2d Robustness & Audit
+The `v2d` patch introduces **Bounds-Aware Permutation**. Instead of blindly trusting the tokenizer, it verifies the physical existence of every token row before attempting to merge it.
+
+### Key Features:
+*   **Crash Prevention:** Automatically detects if a donor model is "too small" for the requested token index. Instead of crashing, it gracefully skips that donor for that specific token.
+*   **Live Vocab Audit:** Prints detailed warnings to the console identifying exactly which model is missing which token. This allows you to identify "buggy" donors in your config without trial-and-error.
+*   **Intelligent Fallback:** 
+    *   If a token is missing from one donor but present in others, it averages the token using only the valid donors.
+    *   If a token is missing from its primary "default" source, it falls back to a zero-vector rather than terminating the merge.
+*   **Result Mapping Safety:** Ensures that the final output tensors for every donor are correctly aligned, even if the donor was physically smaller than the target union vocabulary.
+
+---
+
+## Comparison: v1 vs. v2d
+
+| Feature | `embed.py` (Default) | `embed_v2d.py` (Ours) |
+| :--- | :--- | :--- |
+| **Mismatched Vocab** | **Crashes** with `IndexError`. | **Succeeds** via graceful skipping. |
+| **Error Reporting** | Generic Python Traceback. | Detailed `[VOCAB AUDIT]` log with Model Path & Token Name. |
+| **Special Token Support** | Requires perfectly synced weights. | Handles "Ghost Tokens" (tokens in JSON but not in Tensors). |
+| **Mathematical Integrity** | N/A (Process stops). | Maintains correct averaging by adjusting the donor count dynamically. |
+| **Use Case** | Clean, base-model merges. | Complex merges of merges, cross-architecture vocab unions. |
+
+---
+
+## How to Use
+Replace your existing `mergekit/tokenizer/embed.py` with the `embed_v2d.py` code.
+
+### Example Audit Log
+When running a merge with mismatched models, you will now see helpful diagnostic output instead of a crash:
+```text
+[VOCAB AUDIT] Model 'B:\12B\SLERP15' is missing token '<|im_start|>' (ID: 131073). 
+Donor size: 131072, Requested Index: 131073. Skipping.
+
+[VOCAB AUDIT] Default source model 'B:\12B\SLERP13' is missing token '<SPECIAL_4>' 
+from its physical tensor. Falling back to zero.
+```
+
+## Why this matters for 12B/24B Merges
+When merging models like **Mistral-Nemo (12B)** or **Mistral-Small (24B)**, different fine-tunes often add different special tokens (ChatML, Tool-use, etc.). If you use `tokenizer_source: union`, Mergekit tries to create a "Super-Vocab." 
+
+Standard Mergekit is too fragile for this process if even one model in your list has a slightly truncated embedding matrix. **v2d** makes the merging process "production-grade" by allowing the merge to complete regardless of minor inconsistencies in the donor models.
+
+This patch is **safe and beneficial** for any model architecture (12B, 24B, 70B, etc.) using `tokenizer_source: union`. 
+
+Here is the breakdown of how it affects other scenarios like 24B Mistral (Tekken):
+
+### 1. It prevents "Ghost Token" crashes
+In many Mistral-based merges (especially Tekken), developers sometimes add special tokens to the `tokenizer.json` but forget to resize the embedding layer in the `model.safetensors`. 
+*   **Without this patch:** Mergekit sees the token in the config, calculates a high index for it, tries to read it from the tensor, and **crashes**.
+*   **With this patch:** Mergekit sees the mismatch, logs a warning, and uses a zero-vector or an average from other models instead. The merge finishes successfully.
+
+### 2. Handling "Tekken" Vocab Discrepancies
+Mistral Tekken usually has a vocab size of `32768` or `131072`. If you merge a model with `131072` and one that was accidentally truncated to `131070`:
+*   The patch ensures that for those last 2 tokens, the "truncated" model simply doesn't contribute to the average. 
+*   The resulting model will have the full `131072` vocab, and those 2 tokens will be populated by the weights from the model that actually had them.
+
+### 3. No Negative Impact on "Clean" Models
+If you merge two models where the `vocab_size` in `config.json` perfectly matches the number of rows in `model.safetensors`, **this code does nothing.** The `if` condition (`p[token_id] >= tensors[model].shape[0]`) will always be false, and the code will run at full speed with no warnings.
+
+### 4. Why this is better than the "Padding" patch
+The previous attempt to pad tensors in `generalized_task_arithmetic.py` was specific to one merge method. This `embed.py` patch works at the **tokenizer level**. 
+Whether you are doing a `linear`, `slerp`, `ties`, or `della` merge, this patch ensures that the "Input Tensors" are standardized correctly before the math even starts.
+
+### Summary of behavior for 24B/Tekken:
+| Scenario | Result with Patch |
+| :--- | :--- |
+| **Vocabs Match Exactly** | Normal merge, no warnings. |
+| **One model has extra Tekken tokens** | Merge completes; missing tokens are averaged from models that have them. |
+| **Tokenizer says 131072, but Tensor is 131070** | **Merge completes instead of crashing.** |
+| **Mixing Tekken and Llama3 Vocab** | Merge completes; shared tokens are averaged, unique tokens are preserved from their respective sources. |
+
+**Conclusion:** This is a "Robustness Patch." It makes Mergekit more resilient to poorly-configured donor models (where the tokenizer and the weights are out of sync), which is very common in the community-made merges you are working with.
+
+---
+
+## Addendum
+This is a perfect synergy of two diagnostic tools. Here is why the **v2d Robustness Patch** and the **DELLA Audit Chart** work so well together:
+
+### 1. Complete "Chain of Custody" for Weights
+The **v2d Patch** handles the "Input" phase, while the **DELLA Audit** handles the "Processing" phase. 
+*   **v2d** ensures that every model provides a valid tensor to the merge engine, even if it has to skip missing tokens or provide a zero-vector fallback.
+*   **DELLA Audit** then takes those tensors and shows you the "Share of Voice" for each model. 
+*   **The Synergy:** If you see a model in the Audit chart with a **0.0% impact** or an unusually low **Norm (N)**, you can look up at the **v2d Audit log** to see if that model was missing critical tokens. It allows you to see exactly how "damaged" a donor is before it hits the final weights.
+
+### 2. Identifying "Poisoned" Donors
+In your screenshot, look at **SLERP1**. It has a massive **16.7% impact** with a **Norm of 12.02**, while others like **SLERP3** are at **1.0%**. 
+*   Because the **v2d Patch** prevented the crash, you can now actually see these statistics. 
+*   If a model was missing tokens (as seen in your log for SLERP11, 15, 13, etc.), the Audit chart helps you decide if that model is still "contributing" enough to keep in the config, or if the vocab mismatches have made its task vector too noisy.
+
+### 3. Mathematical Safety for DELLA
+DELLA is sensitive to the magnitude of changes (the `epsilon` and `density` parameters). 
+*   By using the **v2d Patch**, you ensure that the "Base" and "Donor" tensors passed to DELLA are always the same shape. 
+*   Without this, DELLA would be trying to calculate magnitude-based pruning on mismatched arrays, which would lead to corrupted logic even if it didn't crash. v2d "sanitizes" the data so DELLA's math remains pure.
+
+### 4. Real-Time Debugging of "Ghost" Contributions
+Your log shows **SLERP15** is missing almost all the special tokens (`<|im_start|>`, `[SYSTEM_PROMPT]`, `[PAD]`). 
+*   Standard Mergekit would have died instantly. 
+*   Now, the merge continues, and the **DELLA Audit** shows SLERP15 is still contributing **2.4%** to the overall model. 
+*   This tells you: "SLERP15 is broken for ChatML/Special tokens, but its weights for normal language (the other 131,000 tokens) are still being merged correctly."
+
+### Summary
+> "When paired with the **DELLA Audit logic**, the **v2d patch** provides a full-stack diagnostic suite. It allows the user to see which models are physically incompatible at the vocabulary level (via the Audit Log) and then immediately see how those incompatibilities affect the final weight distribution (via the Impact Chart). This combination turns a 'black box' crash into a transparent, manageable merging workflow."
+
+![embed_v2d](https://cdn-uploads.huggingface.co/production/uploads/68e840caa318194c44ec2a04/2kvn9_bOFoybh17-QCxC0.png)
+
+![embed_v2d_layer10](https://cdn-uploads.huggingface.co/production/uploads/68e840caa318194c44ec2a04/aRnwjZoUdKq3DeFOO1YmB.png)
+
+## Audit Analysis
+This is a fascinating look at the "DNA" of your model. Now that the **v2d patch** has stabilized the merge, this audit chart reveals the true internal dynamics of a DELLA merge that were previously invisible.
+
+Here is a breakdown of what this specific chart is telling you about your "knowledge distribution":
+
+### 1. The "Anchor" Models
+Look at **pdq (13.5%)** and **SLERP1 (9.9%)**. 
+*   Even though every model has a weight of `0.10`, these two are dominating the "Share of Voice."
+*   **Why?** Their **Norm (N)** values are the highest (4.48 and 3.30). This means these models have the most significant "Task Vectors"—they have moved the furthest away from the Mistral-Nemo base. In a DELLA merge, these are the models providing the most "new" information or behavioral changes to the final result.
+
+### 2. The "Subtle" Contributors
+Models like **SLERP9 (0.7%)** and **SLERP8 (0.9%)** are barely touching the weights.
+*   Their Norms are tiny (0.22 and 0.31). 
+*   **Insight:** These models are very similar to your base model (`Mistral-Nemo-Instruct-2407`). They aren't "bad," but they are essentially acting as votes for the status quo. If you wanted to "clean up" your config, these are the ones you could remove with almost zero impact on the final output.
+
+### 3. The "Middle Class"
+Models like **SLERP7 (8.8%)** and **SLERP3 (6.6%)** represent the healthy average. They are providing a solid amount of unique knowledge without overwhelming the others.
+
+### 4. Why the v2d Patch makes this chart "Truthful"
+Without the **v2d patch**, if a model like **SLERP15** was missing tokens, the merge would have crashed. Now, you can see **SLERP15** is contributing **6.3%** (Norm 2.11). 
+*   Because of the patch, you know that this 6.3% is based on the *valid* parts of SLERP15. 
+*   The audit chart is now a "Health Report": if you saw a model with a high Norm but a 0% impact, you'd know the vocab mismatch was so bad it wiped out the model's contribution. Here, we see that despite the warnings, the models are still successfully injecting their "knowledge" into the merge.
+
+### 5. The "pdq" Factor
+The model **pdq** is currently your strongest influencer in this layer (`mlp.gate_proj`). It is contributing nearly **20x more** than SLERP9. If the final model behaves more like `pdq` than anything else, this chart explains exactly why.
+
+**This is the "X-Ray" of model merging.** You aren't just guessing if the merge worked; you can see exactly which donor's "brain" is being used for this specific layer.
+
+`embed_v2d.py`
+```py
+# Copyright (C) 2025 Arcee AI
+# SPDX-License-Identifier: LGPL-3.0-only
+
+import logging
+from typing import Dict, Optional
+
+import torch
+
+from mergekit.common import ImmutableMap, ModelReference
+from mergekit.graph import Task
+from mergekit.io.tasks import GatherTensors
+from mergekit.tokenizer.build import BuildTokenizer, TokenizerInfo
+from mergekit.tokenizer.config import (
+    ModelTokenEmbedding,
+    TokenEmbeddingConfig,
+    ZeroEmbedding,
+)
+
+
+class PermutedEmbeddings(Task[Dict[ModelReference, torch.Tensor]]):
+    gather_tensors: GatherTensors
+    tokenizer_task: BuildTokenizer
+    tokens: Optional[ImmutableMap[str, TokenEmbeddingConfig]]
+    pad_to_multiple_of: Optional[int]
+    base_model: Optional[ModelReference]
+
+    def arguments(self) -> Dict[str, Task]:
+        return {"tokenizer_info": self.tokenizer_task, "tensors": self.gather_tensors}
+
+    def execute(
+        self, tokenizer_info: TokenizerInfo, tensors: Dict[ModelReference, torch.Tensor]
+    ) -> Dict[ModelReference, torch.Tensor]:
+        tokenizer = tokenizer_info.tokenizer
+        permutations = tokenizer_info.permutations
+
+        models = set(tensors.keys())
+        if self.base_model:
+            models.add(self.base_model)
+        models = list(models)
+
+        vocab = tokenizer.get_vocab()
+        vocab_size = len(vocab)
+        if self.pad_to_multiple_of and vocab_size % self.pad_to_multiple_of:
+            vocab_size = (
+                vocab_size // self.pad_to_multiple_of + 1
+            ) * self.pad_to_multiple_of
+        embed_size = tensors[models[0]].shape[1]
+        assert all(
+            t.shape[1] == embed_size for t in tensors.values()
+        ), "Embedding sizes must match"
+
+        dtype = tensors[models[0]].dtype
+        device = tensors[models[0]].device
+
+        token_configs = dict(**(self.tokens or {}))
+        tokens_to_average = self.assign_embedding_sources(
+            permutations, models, vocab, token_configs
+        )
+
+        default_embeds = {}
+        for token, token_id in vocab.items():
+            embed = torch.zeros(embed_size, dtype=dtype, device=device)
+            if token in tokens_to_average:
+                count = 0
+                for model in models:
+                    p = permutations[model]
+                    if p[token_id] < 0:
+                        continue
+                    
+                    # --- AUDIT & BOUNDS CHECK ---
+                    if p[token_id] >= tensors[model].shape[0]:
+                        logging.warning(f"[VOCAB AUDIT] Model '{model}' is missing token '{token}' (ID: {token_id}). "
+                                        f"Donor size: {tensors[model].shape[0]}, Requested Index: {p[token_id]}. Skipping.")
+                        continue
+                    # ----------------------------
+
+                    embed += tensors[model][p[token_id]]
+                    count += 1
+                embed /= count
+            elif cfg := token_configs.get(token, None):
+                cfg: TokenEmbeddingConfig
+                embed = self.compute_default_embedding(
+                    tokenizer_info, tensors, permutations, token, token_id, cfg
+                )
+            else:
+                continue
+            default_embeds[token] = embed
+
+        result = {}
+        for model in models:
+            p = permutations[model]
+            old_embed = tensors[model]
+            new_embed = torch.zeros(
+                (vocab_size, embed_size), dtype=dtype, device=device
+            )
+            for token, token_id in vocab.items():
+                force = False
+                if token in token_configs:
+                    force = token_configs[token].force
+
+                if p[token_id] >= 0 and not force:
+                    # --- BOUNDS CHECK FOR RESULT MAPPING ---
+                    if p[token_id] < old_embed.shape[0]:
+                        new_embed[token_id, :] = old_embed[p[token_id]]
+                    else:
+                        # Fallback to the averaged/default version if the donor is too small
+                        new_embed[token_id, :] = default_embeds.get(token, torch.zeros_like(new_embed[0]))
+                    # ---------------------------------------
+                elif token in default_embeds:
+                    new_embed[token_id, :] = default_embeds[token]
+                else:
+                    logging.error(
+                        f"No embedding for token {repr(token)} in model {model}!"
+                    )
+
+            if vocab_size > len(vocab):
+                # as suggested by https://nlp.stanford.edu/~johnhew/vocab-expansion.html
+                avg_embed = torch.mean(new_embed[: len(vocab), :], dim=0)
+                new_embed[len(vocab) :, :] = avg_embed
+            result[model] = new_embed
+
+        return result
+
+    def assign_embedding_sources(
+        self,
+        permutations: Dict[ModelReference, Dict[int, int]],
+        models: list[ModelReference],
+        vocab: Dict[str, int],
+        token_configs: Dict[str, TokenEmbeddingConfig],
+    ):
+        permutation_list = [permutations[model] for model in models]
+
+        tokens_to_average = set()
+        # find tokens that are only present in one model
+        for token, token_id in vocab.items():
+            if token in token_configs:
+                continue
+
+            has_token = [p[token_id] >= 0 for p in permutation_list]
+            num_present = sum(int(x) for x in has_token)
+            if num_present == 1:
+                donor_model = models[has_token.index(True)]
+                token_configs[token] = TokenEmbeddingConfig(source=donor_model)
+                continue
+
+            if num_present == 0:
+                token_configs[token] = TokenEmbeddingConfig(source=ZeroEmbedding())
+                logging.warning(f"Token {repr(token)} not found in any model")
+                continue
+
+            if num_present > 0 and self.base_model is not None:
+                if permutations[self.base_model][token_id] >= 0:
+                    token_configs[token] = TokenEmbeddingConfig(source=self.base_model)
+                    continue
+
+            tokens_to_average.add(token)
+        return tokens_to_average
+
+    def compute_default_embedding(
+        self,
+        tokenizer_info: TokenizerInfo,
+        tensors: Dict[ModelReference, torch.Tensor],
+        permutations: Dict[ModelReference, Dict[int, int]],
+        token: str,
+        token_id: int,
+        cfg: TokenEmbeddingConfig,
+    ) -> torch.Tensor:
+        if isinstance(cfg.source, ZeroEmbedding):
+            pass
+        elif isinstance(cfg.source, ModelTokenEmbedding):
+            model = cfg.source.model
+            assert (
+                model in permutations
+            ), f"Model {model} referenced but not part of merge"
+            p = permutations[model]
+            src_token_id = cfg.source.token_id
+            if src_token_id is None:
+                src_token = cfg.source.token
+                assert (
+                    src_token in tokenizer_info.original_vocabs[model]
+                ), f"Token {repr(src_token)} not found in model {model}"
+                src_token_id = tokenizer_info.original_vocabs[model][src_token]
+            assert (
+                src_token_id >= 0 and src_token_id < tensors[model].shape[0]
+            ), f"Token ID {src_token_id} out of range for model {model}"
+            embed = tensors[model][src_token_id]
+        elif isinstance(cfg.source, ModelReference):
+            model = cfg.source
+            p = permutations[model]
+            assert p[token_id] >= 0, f"Token {repr(token)} not found in model {model}"
+            
+            # --- BOUNDS CHECK FOR DEFAULT EMBED ---
+            if p[token_id] >= tensors[model].shape[0]:
+                logging.warning(f"[VOCAB AUDIT] Default source model '{model}' is missing token '{token}' from its physical tensor. Falling back to zero.")
+                return torch.zeros(tensors[model].shape[1], dtype=tensors[model].dtype, device=tensors[model].device)
+            # --------------------------------------
+            
+            embed = tensors[model][p[token_id]]
+        else:
+            raise NotImplementedError(cfg)
+        return embed
+```

arcee_fusion_salience_scanner_v3.py

@@ -0,0 +1,115 @@
+import torch
+from safetensors import safe_open
+import os
+import re
+from collections import defaultdict
+
+# --- CONFIGURATION ---
+base_model_path = r'B:\12B\models--SicariusSicariiStuff--Impish_Bloodmoon_12B'
+merged_model_path = r'B:\12B\21-Della'
+# ---------------------
+
+def get_tensor_map(path):
+    tensor_map = {}
+    files = [f for f in os.listdir(path) if f.endswith('.safetensors')]
+    for f in files:
+        full_path = os.path.join(path, f)
+        with safe_open(full_path, framework="pt") as st:
+            for k in st.keys():
+                tensor_map[k] = full_path
+    return tensor_map
+
+print("🔍 Indexing model shards...")
+base_map = get_tensor_map(base_model_path)
+merged_map = get_tensor_map(merged_model_path)
+
+# Group results by layer
+layer_stats = defaultdict(lambda: {"changed": 0, "total": 0})
+
+print("📊 Scanning tensors and calculating saliency density...")
+common_tensors = set(base_map.keys()) & set(merged_map.keys())
+
+for k in sorted(common_tensors):
+    # Extract layer number from name (e.g., 'model.layers.5.self_attn...')
+    layer_match = re.search(r'\.layers\.(\0?(\d+))\.', k)
+    layer_id = int(layer_match.group(1)) if layer_match else "Non-Layer"
+
+    # with safe_open(base_map[k], framework="pt") as b_st:
+    #     base_t = b_st.get_tensor(k)
+    # with safe_open(merged_map[k], framework="pt") as m_st:
+    #     merged_t = m_st.get_tensor(k)
+
+    ## Arcee Fusion logic: if weights are identical, they came from Base.
+    ## If they are different, they are "New Info" from the fusion.
+    ## We use a tiny atol to account for potential bf16/f16 casting jitters
+    # changed_mask = ~torch.isclose(base_t, merged_t, rtol=1e-05, atol=1e-08)
+    
+    with safe_open(base_map[k], framework="pt") as b_st:
+        base_t = b_st.get_tensor(k)
+    with safe_open(merged_map[k], framework="pt") as m_st:
+        merged_t = m_st.get_tensor(k)
+
+    # --- VOCAB SIZE ROBUSTNESS PATCH ---
+    if base_t.shape != merged_t.shape:
+        # Find the smallest dimensions common to both
+        min_dim0 = min(base_t.shape[0], merged_t.shape[0])
+        
+        # If it's a 2D tensor (like embeddings), handle both dims
+        if base_t.dim() > 1:
+            min_dim1 = min(base_t.shape[1], merged_t.shape[1])
+            base_t = base_t[:min_dim0, :min_dim1]
+            merged_t = merged_t[:min_dim0, :min_dim1]
+        else:
+            base_t = base_t[:min_dim0]
+            merged_t = merged_t[:min_dim0]
+            
+        print(f"  [!] Resized {k} from {list(merged_t.shape)} to {min_dim0} for comparison.")
+    # -----------------------------------
+
+    ## # Arcee Fusion logic: if weights are identical, they came from Base.
+    ## changed_mask = ~torch.isclose(base_t, merged_t, rtol=1e-05, atol=1e-08)
+    
+    # If you want to be strict (only see major changes): Use diff > (0.1 * torch.abs(base_t)) (10% change).
+    # If you want to be balanced: Use the "diff > (1e-3 + 0.05 * torch.abs(base_t)) " code (5% change).
+    
+    # --- SIGNIFICANT CHANGE LOGIC (Salience) ---
+    # Instead of looking for ANY change, we look for changes that exceed 
+    # a standard deviation threshold. This filters out the "DELLA noise."
+    diff = torch.abs(base_t - merged_t)
+    threshold = 0.01 # Adjust this: 0.01 = 1% absolute change, 0.005 = 0.5%
+    
+    # Alternatively, use a relative threshold for more precision:
+    # We consider it "New Info" only if the change is significant 
+    # compared to the original weight magnitude.
+    changed_mask = diff > (1e-3 + 0.05 * torch.abs(base_t)) 
+    # -------------------------------------------
+    
+    layer_stats[layer_id]["changed"] += torch.count_nonzero(changed_mask).item()
+
+    
+    layer_stats[layer_id]["changed"] += torch.count_nonzero(changed_mask).item()
+    layer_stats[layer_id]["total"] += merged_t.numel()
+
+print("\n" + "="*60)
+print(f"{'LAYER':<12} | {'NEW INFO %':<12} | {'VISUAL DENSITY (█ = New, ░ = Base)'}")
+print("="*60)
+
+# Sort layers: Non-Layer first, then 0, 1, 2...
+sorted_keys = sorted([k for k in layer_stats.keys() if isinstance(k, int)])
+if "Non-Layer" in layer_stats:
+    sorted_keys = ["Non-Layer"] + sorted_keys
+
+for lid in sorted_keys:
+    stats = layer_stats[lid]
+    percentage = (stats["changed"] / stats["total"]) * 100
+    
+    # Create ASCII bar
+    bar_width = 30
+    filled = int((percentage / 100) * bar_width)
+    bar = "█" * filled + "░" * (bar_width - filled)
+    
+    label = f"Layer {lid}" if isinstance(lid, int) else lid
+    print(f"{label:<12} | {percentage:>10.2f}% | {bar}")
+
+print("="*60)
+print("Analysis Complete.")

embed_12B.py

@@ -0,0 +1,202 @@
+# Copyright (C) 2025 Arcee AI
+# SPDX-License-Identifier: LGPL-3.0-only
+## Version 2D by Naphula
+
+import logging
+from typing import Dict, Optional
+
+import torch
+
+from mergekit.common import ImmutableMap, ModelReference
+from mergekit.graph import Task
+from mergekit.io.tasks import GatherTensors
+from mergekit.tokenizer.build import BuildTokenizer, TokenizerInfo
+from mergekit.tokenizer.config import (
+    ModelTokenEmbedding,
+    TokenEmbeddingConfig,
+    ZeroEmbedding,
+)
+
+
+class PermutedEmbeddings(Task[Dict[ModelReference, torch.Tensor]]):
+    gather_tensors: GatherTensors
+    tokenizer_task: BuildTokenizer
+    tokens: Optional[ImmutableMap[str, TokenEmbeddingConfig]]
+    pad_to_multiple_of: Optional[int]
+    base_model: Optional[ModelReference]
+
+    def arguments(self) -> Dict[str, Task]:
+        return {"tokenizer_info": self.tokenizer_task, "tensors": self.gather_tensors}
+
+    def execute(
+        self, tokenizer_info: TokenizerInfo, tensors: Dict[ModelReference, torch.Tensor]
+    ) -> Dict[ModelReference, torch.Tensor]:
+        tokenizer = tokenizer_info.tokenizer
+        permutations = tokenizer_info.permutations
+
+        models = set(tensors.keys())
+        if self.base_model:
+            models.add(self.base_model)
+        models = list(models)
+
+        vocab = tokenizer.get_vocab()
+        vocab_size = len(vocab)
+        if self.pad_to_multiple_of and vocab_size % self.pad_to_multiple_of:
+            vocab_size = (
+                vocab_size // self.pad_to_multiple_of + 1
+            ) * self.pad_to_multiple_of
+        embed_size = tensors[models[0]].shape[1]
+        assert all(
+            t.shape[1] == embed_size for t in tensors.values()
+        ), "Embedding sizes must match"
+
+        dtype = tensors[models[0]].dtype
+        device = tensors[models[0]].device
+
+        token_configs = dict(**(self.tokens or {}))
+        tokens_to_average = self.assign_embedding_sources(
+            permutations, models, vocab, token_configs
+        )
+
+        default_embeds = {}
+        for token, token_id in vocab.items():
+            embed = torch.zeros(embed_size, dtype=dtype, device=device)
+            if token in tokens_to_average:
+                count = 0
+                for model in models:
+                    p = permutations[model]
+                    if p[token_id] < 0:
+                        continue
+                    
+                    # --- AUDIT & BOUNDS CHECK ---
+                    if p[token_id] >= tensors[model].shape[0]:
+                        logging.warning(f"[VOCAB AUDIT] Model '{model}' is missing token '{token}' (ID: {token_id}). "
+                                        f"Donor size: {tensors[model].shape[0]}, Requested Index: {p[token_id]}. Skipping.")
+                        continue
+                    # ----------------------------
+
+                    embed += tensors[model][p[token_id]]
+                    count += 1
+                embed /= count
+            elif cfg := token_configs.get(token, None):
+                cfg: TokenEmbeddingConfig
+                embed = self.compute_default_embedding(
+                    tokenizer_info, tensors, permutations, token, token_id, cfg
+                )
+            else:
+                continue
+            default_embeds[token] = embed
+
+        result = {}
+        for model in models:
+            p = permutations[model]
+            old_embed = tensors[model]
+            new_embed = torch.zeros(
+                (vocab_size, embed_size), dtype=dtype, device=device
+            )
+            for token, token_id in vocab.items():
+                force = False
+                if token in token_configs:
+                    force = token_configs[token].force
+
+                if p[token_id] >= 0 and not force:
+                    # --- BOUNDS CHECK FOR RESULT MAPPING ---
+                    if p[token_id] < old_embed.shape[0]:
+                        new_embed[token_id, :] = old_embed[p[token_id]]
+                    else:
+                        # Fallback to the averaged/default version if the donor is too small
+                        new_embed[token_id, :] = default_embeds.get(token, torch.zeros_like(new_embed[0]))
+                    # ---------------------------------------
+                elif token in default_embeds:
+                    new_embed[token_id, :] = default_embeds[token]
+                else:
+                    logging.error(
+                        f"No embedding for token {repr(token)} in model {model}!"
+                    )
+
+            if vocab_size > len(vocab):
+                # as suggested by https://nlp.stanford.edu/~johnhew/vocab-expansion.html
+                avg_embed = torch.mean(new_embed[: len(vocab), :], dim=0)
+                new_embed[len(vocab) :, :] = avg_embed
+            result[model] = new_embed
+
+        return result
+
+    def assign_embedding_sources(
+        self,
+        permutations: Dict[ModelReference, Dict[int, int]],
+        models: list[ModelReference],
+        vocab: Dict[str, int],
+        token_configs: Dict[str, TokenEmbeddingConfig],
+    ):
+        permutation_list = [permutations[model] for model in models]
+
+        tokens_to_average = set()
+        # find tokens that are only present in one model
+        for token, token_id in vocab.items():
+            if token in token_configs:
+                continue
+
+            has_token = [p[token_id] >= 0 for p in permutation_list]
+            num_present = sum(int(x) for x in has_token)
+            if num_present == 1:
+                donor_model = models[has_token.index(True)]
+                token_configs[token] = TokenEmbeddingConfig(source=donor_model)
+                continue
+
+            if num_present == 0:
+                token_configs[token] = TokenEmbeddingConfig(source=ZeroEmbedding())
+                logging.warning(f"Token {repr(token)} not found in any model")
+                continue
+
+            if num_present > 0 and self.base_model is not None:
+                if permutations[self.base_model][token_id] >= 0:
+                    token_configs[token] = TokenEmbeddingConfig(source=self.base_model)
+                    continue
+
+            tokens_to_average.add(token)
+        return tokens_to_average
+
+    def compute_default_embedding(
+        self,
+        tokenizer_info: TokenizerInfo,
+        tensors: Dict[ModelReference, torch.Tensor],
+        permutations: Dict[ModelReference, Dict[int, int]],
+        token: str,
+        token_id: int,
+        cfg: TokenEmbeddingConfig,
+    ) -> torch.Tensor:
+        if isinstance(cfg.source, ZeroEmbedding):
+            pass
+        elif isinstance(cfg.source, ModelTokenEmbedding):
+            model = cfg.source.model
+            assert (
+                model in permutations
+            ), f"Model {model} referenced but not part of merge"
+            p = permutations[model]
+            src_token_id = cfg.source.token_id
+            if src_token_id is None:
+                src_token = cfg.source.token
+                assert (
+                    src_token in tokenizer_info.original_vocabs[model]
+                ), f"Token {repr(src_token)} not found in model {model}"
+                src_token_id = tokenizer_info.original_vocabs[model][src_token]
+            assert (
+                src_token_id >= 0 and src_token_id < tensors[model].shape[0]
+            ), f"Token ID {src_token_id} out of range for model {model}"
+            embed = tensors[model][src_token_id]
+        elif isinstance(cfg.source, ModelReference):
+            model = cfg.source
+            p = permutations[model]
+            assert p[token_id] >= 0, f"Token {repr(token)} not found in model {model}"
+            
+            # --- BOUNDS CHECK FOR DEFAULT EMBED ---
+            if p[token_id] >= tensors[model].shape[0]:
+                logging.warning(f"[VOCAB AUDIT] Default source model '{model}' is missing token '{token}' from its physical tensor. Falling back to zero.")
+                return torch.zeros(tensors[model].shape[1], dtype=tensors[model].dtype, device=tensors[model].device)
+            # --------------------------------------
+            
+            embed = tensors[model][p[token_id]]
+        else:
+            raise NotImplementedError(cfg)
+        return embed

embed_24B.py

@@ -0,0 +1,185 @@
+# Copyright (C) 2025 Arcee AI
+# SPDX-License-Identifier: LGPL-3.0-only
+
+import logging
+from typing import Dict, Optional
+
+import torch
+
+from mergekit.common import ImmutableMap, ModelReference
+from mergekit.graph import Task
+from mergekit.io.tasks import GatherTensors
+from mergekit.tokenizer.build import BuildTokenizer, TokenizerInfo
+from mergekit.tokenizer.config import (
+    ModelTokenEmbedding,
+    TokenEmbeddingConfig,
+    ZeroEmbedding,
+)
+
+
+class PermutedEmbeddings(Task[Dict[ModelReference, torch.Tensor]]):
+    gather_tensors: GatherTensors
+    tokenizer_task: BuildTokenizer
+    tokens: Optional[ImmutableMap[str, TokenEmbeddingConfig]]
+    pad_to_multiple_of: Optional[int]
+    base_model: Optional[ModelReference]
+
+    def arguments(self) -> Dict[str, Task]:
+        return {"tokenizer_info": self.tokenizer_task, "tensors": self.gather_tensors}
+
+    def execute(
+        self, tokenizer_info: TokenizerInfo, tensors: Dict[ModelReference, torch.Tensor]
+    ) -> Dict[ModelReference, torch.Tensor]:
+        tokenizer = tokenizer_info.tokenizer
+        permutations = tokenizer_info.permutations
+
+        models = set(tensors.keys())
+        if self.base_model:
+            models.add(self.base_model)
+        models = list(models)
+
+        vocab = tokenizer.get_vocab()
+        vocab_size = len(vocab)
+        if self.pad_to_multiple_of and vocab_size % self.pad_to_multiple_of:
+            vocab_size = (
+                vocab_size // self.pad_to_multiple_of + 1
+            ) * self.pad_to_multiple_of
+        embed_size = tensors[models[0]].shape[1]
+        assert all(
+            t.shape[1] == embed_size for t in tensors.values()
+        ), "Embedding sizes must match"
+
+        dtype = tensors[models[0]].dtype
+        device = tensors[models[0]].device
+
+        token_configs = dict(**(self.tokens or {}))
+        tokens_to_average = self.assign_embedding_sources(
+            permutations, models, vocab, token_configs
+        )
+
+        default_embeds = {}
+        for token, token_id in vocab.items():
+            embed = torch.zeros(embed_size, dtype=dtype, device=device)
+            if token in tokens_to_average:
+                count = 0
+                for model in models:
+                    p = permutations[model]
+                    if p[token_id] < 0:
+                        continue
+                    embed += tensors[model][p[token_id]]
+                    count += 1
+                embed /= count
+            elif cfg := token_configs.get(token, None):
+                cfg: TokenEmbeddingConfig
+                embed = self.compute_default_embedding(
+                    tokenizer_info, tensors, permutations, token, token_id, cfg,
+                    embed_size, dtype, device  
+                )
+            else:
+                continue
+            default_embeds[token] = embed
+
+        result = {}
+        for model in models:
+            p = permutations[model]
+            old_embed = tensors[model]
+            new_embed = torch.zeros(
+                (vocab_size, embed_size), dtype=dtype, device=device
+            )
+            for token, token_id in vocab.items():
+                force = False
+                if token in token_configs:
+                    force = token_configs[token].force
+
+                if p[token_id] >= 0 and not force:
+                    new_embed[token_id, :] = old_embed[p[token_id]]
+                elif token in default_embeds:
+                    new_embed[token_id, :] = default_embeds[token]
+                else:
+                    logging.error(
+                        f"No embedding for token {repr(token)} in model {model}!"
+                    )
+
+            if vocab_size > len(vocab):
+                # as suggested by https://nlp.stanford.edu/~johnhew/vocab-expansion.html
+                avg_embed = torch.mean(new_embed[: len(vocab), :], dim=0)
+                new_embed[len(vocab) :, :] = avg_embed
+            result[model] = new_embed
+
+        return result
+
+    def assign_embedding_sources(
+        self,
+        permutations: Dict[ModelReference, Dict[int, int]],
+        models: list[ModelReference],
+        vocab: Dict[str, int],
+        token_configs: Dict[str, TokenEmbeddingConfig],
+    ):
+        permutation_list = [permutations[model] for model in models]
+
+        tokens_to_average = set()
+        # find tokens that are only present in one model
+        for token, token_id in vocab.items():
+            if token in token_configs:
+                continue
+
+            has_token = [p[token_id] >= 0 for p in permutation_list]
+            num_present = sum(int(x) for x in has_token)
+            if num_present == 1:
+                donor_model = models[has_token.index(True)]
+                token_configs[token] = TokenEmbeddingConfig(source=donor_model)
+                continue
+
+            if num_present == 0:
+                token_configs[token] = TokenEmbeddingConfig(source=ZeroEmbedding(kind="zero"))
+                logging.warning(f"Token {repr(token)} not found in any model")
+                continue
+
+            if num_present > 0 and self.base_model is not None:
+                if permutations[self.base_model][token_id] >= 0:
+                    token_configs[token] = TokenEmbeddingConfig(source=self.base_model)
+                    continue
+
+            tokens_to_average.add(token)
+        return tokens_to_average
+
+    def compute_default_embedding(
+        self,
+        tokenizer_info: TokenizerInfo,
+        tensors: Dict[ModelReference, torch.Tensor],
+        permutations: Dict[ModelReference, Dict[int, int]],
+        token: str,
+        token_id: int,
+        cfg: TokenEmbeddingConfig,  
+        embed_size: int,  
+        dtype: torch.dtype,  
+        device: torch.device,
+    ) -> torch.Tensor:
+        if isinstance(cfg.source, ZeroEmbedding):
+            embed = torch.zeros(embed_size, dtype=dtype, device=device)
+            pass
+        elif isinstance(cfg.source, ModelTokenEmbedding):
+            model = cfg.source.model
+            assert (
+                model in permutations
+            ), f"Model {model} referenced but not part of merge"
+            p = permutations[model]
+            src_token_id = cfg.source.token_id
+            if src_token_id is None:
+                src_token = cfg.source.token
+                assert (
+                    src_token in tokenizer_info.original_vocabs[model]
+                ), f"Token {repr(src_token)} not found in model {model}"
+                src_token_id = tokenizer_info.original_vocabs[model][src_token]
+            assert (
+                src_token_id >= 0 and src_token_id < tensors[model].shape[0]
+            ), f"Token ID {src_token_id} out of range for model {model}"
+            embed = tensors[model][src_token_id]
+        elif isinstance(cfg.source, ModelReference):
+            model = cfg.source
+            p = permutations[model]
+            assert p[token_id] >= 0, f"Token {repr(token)} not found in model {model}"
+            embed = tensors[model][p[token_id]]
+        else:
+            raise NotImplementedError(cfg)
+        return embed

model_tools.md

@@ -8,7 +8,7 @@ pinned: false
 ---
 
 # Model Tools by Naphula
-Tools to enhance LLM quantizations and merging. Merge and audit large language models with low VRAM.
+Tools to enhance LLM quantizations and merging. Merge and audit large language models on low VRAM GPUs.
 
 # [graph_v18.py](https://huggingface.co/spaces/Naphula/model_tools/blob/main/graph_v18.py)
 - Merge models in minutes instead of hours on low VRAM. For a 3060/3060 Ti user: This script enables functionality that is otherwise impossible (merging 70B models or large 7B merges with `--cuda`) without OOM. [More details here](https://huggingface.co/spaces/Naphula/model_tools/blob/main/mergekit_low-VRAM-graph_patch.md)
@@ -17,6 +17,10 @@ Tools to enhance LLM quantizations and merging. Merge and audit large language m
 # config.py
 - Simply replace line 13 | BEFORE `ScalarOrGradient: TypeAlias = Union[float, List[float]]` → AFTER `ScalarOrGradient: TypeAlias = Union[float, List[float], str, bool]` | to allow for custom filepath strings within parameter settings.
 
+# [embed_12B.py](https://huggingface.co/spaces/Naphula/model_tools/blob/main/embed_12B.py) and [embed_24B.py](https://huggingface.co/spaces/Naphula/model_tools/blob/main/embed_24B.py)
+- This is an alternate solution in cases where `--fix-mistral-regex` and `tokensurgeon` fail, such as `della` or `passthrough` merges between models with mismatched `vocab_size`. Read [the guide](https://huggingface.co/spaces/Naphula/model_tools/blob/main/Mergekit-Robustness-Patch-embed_v2.md) here, download either file and save it as `mergekit-main\mergekit\tokenizer\embed.py`. Attached is one for Mistral Nemo 12B (v2d), and another for Mistral Small 24B (v2a).
+- I noticed that sometimes the default `embed.py` works best so keep a copy of that too, and if it fails for some reason try the 12B or 24B version.
+
 # [enable_fix_mistral_regex_true.md](https://huggingface.co/spaces/Naphula/model_tools/blob/main/enable_fix_mistral_regex_true.md)
 - Merge models with extreme tokenizer incompatibility. Requires modifying the `mergekit.yaml` `tokenizer` section and adding `--fix-mistral-regex` to your merge commands. (Note: Do not use `token_surgeon.py`, `gen_id_patcher.py`, or `vocab_id_patcher.py` with this, they are obsolete now.) Configured for MN 12B by default. Follow the steps in this guide to modify these scripts:
 - `mergekit/merge.py`
@@ -55,6 +59,7 @@ Tools to enhance LLM quantizations and merging. Merge and audit large language m
 
 # [arcee_fusion_salience_scanner.py](https://huggingface.co/spaces/Naphula/model_tools/blob/main/arcee_fusion_salience_scanner.py)
 - Scan the salience % of your arcee_fusion merges. The default `tukey_fence` value is 1.5 which results in 12.5% salience, but [this can be adjusted (see guide here)](modify_arcee_fusion_tukey_fence_parameter.md).
+- Updated version here [arcee_fusion_salience_scanner_v3.py](https://huggingface.co/spaces/Naphula/model_tools/blob/main/arcee_fusion_salience_scanner_v3.py)
 
 # [eos_scanner.py](https://huggingface.co/spaces/Naphula/model_tools/blob/main/eos_scanner.py)
 - Updated! This tool scans the tokenizer jsons to detect any mismatches with EOS tokens, which cause early termination bugs. You can then use the [gen_id_patcher.py](https://huggingface.co/spaces/Naphula/model_tools/blob/main/gen_id_patcher.py) and [vocab_id_patcher.py](https://huggingface.co/spaces/Naphula/model_tools/blob/main/vocab_id_patcher.py), or the [chatml_to_mistral.py](https://huggingface.co/spaces/Naphula/model_tools/blob/main/chatml_to_mistral.py) to patch missing `generation_config.json` files for EOS token. See [this post](https://huggingface.co/Naphula/Q0_Bench/discussions/1?not-for-all-audiences=true#6987717c762f0a45f672e250) as well as the [EOS Scanner ReadMe](https://huggingface.co/spaces/Naphula/model_tools/blob/main/eos_scanner_readme.md) for more info.