Generalize quantization APIs for FP8/FP4/.. recipes by ksivaman · Pull Request #2256 · NVIDIA/TransformerEngine

ksivaman · 2025-10-10T04:58:08Z

Description

#2177 and #2254 add support for the NVFP4 recipe. This PR is a follow-up to generalize the user-facing APIs to be agnostic of the recipe dtype and deprecate the old functions/arguments that explicitly mention FP8.

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

Changes

fp8_autocast → autocast for PyTorch and Jax.
fp8_model_init → quantized_model_init for PyTorch.
Deprecate fp8 specific arguments to PyTorch's make_graphed_callables for more general versions.
Introduce some utility functions into pytorch extensions public API (e.g. is_fp8_available). This PR starts off the process but will likely not be exhaustive. We want to expose as many necessary utils as possible so that users are discouraged to utilize internal TE functionality.

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

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

ksivaman · 2025-10-10T05:25:16Z

/te-ci pytorch

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

ksivaman · 2025-10-13T01:02:01Z

/te-ci L0 L1 L2 L3

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

ksivaman · 2025-10-13T17:23:09Z

/te-ci L0 L1

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

ksivaman · 2025-10-13T23:46:21Z

For the thunder CI failure: Lightning-AI/lightning-thunder#2635

ksivaman · 2025-10-13T23:46:37Z

/te-ci L0 L1

timmoon10

I don't see any blocking issues, but there are still some minor improvements we could make. Better to put the effort in now to get things exactly right, and there's less risk we'll need to make changes in the future.

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

ksivaman · 2025-10-14T02:24:51Z

/te-ci L0 L1

timmoon10 · 2025-10-14T17:34:05Z



-def is_fp8_available() -> Tuple[bool, str]:
+def is_fp8_available(return_reason: bool = False) -> Union[bool, Tuple[bool, str]]:


I think this design is a big improvement, especially since if is_fp8_available(): ... does what you expect and since the pattern is consistent between recipes. I do get some code smells from the variable return type. The alternative would be to expose check_fp8_support in the public API and use that when we want reason strings, but that also means expanding the functions we are committing to maintain. I don't have a strong opinion on which approach is better.

Yeah I went with lower public maintenance while still having the option to do return reason.

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

ksivaman · 2025-10-14T18:20:44Z

/te-ci L0 L1

Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

timmoon10 · 2025-10-14T18:36:41Z

/te-ci L0 L1

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

ksivaman · 2025-10-14T20:05:20Z

/te-ci L0 L1

ksivaman added 5 commits October 9, 2025 22:18

Initial API change
01cc890
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Merge branch 'main' into nvfp4_cleanup
5d584a2

Change all imports and api
c1d30ef
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

format
6cbbc80
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

fixes
08eadf0
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman requested review from ptrendx and timmoon10 October 10, 2025 04:58

ksivaman added the 2.9.0 label Oct 10, 2025

ksivaman marked this pull request as draft October 10, 2025 04:58

ksivaman added 4 commits October 9, 2025 21:59

Merge branch 'main' into nvfp4_cleanup
52ea9f4

fix typo
cc38278
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

fix recipe tets
30b6af8
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

fix more tests
151e1c8
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman changed the title ~~[PyTorch] Generalized quantize API for FP8/FP4~~Generalized quantize API for FP8/FP4Oct 10, 2025

ksivaman added 2 commits October 10, 2025 18:08

fix docs, tests, and make Jax change as well
5b9b6f0
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Change internal uses of fp8_autocast
8ab201b
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

timmoon10 reviewed Oct 10, 2025
View reviewed changes

ksivaman added 4 commits October 12, 2025 23:20

Address nits
b20a7e4
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

rename file
ebdfcee
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

CG function, and small test fixes
3821c3c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Change instances of make_graphed_callables internally
a23e353
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman changed the title ~~Generalized quantize API for FP8/FP4~~Generalized quantization APIs for FP8/FP4/.. recipesOct 13, 2025

ksivaman changed the title ~~Generalized quantization APIs for FP8/FP4/.. recipes~~Generalize quantization APIs for FP8/FP4/.. recipesOct 13, 2025

ksivaman marked this pull request as ready for review October 13, 2025 00:43

ksivaman added 2 commits October 13, 2025 10:12

Merge branch 'main' into nvfp4_cleanup
7fa3ff6

Fix distributed tests
6847f6e
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

timmoon10 reviewed Oct 13, 2025
View reviewed changes

Comment threadtransformer_engine/pytorch/graph.py Outdated
Comment threadtransformer_engine/pytorch/graph.py Outdated
Comment threadtransformer_engine/pytorch/graph.py Outdated

ksivaman added 3 commits October 13, 2025 15:58

Merge branch 'main' into nvfp4_cleanup
6e2a4e5

fixes
47a4179
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

fix tests
06249b1
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

timmoon10 reviewed Oct 13, 2025
View reviewed changes

Comment threadtransformer_engine/pytorch/fp8.py

timmoon10 previously approved these changes Oct 14, 2025
View reviewed changes

Comment threadtransformer_engine/pytorch/utils.py Outdated

Better docs and better api
c43d89c
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman dismissed timmoon10’s stale review via c43d89cOctober 14, 2025 01:47

Merge branch 'main' into nvfp4_cleanup
d6b9c51

timmoon10 previously approved these changes Oct 14, 2025
View reviewed changes

Merge branch 'main' into nvfp4_cleanup
793feae
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman dismissed timmoon10’s stale review via 793feaeOctober 14, 2025 17:59

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

timmoon10 reviewed Oct 14, 2025
View reviewed changes

Comment threadtests/pytorch/distributed/run_numerics_exact.py
Comment threadtests/pytorch/distributed/run_gemm_with_overlap.py Outdated

Apply suggestions from code review
d32fdc0
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

timmoon10 previously approved these changes Oct 14, 2025
View reviewed changes

fix nvfp4 test
5fc015b
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman dismissed timmoon10’s stale review via 5fc015bOctober 14, 2025 18:54

timmoon10 approved these changes Oct 14, 2025
View reviewed changes

ksivaman merged commit 85a9199 into NVIDIA:mainOct 14, 2025
13 of 15 checks passed

timmoon10 mentioned this pull request Oct 22, 2025
[PyTorch] FSDP2 Support for TE with FP8 Model Init #2245
Merged
13 tasks

pggPL mentioned this pull request Oct 24, 2025
Docs fix for 2.9 #2302
Closed

phu0ngng mentioned this pull request Oct 31, 2025
[JAX] Quickstart documentation #2310
Merged
11 tasks

ksivaman mentioned this pull request Nov 25, 2025
Fix execution with Transformer Engine huggingface/accelerate#3852
Merged
5 tasks



		def is_fp8_available() -> Tuple[bool, str]:
		def is_fp8_available(return_reason: bool = False) -> Union[bool, Tuple[bool, str]]:

Uh oh!

Conversation

ksivaman commented Oct 10, 2025• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

ksivaman commented Oct 10, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ksivaman commented Oct 13, 2025

Uh oh!

ksivaman commented Oct 13, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ksivaman commented Oct 13, 2025

Uh oh!

ksivaman commented Oct 13, 2025

Uh oh!

Uh oh!

timmoon10 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ksivaman commented Oct 14, 2025

Uh oh!

timmoon10Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

ksivamanOct 14, 2025

Choose a reason for hiding this comment

Uh oh!

ksivaman commented Oct 14, 2025

Uh oh!

Uh oh!

Uh oh!

timmoon10 commented Oct 14, 2025

Uh oh!

ksivaman commented Oct 14, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ksivaman commented Oct 10, 2025•
edited
Loading