Repository navigation
Dynamo exporter patches - #27
titaiwangms wants to merge 10 commits into
Conversation
There was a problem hiding this comment.
Pull Request Overview
This PR introduces comprehensive patches for PyTorch Dynamo export functionality, specifically targeting compatibility issues with transformers and torch export mechanisms. The changes enable better support for dynamic shapes, caching systems, and model serialization during the export process.
- Adds exclusions to ruff linting configuration for patch files
- Implements core patching infrastructure with context manager support
- Provides extensive patches for transformers models including cache handling, attention mechanisms, and rotary embeddings
- Adds torch-specific patches for shape inference and symbolic computation
Reviewed Changes
Copilot reviewed 5 out of 6 changed files in this pull request and generated 10 comments.
Show a summary per file
| File | Description |
|---|---|
| pyproject.toml | Excludes patch files from ruff linting |
| optimum/torch_export_patches/_treenode_registry.py | Implements tree node registration and serialization for cache classes |
| optimum/torch_export_patches/_patch_transformers.py | Comprehensive transformers patches for export compatibility |
| optimum/torch_export_patches/_patch_torch.py | Core torch patches for shape inference and symbolic computation |
| optimum/torch_export_patches/_core.py | Main patching infrastructure with context manager |
|
|
||
| # check | ||
| if f_check: | ||
| raise NotImplementedError( |
There was a problem hiding this comment.
The f_check parameter is documented as optional but when provided, it always raises NotImplementedError. This makes the parameter unusable. Either implement the check functionality or remove the parameter.
| import packaging.version as pv | ||
| import torch | ||
| import transformers | ||
| from transformers.modeling_attn_mask_utils import AttentionMaskConverter |
There was a problem hiding this comment.
This import should be wrapped in a try-catch block as it may not exist in all transformers versions, similar to how other conditional imports are handled in this file.
| import torch | ||
| import transformers | ||
| from transformers.modeling_attn_mask_utils import AttentionMaskConverter | ||
| from transformers.cache_utils import StaticCache, Cache, DynamicCache |
There was a problem hiding this comment.
This import should be wrapped in a try-catch block as these classes may not exist in all transformers versions, similar to the pattern used for MambaCache import.
| if verbose: | ||
| print("[torch_export_patches] done patching") | ||
|
|
||
| try: |
There was a problem hiding this comment.
The variable 'fct_callable' is referenced but never defined in this code path. This will cause a NameError.
| try: | |
| try: | |
| fct_callable = None # Placeholder definition; replace with actual callable as needed |
|
You add a readme which explains how to use the patches when exporting a model. |
|
I would definitely treat this as experimental before exposing as public apis when they are more stable. I think we should (1) evolve the patches here (2) upstream as much as we can (3) use them in the dynamo export path of optimum (4) decide on a public api for the patches (if that’s what we still need) with HF |
Just wanted to add that I think (4) is required if (3) is done. The patches are more than just modeling code rewrites, and we don't want to force our users to rely on the dynamo exporter here. |
| ) | ||
|
|
||
| from .serialization.transformers_impl import ( | ||
| __dict__ as dtr, |
There was a problem hiding this comment.
dtr contains the list of functions in one submodule, then the function looks into all of the them to see if it is a patch or not.
There was a problem hiding this comment.
Is it possible to register explicitly?
There was a problem hiding this comment.
We can. The patches need to be removed after the export if a user wants to use it to compute discrepancies. So I chose to restore everything back to its original state including serialization functions.
There was a problem hiding this comment.
I agree restoring is preferable. I support that is independent of how we discover the patches here? With the context manager we should be able to restore definitively I suppose
224801d to
59106c3
Compare
There was a problem hiding this comment.
Maybe call the module just patching, for simplicity, and move it into the onnx exporter directory?
There was a problem hiding this comment.
We can then do
with optimum.exporters.onnx.patching.patch_xxx():
…
There was a problem hiding this comment.
I wanted to keep onnx-diagnostic patches independently, so we can have more flexibility when it comes to TS deprecation and exporter updates.
There was a problem hiding this comment.
Could you say more about the flexibility? What flexibility do we loose or gain with either option?
There was a problem hiding this comment.
For me, it's like building new things under the legacy folders (for example, building dynamo exporter under torchscript exporter folders.). When TS exporter is removed, we can just delete onnx/exporters without further checking it. Same logic goes to model_patcher.py, mixing them up increases the maintain complexity.
There was a problem hiding this comment.
Sounds good thanks. I do want to make sure the api we create can stay for a long time so we need to be careful about naming and where we expose them.
|
I suppose some of the transformer patches can go into https://github.com/huggingface/optimum-onnx/blob/main/optimum/exporters/onnx/model_patcher.py? |
|
could it be that this is too much 😅 is this all to enable dynamo export ? |
|
To enable torch.export.export. some of this code should be moved to transformers (cache serialization), some code are necessaru until torch implements a fix for it (vmap). Some others are rewriting of control flows (again until torch has better support). |
|
In addition to what I said, I created a couple of PR to address the necessary rewriting in transformers:
|
|
The PR is ready. This is a first draft that we can use on our benchmark (we haven't exposed patches anyway.). But torch.export-based (PyTorch 2) ONNX exporter won't work on any transformers model without these patches, because not a single Cache is registered in transformers. In addition, the patches also rewrite modeling code to be export-friendly (e.g. control flows). PS: This PR does not enable torch.export-based ONNX exporter, but only patches (context manager) |
IlyasMoutawwakil
left a comment
IlyasMoutawwakil
left a comment
There was a problem hiding this comment.
Hi, I couldn't fully understand what's being patched by some methods, a lot of patches are not documented enough and don't seem related to specific issues, would appreciate if you linked to those issues/discussion. Also the PR doesn't seem to test these patches on actual transformers models.
I think we agree that there is a price paid with every patch suggested, a breaking change in every transformers/diffusers/torch version that changes a line in those patched methods/classes, especially internal/private ones.
Can you explain how this api is gonna be used in the optimum exporter ?
For now I would appreciate it if you add more comments on top of every patched function explaining why it is necessary. And also enabled some exporters testing of these patches.
| @@ -0,0 +1,44 @@ | |||
| name: ONNX / Python - Test | |||
There was a problem hiding this comment.
please use a different name, we already have a workflow with this name so it's confusing.
| strategy: | ||
| fail-fast: false | ||
| matrix: | ||
| python-version: ['3.10'] |
There was a problem hiding this comment.
3.9 is the minimum supported
| return make_dynamic_cache(list(zip(self.key_cache, self.value_cache))) | ||
|
|
||
|
|
||
| if pv.Version(transformers.__version__) > pv.Version("4.49.99999"): |
There was a problem hiding this comment.
we have the utilisty function is_tranformers_version(op, version)
| def _has_transformers(version: str) -> bool: | ||
| "Returns True if transformers version is higher." | ||
| return pv.Version(transformers.__version__) >= pv.Version(version) |
There was a problem hiding this comment.
you can use is_transformers_version(op, version) from optimum.utils
| return False | ||
|
|
||
| try: | ||
| return torch.compiler.is_exporting() |
| except Exception: | ||
| return False | ||
|
|
||
| if patch_masking_utils: |
There was a problem hiding this comment.
we already have an sdpa mask function that doesn't use vmap and supports all features of sdpa_mask, I would prefer using the registration api to register that as it seems to me less error prone than patching torch and transformers. wdyt ?
There was a problem hiding this comment.
Could you share the api for doing this? Is it via attn_implementation=?
There was a problem hiding this comment.
here is the sdpa without vmap function https://github.com/huggingface/optimum-onnx/blob/main/optimum/exporters/onnx/model_patcher.py#L312, it relies on simple tensor operations to create the custom mask.
we register it in https://github.com/huggingface/optimum-onnx/blob/main/optimum/exporters/onnx/model_patcher.py#L564 and then unregister it after the export.
| # mask = | ||
| # torch.where(mask, torch.tensor(0.0, device=mask.device, dtype=dtype), min_dtype) | ||
| mask = (~mask).to(dtype) * min_dtype |
There was a problem hiding this comment.
is the multiplication any more efficient than where ?
| return mask | ||
|
|
||
|
|
||
| def _patch_make_causal_mask( |
There was a problem hiding this comment.
can you elaborate on why this is patched please
|
Hopefully, more of these patches could make it upstream... Otherwise risk becoming outdated quick :( |
Hi @IlyasMoutawwakil, thanks for the review! These patches have been around since last year, so a few legacy ones are still present. The main goals are:
Our plan is to wrap this as a context manager around torch.onnx.export, which will become the core of the optimum-onnx exporter. We are still documenting some of these patches, and testing whether they are still needed, and we’ll polish them in our repos and expose them as an API once finalized. At that point, the optimum-onnx exporter will directly consume it. |
This PR aim to expose a public API to patch torch.export-based (PyTorch 2.0) ONNX exporter.
I plan to land a minimal code to enable the patches. Maybe it's also a good opportunity to review/check whether some of the patches are already outdated.
7/30 added: We can land this with a private API, and then use it on benchmark to make sure it works. Finally, reorganize the folders and parameters, and then publish it.