[Core][PyTorch] NVFP4 recipe by ksivaman · Pull Request #2177 · NVIDIA/TransformerEngine

ksivaman · 2025-09-16T23:53:22Z

Description

Add support for the NVFP4 training recipe.

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

Adds support for NVFP4 block scaling training recipe, including the required kernels and framework infrastructure.

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> Co-authored-by: Frank Sun <frsun@nvidia.com> Co-authored-by: Oleg Goncharov <ogoncharov@nvidia.com> Co-authored-by: Zhongbo Zhu <zhongboz@nvidia.com> Co-authored-by: Evgeny Tsykunov <etsykunov@nvidia.com> Co-authored-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Teddy Do <tdophung@nvidia.com>

for more information, see https://pre-commit.ci

Copilot

Pull Request Overview

This PR adds support for NVFP4 (4-bit floating point) training recipe, implementing a complete quantization infrastructure for NVFP4 block scaling training. The implementation includes quantization kernels, tensor classes, distributed operations, and integration with the existing transformer engine framework.

Key changes:

Implements NVFP4 tensor class and quantizer with block scaling support
Adds NVFP4-specific CUDA kernels for quantization and transpose operations
Integrates NVFP4 support into distributed operations (all-gather)
Updates module infrastructure to support NVFP4 quantization

Reviewed Changes

Copilot reviewed 82 out of 82 changed files in this pull request and generated 10 comments.

Show a summary per file

File	Description
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Core NVFP4 tensor implementation with quantization logic
transformer_engine/pytorch/fp8.py	NVFP4 recipe state management and quantizer creation
transformer_engine/pytorch/distributed.py	NVFP4 all-gather operations for distributed training
transformer_engine/pytorch/csrc/quantizer.cpp	C++ NVFP4 quantizer implementation
transformer_engine/common/util/nvfp4_transpose.cuh	CUDA kernels for NVFP4 quantization and transpose
transformer_engine/pytorch/experimental/	Experimental middleware for NVFP4 reference implementations

Comments suppressed due to low confidence (1)

transformer_engine/pytorch/tensor/nvfp4_tensor.py:1

The TODO comment indicates missing documentation. Please add a comprehensive module docstring that explains the NVFP4 tensor implementation, its purpose, and usage patterns.

# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

jberchtold-nvidia · 2025-09-17T15:53:15Z

+ if (use_2d_quantization && !return_identity) {
+ return;
+ }
+


This function is missing shape checks on the input tensors

Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

Signed-off-by: Tim Moon <tmoon@nvidia.com>

for more information, see https://pre-commit.ci

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

timmoon10 · 2025-09-18T18:15:05Z

/te-ci L1

This reverts commit f8b2a2d. Signed-off-by: Tim Moon <tmoon@nvidia.com>

timmoon10 · 2025-09-19T04:23:52Z

/te-ci L1

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

Signed-off-by: Tim Moon <tmoon@nvidia.com>

for more information, see https://pre-commit.ci

timmoon10 · 2025-09-27T02:11:19Z

/te-ci L1

Signed-off-by: Tim Moon <tmoon@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

ksivaman · 2025-09-29T17:27:48Z

/te-ci L1

