Vocab mismatch fix for cross-arch merging

td_fuse/merge.py

@@ -39,6 +39,7 @@ from .canary import inject_canary, test_all_canaries
 from .transport import (
     setup_tm_repo,
     load_calibration_data,
+    retokenize_calibration,
     extract_activations,
     compute_transport_plans,
     fuse_weights,
@@ -662,6 +663,7 @@ def run_single_merge(
     protection: MergeProtection,
     residual_bank: ResidualBank = None,
     calibration_data: list = None,
+    calibration_raw_texts: list = None,
     baseline_perplexity: float = None,
     merged_sources: list = None,
 ) -> dict:
@@ -717,14 +719,26 @@ def run_single_merge(
     print(f"\n[merge] Step 3/10: Loading calibration data..."); sys.stdout.flush()
     step_t = time.time()
     if calibration_data is None:
-        calibration_data = load_calibration_data(cfg, target_tokenizer)
+        calibration_data, calibration_raw_texts = load_calibration_data(cfg, target_tokenizer)
     print(f"[merge] Step 3/10 done in {time.time()-step_t:.0f}s"); sys.stdout.flush()
 
     # --- Step 4: Extract activations ---
     print(f"\n[merge] Step 4/10: Extracting activations (both models)..."); sys.stdout.flush()
     step_t = time.time()
-    print(f"[merge] Extracting source activations...")
-    source_activations = extract_activations(source_model, calibration_data)
+    # Check if source model has a different vocabulary size than target.
+    source_vocab_size = source_model.config.vocab_size if hasattr(source_model.config, 'vocab_size') else None
+    target_vocab_size = target_model.config.vocab_size if hasattr(target_model.config, 'vocab_size') else None
+    print(f"[merge] Vocab sizes -- target: {target_vocab_size}, source: {source_vocab_size}")
+
+    if source_vocab_size and target_vocab_size and source_vocab_size != target_vocab_size:
+        print(f"[merge] VOCAB MISMATCH detected! Re-tokenizing calibration data for {source_config.name}...")
+        source_calibration = retokenize_calibration(calibration_raw_texts, source_tokenizer, cfg)
+        print(f"[merge] Extracting source activations (with source-tokenized data)...")
+        source_activations = extract_activations(source_model, source_calibration)
+        del source_calibration
+    else:
+        print(f"[merge] Extracting source activations...")
+        source_activations = extract_activations(source_model, calibration_data)
 
     print(f"[merge] Extracting target activations...")
     pre_merge_target_activations = extract_activations(target_model, calibration_data)
@@ -1101,7 +1115,7 @@ def run_pipeline(
     print(f"[pipeline] Baseline perplexity: {baseline_ppl:.2f}")
 
     # --- Load calibration data once ---
-    calibration_data = load_calibration_data(cfg, target_tokenizer)
+    calibration_data, calibration_raw_texts = load_calibration_data(cfg, target_tokenizer)
 
     # --- Initialize merge protection + residual bank ---
     protection = MergeProtection(cfg)
@@ -1138,6 +1152,7 @@ def run_pipeline(
             protection,
             residual_bank=residual_bank,
             calibration_data=calibration_data,
+            calibration_raw_texts=calibration_raw_texts,
             baseline_perplexity=baseline_ppl,
             merged_sources=merged_sources,
         )

td_fuse/transport.py

@@ -105,7 +105,7 @@ def setup_tm_repo(cfg: MergeConfig):
         print(f"[transport] Added T&M core to path: {core_path}")
 
 
-def load_calibration_data(cfg: MergeConfig, tokenizer: AutoTokenizer) -> list:
+def load_calibration_data(cfg: MergeConfig, tokenizer: AutoTokenizer) -> tuple:
     """
     Load calibration data for activation extraction.
 
@@ -118,6 +118,7 @@ def load_calibration_data(cfg: MergeConfig, tokenizer: AutoTokenizer) -> list:
     print(f"[transport] Loading calibration data ({cfg.calibration_samples} samples)...")
 
     samples = []
+    raw_texts = []  # Store raw text for cross-vocab re-tokenization
 
     # --- Pile: general text (600 samples) ---
     try:
@@ -140,6 +141,7 @@ def load_calibration_data(cfg: MergeConfig, tokenizer: AutoTokenizer) -> list:
                     return_tensors="pt",
                 )
                 samples.append(tokens)
+                raw_texts.append(text)
                 count += 1
                 if count % 100 == 0:
                     print(f"  Pile: {count}/600 samples loaded...")
@@ -171,6 +173,7 @@ def load_calibration_data(cfg: MergeConfig, tokenizer: AutoTokenizer) -> list:
                         return_tensors="pt",
                     )
                     samples.append(tokens)
+                    raw_texts.append(str(text))
                     count += 1
                     if count % 100 == 0:
                         print(f"  neuralmagic: {count}/{remaining} samples loaded...")
@@ -182,6 +185,32 @@ def load_calibration_data(cfg: MergeConfig, tokenizer: AutoTokenizer) -> list:
     tracker.done()
     print(f"[transport] Total calibration samples: {len(samples)}")
     sys.stdout.flush()
+    return samples, raw_texts
+
+
+def retokenize_calibration(raw_texts: list, tokenizer: AutoTokenizer, cfg: MergeConfig) -> list:
+    """
+    Re-tokenize calibration texts with a different tokenizer.
+
+    Used when the source model has a different vocabulary than the target.
+    For example, Llama (128K vocab) vs Qwen (152K vocab).
+    """
+    print(f"[transport] Re-tokenizing {len(raw_texts)} samples for source model vocabulary...")
+    sys.stdout.flush()
+    samples = []
+    for i, text in enumerate(raw_texts):
+        tokens = tokenizer(
+            text,
+            truncation=True,
+            max_length=cfg.calibration_seq_len,
+            return_tensors="pt",
+        )
+        samples.append(tokens)
+        if (i + 1) % 500 == 0:
+            print(f"  Re-tokenized {i + 1}/{len(raw_texts)} samples...")
+            sys.stdout.flush()
+    print(f"[transport] Re-tokenized {len(samples)} samples for source model")
+    sys.stdout.flush()
     return samples