Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Pytorch] Dynamo ONNX export support #1497

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

pggPL
Copy link
Collaborator

@pggPL pggPL commented Feb 19, 2025

Description

This PR adds torch onnx dynamo-based export support. This draft is published to discuss the direction of the progress, it is not functional yet.

Current idea:

  • define onnx_forward in modules, which is much simpler version of forward, which will be run in onnx_export mode. Idea: it was hard to make dynamo work with original layers and inference flow is much much simpler, so this is not painful,
  • I implement many parts in onnx_export using torch modules, which work with onnx export - like layer norm,
  • I implement onnx gemm and onnx quantization functions in Python and I register them in torch.ops,
  • I put mapping torch.ops call <-> symbolic function inte_translation_table, which need to be passed to torch.onnx.export(dynamo=True, custom_translation_table=te_translation_table).
  • torch.onnx.export(dynamo=True) seem to be working for some modules, torch.onnx.dynamo_export will not work, because it has unsupported features like generators, which are essential part of the TE layers,

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@borisfom
Copy link

I don’t like onnx_forward idea, we used similar design for Nemo and it’s not something I am proud of …
Now if we do indeed have a need to use simplified versions for forward for ONNX export, I would suggest using monkeypatch techniques to temporarily modify exported objects just for ONNX export.
Putting is_in_onnx_export_mode conditional blocks permanently in the code may cause all sort of issues with other types of export clients may need : jit.script(), torch.compile(), torch.export()

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants