Skip to content

Fix import of TE Recipe - #2635

Merged
t-vi merged 2 commits into
Lightning-AI:mainfrom
ksivaman:fix_internal_import
Oct 13, 2025
Merged

Fix import of TE Recipe#2635
t-vi merged 2 commits into
Lightning-AI:mainfrom
ksivaman:fix_internal_import

Conversation

@ksivaman

Copy link
Copy Markdown
Contributor
Before submitting
  • Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
  • Did you read the contributor guideline, Pull Request section?
  • Did you make sure to update the docs?
  • Did you write any new necessary tests?

What does this PR do?

Thunder uses internal TE API which is being refactored in NVIDIA/TransformerEngine#2259. The TE PR ensures full backwards compatibility, however, the Recipe class was not defined in fp8.py to begin with and this only works because it is imported indirectly in it. This PR imports it directly in order to not break integration immediately.

Future work

Use of internal TE API should be minimized, and imports from fp8.py should be changed as it will be eventually removed (currently renamed to quantization.py starting TE v2.9).

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@nvMelissa

Copy link
Copy Markdown
Collaborator

Adding @riccardofelluga as reviewer. Please take a look

@t-vi t-vi left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you @ksivaman

@t-vi
t-vi merged commit 2e09afe into Lightning-AI:main Oct 13, 2025
48 of 50 checks passed
Comment on lines +33 to -36
from transformer_engine.common.recipe import Recipe
from transformer_engine.pytorch.fp8 import (
_amax_and_scale_update,
get_fp8_max,
Recipe,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the update! LGTM

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.

4 participants