exir: reuse the graph's fake mode for constant-only submodules

[psiddh]

Problem

Exporting DINOv2 (or any model that partitions into a constant-only subgraph) fails in to_edge_transform_and_lower():

Exception: An error occurred when running the 'FuseBatchNormPass' pass after the following passes: []
  AssertionError: fake mode (...) from fake tensor input 0
                  doesn't match mode (...) from fake tensor input 6

Reported for DINOv2 on WebGPU. The error is doubly misleading: it names a batch-norm pass on a ViT with no batch norm, and the actual cause is in exir/pass_base.py, not in the pass or the backend.

Root cause

_ExportPassBase.call picks the FakeTensorMode to retrace under by scanning self.inputs(graph_module):

fake_tensor_mode = None
for i in inputs:
    if isinstance(i, FakeTensor):
        ...
        fake_tensor_mode = i.fake_mode
if fake_tensor_mode is None:
    fake_tensor_mode = FakeTensorMode(allow_non_fake_inputs=True)

extract_input unwraps a placeholder's FakeTensor to its materialized .constant when one exists, returning a real tensor. A submodule whose placeholders are all lifted tensor constants therefore yields no FakeTensor to scan, and a fresh mode is created — even though the graph already carries exactly one consistent mode.

The retrace then writes vals from the new mode onto the traced nodes while the placeholders keep their originals, leaving two modes in one graph. _get_updated_range_constraints → detect_fake_mode asserts on the mixture.

Instrumenting the failing DINOv2 submodule shows it precisely:

MULTIPLE FAKE MODES: {A: 6, B: 11}
 idx 0-5   placeholder     _lifted_tensor_constant0..5   <- mode A (from export)
 idx 6-7   call_function   aten.arange.start_step        <- mode B (from pass_base:867)

placeholders: 6  extracted: 6  fake among extracted: 0
distinct modes on placeholder meta['val']: 1
_lifted_tensor_constant0: val=FakeTensor .constant_set=True

FuseBatchNormPass appears only because it is the first pass in the Vulkan pipeline that calls super().call() to retrace (fuse_batch_norm.py, "To regenerate metadata and shape information, retrace module"). Any pass in that position would fail identically.

Fix

Fall back to the mode already recorded on the placeholders before creating a new one. Strictly narrowing: graphs where the input scan already found a mode are untouched, and a new mode is still created when the placeholders carry no FakeTensor or disagree on one.

Verification

• dinov2_vits14 exports through WebGPUPartitioner to an 84 MB .pte with 27 VulkanBackend delegates; the 4 remaining CPU ops (floor, index, mul, any) match the partitioner's own skip list.
• Reproduces without the fix and passes with it, using the reporter's unmodified script.
• Two regression tests added; they fail (AssertionError: 2 != 1) with the fallback stubbed out.
• exir/tests/test_pass_infra.py: 19/19 pass.
• Exporting DINOv2 through XNNPACK yields byte-identical .pte files with the fix enabled and disabled, confirming unaffected paths are unperturbed.

Not covered: the Vulkan .pte has not been executed (no VulkanBackend in the test runtime), so runtime numerics for DINOv2 on Vulkan/WebGPU remain unverified and are out of scope for this change.

Authored with assistance from Claude Code.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22891

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 5 New Failures

As of commit 276ed55 with merge base b03dec2 ():

NEW FAILURES - The following jobs have failed:

• Build documentation / build (buck2) / Build doc (gh)
  Could not load credentials from any providers
• pull / test-lora-linux / linux-job (gh)
  Could not load credentials from any providers
• pull / test-mediatek-models-linux / linux-job (gh)
  Could not load credentials from any providers
• pull / unittest / windows / windows-job (gh)
  executorch
• pull / unittest-buck / macos / macos-job (gh)
  ##[error]fatal: Fetched in submodule path 'kernels/optimized/third-party/eigen', but it did not contain 729443409942a1816ddf74b95224003b83f4925c. Direct fetching of that commit failed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[Copilot]

🟢 Approval recommended

The change is narrowly scoped, preserves existing behavior when a mode is already found via inputs(), and includes targeted regression tests for the reported failure mode.

Pull request overview

This PR fixes an EXIR export-pass replay failure for submodules whose FX graphs contain only constant placeholders by reusing the existing FakeTensorMode already recorded on placeholder node metadata, instead of minting a new mode during retrace.

Changes:

• Add _ExportPassBase._get_fake_mode_from_placeholders() to reuse a single placeholder FakeTensorMode when inputs() yields no FakeTensor (e.g., all placeholders were unwrapped to .constant).
• Update _ExportPassBase.call() to fall back to the placeholder mode before creating a fresh FakeTensorMode.
• Add regression tests covering the constant-only placeholder case and the “placeholders disagree on mode” guard.

File summaries

File	Description
exir/pass_base.py	Reuses placeholder FakeTensorMode when inputs() returns no FakeTensor, preventing mixed fake modes after retrace.
exir/tests/test_pass_infra.py	Adds regression tests ensuring constant-only graphs keep a single fake mode and that mode reuse is skipped when placeholders disagree.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 0
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[digantdesai]