* Add NVFP4 recipe Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Co-authored-by: Frank Sun <frsun@nvidia.com> Co-authored-by: Oleg Goncharov <ogoncharov@nvidia.com> Co-authored-by: Zhongbo Zhu <zhongboz@nvidia.com> Co-authored-by: Evgeny Tsykunov <etsykunov@nvidia.com> Co-authored-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Teddy Do <tdophung@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add MathDx dependency to GitHub builds Signed-off-by: Tim Moon <tmoon@nvidia.com> * Suggestions from GitHub Copilot Signed-off-by: Tim Moon <tmoon@nvidia.com> * Move 2x shape logic from core to PyTorch Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Fix compilation errors with CUDA 12.1 Signed-off-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * SM 70 is not supported in CUDA 13 Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Typo Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Revert "Move 2x shape logic from core to PyTorch" This reverts commit f8b2a2d. Signed-off-by: Tim Moon <tmoon@nvidia.com> * Added dequantize kernel for FP4 Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix linter warning Signed-off-by: Tim Moon <tmoon@nvidia.com> * Add NVFP4 support with fusible ops Use logical tensor dims for PyTorch NVFP4 tensors. Temporarily add unfused dequantize impl. Fix bug where NVFP4 recipe was not configurable. Signed-off-by: Tim Moon <tmoon@nvidia.com> * Fix logic for 2x shapes and move to PyTorch Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Fix CG test model config Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Debug NVFP4 tensor size function Signed-off-by: Tim Moon <tmoon@nvidia.com> * Proper handling of the RNG state Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Test SR properly Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix workspace size for GEMM heuristic. Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Fix compile error in C++ NVFP4 test Some some numeric errors when blocks are all zero. Signed-off-by: Tim Moon <tmoon@nvidia.com> * fix distrbuted test problem shape Signed-off-by: zhongboz <zhongboz@nvidia.com> * proper assert dim for low precision AG TP Signed-off-by: zhongboz <zhongboz@nvidia.com> * clean up duplicated code in nvfp4_utils.cuh Signed-off-by: zhongboz <zhongboz@nvidia.com> * lint Signed-off-by: zhongboz <zhongboz@nvidia.com> * pylint: disable=unused-argument Signed-off-by: zhongboz <zhongboz@nvidia.com> * `nvte_cublas_gemm_v2` to take alpha pointer (#12) * make nvte_cublas_gemm_v2 to take alpha/beta pointers Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * users are expected to pass a valid C_tensor Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * typos Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * API to have const float* alpha Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * Minor tweaks Support arbitrary beta scales. Increase workspace to be aligned to 128 bytes. Signed-off-by: Tim Moon <tmoon@nvidia.com> * Debug IMA with alpha pointer Signed-off-by: Tim Moon <tmoon@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Support fused amax kernels with NVFP4 quantization Signed-off-by: Tim Moon <tmoon@nvidia.com> * Disable fused amax with cuDNN LayerNorm kernel Signed-off-by: Tim Moon <tmoon@nvidia.com> * Add NVFP4 cases to distributed tests for TE ops Signed-off-by: Tim Moon <tmoon@nvidia.com> * Change assert to NVTE_CHECK in the hadamard cast fusion Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix compile error Signed-off-by: Tim Moon <tmoon@nvidia.com> * Use global thread IDs for Philox subsequences Signed-off-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add shape checks for NVFP4 cast kernels Signed-off-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Do not fuse amax if cuDNN normalization is forced by envvar Signed-off-by: Przemek Tredak <ptredak@nvidia.com> --------- Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Przemek Tredak <ptredak@nvidia.com> Signed-off-by: zhongboz <zhongboz@nvidia.com> Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> Co-authored-by: Frank Sun <frsun@nvidia.com> Co-authored-by: Oleg Goncharov <ogoncharov@nvidia.com> Co-authored-by: Zhongbo Zhu <zhongboz@nvidia.com> Co-authored-by: Evgeny Tsykunov <etsykunov@nvidia.com> Co-authored-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Teddy Do <tdophung@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Przemek Tredak <ptredak@nvidia.com> Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>

yash3056 · 2025-10-06T11:38:49Z

does nvfp4 only work with sm100 or all arch above sm100?

- Adds NVFP4 pretraining research paper with PR NVIDIA#2177 reference Signed-off-by: Santosh Bhavani <santosh.bhavani@live.com>

* Enhance Latest News section with recent TE and FP8 developments - Adds NVFP4 pretraining research paper with PR #2177 reference Signed-off-by: Santosh Bhavani <santosh.bhavani@live.com> * update nvfp4 reference Signed-off-by: Santosh Bhavani <santosh.bhavani@live.com> * Update README.rst Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> --------- Signed-off-by: Santosh Bhavani <santosh.bhavani@live.com> Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman requested review from ptrendx and timmoon10 September 16, 2025 23:53

ksivaman self-assigned this Sep 16, 2025

ksivaman added the 2.8.0 label Sep 16, 2025

[pre-commit.ci] auto fixes from pre-commit.com hooks
d5919e1
for more information, see https://pre-commit.ci

ptrendx requested a review from Copilot September 17, 2025 00:05

CopilotAI reviewed Sep 17, 2025
View reviewed changes

phu0ngng previously requested changes Sep 17, 2025
View reviewed changes

Comment threadtransformer_engine/common/common.h Outdated

jberchtold-nvidia requested changes Sep 17, 2025
View reviewed changes

ptrendx mentioned this pull request Sep 17, 2025
nvfp4 training with hints Megatron #2171
Closed

timmoon10 added 3 commits September 17, 2025 18:11

Add MathDx dependency to GitHub builds
49e939b
Signed-off-by: Tim Moon <tmoon@nvidia.com>

Suggestions from GitHub Copilot
689c66e
Signed-off-by: Tim Moon <tmoon@nvidia.com>

Merge branch 'main' into nvfp4_recipe
afd411d
Signed-off-by: Tim Moon <tmoon@nvidia.com>

jberchtold-nvidia reviewed Sep 17, 2025
View reviewed changes

Comment threadtransformer_engine/pytorch/csrc/quantizer.cpp

ptrendx previously requested changes Sep 17, 2025
View reviewed changes

Comment threadtransformer_engine/common/common.h Outdated

ksivamanand others added 4 commits September 18, 2025 02:18

Move 2x shape logic from core to PyTorch
f8b2a2d
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Fix compilation errors with CUDA 12.1
6427c09
Signed-off-by: Tim Moon <tmoon@nvidia.com>

Merge branch 'main' into nvfp4_recipe
d8ae3b0
Signed-off-by: Tim Moon <tmoon@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks
279b013
for more information, see https://pre-commit.ci

timmoon10 reviewed Sep 18, 2025
View reviewed changes

Comment threadbuild_tools/utils.py Outdated

SM 70 is not supported in CUDA 13
1465744
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

timmoon10 reviewed Sep 18, 2025
View reviewed changes

Comment threadbuild_tools/utils.py Outdated

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

timmoon10 mentioned this pull request Sep 18, 2025
[PyTorch] Quantizer as API + CustomRecipe #2039
Merged
11 tasks

Revert "Move 2x shape logic from core to PyTorch"
6430896
This reverts commit f8b2a2d. Signed-off-by: Tim Moon <tmoon@nvidia.com>

timmoon10 force-pushed the nvfp4_recipe branch from c18c6c6 to 6430896Compare September 19, 2025 04:22

Merge branch 'main' into nvfp4_recipe
788ee72

ptrendxand others added 5 commits September 26, 2025 16:56

Change assert to NVTE_CHECK in the hadamard cast fusion
8550cb7
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

Fix compile error
d01a658
Signed-off-by: Tim Moon <tmoon@nvidia.com>

Use global thread IDs for Philox subsequences
e519ea4
Signed-off-by: Tim Moon <tmoon@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks
8be0a74
for more information, see https://pre-commit.ci

Merge branch 'main' into nvfp4_recipe
6d2a106

timmoon10and others added 5 commits September 28, 2025 04:54

Add shape checks for NVFP4 cast kernels
a0cb86d
Signed-off-by: Tim Moon <tmoon@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks
b396119
for more information, see https://pre-commit.ci

Merge branch 'main' into nvfp4_recipe
f549d78

Do not fuse amax if cuDNN normalization is forced by envvar
eaa5c46
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

Merge branch 'main' into nvfp4_recipe
115dbfc

timmoon10 approved these changes Sep 29, 2025
View reviewed changes

jberchtold-nvidia approved these changes Sep 29, 2025
View reviewed changes

ksivaman merged commit 3f5b475 into NVIDIA:mainSep 29, 2025
11 of 13 checks passed

ksivaman mentioned this pull request Sep 30, 2025
Fix hang during debug build #2221
Merged
13 tasks

timmoon10 mentioned this pull request Oct 1, 2025
is there a way to use NVTE_FWD_NVFP4_BWD_MXFP8_SCALING? #2102
Closed

ksivaman mentioned this pull request Oct 13, 2025
Generalize quantization APIs for FP8/FP4/.. recipes #2256
Merged
13 tasks

kooshi mentioned this pull request Oct 13, 2025
FP4 Training #1701
Closed

ksivaman mentioned this pull request Oct 16, 2025
Wheels for cuda 13 #2278
Merged
13 tasks

ehartford mentioned this pull request Nov 1, 2025
nvfp4 ROCm/TransformerEngine#358
Open

timmoon10 mentioned this pull request Apr 30, 2026
[Core] Remove unused NVFP4 quantize kernel #2946
Merged
13 tasks

hungryGeek16 mentioned this pull request May 31, 2026
fix unfused padding causal sdpa #3063
Open

sbhavani mentioned this pull request Jul 2, 2026
Release training examples for https://www.arxiv.org/pdf/2509.25149 #2389
Open

Uh oh!

Conversation

ksivaman commented Sep 16, 2025

Description

Type of change

Changes

Checklist:

Uh oh!

CopilotAI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jberchtold-nvidiaSep 17, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

timmoon10 commented Sep 18, 2025

Uh oh!

timmoon10 commented Sep 19, 2025

Uh oh!

timmoon10 commented Sep 27, 2025

Uh oh!

ksivaman commented Sep 29, 2025

Uh oh!

Uh oh!

yash3056 commented Oct 6, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

12 participants