Allow override inputs to export recipe #37508

guangy10 · 2025-04-15T00:49:16Z

What does this PR do?

Enable dynamism at seq_len dim in order to utilize parallel prefill in the executorch runtime. In this PR,

allow caller side to override example inputs, dynamic shapes and strict flag, but keep the default unchanged for BC
add unit test to cover export with dynamic shapes, and strict False as it's the mainstream in latest version of torch.export
make the unit test non-slow, to avoid being skipped on PRs and causing regressions
add test for HybirdCache

Tests
pytest tests/utils/test_cache_utils.py -vv -s -k cache_exportability

collected 23 items / 20 deselected / 3 selected

tests/utils/test_cache_utils.py::CacheExportIntegrationTest::test_dynamic_cache_exportability PASSED                                                             [ 33%]
tests/utils/test_cache_utils.py::CacheExportIntegrationTest::test_hybrid_cache_exportability PASSED                                                              [ 66%]
tests/utils/test_cache_utils.py::CacheExportIntegrationTest::test_static_cache_exportability PASSED                                                              [100%]

Before submitting

Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case. Needed by the downstream in optimum-executorch. Add export dynamic shape for causal lm optimum-executorch#53
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you add new tests? Yes

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@ArthurZucker @ydshieh

CC: @tugsbayasgalan

github-actions · 2025-04-15T00:49:29Z

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the Ready for review button (at the bottom of the PR page). This will assign reviewers and trigger CI.

ydshieh · 2025-04-15T06:55:46Z

src/transformers/integrations/executorch.py

+        dynamic_shapes = (
+            dynamic_shapes
+            if dynamic_shapes is not None
+            else {"input_ids": {1: torch.export.Dim.AUTO}, "cache_position": None}


Could you explain a bit what is the value like torch.export.Dim.AUTO means, and the one in this guide

# Create a dynamic batch size batch = Dim("batch") # Specify that the first dimension of each input is that batch size dynamic_shapes = {"x1": {0: batch}, "x2": {0: batch}}

IIRC, the key 0 or 1 means which dimension, but don't know torch.export.Dim.AUTO or Dim("batch").

key 0 or 1 means which dimension

Correct.

torch.export.Dim.AUTO or Dim("batch")

It's new feature introduced since 2.6.0, smarter than the traditional way of specifying the dynamic range, e.g Dim("seq_len", min, max). @tugsbayasgalan can explain more. @ydshieh if you think Dim("seq_len", min, max) is easier to understand, we can switch to explicit dynamism.

If you can find a doc that explains what torch.export.Dim.AUTO does and why it's smarter, adding the link along with a a short comment in the source code here, that would be nice.

When I look doc or source, I can't find any info.

But if you can't find neither, it's fine. In this case, maybe we can reflect this situation to pytorch team.

I will leave it for @tugsbayasgalan to chime in

Dim.AUTO works by refining the range as you encounter shape related constraints (possibly specializing to static integer). So in that sense, Dim.AUTO is less likely to run into shape related errors and doesn't require user code change as it will just respect user code intention. If you really want to keep seq_len to be dynamic, you should use Dim.DYNAMIC which will error when we specialize. For official doc, cc @pianpwk

src/transformers/integrations/executorch.py

SunMarc

Thanks ! left a few comments

src/transformers/integrations/executorch.py

HuggingFaceDocBuilderDev · 2025-04-15T13:08:00Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

ydshieh · 2025-04-15T18:12:47Z

Other than the nit comments I left, just a question: so this dynamism is necessary if the exported model want to be used to call prompts (i.e. prefill stage) as well as to use the generation steps (here seq=1)?

guangy10 · 2025-04-15T21:00:57Z

tests/utils/test_cache_utils.py

+        input_ids = tokenizer("Here's everything I know", return_tensors="pt").input_ids
+        dynamic_shapes = {"input_ids": {1: torch.export.Dim.AUTO}, "cache_position": None}
+        exported_program = convert_and_export_with_cache(
+            model, example_input_ids=input_ids, dynamic_shapes=dynamic_shapes
+        )


I'm export this model with a non-specialized input_ids and the dynamic shape on "seq_len" dim. @tugsbayasgalan Is there a way I can assert the exported program do have the dim set to be dynamic, and inspect what the dynamic range is?

ydshieh · 2025-04-16T08:03:25Z

src/transformers/integrations/executorch.py

+            if dynamic_shapes is not None:
+                logging.warning("Dynamic shapes spec will be ingored for < 2.6.0.")
+            if strict is not None:
+                logging.warning("strict flag spec will be ingored for < 2.6.0.")


Hmm, first, I think we should have

if is_torch_greater_or_equal("2.6.0"):
...
elif is_torch_greater_or_equal("2.5.0"):
...
else:
...
right ...?

If so, the 2 new messages have to be in both elif and else.

Also, let's rephrase it as

Dynamic shapes spec will be ignored by convert_and_export_with_cache for < 2.6.0.

@ydshieh That could be one option. Alternatively, we can simplify it by just not use dynamic shapes for torch<2.6 including 2.5, which is BC.
Additionally, in torch 2.5 though it supports dynamic shapes, but no Dim.AUTO. As Tugsuu mentioned above, w/o Dim.AUTO it is more likely to run into shape related errors. Hence I'd rather keep things BC by not using dynamic shapes for torch<2.6. WDYT?

src/transformers/integrations/executorch.py

ydshieh · 2025-04-16T08:05:55Z

just a if/else v.s. if/elif/else question and 2 nits, but overall LGTM.

SunMarc · 2025-04-16T14:58:35Z

src/transformers/integrations/executorch.py

+            if strict is not None:
+                logging.warning("strict flag spec will be ingored for < 2.6.0.")


let's default strict to True. Does it require 2.6 to set strict to False ?

I don't have strong option to default True or False, but I think the export team wants to promote default to False in newer version of torch. cc: @tugsbayasgalan

guangy10 · 2025-04-16T18:08:39Z

just a if/else v.s. if/elif/else question and 2 nits, but overall LGTM.

Shared my thoughts here: #37508 (comment). Just in case you missed it. Let me know you thoughts

guangy10 · 2025-04-22T22:57:07Z

More comments on this PR? Can we merge it if no?

guangy10 · 2025-04-28T19:58:55Z

I will rebase this PR since #37728 has been merged

ydshieh

Thank you for the iterations. LGTM, torched test is passing and fast to run 👍

I will wait until the end of the day before merge to see if @SunMarc has any comment.

SunMarc

LGTM ! Please fix the conflits and we are good to merge

guangy10 · 2025-04-29T23:43:40Z

LGTM ! Please fix the conflits and we are good to merge

Couple updates:

Added a non-slow test to cover export for hybird cache
Rebased to latest trunk

@SunMarc @ydshieh Let me know if there are new comments

ydshieh · 2025-04-30T08:19:32Z

Thanks

Add option to specify dynamic shapes during export Co-authored-by: Guang Yang <[email protected]>

github-actions bot marked this pull request as draft April 15, 2025 00:49

guangy10 marked this pull request as ready for review April 15, 2025 00:57

github-actions bot requested review from MekkCyber and SunMarc April 15, 2025 00:57

ydshieh reviewed Apr 15, 2025

View reviewed changes

SunMarc reviewed Apr 15, 2025

View reviewed changes

src/transformers/integrations/executorch.py Outdated Show resolved Hide resolved

src/transformers/integrations/executorch.py Show resolved Hide resolved

guangy10 force-pushed the export_w_dynamism branch 2 times, most recently from 10e7489 to 772606e Compare April 15, 2025 20:58

guangy10 commented Apr 15, 2025

View reviewed changes

guangy10 force-pushed the export_w_dynamism branch from 772606e to 1451ea8 Compare April 16, 2025 00:57

guangy10 changed the title ~~Add option to specify dynamic shapes during export~~ Allow override dynamic shapes and strict in export recipe Apr 16, 2025

guangy10 changed the title ~~Allow override dynamic shapes and strict in export recipe~~ Allow override inputs to export recipe Apr 16, 2025

ydshieh reviewed Apr 16, 2025

View reviewed changes

src/transformers/integrations/executorch.py Outdated Show resolved Hide resolved

SunMarc reviewed Apr 16, 2025

View reviewed changes

guangy10 force-pushed the export_w_dynamism branch from 1451ea8 to 174c382 Compare April 16, 2025 18:11

guangy10 force-pushed the export_w_dynamism branch 2 times, most recently from 5390845 to c2bd31f Compare April 29, 2025 03:32

ydshieh approved these changes Apr 29, 2025

View reviewed changes

SunMarc approved these changes Apr 29, 2025

View reviewed changes

Add option to specify dynamic shapes during export

cbf2207

guangy10 force-pushed the export_w_dynamism branch from bee41b7 to cbf2207 Compare April 29, 2025 23:41

ydshieh merged commit a572744 into huggingface:main Apr 30, 2025
13 checks passed

zucchini-nlp pushed a commit to zucchini-nlp/transformers that referenced this pull request May 14, 2025

Allow override inputs to export recipe (huggingface#37508)

b7fe854

Add option to specify dynamic shapes during export Co-authored-by: Guang Yang <[email protected]>

		if strict is not None:
		logging.warning("strict flag spec will be ingored for < 2.6.0.")

Allow override inputs to export recipe #37508

Allow override inputs to export recipe #37508

Uh oh!

Conversation

guangy10 commented Apr 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

github-actions bot commented Apr 15, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ydshieh Apr 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

SunMarc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Apr 15, 2025

Uh oh!

ydshieh commented Apr 15, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

guangy10 Apr 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ydshieh commented Apr 16, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

guangy10 commented Apr 16, 2025

Uh oh!

guangy10 commented Apr 22, 2025

Uh oh!

guangy10 commented Apr 28, 2025

Uh oh!

ydshieh left a comment

Choose a reason for hiding this comment

Uh oh!

SunMarc left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

guangy10 commented Apr 29, 2025

Uh oh!

Uh oh!

ydshieh commented Apr 30, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

guangy10 commented Apr 15, 2025 •

edited

Loading

ydshieh Apr 15, 2025 •

edited

Loading

guangy10 Apr 16, 2025 •

edited

Loading

SunMarc left a comment •

edited

Loading