how did this work for xnnpack?

Not covered: the Vulkan .pte has not been executed (no VulkanBackend in the test runtime), so runtime numerics for DINOv2 on Vulkan/WebGPU remain unverified and are out of scope for this change.

And does this now work on both vulkan and wgpu e2e?

[digantdesai]

@claude review this PR, and verify if there are any existing assumption around fakemode == none which can be broken by this.

[digantdesai]

Also in a different PR, do you want to add dinov2 on wgpu and/or vk as a model level test ci? perhaps on a trunk or something?

[claude]

Claude finished @digantdesai's task in 3m 55s —— View job

---

Review: reuse graph fake mode for constant-only submodules

Task

• Read the diff (exir/pass_base.py, exir/tests/test_pass_infra.py)
• Trace how fake_tensor_mode is consumed downstream of call()
• Hunt for existing fake_mode is None assumptions the fallback could break
• Post findings

---

The fix itself is correct and the diagnosis is accurate: extract_input() (pass_base.py:612-614) unwraps a placeholder's FakeTensor to its real .constant, so a constant-only submodule leaves the inputs() scan (:881-886) with no FakeTensor, and the pre-existing code minted a fresh mode that then diverged from the placeholders. Reusing the placeholder mode is the right correction.

But you asked specifically about assumptions around fake_tensor_mode is None, and there are two behavioral shifts worth calling out — the fix doesn't just supply a mode, it moves affected graphs from the if branch into the else branch at :894-899:

if fake_tensor_mode is None:
    fake_tensor_mode = FakeTensorMode(allow_non_fake_inputs=True)
    dispatcher_mode = nullcontext()          # <- old path for const-only graphs
else:
    fake_tensor_mode.allow_non_fake_inputs = True
    dispatcher_mode = enable_python_dispatcher()   # <- new path for const-only graphs

1. enable_python_dispatcher() now runs where nullcontext() used to.
Constant-only graphs previously retraced without the python dispatcher; they now retrace with it. This is almost certainly more correct (it's the same context every non-empty-mode graph already uses), and it's what makes the modes line up. Just flag that the change is not purely "supply a mode" — the dispatch context flips too. Worth a one-line mention in the PR description so a reviewer isn't surprised, and worth confirming no constant-only pass relied on the no-dispatcher behavior (I found none in-tree).

2. The reused mode is mutated in place: fake_tensor_mode.allow_non_fake_inputs = True.
This is the assumption most likely to bite. In the old None path the mutated object was a throwaway. Now it's the same FakeTensorMode instance carried by the exported program's placeholder tensors, so flipping allow_non_fake_inputs persists on the graph's shared mode after the pass returns. In practice this is benign (it only loosens the mode — permitting real tensors — and export-created modes generally already have it set), but it is a real side effect on shared state that didn't exist before. If you want to be strictly non-invasive you could snapshot and restore it, but I wouldn't block on it — just be aware it's not as hermetic as "graphs are untouched."

Assumptions I checked that are NOT broken:

• extract_input()'s assert self.fake_tensor_mode is not None (:616) fires during inputs(), which runs at :878 before the new code — unaffected.
• Graphs where the inputs() scan already finds a mode never reach the new branch — genuinely untouched (consistent with the byte-identical XNNPACK .pte claim).
• Graphs with no fake placeholders (all real, or none) still get None from the helper and fall through to the fresh-mode + nullcontext() path — old behavior preserved.
• detect_fake_mode / _get_updated_range_constraints and _fake_program.py:28-38 (which requires a single detectable mode) are the very consumers this fix satisfies — the change moves them from "asserts on mixed modes" to "sees one mode." Strictly an improvement.
• Re-entering an already-existing mode via with fake_tensor_mode is already the established pattern in the pre-existing else branch, so no new re-entrancy risk.

Minor / nits:

• _get_fake_mode_from_placeholders iterates all nodes but only inspects placeholders. Since placeholders lead the graph, you could break after the first non-placeholder op — micro-optimization, optional.
• The helper's "disagree → return None → mint a fresh mode" behavior means a genuinely mixed-mode graph silently gets a third mode rather than raising. That matches the old fall-through so it's fine, but note it won't surface a pre-existing corruption — it just won't make it worse.

Re @digantdesai's inline question ("how did this work for xnnpack?"): XNNPACK partitioning doesn't produce the all-constant-placeholder submodule that trips this, so its passes' inputs() scan always finds a FakeTensor and never reaches the new branch — which is exactly why the XNNPACK .pte is byte-identical with and without the fix. On the e2e question: per the PR's own "Not covered" note, the Vulkan/WebGPU .pte is produced but not executed (no VulkanBackend in the test runtime), so end-to-end numerics on Vulkan and WebGPU are still unverified.

Recommendation: logic is sound and narrowly scoped. Before merge I'd suggest (a) a sentence in the description noting the dispatcher-context flip, and (b) a decision on whether to restore allow_non_fake_inputs on the shared mode rather than mutate it permanently. Neither is a correctness blocker.

Note: I reviewed statically; I did not run test_pass_infra.py in this environment.