fix(moe): avoid incorrect Granite grouped expert execution on MPS

docs/source/content/drivers.md

@@ -48,6 +48,96 @@ logits, cache = bridge.run_with_cache("Hello, world")
 
 There is also `TransformerBridge.boot_native(config)`, which builds a small randomly-initialized TL-native model on the same driver — no HuggingFace Hub call — useful for tests and toy-model experiments.
 
+#### Granite MoE Hybrid grouped experts on Apple MPS
+
+The local transformers backend applies a narrow correctness fallback for the exact
+HF `GraniteMoeHybridExperts` class when all of these conditions hold:
+
+- PyTorch release **2.11.0** and transformers release **5.13.0**, including local
+  version suffixes but excluding prerelease and development versions.
+- The live expert input is on **MPS**.
+- The expert config selects **`grouped_mm`**.
+
+In this verified environment, grouped expert matrix products can be incorrect or
+unstable. The HF grouped fallback also uses nonzero-offset `torch.mm(..., out=...)`
+slice writes that reproduce incorrect results, so disabling only native grouped
+GEMM is insufficient. The bridge selects HF's original ordinary eager expert
+forward instead, emitting a `RuntimeWarning` about the correctness fallback and
+its possible performance cost. Ordinary eager execution may be slower; no
+throughput guarantee is made.
+
+Dispatch is installed on the expert instance when its containing `MoEBridge`
+runs, and checks the live input device and expert policy on every expert call.
+CPU-to-MPS and MPS-to-CPU migration therefore follow the current device. Shared
+HF configs, global dispatch tables, weights, routing indices, and routing weights
+are not mutated. Expert module hooks still execute, including output edits, and
+input, routing-weight, and expert-parameter gradients remain connected.
+
+The fallback does not serialize forwards or mutate shared configuration while
+other layers run. This avoids a shared-policy race, but does not guarantee that
+arbitrary concurrent model calls, device migration, config edits, or hook edits
+are safe. Direct HF execution before the containing bridge's first call is not
+protected; after installation, calls to that same expert instance use the local
+wrapper. Custom instance forwards are left unchanged.
+
+CPU/CUDA inputs, explicit eager execution, other expert policies, and other
+expert classes retain their existing dispatch. The observed defect and regression
+validation cover only the versions above, Apple MPS, and float32. Older/newer
+releases are neither declared affected nor declared fixed, and CUDA behavior has
+not been device-validated for this fallback. Unverified versions receive no
+workaround. MPS remains opt-in via `TRANSFORMERLENS_ALLOW_MPS=1`; this fallback is
+not a general guarantee of MPS correctness.
+
+##### Contributor diagnostics
+
+From the repository root, the offline regression selection requires no Hub
+model download and checks independent expert-oracle agreement, repeatability,
+device migration, hooks, gradients, and failure recovery:
+
+```bash
+TRANSFORMERLENS_ALLOW_MPS=1 uv run pytest \
+  tests/mps/test_granite_grouped_expert_dispatch.py -q
+uv run pytest \
+  tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py -q
+TRANSFORMERLENS_ALLOW_MPS=1 uv run pytest \
+  tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py -q
+```
+
+The primitive diagnostic below is independent of TransformerLens and HF expert
+routing. It reports error rather than asserting that an upstream kernel must
+remain defective. Run it on an MPS-capable host and record the installed versions
+when evaluating an upstream fix:
+
+```python
+import platform
+
+import torch
+import transformers
+
+if not torch.backends.mps.is_available():
+    raise RuntimeError("This diagnostic requires MPS")
+
+print(platform.system(), platform.release(), platform.machine())
+print("torch:", torch.__version__, "transformers:", transformers.__version__)
+generator = torch.Generator().manual_seed(0)
+inputs = torch.randn(10, 64, generator=generator)
+weights = torch.randn(4, 64, 32, generator=generator)
+expected = inputs[5:] @ weights[3]
+x, w = inputs.to("mps"), weights.to("mps")
+with torch.no_grad():
+    output = torch.zeros(10, 32, device="mps")
+    torch.mm(x[5:], w[3], out=output[5:])
+    assigned = torch.zeros_like(output)
+    assigned[5:] = x[5:] @ w[3]
+    torch.mps.synchronize()
+    print("out-slice max error:", (output[5:].cpu() - expected).abs().max().item())
+    print("assignment max error:", (assigned[5:].cpu() - expected).abs().max().item())
+```
+
+No dependency pin change or confirmed upstream fixed version is implied by this
+workaround. Revalidate both the primitive and full expert regressions before
+changing its scope.
+
 ### vLLM — high-throughput capture + declarative interventions
 
 `RemoteBridge.boot_vllm` constructs a vLLM engine (PagedAttention, `torch.compile`, CUDA graphs) and installs capture hooks inside the worker *before* compilation. Activations come back over `collective_rpc` and replay through the bridge's HookPoint tree. This is the throughput path for SAE/probe data collection; unlike observation-only tools (vllm-lens), each hook also applies an affine transform `output = output * scale + bias`, so declarative interventions (`suppress` / `scale` / `add` / `set`) propagate to downstream layers.

