feat(tuner): support shared 2D LoRA adapters on SwitchLinear

PEFT target_parameters fine-tunes (e.g. for GPT-OSS MoE) emit a single
low-rank update that is logically shared across all experts:
  lora_A.weight : (r, hidden_size)
  lora_B.weight : (out_features, r)

LoRASwitchLinear.from_base expects per-expert tensors of shape
(num_experts, r, in) / (num_experts, out, r), so loading a shared
2-D adapter into the MoE layer used to crash at fuse() time on the
matmul lora_b @ lora_a.

Detect 2-D adapter tensors in both fuse() and __call__():
- fuse(): broadcast lora_a / lora_b across the experts axis before
  the existing (lora_b @ lora_a) outer product, so every expert gets
  the same low-rank delta added to its base weight.
- __call__(): when both tensors are 2-D, do a single dense matmul
  instead of gather_mm — the dense form is correct for shared
  adapters and avoids gathering identical rows num_experts times.

The pre-existing per-expert layout (3-D tensors) continues to work
through the unchanged code paths.

mlx_lm/tuner/lora.py

@@ -135,8 +135,25 @@ class LoRASwitchLinear(nn.Module):
         num_experts, output_dims, input_dims = weight.shape
         fused_linear = SwitchLinear(input_dims, output_dims, num_experts, bias=bias)
 
-        lora_b = self.scale * self.lora_b
-        lora_a = self.lora_a.reshape(num_experts, -1, input_dims)
+        # Support two adapter formats:
+        #   (a) per-expert (mlx-lm native): lora_a (E, r, D), lora_b (E, out, r).
+        #   (b) shared / PEFT target_parameters: lora_a (r, D), lora_b (out, r).
+        #       Same low-rank update added to every expert -> broadcast to (E, ...).
+        lora_a = self.lora_a
+        lora_b = self.lora_b
+        if lora_a.ndim == 2:
+            # (r, D) -> (E, r, D); same matrix replicated across experts.
+            lora_a = mx.broadcast_to(
+                lora_a[None, :, :], (num_experts, lora_a.shape[0], lora_a.shape[1])
+            )
+        else:
+            lora_a = lora_a.reshape(num_experts, -1, input_dims)
+        if lora_b.ndim == 2:
+            lora_b = mx.broadcast_to(
+                lora_b[None, :, :], (num_experts, lora_b.shape[0], lora_b.shape[1])
+            )
+
+        lora_b = self.scale * lora_b
         fused_linear.weight = weight + (lora_b @ lora_a).astype(weight.dtype)
         if bias:
             fused_linear.bias = linear.bias
@@ -180,18 +197,24 @@ class LoRASwitchLinear(nn.Module):
 
     def __call__(self, x, indices, sorted_indices=False):
         y = self.linear(x, indices, sorted_indices=sorted_indices)
-        z = mx.gather_mm(
-            self.dropout(x),
-            self.lora_a.swapaxes(-1, -2),
-            rhs_indices=indices,
-            sorted_indices=sorted_indices,
-        )
-        z = mx.gather_mm(
-            z,
-            self.lora_b.swapaxes(-1, -2),
-            rhs_indices=indices,
-            sorted_indices=sorted_indices,
-        )
+        # Shared 2D adapter (PEFT target_parameters) -> dense matmul, broadcast to all routed experts.
+        if self.lora_a.ndim == 2 and self.lora_b.ndim == 2:
+            xd = self.dropout(x)
+            z = xd @ self.lora_a.swapaxes(-1, -2)
+            z = z @ self.lora_b.swapaxes(-1, -2)
+        else:
+            z = mx.gather_mm(
+                self.dropout(x),
+                self.lora_a.swapaxes(-1, -2),
+                rhs_indices=indices,
+                sorted_indices=sorted_indices,
+            )
+            z = mx.gather_mm(
+                z,
+                self.lora_b.swapaxes(-1, -2),
+                rhs_indices=indices,
+                sorted_indices=sorted_indices,
+            )
         return y + (self.scale * z).astype(x.dtype)
 
 

tests/test_lora_switch_linear_shared.py

@@ -0,0 +1,66 @@
+"""Tests: LoRASwitchLinear supports PEFT shared 2D lora_a/lora_b (target_parameters)."""
+import unittest
+
+import mlx.core as mx
+
+from mlx_lm.models.switch_layers import SwitchLinear
+from mlx_lm.tuner.lora import LoRASwitchLinear
+
+
+class TestLoRASwitchLinearSharedAdapter(unittest.TestCase):
+    E = 4
+    D = 16
+    H = 16
+    r = 4
+
+    def _base(self):
+        return SwitchLinear(self.D, self.H, self.E, bias=False)
+
+    def test_fuse_per_expert(self):
+        """Original (E, r, D) / (E, out, r) layout still works."""
+        base = self._base()
+        lora = LoRASwitchLinear.from_base(base, r=self.r, scale=1.0)
+        lora.lora_a = mx.random.normal(shape=(self.E, self.r, self.D))
+        lora.lora_b = mx.random.normal(shape=(self.E, self.H, self.r))
+        fused = lora.fuse()
+        self.assertEqual(fused.weight.shape, (self.E, self.H, self.D))
+
+    def test_fuse_shared_2d(self):
+        """PEFT shared (r, D) / (out, r) layout fuses to same E experts."""
+        base = self._base()
+        lora = LoRASwitchLinear.from_base(base, r=self.r, scale=1.0)
+        a = mx.random.normal(shape=(self.r, self.D))
+        b = mx.random.normal(shape=(self.H, self.r))
+        lora.lora_a = a
+        lora.lora_b = b
+        fused = lora.fuse()
+        self.assertEqual(fused.weight.shape, (self.E, self.H, self.D))
+        # Each expert receives the SAME (B @ A) update on top of its base weight.
+        delta = (b @ a).astype(fused.weight.dtype)
+        for i in range(self.E):
+            self.assertTrue(
+                mx.allclose(
+                    fused.weight[i] - base.weight[i], delta, atol=1e-4
+                ).item()
+            )
+
+    def test_call_shared_2d_matches_fuse_then_call(self):
+        """Forward pass with shared adapter equals fuse-then-call equivalent."""
+        base = self._base()
+        lora = LoRASwitchLinear.from_base(base, r=self.r, scale=1.0)
+        a = mx.random.normal(shape=(self.r, self.D))
+        b = mx.random.normal(shape=(self.H, self.r))
+        lora.lora_a = a
+        lora.lora_b = b
+
+        x = mx.random.normal(shape=(2, 1, self.D))
+        indices = mx.array([[0], [1]], dtype=mx.uint32)
+
+        y_lora = lora(x, indices)
+        fused = lora.fuse()
+        y_fused = fused(x, indices)
+        self.assertTrue(mx.allclose(y_lora, y_fused, atol=1e-3).item())
+
+
+if __name__ == "__main__":
+    unittest.main()