feat(moe): add orthogonal initialization for gate parameters

mergekit/architecture/json_definitions.py

@@ -168,9 +168,9 @@ def _load_architecture_json(text: str) -> ModelArchitecture:
         raise RuntimeError(f"Unexpected architecture kind: {data['kind']}")
 
 
-def _load_all_architectures() -> (
-    Tuple[List[ModelArchitecture], Dict[str, List[ModelArchitecture]]]
-):
+def _load_all_architectures() -> Tuple[
+    List[ModelArchitecture], Dict[str, List[ModelArchitecture]]
+]:
     architectures: List[ModelArchitecture] = []
     for f in importlib.resources.files(mergekit._data.architectures).iterdir():
         if f.is_file() and f.name.lower().endswith(".json"):

mergekit/common.py

@@ -153,9 +153,9 @@ def local_path(
         self, cache_dir: Optional[str] = None, ignore_lora: bool = False
     ) -> str:
         if not ignore_lora:
-            assert (
-                self.lora is None
-            ), "LoRA not merged - use .merged() to get a local path"
+            assert self.lora is None, (
+                "LoRA not merged - use .merged() to get a local path"
+            )
 
         path = self.model.path
         if not os.path.exists(path):
@@ -400,7 +400,7 @@ def get_torch_accelerator_count(accelerator_name: Optional[str] = None):
     if accelerator_name is not None:
         accelerator = torch.device(accelerator_name)
         # if user passes the device index in `accelerator_name`, then 1
-        if accelerator.index != None:
+        if accelerator.index is not None:
             return 1
         torch_accelerator_module = getattr(torch, accelerator.type)
     else:

mergekit/evo/actors.py

@@ -275,9 +275,9 @@ def evaluate(self, genotype: torch.Tensor) -> dict:
 
         model = self.model.model
         if vllm is not None and isinstance(model, vllm.LLM):
-            assert (
-                model.llm_engine.parallel_config.world_size == 1
-            ), "Must be single GPU"
+            assert model.llm_engine.parallel_config.world_size == 1, (
+                "Must be single GPU"
+            )
             engine = model.llm_engine
             if hasattr(engine, "model_executor"):
                 worker = engine.model_executor.worker

mergekit/evo/genome.py

@@ -46,9 +46,9 @@ def validate(self):
 
         if self.merge_method == "slerp":
             assert not self.smooth, "smooth is not supported for slerp merge method"
-            assert (
-                not self.filters
-            ), "tensor name filtering is not supported for slerp merge method"
+            assert not self.filters, (
+                "tensor name filtering is not supported for slerp merge method"
+            )
 
         return self
 

mergekit/graph.py

@@ -167,9 +167,9 @@ def add_task(self, task: Task, recursive: bool = True) -> "TaskHandle":
         _ti_key = (type(task), id(task))
         if _ti_key in self._type_id_to_index:
             index = self._type_id_to_index[_ti_key]
-            assert (
-                self.tasks[index] == task
-            ), "Task modified after being added to universe"
+            assert self.tasks[index] == task, (
+                "Task modified after being added to universe"
+            )
             return TaskHandle(self, index)
 
         index = self.task_to_index.setdefault(task, len(self.tasks))
