[janmenjayap]

---

Description

Fixes #1886

Granite MoE Hybrid expert execution can produce incorrect or unstable outputs on Apple MPS when HuggingFace selects grouped_mm in the verified torch 2.11.0 / transformers 5.13.0 environment. This can make the existing SSM eager/fused parity regression fail even though the first substantial divergence occurs in the MoE expert output rather than the SSM scan.

A standalone fixed-data diagnostic isolates incorrect nonzero-offset torch.mm(..., out=output[5:]) writes. HF's grouped fallback also uses output-slice writes, so disabling native grouped execution alone is insufficient.

---

Changes

• Add an instance-local dispatch wrapper for the exact HF GraniteMoeHybridExperts class, installed through MoEBridge.forward.
• Select HF's original ordinary eager expert forward only for live MPS inputs with grouped_mm on torch release 2.11.0 and transformers release 5.13.0. Local version suffixes are accepted; prerelease/development versions are not.
• Leave shared HF configs and global dispatch tables untouched. No production activation transfers to CPU, routing reimplementation, dependency changes, or global kernel patches are introduced.
• Preserve CPU/CUDA dispatch, explicit eager execution, other policies/classes, and custom instance forwards.
• Add independent expert-oracle and repeatability tests, CPU-to-MPS-to-CPU migration coverage, input/routing-weight/expert-parameter gradient comparisons, observer/output-editing hook coverage, and failure-recovery/state-preservation tests.
• Document the scope, warning, performance cost, concurrency boundary, and primitive diagnostic in the execution-backends guide.

The existing Granite MoE Hybrid MPS scan-parity assertion passes through the corrected production path without a changed threshold, new skip/xfail, or forced eager policy in its fixture.

---

Scope and limitations

• Verified locally on Darwin 23.6.0 arm64, Python 3.12.14, torch 2.11.0, transformers 5.13.0, float32, Apple MPS.
• The fallback deliberately emits a RuntimeWarning; eager expert execution may be slower.
• Dispatch follows the live input device and policy. Avoiding shared-config mutation removes a policy race, but does not guarantee 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. Once installed, the same expert instance retains the wrapper.
• Older/newer dependency versions are neither declared affected nor declared fixed. CUDA and other dependency versions were not device-validated.
• This workaround does not establish the internal root cause of native grouped execution or imply a confirmed upstream fixed version.

---

Relationship to other work

This is separate from #1878, which fixes Granite embedding scaling in cache residual decomposition and component-level direct logit attribution without changing model forward execution. That analysis correction does not repair incorrect expert matrix products.

It is also independent of GPT-2 contextual decomposition.

A potentially related upstream report is pytorch/pytorch#193487, concerning allocator-dependent incorrect MPS matmul results. Its reproduction differs from the small nonzero-offset output-write reproduction here; a shared root cause has not been established. No upstream report or comment has been posted for this work.

---

Type of change

• Bug fix (non-breaking change which fixes an issue)
• This change requires a documentation update

---

Validation

Before and after

On baseline a8bc19e7, the CPU scan-parity case passed and the MPS case failed with relative difference 1.0326683320021979 (required < 1e-4). On this branch, the original test passes without weakened assertions.

The standalone primitive diagnostic measured:

MPS out-slice max error: 29.132984161376953
MPS slice-assignment max error: 4.76837158203125e-06

Focused regressions

TRANSFORMERLENS_ALLOW_MPS=1 uv run pytest \
  tests/mps/test_granite_grouped_expert_dispatch.py \
  tests/unit/model_bridge/generalized_components/test_moe_expert_dispatch.py \
  tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py \
  -q --tb=short

Result: 49 passed, with the intended fallback warnings and other dependency/model warnings reported.

Standard checks

Full unit tier: 7322 passed, 58 skipped, 57 deselected, 3 xfailed
Docstring tier: 13 passed, 24 skipped
Acceptance tier: 163 passed, 69 skipped, 11 deselected
Mypy: no issues found in 391 source files
Docs: build succeeded, 19 warnings

make format, make check-format, and git diff --check passed. The skip/xfail counts above are suite outcomes, not new markers added to bypass this defect.

Full PR validation remains incomplete. make test-pr completed unit, docstring, and acceptance tiers, then hit the 30-minute command limit during integration testing. No integration failure had been reported before interruption, but the integration tier is not claimed complete or passing. The docs build is not claimed warning-free. CUDA was not tested.

---

Checklist

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

The warning checkbox is intentionally unchecked because the affected MPS path emits an explicit correctness-fallback warning. Completing the full integration tier remains a validation follow-up before merge.

---

[jlarson4]

Hi @janmenjayap, the idea behind your fix is sound, but there are some fundamental flaws that need to be resolved before this can be merged. Please take a look at my comments below, and let me know if you have any questions

[jlarson4]

This gates on torch 2.11.0 and transformers 5.13.0, but the bad slice writes are not tied to those versions. On macOS 26.6 with exactly those packages, the primitive, native grouped_mm, and HF's fallback are all correct and the parity test passes without this change, yet the gate still fires, so every Granite Hybrid forward on MPS takes the eager loop and receive a false warning saying its products are wrong. Please key the fallback on a direct check of the slice-write primitive rather than on package versions.

