Skip to content

Fix crashes, wrong outputs and lowering failures in several backends and in exir - #24222

Open
shoumikhin wants to merge 22 commits into
pytorch:mainfrom
shoumikhin:fix-lerobot-export-bugs
Open

shoumikhin wants to merge 22 commits into
pytorch:mainfrom
shoumikhin:fix-lerobot-export-bugs

Conversation

@shoumikhin

@shoumikhin shoumikhin commented Oct 28, 2026 •

Copy link
Copy Markdown
Contributor

Bug fixes found while exporting models to ExecuTorch backends, one commit per fix. AOTI (AOTInductor) is the compiler behind the CUDA and Metal backends.

Fixes

  1. AOTI: the missing-kernel check also catches scatter and index_put fallbacks at lowering.
  2. OpenVINO 2026.4: init no longer frees model data the compiled model still reads.
  3. Qualcomm: DecomposeWrapWithAutocast matches submodule inputs by position and inlines nested regions first.
  4. AOTI: detach_copy maps to alias.
  5. CUDA: strided method outputs are copied, not rejected.
  6. exir: constant_prop_pass no longer folds tensors or buffers that schema-marked in-place writes change, directly or through views.
  7. MLX: tensors nested in positional arguments are checked for unsupported dtypes.
  8. MLX: placeholders of one shared weight in a subgraph share a constant slot and load.
  9. Core ML: ops on bfloat16 tensors stay on the CPU.
  10. Core ML: ops on tensors above rank 5 stay on the CPU, except layout chains coremltools rewrites.
  11. Vulkan: programs with a constant of more than 4 dims load.
  12. Arm: the TOSA backend keeps unused delegate user inputs. Dropping them gave wrong VGF outputs. Ethos-U is not fixed.
  13. NXP: the quantizer accepts non-strict torch.export programs.
  14. XNNPACK: the quantizer skips bfloat16 tensors. cat now follows module scoping too.
  15. exir: to_backend keeps the output node last when it fuses a partition.
  16. exir: to_edge lowers torch._check_tensor_all and torch.sym_ite. Partitioned programs also need fix 15.
  17. CUDA: adds the aoti_torch_as_strided shim.
  18. exir: GroupBasedPartitioner no longer merges partitions into mutually dependent ones.
  19. Qualcomm: unquantized bfloat16 nodes stay on the CPU.
  20. Cortex-M: MatmulToBmmPass handles a constant operand.
  21. exir: ExportPass gives tensor constants static shapes.
  22. MLX: an out-of-place slice update can read a constant base.

For release/1.6: fix 6 needs #24224, fixes 16 and 18 need #24247, fix 19's test needs #23915, and fix 21's test needs #23809.

Testing

Each fix adds or extends a test that fails without it. CUDA and Vulkan ran on an A100 GPU, OpenVINO on an x86 CPU with 2026.4.1 and 2026.3.0, Qualcomm with a QNN SDK build, Core ML and MLX on macOS, Arm with the TOSA reference model and the VGF emulation layer, and NXP with the eIQ Neutron SDK. Metal was not run. GitHub CI does not run the tests for fixes 6 and 11, and skips fix 19's test without the QNN SDK. All three ran by hand.

Copilot AI balanced review requested due to automatic review settings October 28, 2026 03:40
@pytorch-bot

pytorch-bot Bot commented Oct 28, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

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

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

⏳ No Failures, 115 Pending

As of commit 3119872 with merge base 0db7a88 (image):
💚 Looks good so far! There are no failures yet. 💚

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

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 28, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

@shoumikhin

Copy link
Copy Markdown
Contributor Author

@claude /pr-review

@claude

claude Bot commented Oct 28, 2026 •

Copy link
Copy Markdown

Claude finished @shoumikhin's task in 22m 42s —— View job


Approve: nothing blocks. The main open point is that on OpenVINO before 2026.4, the delegate blob now stays in memory next to OpenVINO's own copy for the whole life of the method. Three smaller robustness and test gaps are noted below.

PR Review: #24222

Summary

The PR does two things:

  • AOTI lowering now reports missing scatter and index_put fallback kernels.
  • OpenvinoBackend::init no longer frees a blob that the compiled model may still read.

Nothing blocks merging. The OpenVINO memory cost on versions before 2026.4 is worth settling, and three smaller points follow. The PyTorch Inductor and OpenVINO sources are not in this checkout, so I mark statements about their internals as low confidence. I ran no tests.

