MPS Backend issue - torch._export.verifier.SpecViolationError: Tensor should not be used in dim_order mode

[nil-is-all]

🐛 Describe the bug

I am trying to export MobileNetV3_small model in all 3 backends - CoreML, XNNPack, MPS in Fp32. The export.py script is inspired from this available script for DL3 model for Xnnpack backend.

When I remove the MPS parts, the export script runs without errors.

My export.py script:

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torchvision.models as models
from executorch.backends.apple.coreml.partition import CoreMLPartitioner
from executorch.backends.apple.mps.partition import MPSPartitioner
from executorch.exir.backend.backend_details import CompileSpec #added to handle MPS backend error missing CompileSpec
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
from executorch.exir import to_edge_transform_and_lower
import ssl
import certifi


def main() -> None:
    # change model to MV3
    # model = models.segmentation.deeplabv3_resnet101(weights="DEFAULT").eval()
    model = models.mobilenet_v3_small(weights="DEFAULT").eval()
    sample_inputs = (torch.randn(1, 3, 224, 224),)

    et_program_coreml = to_edge_transform_and_lower(
        torch.export.export(model, sample_inputs),
        partitioner=[CoreMLPartitioner()],
    ).to_executorch()

    et_program_mps = to_edge_transform_and_lower(
        torch.export.export(model, sample_inputs),
        partitioner=[MPSPartitioner([CompileSpec("use_fp32", bytes([True]))])], #added to handle MPS backend error missing CompileSpec
    ).to_executorch()

    et_program_xnnpack = to_edge_transform_and_lower(
        torch.export.export(model, sample_inputs),
        partitioner=[XnnpackPartitioner()],
    ).to_executorch()

    with open("mv3_coreml_fp32.pte", "wb") as file:
        et_program_coreml.write_to_file(file)
    with open("mv3_mps_fp32.pte", "wb") as file:
        et_program_mps.write_to_file(file)
    with open("mv3_xnnpack_fp32.pte", "wb") as file:
        et_program_xnnpack.write_to_file(file)


if __name__ == "__main__":
    main()

The error stack trace suggests the issue happens when performing _partition_and_lower()

# To see more debug info, please use `graph_module.print_readable()` Traceback (most recent call last): File "/Users/nikhilviswanath/executorch/export.py", line 49, in <module> main() File "/Users/nikhilviswanath/executorch/export.py", line 30, in main et_program_mps = to_edge_transform_and_lower( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/executorch/exir/program/_program.py", line 107, in wrapper return func(self, *args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/executorch/exir/program/_program.py", line 1209, in to_edge_transform_and_lower edge_manager = edge_manager.to_backend({name: curr_partitioner}) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/executorch/exir/program/_program.py", line 107, in wrapper return func(self, *args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/executorch/exir/program/_program.py", line 1471, in to_backend new_edge_programs[name] = to_backend(program, partitioner[name]) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/functools.py", line 912, in wrapper return dispatch(args[0].__class__)(*args, **kw) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/executorch/exir/backend/backend_api.py", line 398, in _ tagged_graph_module = _partition_and_lower( ^^^^^^^^^^^^^^^^^^^^^ File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/executorch/exir/backend/backend_api.py", line 321, in _partition_and_lower partitioned_module = _partition_and_lower_one_graph_module( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/executorch/exir/backend/backend_api.py", line 251, in _partition_and_lower_one_graph_module lowered_submodule = to_backend( ^^^^^^^^^^^ File "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/functools.py", line 912, in wrapper return dispatch(args[0].__class__)(*args, **kw) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/executorch/exir/backend/backend_api.py", line 114, in _ preprocess_result: PreprocessResult = cls.preprocess( ^^^^^^^^^^^^^^^ File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/executorch/backends/apple/mps/mps_preprocess.py", line 90, in preprocess edge_program = _transform(edge_program, DimOrderOpsRevertPass()) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/executorch/exir/program/_program.py", line 224, in _transform return _update_exported_program_graph_module(self, transformed_gm) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/executorch/exir/program/_program.py", line 232, in _update_exported_program_graph_module transformed_ep = ExportedProgram( ^^^^^^^^^^^^^^^^ File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/torch/export/exported_program.py", line 970, in __init__ self.validate() File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/torch/export/exported_program.py", line 1526, in validate self._validate() File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/torch/export/exported_program.py", line 1535, in _validate v().check(self) File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/torch/_export/verifier.py", line 164, in check self._check_graph_module(ep.graph_module) File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/torch/_export/verifier.py", line 242, in _check_graph_module _check_valid_op(node.target) File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/torch/_export/verifier.py", line 223, in _check_valid_op self.check_valid_op(op) File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/executorch/exir/verification/verifier.py", line 272, in check_valid_edge_op _check_valid_dim_order_ops(op, self.use_dim_order) File "/Users/nikhilviswanath/executorch/.venv/lib/python3.12/site-packages/executorch/exir/verification/verifier.py", line 52, in _check_valid_dim_order_ops raise SpecViolationError(f"{op} should not be used in dim_order mode") torch._export.verifier.SpecViolationError: <EdgeOpOverload: aten._to_copy.default>: schema = aten::_to_copy(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, bool non_blocking=False, MemoryFormat? memory_format=None) -> Tensor should not be used in dim_order mode

Versions

N/A

cc @mergennachin @iseeyuan @lucylq @helunwencser @tarun292 @kimishpatel @jackzhxng

[digantdesai]

WIP #10276

[digantdesai]

Merged, let me know if this issue still persists.