Core ML: stop the gather opt-out from failing tied embeddings (pytorch#21969)

Fixes pytorch#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.

Author: john-rocky

backends/apple/coreml/compiler/coreml_preprocess.py

@@ -36,6 +36,37 @@
 logger.setLevel(get_coreml_log_level(default_level=logging.WARNING))
 
 
+def _gathers_sharing_a_weight(mlmodel: ct.models.MLModel) -> List[str]:
+    """
+    Names of gather ops whose table is also consumed by an op of another type.
+
+    Such a constant reaches the compressor under two different configurations at once,
+    which it rejects rather than resolves. Naming these gathers lets them be configured
+    the same way as the op they share the constant with; every other gather keeps the
+    opt-out. Returns nothing for a model whose program is not available, which leaves
+    the behaviour as it was.
+    """
+    program = getattr(mlmodel, "_mil_program", None)
+    if program is None:
+        return []
+
+    consumers: Dict[str, List[Any]] = {}
+    for function in program.functions.values():
+        for op in function.operations:
+            for value in op.inputs.values():
+                for var in value if isinstance(value, (list, tuple)) else [value]:
+                    source = getattr(var, "op", None)
+                    if source is not None and source.op_type == "const":
+                        consumers.setdefault(var.name, []).append(op)
+
+    shared: List[str] = []
+    for ops in consumers.values():
+        gathers = [op for op in ops if op.op_type == "gather"]
+        if gathers and len(gathers) < len(ops):
+            shared += [op.name for op in gathers]
+    return shared
+
+
 class COMPILE_SPEC_KEYS(Enum):
     COMPUTE_UNITS = "compute_units"
     MODEL_TYPE = "model_type"
@@ -579,9 +610,16 @@ def _convert_to_mlmodel(
             logger.warning(
                 "Core ML Backend op_linear_quantizer_config API is experimental"
             )
+            # Gathers are left alone so that embedding tables are not quantized by this
+            # config, except where the table is also some other op's weight. A tied
+            # embedding is one constant feeding both a gather and a linear, and coremltools
+            # refuses to compress a constant its consumers disagree about, so opting the
+            # gather out there does not skip the table, it fails the whole lowering.
+            tied = _gathers_sharing_a_weight(mlmodel)
             config = cto.coreml.OptimizationConfig(
                 global_config=op_linear_quantizer_config,
                 op_type_configs={"gather": None},
+                op_name_configs={name: op_linear_quantizer_config for name in tied},
             )
             mlmodel = cto.coreml.linear_quantize_weights(mlmodel, config=config)
 

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -433,6 +433,46 @@ def forward(self, x):
         self.assertIn("executorch_call_delegate", op_names)
         self.assertNotIn("aten.argmax.default", op_names)
 
+    def test_tied_embedding_is_quantized_with_its_linear(self):
+        """A weight that is both an embedding table and an output projection lowers.
+
+        The quantizer config opts gathers out so that embedding tables are not compressed
+        by it. When the table is tied to a linear's weight, that opt-out leaves one
+        constant with two configurations, which coremltools rejects — so the opt-out has
+        to stop at the tie rather than fail the whole lowering.
+        """
+
+        class Tied(torch.nn.Module):
+            def __init__(self, vocab=512, dim=64, tie=True):
+                super().__init__()
+                self.embedding = torch.nn.Embedding(vocab, dim)
+                self.output = torch.nn.Linear(dim, vocab, bias=False)
+                if tie:
+                    self.output.weight = self.embedding.weight
+
+            def forward(self, ids):
+                return self.output(self.embedding(ids))
+
+        ids = torch.zeros(1, 4, dtype=torch.long)
+        for tie in (True, False):
+            compile_specs = CoreMLBackend.generate_compile_specs(
+                minimum_deployment_target=ct.target.iOS18,
+                compute_precision=ct.precision(ct.precision.FLOAT16.value),
+                op_linear_quantizer_config={
+                    "mode": "linear_symmetric",
+                    "dtype": "int4",
+                    "granularity": "per_channel",
+                },
+            )
+            exported = torch.export.export(Tied(tie=tie).eval(), (ids,), strict=True)
+            delegated = executorch.exir.to_edge_transform_and_lower(
+                exported,
+                partitioner=[CoreMLPartitioner(compile_specs=compile_specs)],
+            )
+            # Reaching a program at all is the assertion: the failure this covers is
+            # raised during lowering, not carried in the result.
+            self.assertIsNotNone(delegated.to_executorch())
+
     def test_deprecation_warning_for_to_backend_workflow(self):
         """
         Test that the deprecated to_edge + to_backend workflow shows a deprecation warning.
@@ -532,5 +572,6 @@ def forward(self, x):
     test_runner.test_take_over_constant_data_false()
     test_runner.test_aten_rand_default_falls_back_to_portable()
     test_runner.test_aten_randn_is_still_delegated()
+    test_runner.test_tied_embedding_is_quantized_with_its_linear()
     test_runner.test_deprecation_warning_for_to_backend_workflow()
     test_runner.test_no_warning_for_to_edge_transform_and_lower_workflow()