Code Quality

  • backends/aoti/aoti_backend.py:208-237: the scatter check rebuilds Inductor's private symbol name instead of checking what Inductor writes.

    • The other three wrappers check a name that Inductor passes in. This one copies the naming logic inside CppWrapperCpu._generate_scatter_fallback (get_c_shim_func_name(...).replace("__", "_") + "_out", lines 220-223). It also spells out all 8 of that method's parameters, so a signature change turns every scatter fallback into a TypeError.
    • The test asserts the same rebuilt string (backends/cuda/tests/test_cuda_export.py:580).
      • If a torch bump changes the name Inductor emits, CI stays green, but the check reports a name the .wrapper.so never calls.
      • Once someone adds that name to get_supported_fallback_kernels, lowering passes and the program crashes on its first run again.
    • Comparing against what the original method writes would catch this drift at lowering time. For example, wrap self.writeline for the duration of the call, and fail if the emitted call is not kernel.
    • Low confidence: one check over the aoti_torch_* calls in the generated wrapper would also cover direct-emission paths that nothing watches yet. The CUDA shim defines no aoti_torch_item_float16 or aoti_torch_item_float64 (backends/cuda/runtime/shims/memory.cpp:451-458), and I believe Inductor's codegen_tensor_item emits those calls directly.

    Fix this →

  • backends/openvino/tests/ops/base_openvino_op_test.py:71-80: the mapped .pte is released only on success.

    • If load_method, execute, or the assert raises, the traceback keeps program and method alive while TemporaryDirectory cleans up.
    • On Windows (the case the comment on line 79 is about), the cleanup then raises PermissionError on top of the real failure.
    • tempfile.TemporaryDirectory(ignore_cleanup_errors=True) avoids this; pyproject.toml requires Python 3.11 or newer, so it is available.

    Fix this →

Testing

  • backends/aoti/aoti_backend.py:251-256: the Metal part of fix 1 has no test.

    • The PR says Metal was not run, and nothing under backends/apple/metal/tests/ uses scatter, index_put, or the missing-kernel error.
    • The patch is on CppWrapperCpu. The CUDA test shows that the CUDA wrapper picks it up, but nothing shows that the MPS wrapper does not override _generate_scatter_fallback or _generate_index_put_fallback. This part is low confidence, since I could not read Inductor.
    • If the MPS wrapper does override them, Metal programs still lower and then crash on their first run, because the .so is opened with RTLD_LAZY (backends/apple/metal/runtime/metal_backend.cpp:347).
    • Adding the Int64Accumulate case to backends/apple/metal/tests/test_metal_backend.py, built with export_model_to_metal (test_modules.py:1218), would cover this.

    Fix this →

Performance

  • backends/openvino/runtime/OpenvinoBackend.cpp:336-337: on OpenVINO before 2026.4, the blob now stays alive next to OpenVINO's own copy for the life of the method.

    • The PR description accepts this cost, but neither the code comment nor the docs mention it.
    • 2025.1 to 2026.3 are still advertised as supported. docs/source/build-run-openvino.md:64 and the runtime-not-found error at OpenvinoBackend.cpp:215 both say openvino>=2025.1.0,<2027.0.0.
    • What the extra copy costs depends on the loader:
      • With the default Module load mode, LoadMode::File (extension/module/module.h:100), it is a heap buffer the size of the compiled model.
      • With the LLM runner's default, MmapUseMlockIgnoreErrors (extension/llm/runner/llm_runner_helper.h:112), the pages stay locked in RAM whenever mlock succeeds.
    • For the examples/openvino/llama and qwen2_5 exports, the blob is the whole quantized model.
    • Freeing only on versions before 2026.4 gives those versions back exactly the behavior they had before this PR. This needs ov_get_openvino_version and ov_version_free declared in OpenvinoApi.h and loaded with the other LOAD_SYMs. At minimum, state the cost in the comment and on the docs page.

    Fix this →

Recommendation

Approve

Nothing here blocks the merge. The one point to settle is the OpenVINO memory cost: either free the blob on versions before 2026.4, or document the cost. That can happen here or in a follow-up.


  • Gather context and fetch the diff (the checkout is shallow, so I compared against the merge base 8ecf0fb)
  • Run three independent reviews, then consolidate and fact-check them
  • Post review (I ran no tests)

@shoumikhin
shoumikhin force-pushed the fix-lerobot-export-bugs branch from 35f7db1 to 55b4840 Compare October 28, 2026 09:44
@shoumikhin
shoumikhin requested a review from psiddh as a code owner October 28, 2026 09:44
Copilot AI balanced review requested due to automatic review settings October 28, 2026 09:44
@shoumikhin shoumikhin changed the title Fix two backend crashes that show up only when a program runs Fix crashes and lowering failures in the AOTI, OpenVINO and Qualcomm backends Oct 28, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

@shoumikhin

Copy link
Copy Markdown
Contributor Author

