diff --git a/backends/apple/coreml/compiler/coreml_preprocess.py b/backends/apple/coreml/compiler/coreml_preprocess.py index ebf9efc3e12..1b838d9a9b7 100644 --- a/backends/apple/coreml/compiler/coreml_preprocess.py +++ b/backends/apple/coreml/compiler/coreml_preprocess.py @@ -579,14 +579,52 @@ 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 = CoreMLBackend.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) return mlmodel + @staticmethod + 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 + @staticmethod def preprocess_model( mlmodel: ct.models.MLModel, model_type: MODEL_TYPE diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py index 0e75d6024e4..78dd114eeda 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/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()