[janmenjayap]

You are right: the same torch/transformers versions behave differently on different hosts, so the versions cannot tell us whether the primitive is broken.

Fixed in commit 7dcd402 (detection and cache) and 3eccd6b (autocast isolation and weight-dtype probing):

• Removed the torch 2.11.0 and transformers 5.13.0 checks and their imports.

• Added _probe_mps_output_slice_writes, which runs a small direct check on its own tensors:

  • writes into a nonzero-offset output slice with torch.mm(..., out=output[5:]);
  • compares the result against an ordinary matmul;
  • checks that the preceding rows are left untouched.

  It runs under no_grad, with autocast disabled and a seeded CPU generator, so the user's RNG state and activations are not affected.

• _mps_output_slice_writes_are_broken caches the result per (device, dtype) behind a lock, so the check (and its device sync) runs once; later calls only take a CPU lock for the cache lookup. If the check fails with a supported error, it raises an actionable error and is not cached, rather than silently treating the host as healthy.

• _mps_experts_forward runs the check only on MPS inputs under the grouped_mm policy, using the dtypes of the batched projection weights. A healthy result keeps the original grouped path with no warning; only a broken result switches to HF's eager forward.

Tests:

• test_dispatch_uses_probe_not_package_versions: fakes both healthy and broken results under different version strings and checks that dispatch follows the probe only.
• test_primitive_probe_checks_products_and_untouched_rows: covers missing writes and writes into the wrong rows, across fp32, fp16 and bf16.
• test_probe_cache_is_per_device_and_dtype, test_concurrent_first_use_probes_once, test_probe_failure_is_diagnostic_and_not_cached and test_unaffected_calls_do_not_probe.
• test_primitive_probe_is_independent_of_ambient_autocast and test_dispatch_checks_both_projection_weight_dtypes.
• On real MPS: test_runtime_probe_matches_direct_slice_write_diagnostic and test_real_expert_dispatch_selects_grouped_or_eager_from_probe.

[jlarson4]

Only GraniteMoeHybridExperts is recognized, but the primitive that fails is shared. Every @use_experts_implementation class routes through the same _grouped_mm, and HF defaults all of them to grouped_mm with no device check. On an affected host a GraniteMoe, Mixtral, or Qwen3-MoE bridge produces the same wrong products with no fallback. Please adjust to recognize the shared dispatch decorator.

[janmenjayap]

Agreed. The failing primitive sits behind HF's shared _grouped_mm, so every @use_experts_implementation class is exposed.

Fixed in commit a786fae:

• Removed the type(experts) is GraniteMoeHybridExperts check.

• _hf_eager_experts_forward recognizes forwards produced by HF's use_experts_implementation wrapper. It checks:

  • the wrapper's code object;
  • that the closure's original_forward matches __wrapped__;
  • that the ExpertsInterface still maps grouped_mm to HF's own grouped_mm_experts_forward.

  It then returns HF's original eager forward unchanged, so routing, expert math and layouts are not reimplemented.

• _guard_mps_experts installs the guard as a bound method on that instance only. Configs and HF's global registry are untouched, and custom instance forwards are left alone.

Tests:

• The four real expert classes, Granite MoE Hybrid, Granite MoE, Mixtral and Qwen3-MoE, are run through test_recovers_real_expert_eager_forward, test_installation_is_idempotent_and_deepcopy_safe and test_guarded_eager_cpu_forward_preserves_hooks_and_gradients.
• The same four classes run on real MPS in test_fixed_route_experts_match_independent_oracle_and_repeat, test_corrected_experts_preserve_input_route_and_parameter_gradients, test_dispatch_follows_cpu_mps_cpu_migration and test_autocast_cold_and_warm_dispatch_matches_raw_hf_policy.
• Cases that must not be recognized: test_custom_grouped_interface_is_not_overridden, test_interface_changed_after_installation_is_respected, test_mismatched_wrapped_forward_is_not_recognized, test_unrelated_experts_and_custom_forwards_are_unchanged and test_ordinary_expert_lists_are_unchanged.
• Shared state: test_shared_decorator_dispatch_preserves_arguments_and_global_registry and test_guarded_shared_config_supports_concurrent_cpu_and_mps_dispatch.

One caveat: HF exposes no explicit marker for this wrapper, so recognition depends on the wrapper's internal structure. If a future transformers release restructures it, recognition returns None and the guard is not installed. The real-class tests above will fail for that release rather than pass silently. I'm happy to switch to a public marker if upstream adds one.

[jlarson4]

The warning is raised inside every expert call, and stacklevel=2 lands in jaxtyping's wrapper, so users see a library frame rather than their own call and, under non-default filters, one line per layer per forward. devices.py already keeps its MPS warnings to once per process with a module flag; consider the same here.

[janmenjayap]

Agreed on both points.