@@ -322,9 +322,9 @@ def build_schedule(
         return ExecutionSchedule(tasks=[], last_use_index={})
 
     universe = targets[0]._universe
-    assert all(
-        t._universe is universe for t in targets
-    ), "All tasks must be from the same universe"
+    assert all(t._universe is universe for t in targets), (
+        "All tasks must be from the same universe"
+    )
 
     dummy_handle = TaskHandle(universe, -1)
     edge_tups: List[Tuple[TaskHandle, TaskHandle]] = []

mergekit/merge_methods/easy_define.py

@@ -42,7 +42,7 @@ def __merge_method(
         (tensor_param.annotation is None)
         or (not hasattr(tensor_param.annotation, "__origin__"))
         or not (
-            tensor_param.annotation.__origin__ == list
+            tensor_param.annotation.__origin__ is list
             and tensor_param.annotation.__args__ == (torch.Tensor,)
         )
     ):
@@ -89,7 +89,7 @@ def __merge_method(
                 )
             elif (
                 hasattr(arg_info.annotation, "__origin__")
-                and arg_info.annotation.__origin__ == list
+                and arg_info.annotation.__origin__ is list
                 and arg_info.annotation.__args__[0] in (float, int)
             ):
                 default_value = arg_info.default

mergekit/moe/config.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: LGPL-3.0-only
 
 import logging
-from typing import List, Optional
+from typing import List, Optional, Literal  # Ensure Literal is imported
 
 from pydantic import BaseModel
 
@@ -29,13 +29,24 @@ class MoEMergeConfig(BaseModel):
 
     base_model: ModelReference
     experts: List[Expert]
-    gate_mode: str = (
-        "hidden"  # possible values: "hidden", "cheap_embed", "random", "uniform_random"
-    )
+
+    # Updated to use Literal for strict validation and added "orthogonal"
+    gate_mode: Literal[
+        "hidden",
+        "cheap_embed",
+        "random",
+        "uniform_random",
+        "orthogonal",
+        "hidden_avg",
+        "hidden_last",
+    ] = "hidden"
+
     # "hidden" uses hidden state vectors for the given prompts for each layer
     # "cheap_embed" uses the average of token embeddings for the prompts, same for each layer
-    # "random" is random
+    # "random" is standard normal distribution (torch.randn)
     # "uniform_random" matches default initialization for torch.nn.Linear
+    # "orthogonal" ensures gate vectors are orthogonal for better expert specialization
+
     dtype: Optional[str] = None
     experts_per_token: int = 2
     shared_experts: Optional[List[Expert]] = None

mergekit/moe/router.py

@@ -88,6 +88,30 @@ def get_gate_params(
         return torch.randn(
             (model_cfg.num_hidden_layers, len(experts), model_cfg.hidden_size)
         )
+    # NEW test Orthogonal Initialization
+    elif mode == "orthogonal":
+        num_layers = model_cfg.num_hidden_layers
+        num_experts = len(experts)
+        hidden_size = model_cfg.hidden_size
+
+        # 1. Determine the target dtype
+        # We check if a specific dtype was requested, otherwise use model default
+        target_dtype = getattr(model_cfg, "torch_dtype", torch.float16)
+
+        # 2. Initialize in float32 for mathematical stability
+        # We create a list of tensors to match how "hidden" mode returns data
+        gate_vecs = []
+        for _ in range(num_layers):
+            layer_gate = torch.empty((num_experts, hidden_size), dtype=torch.float32)
+            torch.nn.init.orthogonal_(layer_gate)
+            # 3. Cast to the target dtype and move to the requested device
+            gate_vecs.append(
+                layer_gate.to(
+                    dtype=target_dtype, device=device if device != "auto" else "cpu"
+                )
+            )
+
+        return gate_vecs
     elif mode == "uniform_random":
         in_features = model_cfg.hidden_size
         scale = math.sqrt(1.0 / in_features)

mergekit/options.py

@@ -178,7 +178,7 @@ def wrapper(*args, **kwargs):
             field_type = ShardSizeParamType()
 
         arg_name = field_name.replace("_", "-")
-        if field_type == bool:
+        if field_type is bool:
             arg_str = f"--{arg_name}/--no-{arg_name}"
         else:
             arg_str = f"--{arg_name}"

mergekit/scripts/tokensurgeon.py

@@ -472,7 +472,7 @@ def build_embedding_matrix(
             new_tokens.append(token)
             stats.to_approximate += 1
 
-    donor_tokenizer = transformers.AutoTokenizer.from_pretrained(
+    transformers.AutoTokenizer.from_pretrained(
         options.donor.model.path,
         revision=options.donor.model.revision,
         trust_remote_code=True,

mergekit/tokenizer/embed.py

@@ -45,9 +45,9 @@ def execute(
                 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"
+        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
@@ -155,20 +155,20 @@ def compute_default_embedding(
             pass
         elif isinstance(cfg.source, ModelTokenEmbedding):
             model = cfg.source.model
-            assert (
-                model in permutations
-            ), f"Model {model} referenced but not part of merge"
+            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}"
+                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}"
+            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

mergekit/tokensurgeon/omp.py

@@ -73,14 +73,10 @@ def batch_omp(
             # Project onto existing basis
             projections = torch.bmm(
                 q[:, :, :t].transpose(1, 2), new_atom.unsqueeze(-1)
-            ).squeeze(
-                -1
-            )  # (B, t)
+            ).squeeze(-1)  # (B, t)
             residual = new_atom - torch.bmm(
                 q[:, :, :t], projections.unsqueeze(-1)
-            ).squeeze(
-                -1
-            )  # (B, D)
+            ).squeeze(-1)  # (B, D)
             norm = torch.clamp(torch.norm(residual, dim=1), min=eps)
             # Update R and Q
             r[:, :t, t] = projections
@@ -218,9 +214,9 @@ def batch_mp_rope(
     B, D_a = targets.shape
     N, _ = points_a.shape
     _, D_b = points_b.shape
-    assert (
-        points_a.shape[0] == points_b.shape[0]
-    ), "Number of points in A and B must match"
+    assert points_a.shape[0] == points_b.shape[0], (
+        "Number of points in A and B must match"
+    )
     device = targets.device
     if k > N:
         raise ValueError(f"Cannot select {k} points from {N} candidates")

mergekit/tokensurgeon/pca.py

@@ -22,9 +22,9 @@ def landmark_pca_approximate(
     num_points, d_a = points_a.shape
     batch_size, _ = targets.shape
     _, d_b = points_b.shape
-    assert (
-        points_a.shape[0] == points_b.shape[0]
-    ), "Number of points in A and B must match"
+    assert points_a.shape[0] == points_b.shape[0], (
+        "Number of points in A and B must match"
+    )
     assert targets.shape == (batch_size, d_a)
 
     effective_dim = min(d_a, d_b)

mergekit/tokensurgeon/rope_helpers.py

@@ -11,7 +11,7 @@ def llama_rope_rotationmat(theta: torch.Tensor) -> torch.Tensor:
         theta: Tensor of shape (..., n_heads, head_dim // 2) representing the angles for the rotation.
     """
     # theta shape: (..., n_heads, head_dim // 2)
-    n_heads = theta.shape[-2]
+    theta.shape[-2]
     head_dim = theta.shape[-1] * 2
     theta_p = torch.cat([theta, theta], dim=-1)
     cos_theta = torch.cos(theta_p)

pyproject.toml

@@ -93,4 +93,9 @@ testpaths = ["tests"]
 
 [dependency-groups]
 test = ["pytest~=8.4.0"]
-dev = ["black~=25.1.0", "isort~=6.0.1", "pre-commit~=4.2.0"]
+dev = [
+"black~=25.1.0",
+ "isort~=6.0.1",
+ "pre-commit~=4.2.0",
+ "ruff>=0.15.0",
+]

tests/common.py

@@ -39,9 +39,9 @@ def run_and_check_merge(
         index_exists = os.path.exists(index_path)
         single_shard_exists = os.path.exists(index_path.replace(".index.json", ""))
         assert index_exists or single_shard_exists, "No model produced by merge"
-        assert os.path.exists(
-            os.path.join(tmpdir, "config.json")
-        ), "No config json produced by merge"
+        assert os.path.exists(os.path.join(tmpdir, "config.json")), (
+            "No config json produced by merge"
+        )
 
         if check_nan:
             # check for NaN in output

tests/test_chat_template.py

@@ -25,9 +25,9 @@ def check_chat_template(model_path: str, needle: Optional[str] = None):
     if needle is None:
         assert not tokenizer.chat_template, "Expected no chat template"
         return
-    assert (
-        tokenizer.chat_template and needle in tokenizer.chat_template
-    ), f"Expected chat template to contain {needle}"
+    assert tokenizer.chat_template and needle in tokenizer.chat_template, (
+        f"Expected chat template to contain {needle}"
+    )
 
 
 class TestChatTemplate: