Skip to content

Add NVTE_KEEP_BACKWARD_UNQUANTIZED #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 44 commits into
base: main
Choose a base branch
from
Open

Add NVTE_KEEP_BACKWARD_UNQUANTIZED #2644

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
44 commits
Select commit Hold shift + click to select a range
3afce1f
Add NVTE_KEEP_BACKWARD_UNQUANTIZED
zianglih Feb 3, 2026
72149be
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 3, 2026
3e6eb64
Merge branch 'main' into keep-bwd
zianglih Feb 3, 2026
927d482
Disable ub and clean up
zianglih Feb 3, 2026
cc85b60
Drop fuser changes
zianglih Feb 3, 2026
fe24f95
Replace use_quantized_bwd with use_fp8_bwd
zianglih Feb 3, 2026
5ca3615
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 3, 2026
5ba7674
Ignore keep_backward_unquantized if delayed scaling
zianglih Feb 3, 2026
02b7b2a
Refactor ignoring NVTE_KEEP_BACKWARD_UNQUANTIZED when delayed scaling…
zianglih Feb 3, 2026
01a7de0
Add back missing ctx.debug
zianglih Feb 3, 2026
bf904aa
Refactor changes under fused
zianglih Feb 3, 2026
b449fc4
Clean up
zianglih Feb 3, 2026
de3acaf
Refactor high-precision overwrite if keep_backward_unquantized
zianglih Feb 3, 2026
fe65d34
Clean up
zianglih Feb 3, 2026
59aaf6b
Drop redundant fp8_recipe_bwd
zianglih Feb 4, 2026
44da625
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 4, 2026
0f58793
Drop redundant ub changes
zianglih Feb 4, 2026
192fbad
Drop more redundant ub changes
zianglih Feb 4, 2026
0dd1268
Drop redundant delayed scaling changes
zianglih Feb 4, 2026
216621d
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 4, 2026
ab8749b
Drop unneeded backwards_needs_fc1_input
zianglih Feb 4, 2026
5881083
Drop and disallow LayerNormMLP implementation
zianglih Feb 4, 2026
431f0c8
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 4, 2026
937e34b
Move interface changes to recipe
zianglih Feb 5, 2026
0d26127
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 5, 2026
0135366
Move ub overrides to fwd
zianglih Feb 5, 2026
1de3c64
Remove duplication
zianglih Feb 5, 2026
04d3543
Simplify use_fp8_bwd logic in bwd
zianglih Feb 5, 2026
454976e
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 5, 2026
f7794c9
Set grad quantizers to none if keep bwd unquantized
zianglih Feb 5, 2026
58db8ea
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 5, 2026
9d0b654
Drop delayed scaling change
zianglih Feb 6, 2026
004cb45
Simplify env var logic
zianglih Feb 9, 2026
9baccfd
Move validation check to recipe
zianglih Feb 9, 2026
207eb5a
Simplify effective_enabled
zianglih Feb 9, 2026
15117b1
Fix inverted assertion logic
zianglih Feb 9, 2026
3fc5270
Simplify changes under ops
zianglih Feb 9, 2026
9201d19
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 9, 2026
1e0f1d2
Simplify ctx.keep_backward_unquantized
zianglih Feb 9, 2026
253873a
Fix missing attribute
zianglih Feb 9, 2026
fd94761
Add unit tests
zianglih Feb 10, 2026
0b2dbf9
Fix bias errors in unit test
zianglih Feb 10, 2026
3643320
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 10, 2026
74c787d
Add more shapes to unit test
zianglih Feb 10, 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 @@ -40,6 +40,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"
NVTE_KEEP_BACKWARD_UNQUANTIZED=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_keep_backward_unquantized.xml $TE_PATH/tests/pytorch/test_keep_backward_unquantized.py || test_fail "test_keep_backward_unquantized.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