[Bug Report] Granite MoE Hybrid grouped experts produce incorrect or unstable outputs on MPS #1886

Summary

On baseline TransformerLens a8bc19e7, this integration regression fails on MPS:

tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py
    ::test_granite_eager_scan_device_correctness[mps]

Although the assertion compares SSM eager and fused execution, intermediate inspection localizes the first substantial discrepancy to the MoE expert output, after the SSM and normalization outputs agree closely with CPU. HF selects grouped_mm in this environment. Repeated identical MPS forwards can produce different expert outputs.

A fixed-data reproduction independent of TransformerLens isolates incorrect nonzero-offset torch.mm(..., out=...) slice writes. HF's grouped fallback uses such writes, so merely disabling native grouped execution is insufficient. Ordinary eager expert execution avoids the observed defect.

The original Granite MPS parity test passes through the implemented correction without changing its assertion, adding skips, or relaxing its threshold. This is an expert execution/device correctness issue, independent of GPT-2 contextual decomposition.

Environment verified

• Baseline TransformerLens: a8bc19e7.
• Python: 3.12.14.
• Platform: Darwin 23.6.0, arm64.
• PyTorch: 2.11.0.
• transformers: 5.13.0.
• Device/dtype: Apple MPS, float32; CPU used as an independent comparison.
• Full-model fixture: locally initialized tiny Granite MoE Hybrid from _build_granite_bridge(device) in the integration test, which sets torch.manual_seed(0).
• Full-model input IDs: [[1, 2, 3, 4, 5]], evaluation mode, no gradients.
• Native expert config: _experts_implementation == "grouped_mm".

Only this dependency/version point is locally verified. Older/newer versions are not declared affected or fixed. CUDA was not tested.

Baseline full-model reproduction

Run on baseline a8bc19e7, not on the corrected branch:

TRANSFORMERLENS_ALLOW_MPS=1 uv run pytest \
  tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py::test_granite_eager_scan_device_correctness \
  -q --tb=short

The fixture is offline and does not download model weights. The controlled pre-fix rerun produced:

CPU: passed
MPS: failed
AssertionError: eager vs fused rel diff on mps: 1.03e+00
Measured relative difference: 1.0326683320021979

Other instrumented runs showed different large discrepancies; this is not a stable small rounding discrepancy or a measured failure-rate estimate.

Standalone primitive reproduction

This diagnostic requires only torch and an available MPS device. It does not depend on TransformerLens, HF expert dispatch, a model download, or an attached local file.

import platform

import torch

if not torch.backends.mps.is_available():
    raise RuntimeError("This diagnostic requires MPS")

print(platform.system(), platform.release(), platform.machine())
print("torch:", torch.__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())

Measured output on the verified host:

Darwin 23.6.0 arm64
torch: 2.11.0
out-slice max error: 29.132984161376953
assignment max error: 4.76837158203125e-06

A separate fixed grouped-matmul diagnostic used CPU generator seed 0, input [10,64], weights [4,64,32], and cumulative offsets [2,5,5,10] (expert 2 receives no tokens). Against concatenated ordinary per-group products:

• CPU grouped maximum absolute error: 0.0.
• Native MPS grouped and HF _grouped_mm_fallback maximum absolute error: 29.132984161376953.
• Ordinary MPS per-group matmul plus concatenation maximum absolute error: 9.5367431640625e-6.

This establishes an upstream MPS primitive correctness problem in the inspected build. Whether native grouped execution has the same internal root cause still requires upstream investigation.

Additional intermediate evidence

One instrumented model run measured MPS-versus-CPU maximum absolute differences of approximately 4.84e-8 at the SSM mixer output and 8.94e-8 at router logits, with identical selected expert indices, followed by approximately 6.19e-3 at the MoE output.

A separate direct expert-capture run observed gate_up_proj shape [4,64,64] and repeated expert outputs differing by approximately 0.838. These measurements came from different instrumentation/configurations and should not be combined into a single experiment.

Implemented local workaround

The correction selects HF's original ordinary eager expert forward through an instance-local dispatch wrapper installed by MoEBridge.forward. It is guarded to:

• Exact HF GraniteMoeHybridExperts class semantics.
• torch release 2.11.0 and transformers release 5.13.0, including local suffixes but excluding prerelease/development versions.
• Live MPS expert input and grouped_mm policy.

It avoids native grouped execution and the grouped fallback, does not copy production activations to CPU, and emits a RuntimeWarning about correctness fallback and possible slower execution. Shared configs and global dispatch tables are not mutated. CPU-to-MPS-to-CPU migration follows the live input device; explicit eager execution and other policies/classes retain their dispatch. Custom instance forwards are left unchanged.

