[RFC][WIP] Introduce PT2 dynamo onnx exporter

[BowenBao]

trafficstars

Hello from PyTorch ONNX exporter team. Since the early announcement of PT2 we have been actively developing the next generation of PyTorch ONNX exporter based on Dynamo, the new Python level JIT compiler.

With the upcoming PyTorch 2.3 release (targeting 4/24/24), we hope to introduce the new exporter as an experimental feature for early beta adoption in Optimum.

This is an initial PR to introduce the code path of calling the new exporter, which is now hidden behind a flag that is by default turned off. The goal is to kick off the discussion and start gathering feedback from the community.

Before submitting

• [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• [ ] Did you make sure to update the documentation with your changes?
• [ ] Did you write any new necessary tests?

[fxmarty]

Hi @BowenBao, thank you for the PR! Yes I think moving to dynamo-based export is a good thing.

One thing that is currently missing in Transformers, Diffusers, TIMM CI is testing that models can be traced with fullgraph=True. I expect the dynamo-based export to fail for some architectures, but it is good to know & we could test that already in optimum.

@giuseppe5 was actually interested in something broader than this, namely generalizing the logic in exporters/onnx to be able to export as well e.g. to ScriptModule, or ExportedProgram. But this is an other discussion.

[thiagocrepaldi]

@fxmarty what is the minimum requirement for us to have an initial version of this PR merged? The dynamo-based ONNX exporter is becoming our default ONNX exporter and it would be great to integrate and perfect it on Optimum framework

[fxmarty]

@thiagocrepaldi In the current state, I think simply adding (slow) tests for dynamo, duplicating https://github.com/huggingface/optimum/blob/7e08a820b65a359a61444abe51df4eb96b26b2e3/tests/exporters/onnx/test_exporters_onnx_cli.py#L384 with a @slow decorator, and using dynamo=True, and running RUN_SLOW=1 pytest tests/exporters/onnx -k "test_exporters_cli_pytorch_cpu" -s -vvvvv and reporting the failing tests / skipping them with a reason why would be enough.

For dynamo export to be the default, having the ability to specify the input_names and output_names would be very helpful. Is there a technical reason why these arguments were not ported to torch.onnx.dynamo_export?

[thiagocrepaldi]

@thiagocrepaldi In the current state, I think simply adding (slow) tests for dynamo, duplicating

https://github.com/huggingface/optimum/blob/7e08a820b65a359a61444abe51df4eb96b26b2e3/tests/exporters/onnx/test_exporters_onnx_cli.py#L384

with a @slow decorator, and using dynamo=True, and running RUN_SLOW=1 pytest tests/exporters/onnx -k "test_exporters_cli_pytorch_cpu" -s -vvvvv and reporting the failing tests / skipping them with a reason why would be enough. For dynamo export to be the default, having the ability to specify the input_names and output_names would be very helpful. Is there a technical reason why these arguments were not ported to torch.onnx.dynamo_export?

I need to check again, but I think dynamo already mimics the original input names in the user model, so the user doesnt have to do it themselves. Dynamo API does not provide such facility, so that would be a development in the pytorch core - not only on the onnx exporter module

[BowenBao]

I'm updating the PR and resolving a few quick infra errors to get reasonable outputs from the unittests. Encountered this error and not sure how to proceed. I recall encountering this kind of error before when running standalone optimum export scripts, it feels like the error message does not describe the error cause accurately. @fxmarty do you have suggestions?

_______ ERROR collecting tests/exporters/onnx/test_exporters_onnx_cli.py _______
/workspace/bowbao/anaconda3/envs/bench/lib/python3.10/site-packages/_pytest/runner.py:341: in from_call
    result: Optional[TResult] = func()
/workspace/bowbao/anaconda3/envs/bench/lib/python3.10/site-packages/_pytest/runner.py:372: in <lambda>
    call = CallInfo.from_call(lambda: list(collector.collect()), "collect")
/workspace/bowbao/anaconda3/envs/bench/lib/python3.10/site-packages/_pytest/python.py:531: in collect
    self._inject_setup_module_fixture()
/workspace/bowbao/anaconda3/envs/bench/lib/python3.10/site-packages/_pytest/python.py:545: in _inject_setup_module_fixture
    self.obj, ("setUpModule", "setup_module")
/workspace/bowbao/anaconda3/envs/bench/lib/python3.10/site-packages/_pytest/python.py:310: in obj
    self._obj = obj = self._getobj()
/workspace/bowbao/anaconda3/envs/bench/lib/python3.10/site-packages/_pytest/python.py:528: in _getobj
    return self._importtestmodule()
/workspace/bowbao/anaconda3/envs/bench/lib/python3.10/site-packages/_pytest/python.py:617: in _importtestmodule
    mod = import_path(self.path, mode=importmode, root=self.config.rootpath)
/workspace/bowbao/anaconda3/envs/bench/lib/python3.10/site-packages/_pytest/pathlib.py:567: in import_path
    importlib.import_module(module_name)
/workspace/bowbao/anaconda3/envs/bench/lib/python3.10/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
<frozen importlib._bootstrap>:1050: in _gcd_import
    ???
<frozen importlib._bootstrap>:1027: in _find_and_load
    ???
<frozen importlib._bootstrap>:1006: in _find_and_load_unlocked
    ???
<frozen importlib._bootstrap>:688: in _load_unlocked
    ???
/workspace/bowbao/anaconda3/envs/bench/lib/python3.10/site-packages/_pytest/assertion/rewrite.py:186: in exec_module
    exec(co, module.__dict__)
tests/exporters/onnx/test_exporters_onnx_cli.py:170: in <module>
    class OnnxCLIExportTestCase(unittest.TestCase):
tests/exporters/onnx/test_exporters_onnx_cli.py:287: in OnnxCLIExportTestCase
    _get_models_to_test(PYTORCH_SENTENCE_TRANSFORMERS_MODEL, library_name="sentence_transformers")
tests/exporters/onnx/test_exporters_onnx_cli.py:57: in _get_models_to_test
    task_config_mapping = TasksManager.get_supported_tasks_for_model_type(
optimum/exporters/tasks.py:1226: in get_supported_tasks_for_model_type
    raise KeyError(
E   KeyError: 'clip is not supported yet with the onnx backend. Only [] are supported. If you want to support onnx please propose a PR or open up an issue.'

[BowenBao]

Found root cause to be https://github.com/huggingface/optimum/blob/548d0ac5a2039c6ae73fda2ff53034365b1637d3/optimum/utils/import_utils.py#L61 returning False when installed package is onnx-weekly instead of onnx.

Looks like similar issue was noted for onnxruntime, and onnxruntime-gpu https://github.com/huggingface/optimum/blob/548d0ac5a2039c6ae73fda2ff53034365b1637d3/optimum/utils/import_utils.py#L63-L64

[fxmarty]

Thank you @BowenBao, feel free to fix in this PR.

[BowenBao]

Got an initial result after some tweaks to enable model patcher for dynamo.

=== 350 failed, 360 passed, 11 skipped, 18766 warnings in 8416.29s (2:20:16) ===

We are also undergoing some changes in exporter to merge a few extra dependencies, will update this PR further once that part is done and triage the errors.

[corehalt]

@BowenBao @fxmarty Hi, is this work still in progress or has been replaced by another PR?

[justinchuby]

I expect a new PR will be created. We will pick this up after pytorch 2.5 release