Thanks for the review. Points 1, 2 and 4 are fixed in the commits that own the code. Point 3 has no code change, see below.

  1. Scatter and index_put check: it no longer rebuilds the kernel name or spells out the method's arguments. It reads the kernel name from the line the fallback writes, so it checks what the generated code really calls. The item path you mention, and view_dtype, also write their calls directly. I left them for a follow-up that checks the built library's symbols instead of hooking one path at a time.
  2. OpenVINO test cleanup: the temporary folder now ignores cleanup errors, so a failing test on Windows reports only its own error.
  3. Metal: the Metal tests need torchao's MPS kernels, which my setup does not have, so I did not add a Metal test. I checked instead that the MPS wrapper does not override the two hooked methods, and that an MPS compile with this check reports both kernels as missing.
  4. OpenVINO memory: init now frees the blob as before when the OpenVINO runtime is older than 2026.4, because those versions copy it during the import. With newer versions the blob is kept until the delegate is destroyed.

@shoumikhin shoumikhin changed the title Fix crashes and lowering failures in the AOTI, OpenVINO and Qualcomm backends Fix crashes and lowering failures in the AOTI, OpenVINO and Qualcomm backends and in constant_prop_pass Oct 28, 2026
@shoumikhin
shoumikhin force-pushed the fix-lerobot-export-bugs branch from 55b4840 to 52e8dcf Compare October 28, 2026 15:04
@shoumikhin
shoumikhin force-pushed the fix-lerobot-export-bugs branch from 52e8dcf to c39d57f Compare October 28, 2026 19:42
Copilot AI balanced review requested due to automatic review settings October 28, 2026 19:42
@shoumikhin shoumikhin changed the title Fix crashes and lowering failures in the AOTI, OpenVINO and Qualcomm backends and in constant_prop_pass Fix crashes, wrong outputs and lowering failures in several backends and in constant_prop_pass Oct 28, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

@github-actions github-actions Bot added ciflow/trunk module: arm Issues related to arm backend labels Oct 28, 2026
@shoumikhin
shoumikhin force-pushed the fix-lerobot-export-bugs branch from c39d57f to d11bd4b Compare October 28, 2026 21:47

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

@shoumikhin shoumikhin changed the title Fix crashes, wrong outputs and lowering failures in several backends and in constant_prop_pass Fix crashes, wrong outputs and lowering failures in several backends and in exir Oct 28, 2026
@shoumikhin

Copy link
Copy Markdown
Contributor Author

@claude /pr-review

@claude

claude Bot commented Oct 28, 2026 •

Copy link
Copy Markdown

There are no blockers, so I recommend approving. The nine issues are four edge cases that individual fixes miss, two possible gaps I'm not confident about, two fixes with incomplete tests, and a host round trip on every call for strided CUDA outputs.

Detailed review

PR Review: #24222

Summary

This PR makes 19 independent fixes for crashes, wrong outputs and lowering failures in exir (constant propagation, the Edge verifier, the decomposition table, GroupBasedPartitioner, to_backend fusion) and in the AOTI/CUDA, Core ML, Arm, MLX, NXP, OpenVINO, Qualcomm, Vulkan and XNNPACK backends. None of the findings blocks merging:

  • Fixes 5, 10, 12 and 18 each miss an edge case.
  • Fixes 6 and 19 may have a gap (low confidence).
  • Fixes 14 and 16 are partly untested.
  • Fix 5 sends strided CUDA outputs through a slow host round trip.

I could not run tests in this job. I also did not read the torch, coremltools, OpenVINO or QNN internals, so anything I say about their behavior comes from ExecuTorch's own code and comments.

Code Quality

  • Fix 12: Ethos-U still silently gives wrong outputs when a delegate user input is unused.
    • backends/arm/tosa/backend.py:344-353 now keeps unused user inputs in the TOSA graph.
    • backends/arm/test/ops/test_copy.py:98-104 xfails the U55/U85 runs: Vela drops those inputs again, but the Ethos-U runtime binds delegate inputs by position.
    • So a model like CopyOverUnusedInput lowers for Ethos-U without complaint and computes on the wrong buffers. Only the output comparison catches it.
    • Fix: fail in EthosUBackend.preprocess (backends/arm/ethosu/backend.py:88) when a delegate user input has no users. That turns this into a lowering error. Fix this →
  • Fix 5: the capture step still fails the call when an output can't be raw-copied.
    • An output can pass warmup and still fail the raw-copy check on the capture step, because that step's inputs differ (for example in shape).
    • When that happens, backends/cuda/runtime/cuda_backend.cpp:1065-1080 makes execute return Error::NotSupported with no outputs.
    • CaptureGuard then disables graphs, so later calls run eagerly and succeed, but this one call fails.
    • The warmup path at lines 1171-1184 handles the same case by disabling graphs and copying the output, with no error. The capture-step behavior goes against the commit's goal of copying such outputs instead of rejecting them.
    • It is rare, but this call could run eagerly once the capture state has unwound. Fix this →
  • Fix 10: the rank > 5 check ignores ops the user skips.
    • _touches_unsupported_high_rank_tensor (backends/apple/coreml/partition/coreml_partitioner.py:84-116, called at line 220) accepts a rank > 5 tensor when every node in its chain is a view, permute or clone.
    • But is_node_supported applies skip_ops_for_coreml_delegation to each node separately (line 269).
    • Example: with skip_ops_for_coreml_delegation=["aten.permute_copy.default"], a view_copy → permute_copy → view_copy chain through rank 6 keeps both views delegated and the permute on CPU. The first view's rank-6 result becomes a delegate output, which coremltools rejects ("Core ML only supports tensors with rank <= 5").
    • Fix: the chain check should treat a node that should_skip_op_for_delegation rejects as unsupported. Fix this →
  • Fix 18: a cached None closure is never invalidated, so later merges can be rejected for no reason.
    • When kept partitions depend on each other, _closure caches None for every partition on its stack (exir/backend/canonical_partitioners/group_partitioner.py:180-183). Its docstring (lines 156-159) says merges never undo that.
    • But the two mutually dependent groups can merge with each other. Take groups {a1, a2} and {b1, b2} with direct edges a1 → b1 and b2 → a2: if nothing else lies between them, _can_merge_partitions accepts the pair, and the cycle is gone.
    • Line 327 drops only truthy closures (if c and c & merged_nodes). The stale None for the merged partition, and for any partition that was on the stack, keeps rejecting every later merge whose downstream reaches them.
    • The partitions stay valid but come out smaller than they need to be.
    • Fix: drop None entries there too (c is None or c & merged_nodes) and correct the docstring.
    • The new test_no_cycle_through_another_partition (exir/backend/test/test_group_partitioner.py:1953-1979) can't catch over-rejection, here or on the compute-weight path: it only uses allows_single_node_partition=True and only checks that fuse_by_partitions does not raise. Asserting the expected partitions would catch it. Fix this →
  • Fix 6 (low confidence): in-place writes inside higher-order ops are not detected.
    • _nodes_written_in_place (exir/passes/constant_prop_pass.py:138-147) skips every node whose target has no _schema. That includes wrap_with_set_grad_enabled and wrap_with_autocast.
    • Suppose training IR wraps with torch.no_grad(): self.buf[1:3].copy_(x) in such an op. Then buf is never marked as written, and a later top-level read like self.buf * 2 still folds to the value from before the write. This is the bug this commit fixes, one level down.
    • _schemaless_aliasing_inputs (exir/passes/insert_write_back_for_buffers_pass.py:68-90) already treats nodes without a schema conservatively.
    • Fix: recurse into the higher-order op's submodules (better), or treat every tensor operand of a HigherOrderOperator call as written.
    • I could not run export here to confirm that the write ends up inside the higher-order op. Fix this →
  • Fix 19 (low confidence): a partly quantized bf16 op is kept from decomposition but never delegated.
    • _in_quantized_region (backends/qualcomm/partition/utils.py:163-168) is true if any input comes from a dq op or any user is a q op. In that case filter_fn (lines 176-179) keeps the node from decomposition.
    • But is_node_supported (backends/qualcomm/partition/qnn_partitioner.py:173-177) rejects the node if any unquantized tensor is bf16.
    • Example: a bf16 linear whose weight comes from torchao.dequantize_affine(..., output_dtype=torch.bfloat16), but whose activation and output are plain bf16. It is preserved as aten.linear, rejected by the partitioner, and left on CPU undecomposed. That most likely ends in the "was not decomposed or delegated" error.
    • Fix: require every unsupported-dtype input to come from a dq op, and the output to be supported or feed only q ops. Then filter_fn matches the partitioner.
    • I could not check what the quantization annotations look like for this pattern at partition time. Fix this →

Testing

  • Fix 16: the sym_ite half has no test.
    • The verifier now allows _PYTHON_SYM_OPS_TO_EXECUTORCH_SYM_OPS (exir/verification/verifier.py:300).
    • But the only new test, test_check_tensor_all (exir/program/test/test_program.py:340), covers _check_tensor_all.
    • The existing test_emit_sym_ite (exir/emit/test/test_emit.py:2540-2543) still passes EdgeCompileConfig(_check_ir_validity=False), so it never reaches the verifier.
    • Fix: remove that config, and the existing test becomes the regression test. Fix this →
  • Fix 14: no test runs in OSS CI, and the cat scoping change has no test at all.
    • _annotate_cat now honors filter_fn (backends/xnnpack/quantizer/xnnpack_quantizer_utils.py:1068-1069).
    • But test_bf16_is_not_quantized only covers the bf16 path. No test scopes a config with set_module_name or set_module_type around a cat.
    • pytest.ini:17 also ignores backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py. Only the internal BUCK target (backends/xnnpack/test/BUCK:150) runs it.
    • The PR description lists only fixes 6 and 11 as skipped by GitHub CI. Either add fix 14 there or run the file in OSS CI. Fix this →

