diff --git a/docs/source/content/getting_started.md b/docs/source/content/getting_started.md index 518daf4393..5160f622bb 100644 --- a/docs/source/content/getting_started.md +++ b/docs/source/content/getting_started.md @@ -65,6 +65,11 @@ Set to `"1"` to opt in to Apple Silicon (MPS) as a target device for model infer export TRANSFORMERLENS_ALLOW_MPS=1 ``` +For Hugging Face MoE experts using the shared grouped dispatcher, `TransformerBridge` +checks the MPS matmul output-slice primitive and uses eager experts only when incorrect +writes are detected. This emits a once-per-process warning and may reduce performance. +Setting `TRANSFORMERLENS_ALLOW_MPS=1` does not disable this correctness safeguard. + ## Huggingface Gated Access For convenience, gated-model access depends only on `HF_TOKEN` above. Once you have set the token and accepted any model-specific agreements on the HuggingFace Hub, gated models load through TransformerLens with no additional configuration. The most popular gated families supported by TransformerLens are the Llama, Mistral/Mixtral, and Gemma models. diff --git a/tests/mps/test_granite_grouped_expert_dispatch.py b/tests/mps/test_granite_grouped_expert_dispatch.py new file mode 100644 index 0000000000..dc0be79f5b --- /dev/null +++ b/tests/mps/test_granite_grouped_expert_dispatch.py @@ -0,0 +1,341 @@ +"""Shared HF expert correctness on MPS with fixed routes and an independent oracle.""" + +import copy +import warnings +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from torch.nn import functional as F +from transformers.integrations import moe as hf_moe +from transformers.models.granitemoe.modeling_granitemoe import GraniteMoeExperts +from transformers.models.granitemoehybrid.modeling_granitemoehybrid import ( + GraniteMoeHybridExperts, +) +from transformers.models.mixtral.modeling_mixtral import MixtralExperts +from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeExperts + +from transformer_lens.model_bridge.generalized_components import moe +from transformer_lens.model_bridge.generalized_components.moe import ( + MoEBridge, + _mps_output_slice_writes_are_broken, +) + +pytestmark = pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available") + + +@pytest.fixture( + params=[GraniteMoeHybridExperts, GraniteMoeExperts, MixtralExperts, Qwen3MoeExperts] +) +def expert_class(request): + return request.param + + +@pytest.fixture(autouse=True) +def isolated_probe_cache(): + moe._MPS_OUTPUT_SLICE_RESULTS.clear() + moe._MPS_EXPERTS_WARNED = False + yield + moe._MPS_OUTPUT_SLICE_RESULTS.clear() + moe._MPS_EXPERTS_WARNED = False + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +def test_runtime_probe_matches_direct_slice_write_diagnostic(dtype): + generator = torch.Generator(device="cpu").manual_seed(0) + inputs = torch.randn(10, 64, generator=generator).to(device="mps", dtype=dtype) + weights = torch.randn(4, 64, 32, generator=generator).to(device="mps", dtype=dtype) + with torch.no_grad(): + expected = inputs[5:] @ weights[3] + output = torch.zeros(10, 32, device="mps", dtype=dtype) + torch.mm(inputs[5:], weights[3], out=output[5:]) + products_match = torch.allclose( + output[5:], expected, rtol=max(1e-5, 2 * torch.finfo(dtype).eps), atol=1e-5 + ) + prefix_intact = torch.count_nonzero(output[:5]).item() == 0 + cpu_rng = torch.get_rng_state().clone() + mps_rng = torch.mps.get_rng_state().clone() + detected = _mps_output_slice_writes_are_broken(inputs.device, dtype) + assert detected == (not (products_match and prefix_intact)) + assert _mps_output_slice_writes_are_broken(inputs.device, dtype) == detected + torch.testing.assert_close(torch.get_rng_state(), cpu_rng, rtol=0, atol=0) + torch.testing.assert_close(torch.mps.get_rng_state(), mps_rng, rtol=0, atol=0) + + +class _FixedRoutes(nn.Module): + def __init__(self, expert_class=GraniteMoeHybridExperts): + super().__init__() + config = SimpleNamespace( + hidden_size=64, + intermediate_size=32, + moe_intermediate_size=32, + num_local_experts=4, + num_experts=4, + hidden_act="silu", + ) + config._experts_implementation = "grouped_mm" + self.experts = expert_class(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(expert_class): + block = _FixedRoutes(expert_class) + 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) + fallback_warnings = [ + warning for warning in caught if "eager experts" in str(warning.message) + ] + assert len(fallback_warnings) == int( + _mps_output_slice_writes_are_broken(first.device, first.dtype) + ) + assert len(observed) == 2 + assert block.experts.config._experts_implementation == "grouped_mm" + finally: + handle.remove() + + +def test_dispatch_follows_cpu_mps_cpu_migration(expert_class): + block = _FixedRoutes(expert_class) + 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(expert_class): + block = _FixedRoutes(expert_class).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(expert_class): + block = _FixedRoutes(expert_class) + 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(expert_class): + block = _FixedRoutes(expert_class).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) + + +@pytest.mark.parametrize("broken", [False, True]) +def test_real_expert_dispatch_selects_grouped_or_eager_from_probe( + monkeypatch, expert_class, broken +): + grouped_calls = [] + + def safe_grouped(inputs, weights, offs): + grouped_calls.append(offs.cpu().tolist()) + products = [] + start = 0 + for index, end in enumerate(offs.cpu().tolist()): + products.append(inputs[start:end] @ weights[index]) + start = end + return torch.cat(products) + + monkeypatch.setattr(hf_moe, "_grouped_mm", safe_grouped) + monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: broken) + registry = dict(hf_moe.ALL_EXPERTS_FUNCTIONS) + block = _FixedRoutes(expert_class) + 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) + with torch.no_grad(), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + actual = bridge(layer_input=inputs.to("mps")) + assert block.experts.forward.__wrapped__ is type(block.experts).forward + torch.testing.assert_close(actual.cpu(), expected, rtol=1e-4, atol=1e-6) + assert len(caught) == int(broken) + assert grouped_calls == ([] if broken else [[2, 5, 5, 10]] * 2) + assert block.experts.config._experts_implementation == "grouped_mm" + assert dict(hf_moe.ALL_EXPERTS_FUNCTIONS) == registry + + +@pytest.mark.parametrize("weight_dtype", [torch.float32, torch.float16]) +def test_autocast_cold_and_warm_dispatch_matches_raw_hf_policy(expert_class, weight_dtype): + block = _FixedRoutes(expert_class).to("mps") + block.experts.to(dtype=weight_dtype) + reference = copy.deepcopy(block) + bridge = MoEBridge("moe") + bridge.set_original_component(block) + inputs = torch.randn(10, 64, generator=torch.Generator().manual_seed(1)).to("mps") + assert not moe._MPS_OUTPUT_SLICE_RESULTS + with torch.no_grad(), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with torch.autocast("mps", dtype=torch.float16): + first = bridge(inputs) + assert torch.is_autocast_enabled("mps") + key = (inputs.device, weight_dtype) + assert set(moe._MPS_OUTPUT_SLICE_RESULTS) == {key} + broken = moe._MPS_OUTPUT_SLICE_RESULTS[key] + reference.experts.config._experts_implementation = "eager" if broken else "grouped_mm" + expected = reference(inputs) + second = bridge(inputs) + assert torch.is_autocast_enabled("mps") + torch.testing.assert_close(first, expected, rtol=0, atol=0) + torch.testing.assert_close(second, first, rtol=0, atol=0) + assert len(caught) == int(broken) + assert set(moe._MPS_OUTPUT_SLICE_RESULTS) == {key} + assert block.experts.config._experts_implementation == "grouped_mm" diff --git a/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py new file mode 100644 index 0000000000..3fd3b680e2 --- /dev/null +++ b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py @@ -0,0 +1,734 @@ +"""Shared HF expert dispatch without shared configuration mutation.""" + +import copy +import sys +import warnings +from concurrent.futures import ThreadPoolExecutor +from functools import wraps +from inspect import signature +from threading import Lock, Thread +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from transformers.integrations import moe as hf_moe +from transformers.integrations.moe import ExpertsInterface, use_experts_implementation +from transformers.models.granitemoe.modeling_granitemoe import GraniteMoeExperts +from transformers.models.granitemoehybrid.modeling_granitemoehybrid import ( + GraniteMoeHybridExperts, +) +from transformers.models.mixtral.modeling_mixtral import MixtralExperts +from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeExperts + +from transformer_lens.model_bridge.generalized_components import moe +from transformer_lens.model_bridge.generalized_components.moe import ( + _guard_mps_experts, + _hf_eager_experts_forward, + _mps_experts_forward, + _mps_output_slice_writes_are_broken, + _probe_mps_output_slice_writes, +) + +_EXPERT_CLASSES = [GraniteMoeHybridExperts, GraniteMoeExperts, MixtralExperts, Qwen3MoeExperts] + + +@pytest.fixture(autouse=True) +def isolated_probe(monkeypatch): + moe._MPS_OUTPUT_SLICE_RESULTS.clear() + moe._MPS_EXPERTS_WARNED = False + monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: True) + yield + moe._MPS_OUTPUT_SLICE_RESULTS.clear() + moe._MPS_EXPERTS_WARNED = False + + +class _MpsInput(torch.Tensor): + @property + def device(self): + return torch.device("mps") + + +class _Experts(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.fail = False + + def eager(self, hidden_states): + if self.fail: + raise ValueError("expert failure") + return "eager" + + @wraps(eager) + def forward(self, hidden_states): + return self.config._experts_implementation + + +@pytest.fixture +def stub_dispatch(monkeypatch): + recognize = moe._hf_eager_experts_forward + + def eager(forward): + if forward is _Experts.forward: + return forward.__wrapped__ + return recognize(forward) + + monkeypatch.setattr(moe, "_hf_eager_experts_forward", eager) + + +def test_dispatch_does_not_mutate_shared_config(stub_dispatch): + config = SimpleNamespace(_experts_implementation="grouped_mm") + experts = [_Experts(config), _Experts(config)] + mps = torch.empty(1).as_subclass(_MpsInput) + cpu = torch.empty(1) + with pytest.warns(RuntimeWarning, match="eager experts"): + with ThreadPoolExecutor(max_workers=2) as pool: + results = list( + pool.map( + lambda pair: _mps_experts_forward(*pair), + [(experts[0], mps), (experts[1], cpu)], + ) + ) + assert results == ["eager", "grouped_mm"] + assert vars(config) == {"_experts_implementation": "grouped_mm"} + + +@pytest.mark.parametrize("policy", ["eager", "batched_mm", None]) +def test_other_policies_are_unchanged(policy): + experts = _Experts(SimpleNamespace(_experts_implementation=policy)) + assert _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == policy + + +def test_exception_does_not_leave_dispatch_state(stub_dispatch): + experts = _Experts(SimpleNamespace(_experts_implementation="grouped_mm")) + experts.fail = True + with pytest.warns(RuntimeWarning), pytest.raises(ValueError, match="expert failure"): + _mps_experts_forward(experts, hidden_states=torch.empty(1).as_subclass(_MpsInput)) + assert experts.config._experts_implementation == "grouped_mm" + experts.fail = False + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + assert _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == "eager" + assert not caught + + +def _block(expert_class=GraniteMoeHybridExperts): + config = SimpleNamespace( + hidden_size=8, + intermediate_size=4, + moe_intermediate_size=4, + num_local_experts=4, + num_experts=4, + hidden_act="silu", + ) + config._experts_implementation = "grouped_mm" + block = nn.Module() + block.experts = expert_class(config) + return block + + +@pytest.mark.parametrize( + "torch_version,hf_version", + [ + ("2.11.0", "5.13.0"), + ("2.11.0+local", "5.13.0+local"), + ("2.11.1", "5.13.0"), + ("2.12.0", "5.13.0"), + ("2.11.0", "5.14.0"), + ("2.11.0rc1", "5.13.0"), + ("2.11.0", "5.13.0.dev0"), + ], +) +@pytest.mark.parametrize("broken", [False, True]) +def test_dispatch_uses_probe_not_package_versions( + monkeypatch, torch_version, hf_version, broken, stub_dispatch +): + import transformers + + monkeypatch.setattr(torch, "__version__", torch_version) + monkeypatch.setattr(transformers, "__version__", hf_version) + probes = [] + + def probe(device, dtype): + probes.append((device, dtype)) + return broken + + monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe) + block = _block() + _guard_mps_experts(block) + assert block.experts.forward.__wrapped__ is type(block.experts).forward + assert not probes + assert block.experts.config._experts_implementation == "grouped_mm" + experts = _Experts(block.experts.config) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + actual = _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) + assert actual == ("eager" if broken else "grouped_mm") + assert len(caught) == int(broken) + assert probes == [(torch.device("mps"), torch.float32)] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("corruption", [None, "missing_write", "prefix_write"]) +def test_primitive_probe_checks_products_and_untouched_rows(monkeypatch, dtype, corruption): + original_mm = torch.mm + + def mm(inputs, weights, *, out): + assert not torch.is_grad_enabled() + assert out.storage_offset() == 5 * 32 + if corruption != "missing_write": + original_mm(inputs, weights, out=out) + if corruption == "prefix_write": + out._base[0, 0] = 1 + return out + + monkeypatch.setattr(torch, "mm", mm) + rng_before = torch.get_rng_state().clone() + assert _probe_mps_output_slice_writes(torch.device("cpu"), dtype) == (corruption is not None) + torch.testing.assert_close(torch.get_rng_state(), rng_before, rtol=0, atol=0) + + +def test_probe_cache_is_per_device_and_dtype(monkeypatch): + calls = [] + + def probe(device, dtype): + calls.append((device, dtype)) + return dtype == torch.float16 + + monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe) + keys = [ + (torch.device("mps:0"), torch.float32), + (torch.device("mps:0"), torch.float16), + (torch.device("mps:1"), torch.float32), + ] + for _ in range(2): + for device, dtype in keys: + assert _mps_output_slice_writes_are_broken(device, dtype) == (dtype == torch.float16) + assert calls == keys + + +def test_concurrent_first_use_probes_once(monkeypatch): + calls = [] + + def probe(device, dtype): + calls.append((device, dtype)) + return False + + monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe) + with ThreadPoolExecutor(max_workers=4) as pool: + results = list( + pool.map( + lambda _: _mps_output_slice_writes_are_broken(torch.device("mps"), torch.float32), + range(16), + ) + ) + assert results == [False] * 16 + assert calls == [(torch.device("mps"), torch.float32)] + + +@pytest.mark.parametrize("error_type", [RuntimeError, NotImplementedError, TypeError]) +def test_probe_failure_is_diagnostic_and_not_cached(monkeypatch, error_type, stub_dispatch): + error = error_type("unsupported diagnostic") + + def probe(device, dtype): + raise error + + monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe) + experts = _Experts(SimpleNamespace(_experts_implementation="grouped_mm")) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with pytest.raises(RuntimeError, match="Select eager experts explicitly") as raised: + _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) + assert raised.value.__cause__ is error + assert not caught + assert not moe._MPS_OUTPUT_SLICE_RESULTS + assert experts.config._experts_implementation == "grouped_mm" + monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: False) + assert _mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == "grouped_mm" + + +def test_unexpected_probe_failure_propagates(monkeypatch): + def probe(device, dtype): + raise ValueError("unexpected diagnostic failure") + + monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe) + with pytest.raises(ValueError, match="unexpected diagnostic failure"): + _mps_output_slice_writes_are_broken(torch.device("mps"), torch.float32) + assert not moe._MPS_OUTPUT_SLICE_RESULTS + + +@pytest.mark.parametrize( + "device,policy", + [ + (device, policy) + for device in ["cpu", "cuda"] + for policy in ["grouped_mm", "eager", "batched_mm", None] + ] + + [("mps", policy) for policy in ["eager", "batched_mm", None]], +) +def test_unaffected_calls_do_not_probe(monkeypatch, policy, device): + class DeviceInput(torch.Tensor): + @property + def device(self): + return torch.device(device) + + def probe(device, dtype): + raise AssertionError("unaffected execution must not probe") + + monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe) + experts = _Experts(SimpleNamespace(_experts_implementation=policy)) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + assert _mps_experts_forward(experts, torch.empty(1).as_subclass(DeviceInput)) == policy + assert not caught + assert not moe._MPS_OUTPUT_SLICE_RESULTS + + +@pytest.mark.parametrize("expert_class", _EXPERT_CLASSES) +def test_installation_is_idempotent_and_deepcopy_safe(monkeypatch, expert_class): + import transformers + + monkeypatch.setattr(torch, "__version__", "2.11.0") + monkeypatch.setattr(transformers, "__version__", "5.13.0") + block = _block(expert_class) + original_signature = signature(block.experts.forward) + _guard_mps_experts(block) + forward = block.experts.forward + assert signature(forward) == original_signature + _guard_mps_experts(block) + assert block.experts.forward is forward + cloned = copy.deepcopy(block) + assert cloned.experts.forward.__self__ is cloned.experts + assert signature(cloned.experts.forward) == original_signature + + +@pytest.mark.parametrize("expert_class", _EXPERT_CLASSES) +def test_guarded_cpu_experts_keep_grouped_execution_without_warning(monkeypatch, expert_class): + import transformers + + monkeypatch.setattr(torch, "__version__", "2.11.0") + monkeypatch.setattr(transformers, "__version__", "5.13.0") + block = _block(expert_class) + with torch.no_grad(): + for parameter in block.experts.parameters(): + parameter.normal_(std=0.1) + inputs = torch.randn(5, 8) + indices = torch.tensor([[0], [0], [1], [3], [3]]) + weights = torch.ones(5, 1) + with torch.no_grad(): + expected = block.experts(inputs, indices, weights) + _guard_mps_experts(block) + with warnings.catch_warnings(record=True) as caught, torch.no_grad(): + warnings.simplefilter("always") + actual = block.experts(inputs, indices, weights) + assert not caught + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert block.experts.config._experts_implementation == "grouped_mm" + + +def test_dispatch_does_not_create_missing_policy(): + class MissingPolicyExperts(_Experts): + def forward(self, hidden_states): + return hidden_states + + experts = MissingPolicyExperts(SimpleNamespace()) + inputs = torch.empty(1).as_subclass(_MpsInput) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + assert _mps_experts_forward(experts, inputs) is inputs + assert not caught + assert vars(experts.config) == {} + + +@pytest.mark.parametrize("expert_class", _EXPERT_CLASSES) +def test_guarded_eager_cpu_forward_preserves_hooks_and_gradients(monkeypatch, expert_class): + import transformers + + monkeypatch.setattr(torch, "__version__", "2.11.0") + monkeypatch.setattr(transformers, "__version__", "5.13.0") + block = _block(expert_class) + block.experts.config._experts_implementation = "eager" + with torch.no_grad(): + for parameter in block.experts.parameters(): + parameter.normal_(std=0.1) + reference = copy.deepcopy(block) + _guard_mps_experts(block) + inputs = torch.randn(5, 8, requires_grad=True) + routes = torch.rand(5, 1, requires_grad=True) + indices = torch.tensor([[0], [0], [1], [3], [3]]) + calls = [] + handle = block.experts.register_forward_hook(lambda module, args, output: calls.append(args)) + try: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + actual = block.experts(inputs, indices, routes) + expected = reference.experts(inputs, indices, routes) + assert not caught + assert len(calls) == 1 + assert calls[0][1] is indices + assert calls[0][2] is routes + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + actual_grads = torch.autograd.grad( + actual.sum(), (inputs, routes, *block.experts.parameters()) + ) + expected_grads = torch.autograd.grad( + expected.sum(), (inputs, routes, *reference.experts.parameters()) + ) + for measured, oracle in zip(actual_grads, expected_grads): + torch.testing.assert_close(measured, oracle, rtol=0, atol=0) + assert block.experts.config._experts_implementation == "eager" + finally: + handle.remove() + + +def test_unrelated_experts_and_custom_forwards_are_unchanged(): + block = nn.Module() + block.experts = _Experts(SimpleNamespace(_experts_implementation="grouped_mm")) + _guard_mps_experts(block) + assert block.experts.forward.__func__ is _Experts.forward + block = _block() + custom = lambda hidden_states: hidden_states + block.experts.forward = custom + _guard_mps_experts(block) + assert block.experts.forward is custom + + +def _decorated_block(interface, config=None): + @use_experts_implementation(experts_interface=interface) + class DecoratedExperts(nn.Module): + def __init__(self, config): + super().__init__() + self.received = None + + def forward(self, hidden_states, token=None, *, payload=None): + self.received = (hidden_states, token, payload) + return "eager" + + block = nn.Module() + block.experts = DecoratedExperts( + config if config is not None else SimpleNamespace(_experts_implementation="grouped_mm") + ) + return block + + +@pytest.mark.parametrize("broken", [False, True]) +@pytest.mark.parametrize("keyword_input", [False, True]) +def test_shared_decorator_dispatch_preserves_arguments_and_global_registry( + monkeypatch, broken, keyword_input +): + def grouped(experts, hidden_states, token=None, *, payload=None): + experts.received = (hidden_states, token, payload) + return "grouped_mm" + + registry = dict(hf_moe.ALL_EXPERTS_FUNCTIONS) + interface = ExpertsInterface() + interface["grouped_mm"] = grouped + monkeypatch.setattr(hf_moe, "grouped_mm_experts_forward", grouped) + monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: broken) + block = _decorated_block(interface) + _guard_mps_experts(block) + assert block.experts.forward.__wrapped__ is type(block.experts).forward + inputs = torch.empty(1).as_subclass(_MpsInput) + token, payload = object(), object() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + if keyword_input: + actual = block.experts(hidden_states=inputs, token=token, payload=payload) + else: + actual = block.experts(inputs, token, payload=payload) + assert actual == ("eager" if broken else "grouped_mm") + assert block.experts.received == (inputs, token, payload) + assert len(caught) == int(broken) + assert vars(block.experts.config) == {"_experts_implementation": "grouped_mm"} + assert dict(hf_moe.ALL_EXPERTS_FUNCTIONS) == registry + + +def test_guarded_shared_config_supports_concurrent_cpu_and_mps_dispatch(monkeypatch): + def grouped(experts, hidden_states): + return "grouped_mm" + + interface = ExpertsInterface() + interface["grouped_mm"] = grouped + monkeypatch.setattr(hf_moe, "grouped_mm_experts_forward", grouped) + config = SimpleNamespace(_experts_implementation="grouped_mm") + blocks = [_decorated_block(interface, config), _decorated_block(interface, config)] + for block in blocks: + _guard_mps_experts(block) + inputs = [torch.empty(1).as_subclass(_MpsInput), torch.empty(1)] + with pytest.warns(RuntimeWarning, match="eager experts"): + with ThreadPoolExecutor(max_workers=2) as pool: + actual = list(pool.map(lambda pair: pair[0].experts(pair[1]), zip(blocks, inputs))) + assert actual == ["eager", "grouped_mm"] + assert all(block.experts.config is config for block in blocks) + assert vars(config) == {"_experts_implementation": "grouped_mm"} + + +@pytest.mark.parametrize("expert_class", _EXPERT_CLASSES) +def test_recovers_real_expert_eager_forward(expert_class): + assert _hf_eager_experts_forward(expert_class.forward) is expert_class.forward.__wrapped__ + + +def test_custom_grouped_interface_is_not_overridden(): + def custom(experts, hidden_states): + return "custom" + + interface = ExpertsInterface() + interface["grouped_mm"] = custom + block = _decorated_block(interface) + forward = block.experts.forward + _guard_mps_experts(block) + assert block.experts.forward == forward + with warnings.catch_warnings(record=True) as caught: + assert block.experts(torch.empty(1).as_subclass(_MpsInput)) == "custom" + assert not caught + assert not moe._MPS_OUTPUT_SLICE_RESULTS + + +def test_interface_changed_after_installation_is_respected(): + interface = ExpertsInterface() + block = _decorated_block(interface) + _guard_mps_experts(block) + assert block.experts.forward.__wrapped__ is type(block.experts).forward + interface["grouped_mm"] = lambda experts, hidden_states: "custom" + with warnings.catch_warnings(record=True) as caught: + assert block.experts(torch.empty(1).as_subclass(_MpsInput)) == "custom" + assert not caught + assert not moe._MPS_OUTPUT_SLICE_RESULTS + + +def test_mismatched_wrapped_forward_is_not_recognized(monkeypatch): + block = _decorated_block(ExpertsInterface()) + forward = type(block.experts).forward + monkeypatch.setattr(forward, "__wrapped__", lambda experts, inputs: inputs) + _guard_mps_experts(block) + assert block.experts.forward.__func__ is forward + + +def test_ordinary_expert_lists_are_unchanged(): + block = nn.Module() + block.experts = nn.ModuleList([nn.Linear(8, 4)]) + forward = block.experts.forward + _guard_mps_experts(block) + assert block.experts.forward == forward + + +@pytest.mark.parametrize("concurrent", [False, True]) +def test_warning_is_once_across_layers_models_and_forwards(concurrent): + blocks = [_decorated_block(ExpertsInterface()) for _ in range(4)] + for block in blocks: + _guard_mps_experts(block) + inputs = torch.empty(1).as_subclass(_MpsInput) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + if concurrent: + with ThreadPoolExecutor(max_workers=8) as pool: + results = list(pool.map(lambda block: block.experts(inputs), blocks * 8)) + else: + results = [block.experts(inputs) for block in blocks * 8] + assert results == ["eager"] * 32 + assert len(caught) == 1 + assert caught[0].category is RuntimeWarning + assert "eager experts" in str(caught[0].message) + assert moe._MPS_EXPERTS_WARNED + + +def test_warning_identifies_external_caller_through_runtime_wrappers(monkeypatch): + modules = [] + + def currentframe(): + frame = sys._getframe(1) + while frame is not None: + modules.append(frame.f_globals.get("__name__", "")) + frame = frame.f_back + return sys._getframe(1) + + monkeypatch.setattr(moe, "currentframe", currentframe) + block = _decorated_block(ExpertsInterface()) + _guard_mps_experts(block) + assert not moe._MPS_EXPERTS_WARNED + inputs = torch.empty(1).as_subclass(_MpsInput) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + caller_line = sys._getframe().f_lineno + 1 + result = block.experts(inputs) + assert result == "eager" + assert len(caught) == 1 + assert caught[0].filename == __file__ + assert caught[0].lineno == caller_line + assert any(module.startswith("jaxtyping.") for module in modules) + assert any(module.startswith("torch.") for module in modules) + + +def test_warning_as_error_propagates_without_consuming_notification(): + block = _decorated_block(ExpertsInterface()) + _guard_mps_experts(block) + inputs = torch.empty(1).as_subclass(_MpsInput) + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) + for _ in range(2): + with pytest.raises(RuntimeWarning, match="eager experts"): + block.experts(inputs) + assert not moe._MPS_EXPERTS_WARNED + assert block.experts.received is None + assert block.experts.config._experts_implementation == "grouped_mm" + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + assert block.experts(inputs) == "eager" + assert block.experts(inputs) == "eager" + assert len(caught) == 1 + assert moe._MPS_EXPERTS_WARNED + + +def test_unaffected_calls_do_not_consume_fallback_notification(monkeypatch): + def grouped(experts, hidden_states): + return "grouped_mm" + + interface = ExpertsInterface() + interface["grouped_mm"] = grouped + monkeypatch.setattr(hf_moe, "grouped_mm_experts_forward", grouped) + monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: False) + block = _decorated_block(interface) + _guard_mps_experts(block) + inputs = torch.empty(1).as_subclass(_MpsInput) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + assert block.experts(torch.empty(1)) == "grouped_mm" + assert block.experts(inputs) == "grouped_mm" + block.experts.config._experts_implementation = "eager" + assert block.experts(inputs) == "eager" + assert not caught + assert not moe._MPS_EXPERTS_WARNED + block.experts.config._experts_implementation = "grouped_mm" + moe._MPS_OUTPUT_SLICE_RESULTS.clear() + monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", lambda device, dtype: True) + with pytest.warns(RuntimeWarning, match="eager experts"): + assert block.experts(inputs) == "eager" + + +def test_fallback_warning_does_not_suppress_unrelated_warnings(): + block = _decorated_block(ExpertsInterface()) + _guard_mps_experts(block) + inputs = torch.empty(1).as_subclass(_MpsInput) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + assert block.experts(inputs) == "eager" + warnings.warn("Unrelated diagnostic", UserWarning) + assert block.experts(inputs) == "eager" + warnings.warn("Unrelated diagnostic", UserWarning) + assert [warning.category for warning in caught] == [RuntimeWarning, UserWarning, UserWarning] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("broken", [False, True]) +def test_primitive_probe_is_independent_of_ambient_autocast(monkeypatch, dtype, broken): + original_mm = torch.mm + + def mm(inputs, weights, *, out): + assert not torch.is_autocast_enabled("cpu") + assert inputs.dtype == weights.dtype == out.dtype == dtype + if broken: + return out + return original_mm(inputs, weights, out=out) + + monkeypatch.setattr(torch, "mm", mm) + with torch.autocast("cpu", dtype=torch.bfloat16): + assert _probe_mps_output_slice_writes(torch.device("cpu"), dtype) == broken + assert torch.is_autocast_enabled("cpu") + + +def test_dispatch_probes_batched_weight_dtype_instead_of_input_dtype(monkeypatch): + block = _decorated_block(ExpertsInterface()) + block.experts.register_parameter( + "gate_up_proj", nn.Parameter(torch.empty(4, 8, 8, dtype=torch.float16)) + ) + block.experts.register_parameter( + "down_proj", nn.Parameter(torch.empty(4, 8, 4, dtype=torch.float16)) + ) + probed = [] + + def probe(device, dtype): + probed.append(dtype) + return True + + monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe) + _guard_mps_experts(block) + with pytest.warns(RuntimeWarning, match="eager experts"): + assert block.experts(torch.empty(1).as_subclass(_MpsInput)) == "eager" + assert probed == [torch.float16] + + +def test_warning_callback_can_reenter_guarded_experts(monkeypatch): + monkeypatch.setattr(moe, "_MPS_EXPERTS_WARNING_LOCK", Lock()) + block = _decorated_block(ExpertsInterface()) + _guard_mps_experts(block) + inputs = torch.empty(1).as_subclass(_MpsInput) + callbacks, results, errors = [], [], [] + + def showwarning(*args, **kwargs): + callbacks.append(block.experts(inputs)) + + def invoke(): + try: + results.append(block.experts(inputs)) + except Exception as error: + errors.append(error) + + with warnings.catch_warnings(): + warnings.simplefilter("always") + monkeypatch.setattr(warnings, "showwarning", showwarning) + worker = Thread(target=invoke, daemon=True) + worker.start() + worker.join(timeout=2) + assert not worker.is_alive(), "warning callback deadlocked on re-entry" + assert not errors + assert callbacks == ["eager"] + assert results == ["eager"] + assert moe._MPS_EXPERTS_WARNED + + +def test_warning_callback_failure_releases_reservation(monkeypatch): + block = _decorated_block(ExpertsInterface()) + _guard_mps_experts(block) + inputs = torch.empty(1).as_subclass(_MpsInput) + + def showwarning(*args, **kwargs): + raise ValueError("warning callback failed") + + with warnings.catch_warnings(), monkeypatch.context() as context: + warnings.simplefilter("always") + context.setattr(warnings, "showwarning", showwarning) + with pytest.raises(ValueError, match="warning callback failed"): + block.experts(inputs) + assert not moe._MPS_EXPERTS_WARNED + with pytest.warns(RuntimeWarning, match="eager experts"): + assert block.experts(inputs) == "eager" + + +def test_dispatch_checks_both_projection_weight_dtypes(monkeypatch): + def grouped(experts, hidden_states): + return "grouped_mm" + + interface = ExpertsInterface() + interface["grouped_mm"] = grouped + monkeypatch.setattr(hf_moe, "grouped_mm_experts_forward", grouped) + block = _decorated_block(interface) + block.experts.register_parameter( + "gate_up_proj", nn.Parameter(torch.empty(4, 8, 8, dtype=torch.float16)) + ) + block.experts.register_parameter( + "down_proj", nn.Parameter(torch.empty(4, 8, 4, dtype=torch.float32)) + ) + probed = [] + + def probe(device, dtype): + probed.append(dtype) + return False + + monkeypatch.setattr(moe, "_probe_mps_output_slice_writes", probe) + _guard_mps_experts(block) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + assert block.experts(torch.empty(1).as_subclass(_MpsInput)) == "grouped_mm" + assert set(probed) == {torch.float16, torch.float32} + assert not caught diff --git a/transformer_lens/model_bridge/generalized_components/moe.py b/transformer_lens/model_bridge/generalized_components/moe.py index ca96426d30..be46ed1f55 100644 --- a/transformer_lens/model_bridge/generalized_components/moe.py +++ b/transformer_lens/model_bridge/generalized_components/moe.py @@ -6,7 +6,12 @@ from __future__ import annotations import logging -from typing import Any, Dict, List, Mapping, Optional, Tuple +import warnings +from functools import wraps +from inspect import currentframe, getclosurevars +from threading import Lock +from types import CodeType, FunctionType, MethodType +from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple import torch from torch import nn @@ -29,6 +34,157 @@ _BATCHED_OUTPUT_PROJECTIONS = frozenset({"down_proj"}) _BATCHED_PROJECTIONS = _BATCHED_INPUT_PROJECTIONS | _BATCHED_OUTPUT_PROJECTIONS +_MPS_OUTPUT_SLICE_RESULTS: Dict[Tuple[torch.device, torch.dtype], bool] = {} +_MPS_OUTPUT_SLICE_LOCK = Lock() +_MPS_EXPERTS_WARNED = False +_MPS_EXPERTS_WARNING_LOCK = Lock() + + +def _warn_mps_experts_once() -> None: + """Warn once at the first caller outside the model and runtime-checking libraries.""" + global _MPS_EXPERTS_WARNED + with _MPS_EXPERTS_WARNING_LOCK: + if _MPS_EXPERTS_WARNED: + return + _MPS_EXPERTS_WARNED = True + frame = None + emitted = False + try: + frame = currentframe() + stacklevel = 1 + while frame is not None: + module = frame.f_globals.get("__name__", "") + if not any( + module == prefix or module.startswith(prefix + ".") + for prefix in ("transformer_lens", "transformers", "torch", "jaxtyping", "beartype") + ): + break + frame = frame.f_back + stacklevel += 1 + warnings.warn( + "Using eager experts on MPS after detecting incorrect " + "matmul output-slice writes; expert execution may be slower.", + RuntimeWarning, + stacklevel=stacklevel, + ) + emitted = True + finally: + del frame + if not emitted: + with _MPS_EXPERTS_WARNING_LOCK: + _MPS_EXPERTS_WARNED = False + + +def _probe_mps_output_slice_writes(device: torch.device, dtype: torch.dtype) -> bool: + """Check for incorrect nonzero-offset matmul output writes using isolated tensors.""" + with torch.no_grad(), torch.autocast(device_type=device.type, enabled=False): + generator = torch.Generator(device="cpu").manual_seed(0) + inputs = torch.randn(10, 64, generator=generator).to(device=device, dtype=dtype) + weights = torch.randn(4, 64, 32, generator=generator).to(device=device, dtype=dtype) + output = torch.zeros(10, 32, device=device, dtype=dtype) + expected = inputs[5:] @ weights[3] + torch.mm(inputs[5:], weights[3], out=output[5:]) + rtol = max(1e-5, 2 * torch.finfo(dtype).eps) + prefix_intact = torch.count_nonzero(output[:5]).item() == 0 + products_match = torch.allclose(output[5:], expected, rtol=rtol, atol=1e-5) + return not (prefix_intact and products_match) + + +def _mps_output_slice_writes_are_broken(device: torch.device, dtype: torch.dtype) -> bool: + """Cache the primitive check once per device and dtype, including concurrent first use.""" + key = (device, dtype) + with _MPS_OUTPUT_SLICE_LOCK: + if key not in _MPS_OUTPUT_SLICE_RESULTS: + try: + _MPS_OUTPUT_SLICE_RESULTS[key] = _probe_mps_output_slice_writes(device, dtype) + except (RuntimeError, TypeError) as error: + raise RuntimeError( + f"Cannot verify MPS matmul output-slice writes on {device} with {dtype}. " + "Select eager experts explicitly or use CPU execution." + ) from error + return _MPS_OUTPUT_SLICE_RESULTS[key] + + +def _mps_experts_forward(experts: nn.Module, *args: Any, **kwargs: Any) -> Any: + """Dispatch locally, without changing the config shared by other expert layers.""" + hidden_states = args[0] if args else kwargs.get("hidden_states") + forward = type(experts).forward + config: Any = getattr(experts, "config", None) + eager = None + if ( + isinstance(hidden_states, torch.Tensor) + and hidden_states.device.type == "mps" + and getattr(config, "_experts_implementation", None) == "grouped_mm" + ): + eager = _hf_eager_experts_forward(forward) + if eager is not None: + dtypes = { + parameter.dtype + for parameter in experts.parameters(recurse=False) + if parameter.ndim == 3 and parameter.is_floating_point() + } or {hidden_states.dtype} + if any( + _mps_output_slice_writes_are_broken(hidden_states.device, dtype) for dtype in dtypes + ): + _warn_mps_experts_once() + return eager(experts, *args, **kwargs) + return forward(experts, *args, **kwargs) + + +def _hf_eager_experts_forward(forward: Any) -> Optional[Callable[..., Any]]: + """Recognize HF's shared grouped dispatcher and recover its original eager forward.""" + from transformers.integrations.moe import ( + ExpertsInterface, + grouped_mm_experts_forward, + use_experts_implementation, + ) + + if not isinstance(forward, FunctionType): + return None + if not any( + forward.__code__ is candidate + for wrapper in use_experts_implementation.__code__.co_consts + if isinstance(wrapper, CodeType) and wrapper.co_name == "wrapper" + for candidate in wrapper.co_consts + if isinstance(candidate, CodeType) and candidate.co_name == "forward" + ): + return None + closure = getclosurevars(forward).nonlocals + original = closure.get("original_forward") + interface = closure.get("experts_interface") + if ( + not callable(original) + or original is not getattr(forward, "__wrapped__", None) + or not isinstance(interface, ExpertsInterface) + or interface.get_interface("grouped_mm", original) is not grouped_mm_experts_forward + ): + return None + return original + + +def _guard_mps_experts(component: nn.Module) -> None: + """Install instance-local dispatch only for HF's shared grouped expert implementation. + + The bound method preserves the forward signature, nn.Module hooks, and deepcopy + semantics. Configs and global dispatch tables remain untouched, including during + concurrent forwards. + """ + experts = getattr(component, "experts", None) + if not isinstance(experts, nn.Module): + return + experts = unwrap_bridge(experts) + forward = type(experts).forward + if getattr(experts.forward, "__func__", None) is not forward: + return + if _hf_eager_experts_forward(forward) is None: + return + + @wraps(forward) + def guarded_forward(experts: nn.Module, *args: Any, **kwargs: Any) -> Any: + return _mps_experts_forward(experts, *args, **kwargs) + + setattr(experts, "forward", MethodType(guarded_forward, experts)) + class UnfoldableMoEParameter(Exception): """A parameter whose relationship to the MoE block's input cannot be established.""" @@ -431,6 +587,7 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: elif "hidden_states" in kwargs: hooked = self.hook_in(kwargs["hidden_states"]) kwargs = {**kwargs, "hidden_states": hooked} + _guard_mps_experts(self.original_component) output = self.original_component(*args, **kwargs) if isinstance(output, tuple): if not output: