Add dynamo ONNX export with fallback#212
Merged
1985312383 merged 2 commits intodatawhalechina:mainfrom Apr 5, 2026
Merged
Conversation
Enable Dynamo-based ONNX export with support for torch.export dynamic_shapes and automatic fallback to the legacy exporter. Adds _build_dynamic_shapes to construct dynamo-friendly dynamic_shapes for sequence and sparse features, updates ONNX export logic to prefer dynamo (injecting dynamic_shapes when dynamic_batch=True), and retries with the legacy exporter on failure. Also updates trainer docstrings to document the new "dynamo" onnx_export_kwargs, adds tests for DIN ONNX export covering sequence features and dynamic-batch behavior, and updates pyproject.toml to include ml-dtypes and onnxscript in the onnx extras.
github-actions
bot
added
enhancement
Add `from __future__ import annotations` to config/generate_api_md.py and expand/clarify API documentation in both English and Chinese. Docs changes clarify the ONNX exporter behavior (dynamo vs legacy fallback and how to force a path), translate and expand SeqTrainer docs (parameters, loss options, logger, eval_step), and add RQVAE-specific methods and return-value details. These updates improve clarity and consistency across the generated API docs.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Pull Request / 拉取请求
What does this PR do? / 这个PR做了什么？
Enable Dynamo-based ONNX export with support for torch.export dynamic_shapes and automatic fallback to the legacy exporter. Adds _build_dynamic_shapes to construct dynamo-friendly dynamic_shapes for sequence and sparse features, updates ONNX export logic to prefer dynamo (injecting dynamic_shapes when dynamic_batch=True), and retries with the legacy exporter on failure. Also updates trainer docstrings to document the new "dynamo" onnx_export_kwargs, adds tests for DIN ONNX export covering sequence features and dynamic-batch behavior, and updates pyproject.toml to include ml-dtypes and onnxscript in the onnx extras.
Type of Change / 变更类型
Related Issues / 相关Issues
Fixes #211
How to Test / 如何测试
# Example of how to test your changes / 如何测试您的更改的示例Checklist / 检查清单
python config/format_code.py) / 代码遵循项目风格（运行了格式化脚本）Additional Notes / 附加说明
Any other information for reviewers / 给审查者的其他信息