tests/mps/test_granite_grouped_expert_dispatch.py

@@ -0,0 +1,227 @@
+"""Granite expert correctness on MPS with fixed routes and an independent oracle."""
+
+import copy
+import warnings
+from types import SimpleNamespace
+
+import pytest
+import torch
+import transformers
+from packaging.version import Version
+from torch import nn
+from torch.nn import functional as F
+from transformers.models.granitemoehybrid.modeling_granitemoehybrid import (
+    GraniteMoeHybridExperts,
+)
+
+from transformer_lens.model_bridge.generalized_components.moe import MoEBridge
+
+pytestmark = pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
+
+
+class _FixedRoutes(nn.Module):
+    def __init__(self):
+        super().__init__()
+        config = SimpleNamespace(
+            hidden_size=64, intermediate_size=32, num_local_experts=4, hidden_act="silu"
+        )
+        config._experts_implementation = "grouped_mm"
+        self.experts = GraniteMoeHybridExperts(config)
+        generator = torch.Generator().manual_seed(0)
+        with torch.no_grad():
+            for parameter in self.experts.parameters():
+                parameter.copy_(torch.randn(parameter.shape, generator=generator) * 0.1)
+        self.register_buffer(
+            "indices", torch.tensor([[0], [0], [1], [1], [1], [3], [3], [3], [3], [3]])
+        )
+        self.register_buffer("weights", torch.linspace(0.2, 1.0, 10).unsqueeze(1))
+
+    def forward(self, layer_input, top_k_weights=None):
+        weights = self.weights if top_k_weights is None else top_k_weights
+        return self.experts(layer_input, self.indices, weights)
+
+
+def _oracle(block, inputs, top_k_weights=None):
+    contributions = []
+    route_weights = block.weights if top_k_weights is None else top_k_weights
+    for token, indices, weights in zip(inputs, block.indices, route_weights):
+        routes = []
+        for index, weight in zip(indices, weights):
+            gate, up = F.linear(token, block.experts.gate_up_proj[index]).chunk(2, dim=-1)
+            routes.append(F.linear(F.silu(gate) * up, block.experts.down_proj[index]) * weight)
+        contributions.append(torch.stack(routes).sum(0))
+    return torch.stack(contributions)
+
+
+def test_fixed_route_experts_match_independent_oracle_and_repeat():
+    block = _FixedRoutes()
+    inputs = torch.randn(10, 64, generator=torch.Generator().manual_seed(1))
+    with torch.no_grad():
+        expected = _oracle(block, inputs)
+    block.to("mps")
+    bridge = MoEBridge("moe")
+    bridge.set_original_component(block)
+    observed = []
+    handle = block.experts.register_forward_hook(
+        lambda module, args, output: observed.append(output)
+    )
+    try:
+        with torch.no_grad(), warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            first = bridge(inputs.to("mps"))
+            second = bridge(layer_input=inputs.to("mps"))
+        assert first.device.type == "mps"
+        torch.testing.assert_close(first.cpu(), expected, rtol=1e-4, atol=1e-6)
+        torch.testing.assert_close(second, first, rtol=0, atol=0)
+        if (
+            Version(torch.__version__).public == "2.11.0"
+            and Version(transformers.__version__).public == "5.13.0"
+        ):
+            assert any("eager Granite" in str(warning.message) for warning in caught)
+        assert len(observed) == 2
+        assert block.experts.config._experts_implementation == "grouped_mm"
+    finally:
+        handle.remove()
+
+
+def test_dispatch_follows_cpu_mps_cpu_migration():
+    block = _FixedRoutes()
+    bridge = MoEBridge("moe")
+    bridge.set_original_component(block)
+    inputs = torch.randn(10, 64, generator=torch.Generator().manual_seed(1))
+    with torch.no_grad(), warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        cpu_before = bridge(inputs)
+    assert not caught
+    expected = _oracle(block, inputs)
+    torch.testing.assert_close(cpu_before, expected, rtol=1e-4, atol=1e-6)
+    bridge.to("mps")
+    with torch.no_grad():
+        metal = bridge(inputs.to("mps"))
+    assert metal.device.type == "mps"
+    torch.testing.assert_close(metal.cpu(), expected, rtol=1e-4, atol=1e-6)
+    bridge.to("cpu")
+    with torch.no_grad(), warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        cpu_after = bridge(inputs)
+    assert not caught
+    torch.testing.assert_close(cpu_after, cpu_before, rtol=0, atol=0)
+    assert block.experts.config._experts_implementation == "grouped_mm"
+
+
+def test_explicit_eager_policy_and_expert_hooks_are_preserved():
+    block = _FixedRoutes().to("mps")
+    block.experts.config._experts_implementation = "eager"
+    bridge = MoEBridge("moe")
+    bridge.set_original_component(block)
+    inputs = torch.randn(10, 64).to("mps")
+    indices, weights = block.indices.clone(), block.weights.clone()
+    calls = []
+
+    def observe(module, args, output):
+        calls.append(args)
+        return output * 2
+
+    handle = block.experts.register_forward_hook(observe)
+    try:
+        with torch.no_grad(), warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            expected = _oracle(block, inputs) * 2
+            actual = bridge(inputs)
+        assert not caught
+        assert len(calls) == 1
+        torch.testing.assert_close(calls[0][0], inputs, rtol=0, atol=0)
+        torch.testing.assert_close(calls[0][1], indices, rtol=0, atol=0)
+        torch.testing.assert_close(calls[0][2], weights, rtol=0, atol=0)
+        torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-6)
+        assert actual.shape == inputs.shape
+        assert block.experts.config._experts_implementation == "eager"
+        torch.testing.assert_close(block.indices, indices, rtol=0, atol=0)
+        torch.testing.assert_close(block.weights, weights, rtol=0, atol=0)
+    finally:
+        handle.remove()
+
+
+def test_corrected_experts_preserve_input_route_and_parameter_gradients():
+    block = _FixedRoutes()
+    block.indices = torch.tensor(
+        [
+            [0, 1],
+            [0, 3],
+            [1, 0],
+            [1, 3],
+            [1, 0],
+            [3, 0],
+            [3, 1],
+            [3, 0],
+            [3, 1],
+            [3, 0],
+        ]
+    )
+    reference = copy.deepcopy(block)
+    block.to("mps")
+    bridge = MoEBridge("moe")
+    bridge.set_original_component(block)
+    data = torch.randn(10, 64, generator=torch.Generator().manual_seed(1))
+    cpu_input = data.clone().requires_grad_()
+    mps_input = data.to("mps").requires_grad_()
+    cpu_routes = torch.linspace(0.2, 1.0, 20).reshape(10, 2).requires_grad_()
+    mps_routes = torch.linspace(0.2, 1.0, 20, device="mps").reshape(10, 2).requires_grad_()
+    calls = []
+
+    def observe(module, args, output):
+        calls.append(args)
+        return output * 2
+
+    handle = block.experts.register_forward_hook(observe)
+    try:
+        expected = _oracle(reference, cpu_input, cpu_routes) * 2
+        actual = bridge(mps_input, top_k_weights=mps_routes)
+    finally:
+        handle.remove()
+    assert len(calls) == 1
+    assert calls[0][0] is mps_input
+    assert calls[0][1] is block.indices
+    assert calls[0][2] is mps_routes
+    assert actual.shape == mps_input.shape
+    torch.testing.assert_close(actual.cpu(), expected, rtol=1e-4, atol=1e-6)
+    cpu_targets = (cpu_input, cpu_routes, *reference.experts.parameters())
+    mps_targets = (mps_input, mps_routes, *block.experts.parameters())
+    cpu_grads = torch.autograd.grad(expected.square().sum(), cpu_targets)
+    mps_grads = torch.autograd.grad(actual.square().sum(), mps_targets)
+    for measured, oracle in zip(mps_grads, cpu_grads):
+        assert measured.device.type == "mps"
+        torch.testing.assert_close(measured.cpu(), oracle, rtol=1e-4, atol=1e-6)
+    assert cpu_grads[1].abs().max() > 0
+    assert all(gradient.abs().max() > 0 for gradient in cpu_grads[2:])
+    torch.testing.assert_close(block.indices.cpu(), reference.indices, rtol=0, atol=0)
+    for parameter, saved in zip(block.experts.parameters(), reference.experts.parameters()):
+        torch.testing.assert_close(parameter.cpu(), saved, rtol=0, atol=0)
+    assert block.experts.config._experts_implementation == "grouped_mm"
+
+
+def test_failed_expert_forward_leaves_policy_and_parameters_intact():
+    block = _FixedRoutes().to("mps")
+    bridge = MoEBridge("moe")
+    bridge.set_original_component(block)
+    inputs = torch.randn(10, 64).to("mps")
+    parameters = [parameter.clone() for parameter in block.experts.parameters()]
+    activation = block.experts.act_fn
+
+    class FailingActivation(nn.Module):
+        def forward(self, value):
+            raise ValueError("expert activation failed")
+
+    block.experts.act_fn = FailingActivation()
+    try:
+        with pytest.raises(ValueError, match="expert activation failed"):
+            bridge(inputs)
+    finally:
+        block.experts.act_fn = activation
+    assert block.experts.config._experts_implementation == "grouped_mm"
+    with torch.no_grad():
+        actual = bridge(inputs)
+        expected = _oracle(block, inputs)
+    torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-6)
+    for parameter, saved in zip(block.experts.parameters(), parameters):
+        torch.testing.assert_close(parameter, saved, rtol=0, atol=0)