Add common test for `torch.export` and fix some vision models #35124

qubvel · 2024-12-06T15:30:11Z

What does this PR do?

Add a common slow test to check if a model can be exported with no issues using torch.export.export

Add an optional test, to enable it please set test_torch_exportable = True flag for model-specific test.
Enable test for vision and video models
Fix most of the vision models

The main fixes include:

Use a compile-compatible LRU cache for models.
Avoid modifying model parameters in the forward pass (e.g. self.param = self.param + x).
Avoid modifying leaf in-place tensors created in the forward pass.
Avoid creating tensors with requires_grad=True in the forward pass.

Testing is not complete, there might be code paths that can't be exported. I did additional testing with specific checkpoints. In most cases, we are safe. The only two situations I found where tests pass but checkpoint export does not pass are:

beit (fixed)
zoedepth (not fixed)

Results

✅ - can be exported with torch.export.export
🔵 - export fixed in this PR
❌ - can't be exported

Vision models

🔵 beit
🔵 bit
🔵 conditional_detr
✅ convnext
✅ convnextv2
✅ cvt
✅ dab_detr
🔵 deformable_detr
✅ deit
✅ depth_anything
✅ depth_pro
🔵 detr
✅ dinat
✅ dinov2
✅ dinov2_with_registers
✅ dit
✅ dpt
✅ efficientnet
🔵 focalnet
✅ glpn
✅ hiera
✅ ijepa
🔵 imagegpt
❌ levit (low usage, won't fix)
✅ mask2former
🔵 maskformer
✅ mobilenet_v1
✅ mobilenet_v2
✅ mobilevit
✅ mobilevitv2
✅ poolformer
✅ pvt
✅ pvt_v2
✅ regnet
✅ resnet
✅ rt_detr
🔵 rt_detr_v2
✅ segformer
🔵 seggpt
❌ superpoint (data-dependent expression)
✅ swiftformer
✅ swin
✅ swinv2
🔵 swin2sr
✅ table_transformer
✅ textnet
✅ upernet
✅ vit
✅ vitdet
✅ vit_mae
✅ vitmatte
✅ vit_msn
✅ vitpose
✅ vitpose_backbone
✅ yolos
❌ zoedept (data-dependent expression, test config pass but checkpoint not)

Video models

✅ timesformer
✅ vivit
✅ videomae

Fixes # (issue)

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
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.
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 write any new necessary tests?

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.

HuggingFaceDocBuilderDev · 2024-12-06T16:08:06Z

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.

qubvel · 2024-12-11T16:52:22Z

@guangy10 please have a look if you have bandwidth! Do you have anything in mind that should be added to the common test on the side of the transformers to ensure the model is exportable and executorch compatible?

ydshieh

Thanks @qubvel for working on this!

I haven't check the changes in models but left a few comments in test_modeling_common.py

tests/test_modeling_common.py

ydshieh

Just left a few tiny comments.

Could we also trigger slow CI for the modified models?

You can rebase on main to use the new way to trigger slow CI.

src/transformers/models/conditional_detr/modeling_conditional_detr.py

src/transformers/models/maskformer/modeling_maskformer.py

src/transformers/models/seggpt/modeling_seggpt.py

src/transformers/models/swin2sr/modeling_swin2sr.py

ydshieh

Make sure run a slow CI 🙏

Thanks for the work.

qubvel · 2024-12-18T16:42:46Z

Ok, sure, what is a new way to trigger slow tests?

Thanks for review 🤗

ydshieh · 2024-12-18T17:14:10Z

Second part in

https://www.notion.so/huggingface2/CI-for-pull-requests-8335bd9217d24d1e8ed1d3a016f39804

or my message in slack

https://huggingface.slack.com/archives/C01NE71C4F7/p1734527174922859

guangy10

@qubvel Great work and thank you for standardizing the test for exportability! Does this PR cover all vision existing vision models in transformers? Is there a plan to set up the same standard for audio models (in separate PRs)?

guangy10 · 2025-01-10T21:04:44Z

tests/test_modeling_common.py

    test_mismatched_shapes = True
    test_missing_keys = True
    test_model_parallel = False
+    test_torch_exportable = False


If vast majority vision models are exportable, should we strategically turn this flag to True? Typically new models are more popular and important than old models, the biggest benefit of reversing the default to True is to "softly enforce" new models to be exportable, naturally growing this path to be the default and hard enforced over time. When I say softly enforce, I mean the new model can still have the option to disable export test if confirmed the failure is not obvious to fix and file a github issue in the backlog. But I do think most of the failures would be common and easy fixable given the work you did in this PR.

Common tests are applied to all modalities, so for now, I suppose we will set it to False before updating for other models.

Including text models right? @qubvel Could you guide me how I can leverage this common test to cover more test models like this: https://github.com/search?q=repo%3Ahuggingface%2Ftransformers%20test_export_static_cache&type=code

It' s fine to have to as we want to reach 100% exportability no?

tests/test_modeling_common.py

guangy10 · 2025-01-10T21:16:49Z

tests/test_modeling_common.py

+                for key in eager_outputs:
+                    is_tested = is_tested or recursively_check(eager_outputs[key], exported_outputs[key])
+                return is_tested
+            return is_tested


A nit for debuggability: Is silent return from here is expected? I think we should explicitly report/catch an error instead because it bypasses all checks due to output type mismatch?

I suppose the main targets are covered, in case no matching types is_tested will be False and the error will be raised below.

guangy10 · 2025-01-10T21:34:29Z

tests/test_modeling_common.py

+                    model,
+                    args=(),
+                    kwargs=inputs_dict,
+                    strict=True,


@gmagogsfm Oh I think I recommended strict=True somewhere. The reason is to ensure the safety. If some python code block is unsupported by torchdynamo and got traced out in non-strict mode, instead of proceed silently defer seeing inconsistent results in the exported graph, a better process is to first surface the error up with strict mode, let the the model author or reviewers to see the source and confirm if it's safe and then turn the strict flag to False to unblock. cc: @ydshieh @qubvel

guangy10 · 2025-01-10T21:43:03Z

tests/test_modeling_common.py

+                with tempfile.TemporaryDirectory() as tmpdirname:
+                    save_path = os.path.join(tmpdirname, "exported_model.pt2")
+                    torch.export.save(exported_model, save_path)
+                    exported_model = torch.export.load(save_path)


@qubvel You may want to add a check to ensure the loaded exported model is identical, to avoid any weird issues due to bugs in (de)serialization as you are checking the export outputs against eager using the loaded artifact.

In a workflow (e.g. to ExecuTorch for inference on-device) where exported model is just an intermediate representation, there is no need to save the exported program to a local fs first. We don't want any bug in (de)serialization of the exported IR affect the testability of the downstream workflow, so it's better to first compare the export and eager outputs, then test the (de)serialization.

Ok, basically, there is almost nothing we can do on our side in case of a deserialization error (just report to a torch team). So, I suppose we can remove this entirely, WDYT?

Sure, we can remove this entirely

guangy10 · 2025-01-10T21:45:46Z

src/transformers/pytorch_utils.py

        raise ValueError(f"Unsupported parallel style value: {style}")
+
+
+def compile_compatible_method_lru_cache(*lru_args, **lru_kwargs):


@qubvel Do you mind elaborate a bit more how this decorator helps with export?

please, see below

guangy10 · 2025-01-10T21:46:38Z

src/transformers/pytorch_utils.py

+def compile_compatible_method_lru_cache(*lru_args, **lru_kwargs):
+    """
+    LRU cache decorator from standard functools library, but with a workaround to disable
+    caching when torchdynamo is compiling. Expected to work with class methods.


what particular caching is disabled when torchdynamo is tracing?

We turning off standard lru_cache when torchdynamo is compiling.

We had this decorator previously only for RT-DETR, I just moved it to make it common:

RT-DETR (afair) caches anchors to avoid it's creation for the same image size.
OmDet-Turbo (zero-shot-object-detection, out of this PR) cashes text label embeddings in order to avoid recomputing them once again in case the same label is passed.

Not that much speedup coming from these optimizations tbh, so I suppose it's fine to turn it off to enable compile/export.

guangy10 · 2025-01-10T22:01:50Z

tests/test_modeling_common.py

+
+        default_config, default_inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+        config = config or default_config
+        inputs_dict = inputs_dict or default_inputs_dict


What are required parameters to be passed to the forward()? I think we should standardize the signature of forward so that we can simplify and standardize developing an ExecuTorch runtime for all vision models? Like what we did for text model in transformers.integrations.executorch where only input_ids and current cache_position are required as the model is always exported with the static_cache. We don't have to maintain backwards compatibility (BC) for all arguments.

It might be different for different vision models, but most of them should have only pixel_values as the required parameter.

guangy10

Looks good! 🚀 🚀

ArthurZucker

Missing one big thing: documentation about supported models for export!
otherwise 🚀 ! Kudos

src/transformers/models/bit/modeling_bit.py

src/transformers/models/maskformer/modeling_maskformer.py

ArthurZucker · 2025-02-03T15:44:27Z

tests/test_modeling_common.py

    test_mismatched_shapes = True
    test_missing_keys = True
    test_model_parallel = False
+    test_torch_exportable = False


It' s fine to have to as we want to reach 100% exportability no?

tests/test_modeling_common.py

guangy10 · 2025-02-04T01:11:06Z

@qubvel Let me know the timeline. I can start working on e2e enablement by connecting these models to optimum-executorch once this PR merged.

qubvel · 2025-02-04T16:53:26Z

Hey @guangy10, going to finish this week!

guangy10 · 2025-02-10T21:38:08Z

Hey @guangy10, going to finish this week!

@qubvel Just a friendly reminder to merge this PR. Let mw know how I can help if there is any blocker.

chrsmcgrr · 2025-02-11T10:20:10Z

Nice work in bringing the torch.export coverage up.

Are the export tests executed with torch==2.6.0?

qubvel · 2025-02-11T10:32:33Z

Hey @chrsmcgrr, yes we have 2.6.0 in CI. I have also run these tests with torch 2.5.0 locally

qubvel · 2025-02-11T11:12:57Z

run-slow: rt_detr_v2, dab_detr, beit, conditional_detr, deformable_detr, detr, swin2sr

github-actions · 2025-02-11T11:14:14Z

This comment contains run-slow, running the specified jobs: This comment contains run-slow, running the specified jobs:

models: ['models/beit', 'models/conditional_detr', 'models/dab_detr', 'models/deformable_detr', 'models/detr', 'models/rt_detr_v2', 'models/swin2sr']
quantizations: [] ...

qubvel · 2025-02-11T11:36:40Z

Test failures are unrelated (fixed in #35654).

Merging this PR to unblock @guangy10 work. I will work on docs in the follow-up PR asap. cc @stevhliu in case you have suggestions on how to organize it better.

stevhliu · 2025-02-11T18:09:25Z

Depending on when the new docs ship (hopefully this week or the next), you can add them here I think :)

If it takes longer, then feel free to add them here and I can move them later!

guangy10 · 2025-02-12T18:46:47Z

Depending on when the new docs ship (hopefully this week or the next), you can add them here I think :)

If it takes longer, then feel free to add them here and I can move them later!

Feel free to subscribe me on the new PR for the doc fix.

guangy10 · 2025-02-12T18:50:19Z

@qubvel While I'm working on lowering these vision models e2e to ExecuTorch in optimum, would you like to start expanding the export coverage over audio models as well? I’d really appreciate the effort!

qubvel · 2025-02-12T22:45:50Z

@guangy10 I'm not sure if I will have bandwidth in the coming few weeks, but maybe someone from the audio team can have a look
cc @eustlb

guangy10 · 2025-02-13T02:19:37Z

@qubvel Thanks for looping in the audio team!

👋 @eustlb, nice to e-meet you! I'm from PyTorch team at Meta. We've been collaborating with 🤗 to expand torch.export coverage on transformer models since last year (FYI there is a parallel efforts focusing on torch.compile coverage). So far we have covered text and vision models, but not on audio models yet. It seems like your team is the right one to talk. I'd happy to share more context in slack.

eustlb · 2025-02-13T10:22:33Z

Hey @guangy10, nice to e-meet you too! It would be a pleasure to help 🤗

…gface#35124) * Add is_torch_greater_or_equal test decorator * Add common test for torch.export * Fix bit * Fix focalnet * Fix imagegpt * Fix seggpt * Fix swin2sr * Enable torch.export test for vision models * Enable test for video models * Remove json * Enable for hiera * Enable for ijepa * Fix detr * Fic conditional_detr * Fix maskformer * Enable test maskformer * Fix test for deformable detr * Fix custom kernels for export in rt-detr and deformable-detr * Enable test for all DPT * Remove custom test for deformable detr * Simplify test to use only kwargs for export * Add comment * Move compile_compatible_method_lru_cache to utils * Fix beit export * Fix deformable detr * Fix copies data2vec<->beit * Fix typos, update test to work with dict * Add seed to the test * Enable test for vit_mae * Fix beit tests * [run-slow] beit, bit, conditional_detr, data2vec, deformable_detr, detr, focalnet, imagegpt, maskformer, rt_detr, seggpt, swin2sr * Add vitpose test * Add textnet test * Add dinov2 with registers * Update tests/test_modeling_common.py * Switch to torch.testing.assert_close * Fix masformer * Remove save-load from test * Add dab_detr * Add depth_pro * Fix and test RT-DETRv2 * Fix dab_detr

qubvel added 12 commits December 5, 2024 22:57

Add is_torch_greater_or_equal test decorator

645b6b2

Add common test for torch.export

98698b0

Fix bit

208f9c4

Fix focalnet

b464e26

Fix imagegpt

a64f10b

Fix seggpt

70fb0da

Fix swin2sr

552a98e

Enable torch.export test for vision models

5a2e048

Enable test for video models

c63b95c

Remove json

7dd1cce

Enable for hiera

6e2e601

Enable for ijepa

4d48819

qubvel added the Vision label Dec 6, 2024

qubvel added 5 commits December 6, 2024 16:35

Fix detr

b60030b

Fic conditional_detr

6f7a0fd

Fix maskformer

665f583

Enable test maskformer

abb883c

Fix test for deformable detr

3391cf0

qubvel requested a review from ydshieh December 17, 2024 22:32

qubvel added run-slow torch export Issues and PRs related to torch.export compatibility labels Dec 17, 2024

ydshieh reviewed Dec 18, 2024

View reviewed changes

ydshieh approved these changes Dec 18, 2024

View reviewed changes

qubvel added 2 commits December 18, 2024 19:03

Merge branch 'main' into test-torch-export

cd569af

Fix custom kernels for export in rt-detr and deformable-detr

c22475b

Merge branch 'main' into test-torch-export

ca4b6ac

guangy10 suggested changes Jan 10, 2025

View reviewed changes

Update tests/test_modeling_common.py

1bc76e9

guangy10 approved these changes Jan 31, 2025

View reviewed changes

ArthurZucker approved these changes Feb 3, 2025

View reviewed changes

guangy10 mentioned this pull request Feb 5, 2025

Support vision transformers huggingface/optimum-executorch#18

Closed

qubvel added 6 commits February 11, 2025 10:02

Merge branch 'main' into test-torch-export

2cc7e6a

Switch to torch.testing.assert_close

b3c17cc

Fix masformer

130bf0a

Remove save-load from test

617403b

Add dab_detr

b9b6aa6

Add depth_pro

b1b7944

Fix and test RT-DETRv2

8749a9a

Fix dab_detr

c6e0f0f

qubvel merged commit f42d46c into huggingface:main Feb 11, 2025
25 of 26 checks passed

		raise ValueError(f"Unsupported parallel style value: {style}")


		def compile_compatible_method_lru_cache(lru_args, *lru_kwargs):

Add common test for torch.export and fix some vision models #35124

Add common test for torch.export and fix some vision models #35124

Uh oh!

Conversation

qubvel commented Dec 6, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Results

Vision models

Video models

Before submitting

Who can review?

Uh oh!

HuggingFaceDocBuilderDev commented Dec 6, 2024

Uh oh!

qubvel commented Dec 11, 2024

Uh oh!

ydshieh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ydshieh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ydshieh left a comment

Choose a reason for hiding this comment

Uh oh!

qubvel commented Dec 18, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ydshieh commented Dec 18, 2024

Uh oh!

guangy10 left a comment

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!

guangy10 Jan 31, 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!

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!

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!

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!

qubvel Jan 31, 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!

Add common test for `torch.export` and fix some vision models #35124

Add common test for `torch.export` and fix some vision models #35124

qubvel commented Dec 6, 2024 •

edited

Loading

qubvel commented Dec 18, 2024 •

edited

Loading

guangy10 Jan 31, 2025 •

edited

Loading

qubvel Jan 31, 2025 •

edited

Loading

qubvel commented Feb 12, 2025 •

edited

Loading