diff --git a/docs/source/content/drivers.md b/docs/source/content/drivers.md index 43918889c9..4652ef19ed 100644 --- a/docs/source/content/drivers.md +++ b/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. 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..6413c2e561 --- /dev/null +++ b/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) 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..46cf128695 --- /dev/null +++ b/tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py @@ -0,0 +1,213 @@ +"""Scoped Granite expert dispatch without shared configuration mutation.""" + +import copy +import warnings +from concurrent.futures import ThreadPoolExecutor +from functools import wraps +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from transformers.models.granitemoehybrid.modeling_granitemoehybrid import ( + GraniteMoeHybridExperts, +) + +from transformer_lens.model_bridge.generalized_components.moe import ( + _granite_mps_experts_forward, + _guard_granite_mps_experts, +) + + +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 + + +def test_dispatch_does_not_mutate_shared_config(): + 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 Granite"): + with ThreadPoolExecutor(max_workers=2) as pool: + results = list( + pool.map( + lambda pair: _granite_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 _granite_mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == policy + + +def test_exception_does_not_leave_dispatch_state(): + experts = _Experts(SimpleNamespace(_experts_implementation="grouped_mm")) + experts.fail = True + with pytest.warns(RuntimeWarning), pytest.raises(ValueError, match="expert failure"): + _granite_mps_experts_forward(experts, hidden_states=torch.empty(1).as_subclass(_MpsInput)) + assert experts.config._experts_implementation == "grouped_mm" + experts.fail = False + with pytest.warns(RuntimeWarning): + assert ( + _granite_mps_experts_forward(experts, torch.empty(1).as_subclass(_MpsInput)) == "eager" + ) + + +def _block(): + config = SimpleNamespace( + hidden_size=8, intermediate_size=4, num_local_experts=4, hidden_act="silu" + ) + config._experts_implementation = "grouped_mm" + block = nn.Module() + block.experts = GraniteMoeHybridExperts(config) + return block + + +@pytest.mark.parametrize( + "torch_version,hf_version,installed", + [ + ("2.11.0", "5.13.0", True), + ("2.11.0+local", "5.13.0+local", True), + ("2.11.1", "5.13.0", False), + ("2.12.0", "5.13.0", False), + ("2.11.0", "5.14.0", False), + ("2.11.0rc1", "5.13.0", False), + ("2.11.0", "5.13.0.dev0", False), + ], +) +def test_installation_version_scope(monkeypatch, torch_version, hf_version, installed): + import transformers + + monkeypatch.setattr(torch, "__version__", torch_version) + monkeypatch.setattr(transformers, "__version__", hf_version) + block = _block() + _guard_granite_mps_experts(block) + assert (block.experts.forward.__func__ is _granite_mps_experts_forward) == installed + assert block.experts.config._experts_implementation == "grouped_mm" + + +def test_installation_is_idempotent_and_deepcopy_safe(monkeypatch): + import transformers + + monkeypatch.setattr(torch, "__version__", "2.11.0") + monkeypatch.setattr(transformers, "__version__", "5.13.0") + block = _block() + _guard_granite_mps_experts(block) + forward = block.experts.forward + _guard_granite_mps_experts(block) + assert block.experts.forward is forward + cloned = copy.deepcopy(block) + assert cloned.experts.forward.__self__ is cloned.experts + + +def test_guarded_cpu_experts_keep_grouped_execution_without_warning(monkeypatch): + import transformers + + monkeypatch.setattr(torch, "__version__", "2.11.0") + monkeypatch.setattr(transformers, "__version__", "5.13.0") + block = _block() + 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_granite_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 _granite_mps_experts_forward(experts, inputs) is inputs + assert not caught + assert vars(experts.config) == {} + + +def test_guarded_eager_cpu_forward_preserves_hooks_and_gradients(monkeypatch): + import transformers + + monkeypatch.setattr(torch, "__version__", "2.11.0") + monkeypatch.setattr(transformers, "__version__", "5.13.0") + block = _block() + 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_granite_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_granite_mps_experts(block) + assert block.experts.forward.__func__ is _Experts.forward + block = _block() + custom = lambda hidden_states: hidden_states + block.experts.forward = custom + _guard_granite_mps_experts(block) + assert block.experts.forward is custom diff --git a/transformer_lens/model_bridge/generalized_components/moe.py b/transformer_lens/model_bridge/generalized_components/moe.py index ca96426d30..dc50c0441d 100644 --- a/transformer_lens/model_bridge/generalized_components/moe.py +++ b/transformer_lens/model_bridge/generalized_components/moe.py @@ -6,9 +6,12 @@ from __future__ import annotations import logging +import warnings +from types import MethodType from typing import Any, Dict, List, Mapping, Optional, Tuple import torch +from packaging.version import Version from torch import nn from transformer_lens.hook_points import HookPoint @@ -30,6 +33,59 @@ _BATCHED_PROJECTIONS = _BATCHED_INPUT_PROJECTIONS | _BATCHED_OUTPUT_PROJECTIONS +def _granite_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) + if ( + isinstance(hidden_states, torch.Tensor) + and hidden_states.device.type == "mps" + and getattr(config, "_experts_implementation", None) == "grouped_mm" + ): + warnings.warn( + "Using eager Granite MoE Hybrid experts on MPS to avoid incorrect grouped " + "matrix products in torch 2.11.0; expert execution may be slower.", + RuntimeWarning, + stacklevel=2, + ) + return getattr(forward, "__wrapped__")(experts, *args, **kwargs) + return forward(experts, *args, **kwargs) + + +def _guard_granite_mps_experts(component: nn.Module) -> None: + """Install instance-local dispatch only for the verified HF expert implementation. + + The bound method preserves 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) + if type(experts).__module__ != ( + "transformers.models.granitemoehybrid.modeling_granitemoehybrid" + ): + return + import transformers + from transformers.models.granitemoehybrid.modeling_granitemoehybrid import ( + GraniteMoeHybridExperts, + ) + + if ( + type(experts) is not GraniteMoeHybridExperts + or Version(torch.__version__).public != "2.11.0" + or Version(transformers.__version__).public != "5.13.0" + ): + return + # Respect custom instance forwards rather than silently replacing them. + if getattr(experts.forward, "__func__", None) is not GraniteMoeHybridExperts.forward: + return + if not callable(getattr(GraniteMoeHybridExperts.forward, "__wrapped__", None)): + raise RuntimeError("Granite eager expert forward is unavailable for the MPS workaround") + setattr(experts, "forward", MethodType(_granite_mps_experts_forward, experts)) + + class UnfoldableMoEParameter(Exception): """A parameter whose relationship to the MoE block's input cannot be established.""" @@ -431,6 +487,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_granite_mps_experts(self.original_component) output = self.original_component(*args, **kwargs) if isinstance(output, tuple): if not output: