feat: Deterministic base-architecture preference for cross-family merges

mergekit/architecture/__init__.py

@@ -86,6 +86,15 @@ def get_architecture_info(
             raise RuntimeError(
                 "Must specify --allow-crimes to attempt to mix different architectures"
             )
+        # Prefer using the base_model's architecture when mixing different families
+        # to ensure the output layout matches the base.
+        if config.base_model is not None:
+            try:
+                idx = models.index(config.base_model)
+                return model_arch_info[idx]
+            except ValueError:
+                # base_model not in referenced models; fall back to first
+                pass
         return model_arch_info[0]
 
     # try to infer from all models

mergekit/io/tensor_writer.py

@@ -39,6 +39,7 @@ def __init__(
         use_async: bool = False,
         max_write_threads: int = 1,
     ) -> None:
+        out_path = os.path.abspath(os.path.expanduser(out_path))
         os.makedirs(out_path, exist_ok=True)
 
         self.out_path = out_path

mergekit/merge_methods/slerp.py

@@ -21,7 +21,7 @@
 class SlerpTask(Task[torch.Tensor]):
     gather_tensors: MergeTensorInput
     base_model: ModelReference
-    t: float
+    t: Optional[float]
     weight_info: WeightInfo
 
     def uses_accelerator(self) -> bool:
@@ -38,6 +38,12 @@ def execute(self, tensors: Dict[ModelReference, torch.Tensor]) -> torch.Tensor:
         elif self.base_model not in tensors:
             raise RuntimeError("Base model not in input tensors")
 
+        # If no interpolation parameter was provided for this tensor, do not attempt to merge;
+        # simply return the base model's weight unchanged. This avoids shape/broadcast errors
+        # when the secondary model has incompatible tensor shapes.
+        if self.t is None:
+            return tensors[self.base_model]
+
         [a, b] = list(tensors.items())
         if a[0] != self.base_model:
             [a, b] = [b, a]
@@ -72,7 +78,7 @@ def reference_url(self):
         return "https://en.wikipedia.org/wiki/Slerp"
 
     def parameters(self) -> List[ConfigParameterDef]:
-        return [ConfigParameterDef(name="t", required=True)]
+        return [ConfigParameterDef(name="t", required=False, default_value=None)]
 
     def make_task(
         self,
@@ -92,7 +98,9 @@ def make_task(
 
 
 def lerp(
-    t: float, v0: Union[np.ndarray, torch.Tensor], v1: Union[np.ndarray, torch.Tensor]
+    t: Optional[float],
+    v0: Union[np.ndarray, torch.Tensor],
+    v1: Union[np.ndarray, torch.Tensor],
 ) -> Union[np.ndarray, torch.Tensor]:
     return (1 - t) * v0 + t * v1
 

mergekit/tokenizer/embed.py

@@ -130,7 +130,9 @@ def assign_embedding_sources(
                 continue
 
             if num_present == 0:
-                token_configs[token] = TokenEmbeddingConfig(source=ZeroEmbedding())
+                token_configs[token] = TokenEmbeddingConfig(
+                    source=ZeroEmbedding(kind="zero")
+                )
                 logging.warning(f"Token {repr(token)} not found in any model")
                 continue
 
@@ -152,7 +154,10 @@ def compute_default_embedding(
         cfg: TokenEmbeddingConfig,
     ) -> torch.Tensor:
         if isinstance(cfg.source, ZeroEmbedding):
-            pass
+            any_tensor = next(iter(tensors.values()))
+            embed = torch.zeros(
+                any_tensor.shape[1], dtype=any_tensor.dtype, device=any_tensor.device
+            )
         elif isinstance(cfg.source, ModelTokenEmbedding):
             model = cfg.source.model
             assert (