Performance

  • Fix 5: strided CUDA outputs take a host round trip on every call.
    • Strided AOTI outputs used to be rejected by _check_tensor_metadata. They now go to the slow path of _copy_slimtensor_to_etensor_impl (backends/cuda/runtime/utils.h:220-288), reached through copy_slimtensor_to_device_etensor_async (lines 341-347) for the usual GPU-resident ETensor.
    • On every execute, that path does:
      • a stream sync;
      • a D2H copy of the view's whole span;
      • a CPU rearrange into a second host buffer;
      • an H2D copy;
      • another sync, with two heap allocations along the way.
    • So the _async copy blocks, and the method also loses CUDA graphs from warmup on (cuda_backend.cpp:1171-1179).
    • Fix: copy on the device.
      • cudaMemcpy2DAsync covers 2-D strided layouts, as at backends/cuda/runtime/cuda_kv_cache.cpp:695.
      • A small kernel covers the rest.
      • SlimTensor::copy_, behind aoti_torch_copy_ (backends/cuda/runtime/shims/memory.cpp:406-424), already handles CUDA-to-CUDA strided copies, but its comment says it is synchronous, so check its cost first.
      • At minimum, log a one-time warning so the slowdown is visible. Fix this →

Recommendation

Approve. Nothing here blocks merging. The two low-confidence items (fixes 6 and 19) need the author's judgment. The rest can be follow-ups, but three are small enough to fold in here: the GroupBasedPartitioner cache fix and the two test changes.


  • Gather context and fetch the diff (42 files, 19 commits, compared against the merge base b831dbe)
  • Run three independent reviews
  • Consolidate and fact-check findings (9 left, none blocking)
  • Post review

View job run

…kends

An AOTI backend (CUDA, Metal) checks at lowering time that every kernel
Inductor leaves to ATen exists in the backend's shim library, so a missing
one fails the export with a clear list. The check watches two code paths
that emit a C shim call, but scatter and index_put fallbacks write their
call directly. A program that needs one of those kernels lowers without an
error, and the process then dies on the first run with an undefined symbol.

For example, Inductor leaves an int64 scatter_add_ to
aoti_torch_cuda_scatter_reduce_two_out, which the CUDA shim does not have.

Watch those two paths too. They do not pass the kernel name to a method
the check can see, so read it from the line they write. This checks the
name the generated code really calls.
Since OpenVINO 2026.4, ov_core_import_model wraps the blob it is given in
an ov::Tensor and the compiled model can keep reading that memory. Earlier
versions read it through a stream during the import only.
OpenvinoBackend::init frees the blob right after the import. When the
program is loaded from a buffer, Free() does nothing, so nothing goes
wrong. When it is loaded from a file, the blob is really freed: with a
memory-mapped file the first execute crashes with a segmentation fault,
even for a two-layer model, and with a file read into memory the model
reads freed memory.

Free the blob in init only when the OpenVINO runtime is older than 2026.4,
as before. With newer versions keep it: ExecuTorch owns it and frees it
after destroy(), so it lives exactly as long as the compiled model.

The op tests loaded every program from a buffer, so they never saw this.
They now load from a file, as most users do.
…gions

torch.export records a `with torch.autocast(...)` block as a
wrap_with_autocast call on a submodule. DecomposeWrapWithAutocast inlines
that submodule into the graph. It had two bugs:

1. It mapped the submodule's placeholders to the operands by name. When
   the region reads inputs or parameters, the placeholders are named like
   arg2_1, arg0_1, arg1_1, not after the operands, so the pass failed with
   a KeyError such as 'arg2_1'. For example:

       with torch.autocast("cpu", enabled=False):
           y = self.linear(x)

2. It only looked at the top graph. A region nested in another region was
   copied out with its get_attr, whose submodule name only means something
   inside the outer submodule. In the top module it resolved to the outer
   body, so the program silently ran the wrong code.

Map placeholders to operands by position, and inline nested regions
before their parent.
Edge turns `x.detach()` into `aten.detach_copy`. Inductor has no lowering
or C shim for it, so an AOTI backend (CUDA, Metal) cannot lower any graph
that calls detach, for example:

    return (x * 2).detach() + 1

fails with "missing fallback kernels: aten::detach_copy".

ReplaceViewCopyWithViewPass already maps copy ops Inductor cannot lower
(slice_copy, select_copy, split_copy) back to views. Map detach_copy to
alias, which is what detach means without autograd. Mapping to detach
itself would still fail Edge verification under the default compile
config, because detach is a non-core view op that the partitioner
preserves.
AOTInductor can return a method output as a view into a larger buffer
with strides that do not match a contiguous tensor, for example a slice
`y[:, :, :14]` of a [1, 50, 32] result. The CUDA backend copies outputs
into the planned ETensor only from contiguous SlimTensors, so execute()
failed with "SlimTensor must be contiguous" although lowering accepted
the program.

