diff --git a/coremltools/_deps/__init__.py b/coremltools/_deps/__init__.py index 77d2d0e3a..63859f7a1 100644 --- a/coremltools/_deps/__init__.py +++ b/coremltools/_deps/__init__.py @@ -163,7 +163,7 @@ def __get_sklearn_version(version): # --------------------------------------------------------------------------------------- _HAS_TORCH = True -_TORCH_MAX_VERSION = "2.8.0" +_TORCH_MAX_VERSION = "2.9.0" _HAS_TORCH_EXPORT_API = False _CT_OPTIMIZE_TORCH_MIN_VERSION = "2.1.0" _IMPORT_CT_OPTIMIZE_TORCH = False diff --git a/coremltools/converters/_converters_entry.py b/coremltools/converters/_converters_entry.py index ff9a26db2..ce646ff31 100644 --- a/coremltools/converters/_converters_entry.py +++ b/coremltools/converters/_converters_entry.py @@ -591,6 +591,18 @@ def forward(self, x): outputs_as_strings, outputs_as_tensor_or_image_types, outputs) + # Starting with torch 2.9, torch.export.export() returns an ExportedProgram + # in the new "TRAINING" IR dialect by default, whereas conversion expects + # the "ATEN" dialect (the prior default). Lower it to ATEN via + # run_decompositions -- the migration torch recommends -- so existing + # convert() calls keep working without a manual run_decompositions() step. + if ( + exact_source == "pytorch" + and _HAS_TORCH_EXPORT_API + and isinstance(model, ExportedProgram) + and model.dialect not in ("ATEN", "EDGE") + ): + model = model.run_decompositions({}) source_dialect = _determine_source_dialect(model, exact_source) exact_target = _determine_target(convert_to, minimum_deployment_target) _validate_conversion_arguments( diff --git a/coremltools/converters/mil/frontend/torch/ops.py b/coremltools/converters/mil/frontend/torch/ops.py index d2470423c..e474c6e0e 100644 --- a/coremltools/converters/mil/frontend/torch/ops.py +++ b/coremltools/converters/mil/frontend/torch/ops.py @@ -3131,12 +3131,18 @@ def _parse_positional_args(context, node) -> Tuple[Var]: @register_torch_op def rms_norm(context, node): - # Parse Inputs - inputs = _get_inputs(context, node, expected=4) + # Parse Inputs. torch.export / ExecuTorch omit the optional weight/eps when + # they use defaults, so only window-style strict counts apply to TorchScript. + inputs = _get_inputs( + context, + node, + expected={TorchFrontend.TORCHSCRIPT: 4}, + min_expected={TorchFrontend.TORCHEXPORT: 2, TorchFrontend.EXECUTORCH: 2}, + ) x = inputs[0] normalized_shape = inputs[1] - weight = inputs[2] - eps = inputs[3] + weight = inputs[2] if len(inputs) > 2 else None + eps = inputs[3] if len(inputs) > 3 else None axes = list(range(-len(normalized_shape.val), 0)) # Store epsilon value to ensure ZeroDivisionError doesn't occur # while computing RMSNorm @@ -8277,14 +8283,28 @@ def remainder(context, node): context.add(remainder_node) -@register_torch_op +@register_torch_op(torch_alias=["hann_window.periodic"]) def hann_window(context, node): - inputs = _get_inputs(context, node, expected=[5, 6]) + inputs = _get_inputs( + context, + node, + expected={TorchFrontend.TORCHSCRIPT: [5, 6]}, + min_expected={TorchFrontend.TORCHEXPORT: 1, TorchFrontend.EXECUTORCH: 1}, + ) if inputs[0].val is None: raise NotImplementedError("variable 'window_length' not supported.") + # TorchScript passes the dtype/layout/device/pin_memory kwargs as graph inputs + # (5 for the default overload, 6 for `.periodic`); torch.export / ExecuTorch + # pass only window_length (and periodic for the `.periodic` overload). So the + # `periodic` flag is the 2nd positional input for the 6-arg TorchScript form or + # any >=2-arg export form. periodic = True - if len(inputs) == 6: + if context.frontend == TorchFrontend.TORCHSCRIPT: + has_periodic = len(inputs) == 6 + else: + has_periodic = len(inputs) >= 2 + if has_periodic: if inputs[1].val is None: raise NotImplementedError("variable 'periodic' not supported.") if not inputs[1].val: diff --git a/coremltools/converters/mil/frontend/torch/test/test_torch_conversion_api.py b/coremltools/converters/mil/frontend/torch/test/test_torch_conversion_api.py index 9b11dc367..6fd782f49 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_torch_conversion_api.py +++ b/coremltools/converters/mil/frontend/torch/test/test_torch_conversion_api.py @@ -22,6 +22,7 @@ from coremltools._deps import ( _HAS_HF, _HAS_TORCH, + _HAS_TORCH_EXPORT_API, _HAS_TORCHAO, MSG_TORCH_NOT_FOUND, MSG_TORCHAO_NOT_FOUND, @@ -571,6 +572,41 @@ def forward(self, x): ) assert isinstance(model._get_mil_internal(), ct.converters.mil.Program) + @staticmethod + @pytest.mark.skipif(not _HAS_TORCH_EXPORT_API, reason="torch.export API not available.") + def test_convert_exported_program_training_dialect(): + # Starting with torch 2.9, torch.export.export() returns an + # ExportedProgram in the new "TRAINING" IR dialect by default. convert() + # should lower it automatically via run_decompositions instead of + # requiring the user to call run_decompositions() themselves. + class Network(torch.nn.Module): + def __init__(self): + super().__init__() + self.hidden = torch.nn.Linear(8, 4) + self.relu = torch.nn.ReLU() + + def forward(self, x): + return self.relu(self.hidden(x)) + + torch_model = Network().eval() + example_input = torch.rand(2, 8) + exported_program = torch.export.export(torch_model, (example_input,)) + + # With torch >= 2.9 the default export dialect is no longer ATEN, so the + # installed torch here actually drives convert()'s auto-lowering path. + # Assert it to guarantee this test exercises that path instead of + # silently passing as a no-op on a torch whose default is still ATEN. + if Version(torch.__version__) >= Version("2.9.0"): + assert exported_program.dialect not in ("ATEN", "EDGE") + + # Convert directly, without a manual run_decompositions() step. + mlmodel = ct.convert( + exported_program, + inputs=[ct.TensorType(name="x", shape=example_input.shape)], + convert_to="mlprogram", + ) + assert isinstance(mlmodel._get_mil_internal(), ct.converters.mil.Program) + @staticmethod @pytest.mark.skipif(ct.utils._macos_version() < (12, 0), reason='Model produces specification 6.') @pytest.mark.parametrize( diff --git a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py index 47a6944c8..6a32cb67e 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py +++ b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py @@ -12566,15 +12566,16 @@ def test_logcumsumexp(self, compute_unit, backend, frontend, axis): class TestHannWindow(TorchBaseTest): @pytest.mark.parametrize( - "compute_unit, backend, window_length, periodic", + "compute_unit, backend, frontend, window_length, periodic", itertools.product( compute_units, backends, + frontends, [1, 3, 6, 10, 12], [True, False], ), ) - def test_hann_window(self, compute_unit, backend, window_length, periodic): + def test_hann_window(self, compute_unit, backend, frontend, window_length, periodic): class HannWindowModel(nn.Module): def forward(self, x): return torch.hann_window(window_length, periodic) @@ -12588,6 +12589,29 @@ def forward(self, x): model, expected_results=torch_out, input_as_shape=False, + frontend=frontend, + backend=backend, + compute_unit=compute_unit, + ) + + +class TestRMSNorm(TorchBaseTest): + @pytest.mark.parametrize( + "compute_unit, backend, frontend, has_weight", + itertools.product(compute_units, backends, frontends, [True, False]), + ) + def test_rms_norm(self, compute_unit, backend, frontend, has_weight): + normalized_shape = (8,) + weight = torch.rand(normalized_shape) if has_weight else None + + class RMSNormModel(nn.Module): + def forward(self, x): + return nn.functional.rms_norm(x, normalized_shape, weight=weight, eps=1e-5) + + self.run_compare_torch( + (2, 4, 8), + RMSNormModel().eval(), + frontend=frontend, backend=backend, compute_unit=compute_unit, ) @@ -13768,6 +13792,23 @@ def test_unfold( if frontend == TorchFrontend.EXECUTORCH: pytest.skip("ExecuTorch produces rank > 5 tensor") + if ( + is_dynamic_hw + and frontend in TORCH_EXPORT_BASED_FRONTENDS + and Version(torch.__version__) >= Version("2.9.0") + ): + # torch 2.9's ExportedProgram.run_decompositions({}) raises + # "NameError: name 'L' is not defined" while interpreting the + # `_guards_fn` submodule that torch 2.9 generates for dynamic-shape + # exports carrying shape guards (here the H/W >= f(kernel, dilation, + # padding, stride) constraint). This is an upstream torch issue, not + # a converter bug; static-shape unfold and all other export ops are + # unaffected. Re-enable once the torch regression is resolved. + pytest.skip( + "torch 2.9 run_decompositions() NameError on _guards_fn " + "for guarded dynamic-shape exports" + ) + if isinstance(kernel_size, int): kernel_size = (kernel_size, kernel_size) if isinstance(dilation, int): diff --git a/reqs/pytorch.pip b/reqs/pytorch.pip index f35449f72..f0712ffc8 100644 --- a/reqs/pytorch.pip +++ b/reqs/pytorch.pip @@ -5,7 +5,7 @@ torchaudio==2.2.0; platform_machine != "arm64" torchvision==0.17.0; platform_machine != "arm64" # Torch dependencies for ARM -torch==2.8.0; platform_machine == "arm64" +torch==2.9.0; platform_machine == "arm64" torchaudio>=2.2.0; platform_machine == "arm64" torchvision>=0.17.0; platform_machine == "arm64" torchsr==1.0.4; platform_machine == "arm64" @@ -13,4 +13,4 @@ torchsr==1.0.4; platform_machine == "arm64" # TODO (rdar://141476729) support a more recent timm timm==0.6.13; platform_machine == "arm64" -torchao==0.12.0; platform_machine == "arm64" and python_version >= '3.10' +torchao==0.14.0; platform_machine == "arm64" and python_version >= '3.10' diff --git a/reqs/test_executorch.pip b/reqs/test_executorch.pip index 7a154e09c..b90839205 100644 --- a/reqs/test_executorch.pip +++ b/reqs/test_executorch.pip @@ -3,5 +3,5 @@ # Warning: Starting from ExecuTorch 0.6.0, coremltools is added as a dependency # so we need to re-install built-from-source coremltools after pip install ExecuTorch -executorch>=0.7.0; platform_machine == "arm64" and python_version >= '3.10' and python_version < '3.13' +executorch>=1.0.0,<1.1.0; platform_machine == "arm64" and python_version >= '3.10' and python_version < '3.13' setuptools<82