diff --git a/mlx_lm/tuner/lora.py b/mlx_lm/tuner/lora.py index 3b5ddef..04ea113 100644 --- a/mlx_lm/tuner/lora.py +++ b/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) diff --git a/tests/test_lora_switch_linear_shared.py b/tests/test_lora_switch_linear_shared.py new file mode 100644 index 0000000..75c4279 --- /dev/null +++ b/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()