The copy already has a slow path that rearranges elements through host
memory when the strides differ from the ETensor's. Let it take
non-contiguous views too: copy the span the view reaches instead of
nbytes(), and check that span against the storage, with checked
arithmetic. Wrapping a SlimTensor without a copy still requires a
contiguous one.

A CUDA graph replay copies each output's raw bytes. For such a strided
output, and also for a contiguous output whose ETensor uses another dim
order, the replay was silently wrong. A raw copy is known to be right
when both are contiguous in the default dim order, so for any other
layout warmup now turns graphs off and the capture step refuses to
replay, so the method runs without a graph.
torch.export keeps in-place ops in the program it returns. In

    t = torch.zeros(6)
    t[1:3] = x[:2]
    return t * 2

`t` and its slice look constant to constant_prop_pass, so the pass turned
them into constants, while `copy_` still writes into them at run time.
Decomposing the program then fails with "Constant ... is mutated in the
forward method". A buffer written in place has the same problem, and
there the folded reads are silently wrong, because buffers_to_mutate
does not list it until the program is decomposed.

Find the tensors an op writes in place (through any argument, including
out= and lists, as insert_write_back_for_buffers_pass already reads
them), directly or through a view, and do not fold them or their reads.
Programs without in-place ops, such as Edge programs, are unchanged.
The MLX partitioner decides which nodes the delegate can take, and part of
that is rejecting nodes whose inputs have a dtype MLX cannot hold, such as
float64. The check only looked at arguments that are a single node, so a
tensor list argument (the inputs of aten.cat, for example) was never
checked.

As a result a cat of float64 tensors was put into the delegate. Its float64
inputs became delegate inputs, and lowering then failed with "placeholder
has unsupported dtype: torch.float64" instead of leaving the cat on the CPU.
Sinusoidal timestep embeddings computed in float64 hit this.

Check every node inside each positional argument, so such a cat stays on
the CPU like any other float64 op.
When a weight is used by several partitions, to_backend copies its
placeholder once per extra user. Every copy keeps the original target, so
one MLX subgraph can get several placeholders for the same weight (a layer
applied at more than one step, for example).

The builder made a separate constant slot for each of them, but the named
slots are keyed by target, so only one of those slots got a name. Loading
the program then failed with "load_constants: no name found for constant
tensor N".

Give placeholders that share a target one constant slot. The weight is then
stored and loaded once per subgraph.
coremltools has no bfloat16 type. The support query that CoreMLPartitioner
asks coremltools only checks the op name, so an op that reads or produces
a bfloat16 tensor was still delegated, and coremltools 9.0 then raised
"KeyError: torch.bfloat16" during lowering. The whole export failed instead
of leaving that op on the CPU. Models that run part of their forward pass
under bfloat16 autocast hit this.

Override support for nodes with a bfloat16 input or output, the same way
the partitioner already skips other ops coremltools cannot convert.

That rejection is per call, but ops_to_not_decompose keeps an op whole as
soon as one of its calls is supported. With attention run once in float32
and once in bfloat16, the bfloat16 call was kept whole, rejected by the
partitioner, and left with no kernel ("Missing out variants"). Keep an op
whole only when every call of it is supported, as the MLX partitioner
does.
A Core ML tensor has at most 5 dimensions. The support query that
CoreMLPartitioner asks coremltools only checks the op name, so an op on a
rank 6 tensor (for example a stack of 5-D tensors) was still delegated, and
coremltools 9.0 then raised "Core ML only supports tensors with rank <= 5"
during lowering. The whole export failed instead of leaving that op on the
CPU.

coremltools does rewrite layout chains through a higher rank, such as the
view, permute, view of a window partition, as long as each step has static
shapes and a single user. So override support for a node when the tensors
above rank 5 it is connected to reach anything other than such a chain of
view, permute and clone ops. Those ops then run on the portable backend,
and window partitions stay delegated as before.
A Vulkan program failed to load when a delegated op used a constant tensor
with more than 4 dims, for example x * c with a 5-D buffer c:

  Tensor.h:790: (sizes_.size() <= 4) is false!

The Vulkan backend stores tensors with more than 4 dims in buffers, and
the partitioner delegates such ops on purpose. 5-D activations already
work. But add_prepack_standard_node, which uploads a constant at load
time, still passed the tensor's sizes, strides and numel as push constants
for buffer storage. Those 4-D push constants assert on a 5-D tensor, and
the nchw_to_buffer shader no longer reads them: it gets everything from
the BufferMetadata UBO, the same way add_staging_to_tensor_node already
works.

Drop the unused push constants for buffer storage.
The TOSA backend skipped every top-level placeholder without users. That is
fine for constants, but a delegate input can also lose all its users during
the Arm pass pipeline. Two examples:

- ConvertFullLikeToFullPass turns zeros_like(y), ones_like(y) and
  full_like(y, v) into full, so y is no longer used.