This avoids a shared-config policy race, but does not claim arbitrary concurrent model calls, device/config changes, or hook edits are safe. Direct HF execution before the containing bridge's first call is not protected; after installation, the same expert instance retains the wrapper.

Local implementation commits:

• 7f23a626: runtime correction and initial regressions.
• 72a38337: migration, hooks, gradients, and failure-recovery regressions.
• caa8e21a: user-facing documentation and primitive diagnostic.

These are local commit references, not public reproduction links. Add the implementation PR link when one exists.

Validation status

On the verified host:

Focused dispatch + MPS + Granite hybrid regressions: 49 passed
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

The recorded skip/xfail counts are suite outcomes, not new markers added by this workaround. Focused regressions exercise repeatability, an independent ordinary-linear expert oracle, migration, explicit eager policy, observer/output-editing hooks, routing/input/expert-weight gradients, state preservation, and recovery after a failed expert activation. The existing Granite MPS scan parity threshold remains unchanged.

Formatting, format checks, and git diff --check passed. The full make test-pr run reached integration after completing the other tiers, but timed out at the 30-minute command limit. Integration is not claimed complete or passing; no failure had been reported before interruption. Docs are not claimed warning-free. CUDA and other dependency versions remain unvalidated.

Relationship to PR #1878

Related PR: #1878 - fix: include Granite embedding scale in residual decomposition.

At the latest filing-preparation check, #1878 is open and unmerged, targeting dev. Its inspected changes propagate embedding_multiplier and scale embedding contributions in ActivationCache residual decomposition/component-level direct logit attribution. They deliberately leave native hooks and model forward execution unchanged. Its MoE regression runs on CPU, and its description does not claim Granite hybrid runtime validation.

This MPS grouped-expert issue is not addressed by the inspected #1878 changes. Cache decomposition scaling cannot repair incorrect expert matrix products, and the primitive reproduction does not depend on ActivationCache. This conclusion is based on its description and complete diff inspected during preparation, not a reproduction run on that PR's head.

Duplicate search and related upstream reports

A bounded GitHub issue search across open/closed issues found no direct TransformerLens duplicate using granite MPS; the returned contextual-decomposition proposal does not itself establish duplicate coverage.

A potentially related open PyTorch report is pytorch/pytorch#193487, "MPS: allocator-state-dependent silently wrong matmul results". It reports incorrect transpose-view matmul after large prior allocations, including on torch 2.11.0. The small nonzero-offset output-write reproduction here differs; shared root cause has not been established. Its version-range and mitigation claims are not independently validated by this report.

Also inspected:

• pytorch/pytorch#116769, closed: a different indexed-matmul correctness reproduction on torch 2.1.2.
• pytorch/pytorch#199882, closed: a native crash on transposed inputs in later nightly builds, rather than this silent output-slice defect.

Searches included MPS mm storage_offset, MPS grouped_mm, and MPS "out=" "offset". This is a bounded search, not proof no duplicate exists. Before opening a separate PyTorch issue, compare the inline primitive reproduction with #193487 or offer it there for maintainer triage. No upstream report or comment has been posted for this work.

Expected behavior and remaining acceptance criteria

Fixed inputs, routes, and weights should yield repeatable expert outputs agreeing with independent eager execution within validated float32 tolerances. SSM scan comparisons should not alter unrelated expert dispatch. Valid CPU/CUDA grouped execution should remain available.

• Establish an independent fixed-data primitive/grouped reproduction.
• Implement the narrow locally verified expert-dispatch correction.
• Add independent fixed-route, repeatability, migration, hook, gradient, and failure-recovery regressions.
• Restore existing Granite MPS scan parity without semantic weakening.
• Verify unchanged CPU dispatch and explicit eager policy.
• Run format/type checks, focused regressions, and a docs build.
• Complete the integration tier and final branch review.
• Push the implementation and link its public PR.
• Publish or attach the primitive reproduction to an appropriate upstream report.
• Confirm native grouped-kernel internals and any shared upstream root cause.
• Validate additional dependency versions before widening the workaround or claiming an upstream fixed version.
• Validate CUDA on hardware before claiming CUDA test coverage.

Scope

This report concerns Granite MoE Hybrid expert execution on the verified MPS environment, not all Granite variants or all MPS grouped matrix multiplication. Keep the correction separate from contextual decomposition and the embedding-decomposition change in #1878. Report the independent primitive defect upstream without inventing affected/fixed version bounds.