Fixed in commit c6ce797, with 3eccd6b moving the emit outside the lock and releasing the reservation on failure:

• _warn_mps_experts_once follows the module-flag pattern in devices.py. A process-wide flag, protected by a lock, allows one notification. The warning itself is emitted outside the lock, so a warning callback can re-enter guarded experts.
• Instead of a fixed stacklevel=2, it walks up the stack past transformer_lens, transformers, torch, jaxtyping and beartype frames. The warning is attributed to the first frame outside those libraries, and this works on Python 3.10 through 3.12.
• The warning fires only when a forward actually switches to eager. Installing the guard, a healthy probe, CPU/CUDA, and an explicit eager policy never use up the notification. If the warning is turned into an error and raised, the notification is released so it is not lost.
• The message no longer names a package version or claims a whole release is broken.

Tests:

• test_warning_is_once_across_layers_models_and_forwards: one warning under an always filter, run both serially and concurrently.
• test_warning_identifies_external_caller_through_runtime_wrappers: checks the exact filename and line, and that real jaxtyping and torch frames were skipped.
• test_warning_as_error_propagates_without_consuming_notification, test_warning_callback_can_reenter_guarded_experts, test_warning_callback_failure_releases_reservation, test_unaffected_calls_do_not_consume_fallback_notification and test_fallback_warning_does_not_suppress_unrelated_warnings.

In the integration run, the single warning points at the bridge call in the test file (fused = bridge(toks)), not at a jaxtyping frame.

[jlarson4]

This subsection sits under the transformers backend in the driver guide, but it is an Apple-specific workaround. Would you please move it to the existing TRANSFORMERLENS_ALLOW_MPS section in getting_started.md? Consider reducing to a short note there, with the diagnostic kept in the test file or the issue, instead of this subsection.

[janmenjayap]

Done in commit 10224e3:

• Removed the "Granite MoE Hybrid grouped experts on Apple MPS" subsection from docs/source/content/drivers.md.
• Added a four-line note under TRANSFORMERLENS_ALLOW_MPS in docs/source/content/getting_started.md. It says that the Bridge checks the MPS output-slice primitive for shared-dispatcher MoE experts, switches to eager only when incorrect writes are detected, emits one warning per process, and may cost performance. It also says that TRANSFORMERLENS_ALLOW_MPS=1 does not disable this safeguard.
• The detailed diagnostic stays in tests/mps/test_granite_grouped_expert_dispatch.py; there is no user-facing script.

[janmenjayap]

Thank you for the careful review, @jlarson4, and especially for testing on a healthy macOS 26.6 host. That test showed the version gate was wrong. Each point is addressed below, with the commit that makes the change on fix/granite-mps-grouped-experts. The branch is now merged with the latest dev, at 0db2dfb3

[janmenjayap]

Latest local results

These were run on the merged head 0db2dfb3 (latest dev merged in, no conflicts), on macOS 14.8.9 (Apple Silicon), Python 3.12.14, torch 2.11.0 and transformers 5.13.0. On this host the probe reports the primitive as broken.

Check	Result
Focused tests: unit dispatch, MPS dispatch and Granite Hybrid integration, with TRANSFORMERLENS_ALLOW_MPS=1	152 passed
make check-format	Passed (1002 files unchanged)
uv run mypy .	No issues in 389 source files
make unit-test	7734 passed, 59 skipped, 57 deselected, 3 xfailed
make docstring-test	13 passed, 24 skipped
make acceptance-test	163 passed, 69 skipped, 11 deselected
make integration-test	1634 passed, 1 failed, 7 skipped, 197 deselected, in two separate local runs (see note below)
CI on 0db2dfb3	All checks pass, including Compatibility Checks (3.10, 3.11, 3.12) and Full Code Coverage Test; MPS Checks skipped (they run only for main)
uv run build-docs	Build succeeded (336 Sphinx warnings), and the same suite passes in CI. A second full local run also ended with exactly one failure; I only kept its summary line, so I can't confirm it was the same test. I haven't yet found what closes the client mid-suite locally.

Because CI skips MPS Checks for PRs into dev, the MPS results above (tests/mps/test_granite_grouped_expert_dispatch.py and the Granite MPS integration case) come from my machine only
The one integration failure was test_mistral3_adapter.py::TestMistral3ForwardEquivalence::test_multimodal_forward_matches_fresh_hf. It could not load the tiny-random/mistral-3 processor because the Hugging Face Hub HTTP client had already been closed (RuntimeError: Cannot send a request, as the client has been closed), 27 minutes into the run. Mistral3 does not use the shared expert dispatcher. Run on its own, the file passes (6 passed). I haven't yet found what closes the client mid-suite.

The original Granite MPS scan-parity regression (test_granite_eager_scan_device_correctness[mps]) passes with its original tolerance. No skips, xfails or # type: ignore were added.

The upstream reports (the PyTorch nonzero-offset out= reproduction and the HF grouped dispatch with no device check) are still pending; I'll link them here once filed.

Thanks again. Please let me know if anything else should change.