[RFC][WIP] Introduce PT2 dynamo onnx exporter #1712

BowenBao · 2024-02-22T22:18:39Z

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 · 2024-02-26T10:18:06Z

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.

optimum/exporters/onnx/convert.py

thiagocrepaldi · 2024-03-12T21:13:45Z

@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 · 2024-03-19T06:17:32Z

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

optimum/tests/exporters/onnx/test_exporters_onnx_cli.py

Line 384 in 7e08a82

def test_exporters_cli_pytorch_cpu(

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 · 2024-03-19T16:58:20Z

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

optimum/tests/exporters/onnx/test_exporters_onnx_cli.py

Line 384 in 7e08a82

def test_exporters_cli_pytorch_cpu(

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 · 2024-03-25T19:59:56Z

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 · 2024-03-25T21:27:53Z

Found root cause to be

optimum/optimum/utils/import_utils.py

Line 61 in 548d0ac

_onnx_available = _is_package_available("onnx")

returning False when installed package is onnx-weekly instead of onnx.

Looks like similar issue was noted for onnxruntime, and onnxruntime-gpu

optimum/optimum/utils/import_utils.py

Lines 63 to 64 in 548d0ac

    
           # importlib.metadata.version seem to not be robust with the ONNX Runtime extensions (`onnxruntime-gpu`, etc.) 
        
           _onnxruntime_available = importlib.util.find_spec("onnxruntime") is not None

fxmarty · 2024-03-26T08:54:42Z

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

BowenBao · 2024-03-29T20:53:22Z

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.

justinchuby · 2024-05-10T20:39:18Z

tests/exporters/onnx/test_exporters_onnx_cli.py

@@ -172,6 +171,12 @@ class OnnxCLIExportTestCase(unittest.TestCase):
    Integration tests ensuring supported models are correctly exported.
    """

+    # TODO: Make this flag configurable.


Do you mean exposing it in the cli?

justinchuby · 2024-05-10T20:48:13Z

optimum/exporters/onnx/__main__.py

@@ -163,6 +164,8 @@ def main_export(
            If True, disables the use of dynamic axes during ONNX export.
        do_constant_folding (bool, defaults to `True`):
            PyTorch-specific argument. If `True`, the PyTorch ONNX export will fold constants into adjacent nodes, if possible.
+        dynamo (bool, default to `False):


Suggested change

dynamo (bool, default to `False):

dynamo (bool, default to `False`):

fxmarty reviewed Feb 26, 2024

View reviewed changes

optimum/exporters/onnx/convert.py Outdated Show resolved Hide resolved

corehalt mentioned this pull request Mar 29, 2024

Cannot export DETR model with torch.export pytorch/pytorch#122934

Open

BowenBao mentioned this pull request Apr 11, 2024

[Tracking] TorchDynamo ONNX exporter issues #1810

Open

74 tasks

BowenBao added 2 commits May 10, 2024 09:20

Initial PR

0a7bec6

WIP: enable dynamo onnx export in tests & group errors

efbefbd

BowenBao force-pushed the dynamo_onnx_export branch from c10f372 to efbefbd Compare May 10, 2024 16:20

justinchuby reviewed May 10, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[RFC][WIP] Introduce PT2 dynamo onnx exporter #1712

[RFC][WIP] Introduce PT2 dynamo onnx exporter #1712

BowenBao commented Feb 22, 2024

fxmarty commented Feb 26, 2024

thiagocrepaldi commented Mar 12, 2024

fxmarty commented Mar 19, 2024

thiagocrepaldi commented Mar 19, 2024

BowenBao commented Mar 25, 2024

BowenBao commented Mar 25, 2024

fxmarty commented Mar 26, 2024

BowenBao commented Mar 29, 2024

justinchuby May 10, 2024

justinchuby May 10, 2024

	dynamo (bool, default to `False):
	dynamo (bool, default to `False`):

[RFC][WIP] Introduce PT2 dynamo onnx exporter #1712

Are you sure you want to change the base?

[RFC][WIP] Introduce PT2 dynamo onnx exporter #1712

Conversation

BowenBao commented Feb 22, 2024

Before submitting

fxmarty commented Feb 26, 2024

thiagocrepaldi commented Mar 12, 2024

fxmarty commented Mar 19, 2024

thiagocrepaldi commented Mar 19, 2024

BowenBao commented Mar 25, 2024

BowenBao commented Mar 25, 2024

fxmarty commented Mar 26, 2024

BowenBao commented Mar 29, 2024

justinchuby May 10, 2024

Choose a reason for hiding this comment

justinchuby May 10, 2024

Choose a reason for hiding this comment