- RemoveNoopPass folds torch.empty(...).copy_(x) into x, so the empty
  tensor, which stays on the CPU, is no longer used.

The runtime still passes every delegate argument, and the VGF backend binds
them by position. With the unused input missing from the graph, inputs after
it were read from the wrong slots, and outputs were written one slot early,
into the last input tensor. The program ran without an error but returned
wrong values (all zeros in the copy example).

Keep unused placeholders that are delegate inputs. The model converter can
still remove dead inputs, which the VGF runtime already handles.

Ethos-U is not fixed by this change: Vela drops the unused input as well,
and the Ethos-U runtime also binds arguments by position. The new test case
is marked as an expected failure there.
NeutronQuantizer did not work on programs from a non-strict torch.export
(the default of torch.export.export), because two places that read the
source_fn_stack node metadata assumed it is there, and a non-strict export
does not record it:

1. get_aten_node_target_partitions skipped every node without it. The same
   conv, relu and linear model got 10 quantize and dequantize nodes with
   strict=True and 2 with strict=False, so almost nothing was delegated.
   The function only uses the metadata to group partitions, and its caller
   flattens the groups anyway.
2. MoveActivationBeforeConcat, which runs before annotation, raised
   KeyError: 'source_fn_stack' on a conv, cat, relu pattern.

Group by the node's target when the metadata is missing, and copy the
metadata to the new activation nodes only when it is there. The new test
quantizes a non-strict conv, cat, relu model, which goes through both.
XNNPACKQuantizer annotated ops that read or produce bfloat16 tensors, such
as linears in a model that casts its activations to bfloat16. The quantize,
dequantize and choose_qparams nodes it produced then ran on bfloat16
tensors. XNNPACK does not take those patterns, so they stayed on the
portable quantized kernels, which do not take bfloat16 and abort the
process at run time ("Expected input to be Float tensor received: 15").

Skip nodes that read or produce a bfloat16 tensor, both in the filter that
every annotator applies and when annotations are propagated to ops that
share their input's quantization parameters. Such ops stay unquantized and
run correctly. The cat annotator did not apply that filter at all; it now
does, like the other annotators, so cats now also respect module and
filter scoping. Otherwise float32 and float16 models quantize as before.
to_backend could fail with "Expected output node, got op=call_function"
when a runtime check outside a partition reads a value the partition
computes, for example torch._check(t.all().item()) where the tensor ops
are delegated and the item() and the check are not.

After fusing the partition, create_submodule_from_nodes re-sorted the graph
with legalize_graph. That sort places a node as soon as its inputs are
ready, so the output node could be placed before check nodes that have no
users. Later code calls graph.output_node(), which requires the output node
to be last.

Use stable_topological_sort instead. It keeps the original order wherever
the dependencies allow it, so the output node stays last. This changes the
order of the remaining top-level nodes of every program that to_backend
touches: they now keep their original order instead of the order the old
sort produced.
A model that uses a tensor-wide check, for example torch._check_tensor_all
or transformers' torch_compilable_check, could not go through to_edge. Two
things stopped it:

1. torch.export turns the check into aten._is_all_true, which is not a core
   ATen op and has no decomposition, so the ATen verifier rejected it.
2. The check then reads the result with item() and compares it through
   torch.sym_ite. The Edge verifier rejected every Python symbolic function
   (torch.sym_ite, sym_max, sym_min, sym_not, sym_float, sym_int), on
   data-dependent values and on shapes alike, even though ExecuTorch lowers
   each of them to its own executorch_prim op and the runtime runs it.

Like other runtime asserts, the check is removed later by to_executorch
unless a backend takes it into a delegate; the problem was that the
program never got that far.

Decompose aten._is_all_true to aten.all in the default decomposition table,
and let the Edge verifier accept the Python symbolic functions that
ExecuTorch maps to prim ops, the same way it already accepts those prim ops.
A program lowered with the CUDA backend could pass lowering and then kill
the process on its first run:

  undefined symbol: aoti_torch_as_strided

Inductor's C++ wrapper calls aoti_torch_as_strided when it allocates a
buffer with a larger shape than its logical shape and then gives it the
logical shape. ExecuTorch's CUDA shim library did not provide it, so the
compiled model could not be loaded.

Add aoti_torch_as_strided with the PyTorch shim's signature. It returns a
view with the given sizes and strides and the tensor's own storage offset,
built on the existing aoti_torch__reinterpret_tensor.
… other

Lowering a dynamically quantized pi0 model to XNNPACK failed with
"Invalid partition, found dependency cycles".

Before merging two partitions, GroupBasedPartitioner checks that no node
outside them is both downstream and upstream of the merged partition. That
check looks at single nodes. But every other partition also becomes a
single node once it is fused. So a path can leave the merged partition,
enter another partition at one node, leave it from a different node that
the first one does not reach, and come back. The two partitions then
depend on each other, and fusing them fails.

