Skip to content

Add NVTE_BACKWARD_MODE=default|unquant|dequant #2644

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.

Sign up for GitHub

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

Open
zianglih wants to merge 51 commits into
base: main
Choose a base branch
from
Open

Add NVTE_BACKWARD_MODE=default|unquant|dequant #2644

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
51 commits
Select commit Hold shift + click to select a range
8cad9ae
Add NVTE_KEEP_BACKWARD_UNQUANTIZED
zianglih Feb 3, 2026
b4f91ef
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 3, 2026
cf8ac99
Disable ub and clean up
zianglih Feb 3, 2026
a90c4d6
Drop fuser changes
zianglih Feb 3, 2026
80917d6
Replace use_quantized_bwd with use_fp8_bwd
zianglih Feb 3, 2026
07cb1df
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 3, 2026
dc5d381
Ignore keep_backward_unquantized if delayed scaling
zianglih Feb 3, 2026
05fb894
Refactor ignoring NVTE_KEEP_BACKWARD_UNQUANTIZED when delayed scaling…
zianglih Feb 3, 2026
f78ffa0
Add back missing ctx.debug
zianglih Feb 3, 2026
91b53e1
Refactor changes under fused
zianglih Feb 3, 2026
4509967
Clean up
zianglih Feb 3, 2026
70be1b5
Refactor high-precision overwrite if keep_backward_unquantized
zianglih Feb 3, 2026
058ad45
Clean up
zianglih Feb 3, 2026
0f53179
Drop redundant fp8_recipe_bwd
zianglih Feb 4, 2026
6f213dc
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 4, 2026
b4836d1
Drop redundant ub changes
zianglih Feb 4, 2026
6b7666a
Drop more redundant ub changes
zianglih Feb 4, 2026
6a10fd1
Drop redundant delayed scaling changes
zianglih Feb 4, 2026
dd9038b
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 4, 2026
c83ed74
Drop unneeded backwards_needs_fc1_input
zianglih Feb 4, 2026
bf632e7
Drop and disallow LayerNormMLP implementation
zianglih Feb 4, 2026
ae939a1
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 4, 2026
25902a2
Move interface changes to recipe
zianglih Feb 5, 2026
f10f0bb
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 5, 2026
074b83f
Move ub overrides to fwd
zianglih Feb 5, 2026
1a504d4
Remove duplication
zianglih Feb 5, 2026
447677b
Simplify use_fp8_bwd logic in bwd
zianglih Feb 5, 2026
81e7feb
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 5, 2026
435859b
Set grad quantizers to none if keep bwd unquantized
zianglih Feb 5, 2026
a191345
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 5, 2026
1fbc22a
Drop delayed scaling change
zianglih Feb 6, 2026
7da1af5
Simplify env var logic
zianglih Feb 9, 2026
442297c
Move validation check to recipe
zianglih Feb 9, 2026
9e9e94f
Simplify effective_enabled
zianglih Feb 9, 2026
e269b85
Fix inverted assertion logic
zianglih Feb 9, 2026
153b1a8
Simplify changes under ops
zianglih Feb 9, 2026
e5eb2ed
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 9, 2026
c233a6d
Simplify ctx.keep_backward_unquantized
zianglih Feb 9, 2026
b967628
Fix missing attribute
zianglih Feb 9, 2026
8ff02a2
Add unit tests
zianglih Feb 10, 2026
01af855
Fix bias errors in unit test
zianglih Feb 10, 2026
7198af2
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 10, 2026
8d985c3
Add more shapes to unit test
zianglih Feb 10, 2026
e35da44
Refator interface to `NVTE_BACKWARD_MODE=default|unquant|dequant`
zianglih Feb 24, 2026
749e3a0
Fix override and clean up
zianglih Feb 25, 2026
03a7fe9
Clean up unit test
zianglih Mar 5, 2026
4ef353f
Clean up unit test
zianglih Mar 5, 2026
c16ba4b
Override `ctx.reduce_and_update_bwd_fp8_tensors = False`
zianglih Mar 9, 2026
27e70bc
Expand unit test
zianglih Mar 9, 2026
6ac9050
Add `test_backward_mode_memory_peak_report`
zianglih Mar 9, 2026
a2b5250
Expand test coverage and fix
zianglih Mar 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions qa/L0_pytorch_unittest/test.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gqa.xml $TE_PATH
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_optimizer.xml $TE_PATH/tests/pytorch/test_fused_optimizer.py || test_fail "test_fused_optimizer.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_multi_tensor.xml $TE_PATH/tests/pytorch/test_multi_tensor.py || test_fail "test_multi_tensor.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/test_fusible_ops.py || test_fail "test_fusible_ops.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_mode.xml $TE_PATH/tests/pytorch/test_backward_mode.py || test_fail "test_backward_mode.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml $TE_PATH/tests/pytorch/test_permutation.py || test_fail "test_permutation.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_entropy.xml $TE_PATH/tests/pytorch/test_parallel_cross_entropy.py || test_fail "test_parallel_cross_entropy.py"
python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py"
Expand Down
Loading