examples/llama: lower to Core ML with to_edge_transform_and_lower

[msluszniak]

Summary

Fixes #19634.

The Core ML branch of _to_edge_and_lower_llama() still used the deprecated export_to_edge() + to_backend() split, and CoreMLPartitioner logs a deprecation warning about it on every invocation.

The split matters beyond the warning. It decomposes the graph before the partitioner runs, so the ops Core ML has its own implementations for are already broken into primitives by the time it sees them. CoreMLPartitioner.ops_to_not_decompose() asks to keep every op Core ML supports, and only to_edge_transform_and_lower honours that request.

This adds _to_edge_and_lower_llama_coreml(), matching the existing xnnpack and mlx helpers, and routes to it when Core ML is the only backend enabled.

Scope

Enabling Core ML together with Vulkan, MPS or QNN still takes the old combined path, which builds a single partitioner list. QNN in particular needs the edge manager for its pass pipeline and model sharding, so it is left alone; migrating it is a separate piece of work.

Etrecord generation comes along for free: to_edge_transform_and_lower takes generate_etrecord directly, so the helper does not need the deepcopy of the edge manager the old path used.

Results

LFM2.5 350M, fp32, max_seq_length=512, ios: 18, compute_units: cpu_and_ne:

	subgraphs	delegated nodes	non-delegated	deprecation warnings
before	1	1156	13	1
after	1	893	13	0

Same partitioning, 263 fewer nodes for Core ML to reassemble. expand_copy drops from 133 to 49 and unsqueeze_copy from 375 to 99.

A note on the original report

The issue also reported a SpecViolationError ("Mutation node aten_index_put_default_N is neither a buffer nor a user input") from the split path on LFM2 hybrid models, where the short-conv conv_state.copy_() decomposes to slice_copy + index_put and only one of the two is recorded as the mutation source. That no longer reproduces on main: the same export now completes on both the old and the new path. Only the deprecation and the extra decomposition remain, which is what this PR addresses.

Test plan

Adds examples/models/lfm2/config/lfm2_coreml_fp32.yaml alongside the existing lfm2_xnnpack_fp32.yaml and lfm2_mlx_4w.yaml, so the Core ML path has a config to run:

python -m extension.llm.export.export_llm \
  --config examples/models/lfm2/config/lfm2_coreml_fp32.yaml \
  +base.model_class=lfm2_5_350m \
  +base.params=examples/models/lfm2/config/lfm2_5_350m_config.json \
  +export.max_seq_length=512 \
  +export.max_context_length=512 \
  +export.output_name=lfm2_coreml.pte

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 20 Awaiting Approval

As of commit de8e0bc with merge base fbd4bbf ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Aarch64 Linux CUDA Wheels (gh)
• Build Aarch64 Linux Wheels (gh)
• Build documentation (gh)
• Build Linux CUDA Wheels (gh)
• Build Linux Wheels (gh)
• Build macOS Wheels (gh)
• Build Presets (gh)
• Build Windows Wheels (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

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.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 5m 7s —— View job

---

Code Review

I read the diff against the surrounding code in export_llama_lib.py, extension/llm/export/builder.py, and exir/program/_program.py. Overall this is a clean, well-scoped change that follows the existing _to_edge_and_lower_llama_{xnnpack,mlx,arm} pattern closely, and the dispatch guard is correct and clearly documented. One substantive concern about ETRecord, plus a couple of minor notes.

1. ETRecord does not actually land on disk in this path (behavior regression vs. the old CoreML path)

The PR body says "Etrecord generation comes along for free," but I don't think that's accurate. The new helper does:

builder_exported.generate_etrecord = generate_etrecord
builder = builder_exported.pt2e_quantize(quantizers).to_edge_transform_and_lower(partitioners)
...
return builder.to_executorch(passes=additional_passes)

That flag flows into to_edge_transform_and_lower (builder.py:525) and, per _program.py:1233-1236 / 1803-1808, the ETRecord is attached to export_program._etrecord. But nothing in the llama flow ever calls .save(...) / get_etrecord().save(...) on it. save_to_pte → save_pte_program only writes the .pte/.ptd (extension/export_util/utils.py:137-143).

Contrast with the old path (export_llama_lib.py:1533-1560), which explicitly writes the file:

generate_etrecord_func(et_record="etrecord.bin", edge_dialect_program=..., executorch_program=builder.export_program)

So a user running --generate_etrecord with CoreML previously got etrecord.bin on disk; with this change they get an ETRecord attached to an in-memory program object that is then discarded. Note the XNNPACK helper has the same gap and even flags it (export_llama_lib.py:1190: # TODO: Enable generating ETRecord with XNNPack and to_edge_transform_and_lower()), which suggests the modern path's ETRecord persistence isn't wired up for this entry point yet.

This is consistent with XNNPACK, so it's defensible to ship — but I'd suggest either (a) toning down the PR description so it doesn't claim ETRecord works, or (b) actually persisting it (e.g. builder.export_program.get_etrecord().save("etrecord.bin") guarded on generate_etrecord) so CoreML users don't silently lose the file they used to get. Fix this →

2. Minor: modelname parameter is unused

The helper accepts modelname but never uses it. The old CoreML path did modelname = f"coreml_{modelname}". The mlx helper also takes it unused, so this matches precedent and the final output name comes from config — but it's dead here. Fine to leave for consistency; just noting it.

3. Minor: test config sets enable_state / preserve_sdpa, which appear inert on this path

lfm2_coreml_fp32.yaml sets enable_state: True and preserve_sdpa: True. Grepping the source, these CoreMLConfig fields are only populated from CLI args in llm_config.py and don't appear to be read anywhere in the export_llm lowering flow — the SDPA/KV-cache source transforms are gated on coreml_ios >= 18 (export_llama_lib.py:2225-2231), not on these flags. So they may be no-ops here (pre-existing, not introduced by this PR). Worth double-checking they actually affect the exported model on this path; if not, they're just noise in the example config.

What looks good

• Dispatch guard coreml.enabled and not (vulkan or mps or qnn) correctly preserves the combined path for the multi-backend / QNN-sharding cases, and the rationale is documented.
• Dropping the deepcopy of the edge manager is a real memory win versus the old etrecord branch.
• Docstring clearly explains why to_edge_transform_and_lower matters (ops_to_not_decompose), which is the substance of the fix.
• The reported node-count reduction (1156 → 893, no partitioning change) is a plausible and meaningful result of honoring ops_to_not_decompose.

The core change is correct and worth landing; the ETRecord point is the one I'd resolve (or at least reword) before merge.
· branch ms/llama-coreml-transform-and-lower