Core ML: stop the gather opt-out from failing tied embeddings

[john-rocky]

Fixes#21856.

What happens today

op_linear_quantizer_config is applied with op_type_configs={"gather": None}, so that
embedding tables are not compressed by the linear quantizer. When a model ties its embedding
table to its output projection — one weight, two consumers — that constant reaches the
compressor configured two ways at once: quantized as a linear weight, exempt as a gather
table. coremltools refuses the conflict rather than resolving it:

ValueError: compression config conflict detected between ops
%aten_embedding_default_...: gather(x=%p_out_weight_to_fp16, ...)
and
%aten_mm_default...: linear(x=..., weight=%p_out_weight_to_fp16, ...)
... has config None while ... has OpLinearQuantizerConfig(mode='LINEAR_SYMMETRIC', ...)

So the opt-out does not skip the table. It fails the whole lowering, and a very common LLM
shape cannot be lowered at all with coreml_quantize set.

The change

Gathers whose constant is also consumed by an op of another type are named in
op_name_configs with the same config as the op they share it with, so the constant is seen
one way. Every other gather keeps the opt-out, and a model whose MIL program is not available
returns no names, which leaves the behaviour exactly as it was.

Checking it

The new test in test_coreml_partitioner.py lowers a two-layer model twice, tied and untied.
Against main the tied case raises the error above; with this change both lower. The untied
case passes either way, which is the point — the opt-out still applies where it always did.

Measured on macOS arm64, coremltools 9.0.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 1 Cancelled Job, 2 Unclassified Failures

As of commit 02f2e1e with merge base 7c3cdbb ():

NEW FAILURE - The following job has failed:

• Apple / build-frameworks-ios / macos-job (gh)
  ##[error]fatal: unable to access 'https://github.com/mreineck/pocketfft/': Failed to connect to github.com port 443 after 12696 ms: Couldn't connect to server

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

• Test ARM Backend / test-arm-vgf / test-backend-linux (arm_vgf_fp, models) / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  ##[error]The operation was canceled.
• Test ARM Backend / test-arm-vgf / test-backend-linux (arm_vgf_fp, operators) / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  ##[error]The operation was canceled.

CANCELLED JOB - The following job was cancelled. Please retry:

• pull / test-arm-backend-public-api-backward-compatibility / linux-job (gh)
  ##[error]The operation was canceled.

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 1m 53s —— View job

---

Code review

• Read coreml_preprocess.py changes
• Read test_coreml_partitioner.py changes
• Analyze correctness & edge cases
• Post review

The change is small, well-scoped, and the reasoning in the comments is excellent — it explains why coremltools rejects the tied constant rather than just what the code does. The graceful degradation (empty list when _mil_program is unavailable → prior behavior) is the right instinct. A few observations, none blocking.

---

Correctness

The detection logic is sound.gathers_sharing_a_weight builds const→consumers and flags a gather only when gathers and len(gathers) < len(ops) — i.e. a non-gather also consumes the same constant. Two gathers sharing one table (len(gathers) == len(ops)) correctly stay opted out, and a gather with no co-consumer stays opted out. That matches the stated intent.

Behavioral trade-off worth calling out (not a bug). The resolution quantizes the tied table (names the gather with op_linear_quantizer_config) rather than leaving the whole constant uncompressed. The original {"gather": None} opt-out exists to keep embedding tables at full precision. For a tied weight that intent now flips: the table gets compressed to int4 along with the projection. This is almost certainly the right call for LLMs (the alternative is a hard failure, and the shared projection benefits from quantization), but it's a silent accuracy change for anyone who tied weights specifically expecting the table to stay uncompressed. Worth a one-line note in the PR description or a logger.info naming the tied gathers.

Minor points

• op_name_configs assumes the co-consumer carries global_config. At line 588–592, tied gathers are assigned op_linear_quantizer_config (the global config). That is correct only because the current OptimizationConfig overrides nothing except gather. If a future change adds another op_type_configs entry (or a per-op override) for the type that shares the table, the gather would be matched to the global config while its co-consumer uses a different one — reintroducing the conflict. A short comment tying the config choice to "same config as the sharing op, given gather is the only type override" would harden this against future edits.

• Nested blocks aren't traversed.for op in function.operations walks the top-level ops of each function; ops inside cond/while sub-blocks aren't visited. Not reachable for the embedding+linear shape here, so fine — just flagging it if this helper ever gets reused more broadly.

• Double-counting is harmless but subtle.consumers[...].append(op) runs per matching input value, so an op referencing the same const in two inputs is counted twice. It can't produce a false positive here (gather's constant is only its x input) and duplicate names collapse in the {name: config ...} dict comprehension, but the len(gathers) < len(ops) comparison is technically counting input-edges, not distinct ops.

Test

• to_executorch() exercises the CoreML preprocess (convert + linear_quantize_weights) on any platform, so this genuinely covers the fix on Linux CI, not just macOS — good. Looping tied/untied and asserting both reach a program is a clean regression guard.
• The name is ..._is_quantized_with_its_linear but the assertion only checks that lowering succeeds, not that the table was actually compressed. Consider asserting on the resulting weight dtype (or at least renaming to ..._lowers) so the test name matches what it verifies.
• ct.precision(ct.precision.FLOAT16.value) (line ~460) is a redundant round-trip; ct.precision.FLOAT16 is equivalent.

---

Overall this is a clean, minimal fix for a real and common failure. The two things I'd most want addressed are surfacing the "tied tables now get quantized" behavior change and tightening the test to match its name; everything else is optional polish.
· coreml-tied-embedding-quantize

[metascroy]

Does this need to be a public API on CoreMLBackend?

Or just a private helper in preproces?

[john-rocky]

You're right — it has exactly one caller, _convert_to_mlmodel in the same file, so nothing needs it public. Renamed to _gathers_sharing_a_weight, matching the class's underscore convention.

[metascroy]

Can you make it a private free function in the file, not a static function on the class?

[john-rocky]

Done in 02f2e1e — it's now a module-level private function. The scan never read anything from the class, so the body is unchanged; _convert_to_mlmodel calls it directly.

[john-rocky]

On the attempt-2 CI failures — none of them reaches this PR's code. android / build-android failed downloading the golden-artifacts snapshot: golden_artifacts_26052718.zip has aged out of S3, so the old curl -sL saved the 404 page and unzip failed; main already pins a fresh snapshot and uses curl -f. The three ARM jobs spent their full 120 minutes waiting for a docker image (executorch-ubuntu-24.04-arm-sdk-1d92118e…) that fork PRs cannot build or push, and were cancelled at the timeout. All four jobs are green on current main, and the free-function push re-runs everything against a fresh merge with main.

[metascroy]

Looks good, thanks for the contribution @john-rocky!