[ONNX] Improve the conversion of `from dynamic axes to shapes` #140488

titaiwangms · 2024-11-13T01:25:56Z

Features:
(1) Add support for tree structure.
(2) Add user warning before axes to shapes conversion
(3) Add suggestion of providing dynamic_shapes when conversion fails

Notes:
(1) input_names is crucial to the conversion, as we don't know the ONNX graph inputs.
(2) min and max are set as default, so LLM has higher chance to fail if users use dynamic_axes in terms of the min/max constraints dependency between attention_mask and sequence_length, etc. (Found in llama-3.2-1B_Instruct)

pytorch-bot · 2024-11-13T01:25:59Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/140488

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

[DomainsOnly] Jobs fail with GLIBC version not found

✅ No Failures

As of commit 123e85d with merge base f98c601 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

torch/onnx/_internal/exporter/_compat.py

titaiwangms · 2024-11-13T20:05:27Z

I have tested this PR on Olive with llama, and it works:

olive capture-onnx-graph -m meta-llama/Llama-3.2-1B-Instruct --torch_dtype float32 -o models/dynamo_export --log_level 0 --use_dynamo_exporter

However, I would say (1) re.sub on torch.export.Dim naming and (2) max=99999 on constraints is kind of over customized. Alternative way is that we don't do these two for users, and ask them to provide the correct information (we would have to change code in Olive side.). But still, max and min are hard to be decided in automation I think.

torch/onnx/_internal/exporter/_compat.py

test/onnx/exporter/test_api.py

torch/onnx/_internal/exporter/_compat.py

test/onnx/exporter/test_api.py

justinchuby · 2024-11-13T21:55:51Z

torch/onnx/_internal/exporter/_compat.py

    raise ValueError("model has no forward method and is not callable")


 def _from_dynamic_axes_to_dynamic_shapes(


Possible to unit test this function a little more? Just to ensure rubustness

justinchuby

Thanks for this PR! Nice use of pytree. LGTM overall; more unit tests should help us be more confident.

Co-authored-by: Justin Chu <[email protected]>

justinchuby · 2024-11-14T17:46:13Z

For test, I suggest reproducing the huggingface model signature and make sure that works.

titaiwangms · 2024-11-14T22:31:56Z

@justinchuby I added some extra tests for pytree, and there are existing small model tests already in test_api.py (they were added when conversion was introduced.). I find it hard to include hf models in the test of _from_dynamic_axes_to_dynamic_shapes, because the function basically asks everything: model, args, kwargs, input_names, etc. Do you have a better way to test this? Otherwise, we might need to build up llm testing...

justinchuby · 2024-11-14T22:45:39Z

We don't need to include the hf models. (And we should not.) I would just replicate its forward signature. The test model doesn't need to do anything except for having the same forward signature as hf models.

justinchuby · 2024-11-15T00:58:10Z

test/onnx/exporter/test_compat.py

+            dynamic_axes=dynamic_axes,
+        )
+
+        # NOTE: torch.export.Dim being an object makes it impossible to compare the objects directly.


Should we implement __eq__ on Dim? Just a thought

I think it's a valid request?

pytorch/torch/export/dynamic_shapes.py

Line 56 in 33191bb

class _Dim(type):

torch/onnx/_internal/exporter/_compat.py

test/onnx/exporter/test_compat.py

justinchuby · 2024-11-15T01:00:39Z

test/onnx/exporter/test_compat.py

+                    "input_ids": torch.randn(2, 16),
+                    "attention_mask": torch.randn(2, 32),
+                    "position_ids": torch.randn(2, 16),
+                    "past_key_values": [


Since 16 entries and 3 entries are effectively the same, we can shorten this test case a little

titaiwangms · 2024-11-15T01:39:49Z

@pytorchbot merge

pytorchmergebot · 2024-11-15T01:41:43Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

…ch#140488) Features: (1) Add support for tree structure. (2) Add user warning before axes to shapes conversion (3) Add suggestion of providing `dynamic_shapes` when conversion fails Notes: (1) `input_names` is crucial to the conversion, as we don't know the ONNX graph inputs. (2) min and max are set as default, so LLM has higher chance to fail if users use `dynamic_axes` in terms of the min/max constraints dependency between `attention_mask` and `sequence_length`, etc. (Found in llama-3.2-1B_Instruct) Pull Request resolved: pytorch#140488 Approved by: https://github.com/justinchuby Co-authored-by: Justin Chu <[email protected]>

improve from dynamic axes to shapes

efe8b0b

titaiwangms requested review from justinchuby, shubhambhokare1 and wschin as code owners November 13, 2024 01:25

pytorch-bot bot added the release notes: onnx torch.onnx related changes that should show up in the release notes label Nov 13, 2024

titaiwangms changed the title ~~[ONNX] Improve from dynamic axes to shapes~~ [ONNX] Improve the conversion of from dynamic axes to shapes Nov 13, 2024

titaiwangms added the topic: improvements topic category label Nov 13, 2024

pytorchbot added the open source label Nov 13, 2024

soulitzer added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Nov 13, 2024