Before accepting a merge, also follow paths through the other partitions:
whenever the merged partition reaches a node of another partition, add
everything that partition reaches as a whole, and reject the merge as soon
as that reaches something upstream of it. Only partitions that will be
returned count, so partitions that are later filtered out (for example
ones without enough compute nodes) do not block a merge. Their compute
node count now comes from one helper shared with the filter.

To stay fast on large graphs, the check only visits the partitions it
actually reaches, and caches what each partition reaches; a merge only
drops the cached entries that reach the merged nodes. On synthetic graphs
of 4k to 16k nodes the partitioner takes about twice as long as before.
The merge loop itself is unchanged and still tries every remaining pair,
so a graph where the check rejects most merges now keeps many partitions
and pays for all their pairs (about 2 s for 1600 partitions on a 4k-node
chain of groups); before, such graphs had few partitions only because
merges that cannot be fused were accepted. If caller-provided node groups
already depend on each other, they are returned as before and fusing them
still reports the cycle; merges whose paths reach such groups are refused.
Because a merge that would need a later merge to become valid is now
rejected, a few graphs end up with a different number of partitions.
The Qualcomm backend has no bfloat16 tensor type (QNN_TENSOR_TYPE_MAP
has no entry for it), but nothing checked data types before a node
visitor was asked to define the node's tensors. An unquantized
bfloat16 node (for example an explicit x.to(torch.bfloat16) region in a
float model) therefore crashed lowering with KeyError: torch.bfloat16
in NodeVisitor.get_data_type instead of staying on CPU.

Add has_unsupported_dtype, true when the node's output or one of its
inputs is an unquantized tensor of a data type the backend cannot
represent, and use it in two places:

- The partitioner skips such a node, so it falls back to CPU like any
  other unsupported node.
- filter_fn does not keep such a node from decomposition when it is
  outside a quantized region. Otherwise an op the backend normally
  lowers whole, such as a bfloat16 linear, would stay on CPU
  undecomposed and fail with "was not decomposed or delegated".

A quantized tensor takes its QNN type from its encoding, so quantized
bfloat16 regions are still delegated as before. A bfloat16 graph input
to a quantized model is still not supported.
MatmulToBmmPass read the operand shapes with ProxyValue.to_tensor(). When
an operand of the matmul is a constant (a parameter, a buffer or a tensor
attribute), ExportPass passes it to call_operator as a plain torch.Tensor,
so CortexMQuantizer crashed in prepare_pt2e for every such matmul:

    AttributeError: 'Tensor' object has no attribute 'to_tensor'

For example softmax(x) @ codebook, where codebook is a tensor attribute.

Read the shapes through .data, which works for both kinds of operand. A
constant operand of a shape the pass rewrites, such as a batched rank-3
codebook used as the right-hand side, now becomes bmm like any other operand
and lowers to cortex_m.quantized_batch_matmul, which already supports a
constant right-hand side. Shapes the pass does not rewrite, such as
rank-3 @ rank-2, no longer crash but stay matmul.
When ExportPass meets a real tensor, for example a get_attr constant, it
makes a fake tensor for the node metadata. Only parameters were faked with
static_shapes=True. Any other tensor that the fake mode had not seen before
got a fresh symbol for every dimension when the fake mode has a ShapeEnv,
which is the case for exported programs.

Such tensors are common: constant folding (for example the Arm quantizer's
ConstantFoldingPass) creates new constant tensors, and the next ExportPass
gives them symbolic sizes like (s45, s96) instead of (9, 2). Later passes
then see symbolic shapes where everything is static. With the LeRobot
diffusion policy, the Arm quantizer then failed in DecomposeVarPass:

    AssertionError: a.node.constant should not be None

A real tensor is constant data with known sizes, so fake all of them with
static shapes, as the existing TODO suggested.
aten.slice_scatter is functional: it returns a new tensor and leaves its
base alone. The MLX builder lowers it to a SliceUpdateNode whose output
slot differs from its base slot, and the runtime handles that case without
writing the base.

Still, exec_slice_update always fetched the base through tensor_ref(),
which is the accessor for tensors an op writes, and which refuses
constants. So a slice_scatter whose base is a buffer or another constant
failed at run time with "tensor_ref: cannot mutate constant tensor".

Read the base with const_tensor_ref(), and use tensor_ref() only for the
in-place form, which is the one that writes it.
@shoumikhin
shoumikhin force-pushed the fix-lerobot-export-bugs branch from d11bd4b to 3119872 Compare October 29, 2026 00:11
Copilot AI balanced review requested due to automatic review settings October 29, 2026 00:11

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

@shoumikhin

Copy link
Copy Markdown
Contributor Author

Please hold off on merging this one for now. I'm adding a last batch of fixes over the next few hours and will comment here when it is final and ready to merge.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/trunk CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. module: arm Issues related to arm backend release notes: backends

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants