From 2387f4db60c80411e1cb8072aee0a59fe060c884 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 30 Jul 2026 17:18:27 +0200 Subject: [PATCH 1/2] [PyTorch] Fix the selective activation checkpointing test tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py fails on all 16 parametrizations on main. The failure is in the test, not in the feature: outputs and all six parameter gradients are bit-exact between the checkpointed and the non-checkpointed path, and checkpointing does save memory. The test asserted `ln_fwd_mem > 6 * sln_fwd_mem`. That threshold was never reachable, because the ratio is fixed by the test's own configuration rather than by anything in TE. Both peaks are derivable and the derivations reproduce the measurements exactly: ln_fwd_mem = layers * (2*s*f + 2*s*h + 2*s) * itemsize sln_fwd_mem = ((layers+1)*s*h + 2*s*f + 2*s) * itemsize For `small` @ 128 this predicts 7876608 B and 1377280 B; measured 7876608 B and 1377280 B. With f = 4h and layers = 12 the ratio is 120/21 = 5.714, matching the measured 5.715. Reaching 6 would need a different model shape, e.g. layers = 16 gives 160/25 = 6.4. The tensor lists in both branches of _LayerNormMLP._forward are unchanged since the test was added, so this is not a regression. - Assert on the memory that recompute actually frees - fc1_out and act_out, derived from the model config - instead of the ratio. The checkpointed peak still holds the transient of one layer, so the expectation covers layers - 1. This keeps the assertion independent of layer count and model shape. - Check outputs and gradients before the memory check. Previously a numerical regression would surface as a memory-ratio failure and the correctness comparison would never run, which is exactly what happens on main today. - Drop `assert ln_bwd_time < sln_bwd_time`. The margin is as low as 13% on an idle GPU (huge @ 128: 11.2 ms vs 12.8 ms), which makes it a CI flake. - Skip parametrizations that do not fit in device memory. large @ 65536 and huge @ 65536 need more than 32 GiB for the non-checkpointed model alone and raise OutOfMemoryError on 48 GiB cards. Verified on RTX 5880 Ada: 12 passed, 4 skipped, was 16 failed. Signed-off-by: Pawel Gadzinski --- .../test_selective_activation_checkpoint.py | 51 +++++++++++++++---- 1 file changed, 42 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py b/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py index 306d0627f5..34aaf32ec9 100644 --- a/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py +++ b/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py @@ -137,10 +137,41 @@ def _param_key(name): return name.split(".")[-1] +def _no_checkpoint_activation_bytes(cfg, seq_size, itemsize): + """Activations LayerNormMLP saves for backward when checkpoint=False. + + Per layer: ln_out and out (seq*hidden each), fc1_out and act_out + (seq*ffn_hidden each), mu and rsigma (seq each). + """ + per_layer = 2 * seq_size * (cfg._ffn_hidden_size + cfg._hidden_size + 1) + return cfg._layers * per_layer * itemsize + + +def _recomputed_activation_bytes(cfg, seq_size, itemsize): + """Activations checkpointing must free: fc1_out and act_out. + + The peak still holds the transient of one layer, so only the remaining + layers count. Keeping this independent of _layers means the assertion + below does not encode the shape of the test models. + """ + return (cfg._layers - 1) * 2 * seq_size * cfg._ffn_hidden_size * itemsize + + @pytest.mark.parametrize("size", config.keys()) @pytest.mark.parametrize("seq_size", seq_sizes) def test_selective_activation_checkpoint(size, seq_size): + itemsize = torch.empty((), dtype=torch.get_default_dtype()).element_size() + no_ckpt_bytes = _no_checkpoint_activation_bytes(config[size], seq_size, itemsize) + + # Both models live in the same process, so budget the non-checkpointed peak twice. + free_bytes, _ = torch.cuda.mem_get_info(device) + if free_bytes < 2 * no_ckpt_bytes: + pytest.skip( + f"needs {2 * no_ckpt_bytes / 2**30:.1f} GiB free device memory, only" + f" {free_bytes / 2**30:.1f} GiB available" + ) + ln_model, sln_model = config[size].build() data = torch.randn((seq_size, config[size]._hidden_size), device=device) @@ -152,15 +183,8 @@ def test_selective_activation_checkpoint(size, seq_size): sln_fwd_out, sln_fwd_time, sln_fwd_mem = _run_fwd(sln_model, data) sln_grads, sln_bwd_time, sln_bwd_mem = _run_bwd(sln_model, sln_fwd_out) - assert ln_fwd_mem > 6 * sln_fwd_mem, ( - "selective activation checkpointing does not reduce forward memory by 6X, only by" - f" {ln_fwd_mem/sln_fwd_mem}!" - ) - assert ln_bwd_time < sln_bwd_time, ( - "selective activation activation checkpointing backward pass is NOT slower than native!" - f" got Native LayerNormMLP Backward Time: {ln_bwd_time} ms and Selective Activation" - f" Checkpointed LayerNormMLP Backward Time: {sln_bwd_time} ms" - ) + # Correctness first, so that a numerical regression is not masked by the + # memory check below. diff = _max_diff(ln_fwd_out, sln_fwd_out) assert diff == 0.0, f"outputs are not equal! maximum difference {diff}" for key in [ @@ -173,3 +197,12 @@ def test_selective_activation_checkpoint(size, seq_size): ]: diff = _max_diff(ln_grads[key], sln_grads[key]) assert diff == 0.0, f"gradients for {key} are not equal! maximum difference: {diff}" + + # Checkpointing recomputes fc1_out and act_out, so it must free at least those. + expected_saving = _recomputed_activation_bytes(config[size], seq_size, itemsize) + saving = ln_fwd_mem - sln_fwd_mem + assert saving >= 0.95 * expected_saving, ( + "selective activation checkpointing did not free the recomputed activations: saved" + f" {saving} B, expected at least {0.95 * expected_saving} B (ln_fwd_mem={ln_fwd_mem}," + f" sln_fwd_mem={sln_fwd_mem})" + ) From da1cadcf5faf67e0b5af5a70d0b939890f5f0dcb Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 30 Jul 2026 17:18:42 +0200 Subject: [PATCH 2/2] [CI] Connect orphaned pytorch test files to qa Six pytorch test files are not referenced by any script in qa/ and have therefore never run in CI: test_qk_norm.py test_float8_current_scaling_exact.py attention/test_cu_seqlens_cache.py test_nvfp4_fsdp2_hooks.py test_fused_router_perf.py layernorm_mlp/test_selective_activation_checkpoint.py `git log -S --all -- qa/` shows none of them was ever added and later removed, and none of the PRs that introduced them touched qa/. They were simply never wired up. All are single-GPU and self-skip on unsupported hardware, so they belong in L0. test_float8_current_scaling_exact.py marks its classes with skipif(not fp8_available), test_nvfp4_fsdp2_hooks.py requires sm_100+, and the one multi-device case in test_cu_seqlens_cache.py checks device_count() first. - Add entries in L0_pytorch_unittest for test_qk_norm.py, test_float8_current_scaling_exact.py, attention/test_cu_seqlens_cache.py and layernorm_mlp/test_selective_activation_checkpoint.py. - Move test_nvfp4_fsdp2_hooks.py into tests/pytorch/nvfp4/, which L0 already runs as a whole directory. It has no local imports, so the move is inert. - Also run attention/test_cu_seqlens_cache.py in L1. Its cross-device test needs two GPUs and would otherwise stay permanently skipped, which would defeat the point of connecting the file - that test is the regression guard for #2728. test_fused_router_perf.py is left out. It is gated behind TE_RUN_PERF_TESTS and, despite its name, has no perf assertions at all - only torch.testing.assert_close on correctness, with timings going to record_property. Ungating the correctness half is worth doing but deserves its own decision. Measured on RTX 5880 Ada: 45 passed, 5 passed, 1 passed + 1 skipped, and 16 skipped respectively, about 13 s of added L0 time. Signed-off-by: Pawel Gadzinski --- qa/L0_pytorch_unittest/test.sh | 4 ++++ qa/L1_pytorch_distributed_unittest/test.sh | 1 + tests/pytorch/{ => nvfp4}/test_nvfp4_fsdp2_hooks.py | 0 3 files changed, 5 insertions(+) rename tests/pytorch/{ => nvfp4}/test_nvfp4_fsdp2_hooks.py (100%) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 39d3e79e62..973077bf4e 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -40,11 +40,14 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_torch_compile.xm python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8blockwisetensor.xml $TE_PATH/tests/pytorch/test_float8blockwisetensor.py || test_fail "test_float8blockwisetensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_scaling_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_scaling_exact.py || test_fail "test_float8_blockwise_scaling_exact.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_gemm_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_gemm_exact.py || test_fail "test_float8_blockwise_gemm_exact.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_current_scaling_exact.xml $TE_PATH/tests/pytorch/test_float8_current_scaling_exact.py || test_fail "test_float8_current_scaling_exact.py" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/test_grouped_tensor.xml $TE_PATH/tests/pytorch/test_grouped_tensor.py || test_fail "test_grouped_tensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gqa.xml $TE_PATH/tests/pytorch/test_gqa.py || test_fail "test_gqa.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_qk_norm.xml $TE_PATH/tests/pytorch/test_qk_norm.py || test_fail "test_qk_norm.py" 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_selective_activation_checkpoint.xml $TE_PATH/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py || test_fail "test_selective_activation_checkpoint.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_distributed_weight.xml $TE_PATH/tests/pytorch/test_distributed_weight.py || test_fail "test_distributed_weight.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_override.xml $TE_PATH/tests/pytorch/test_backward_override.py || test_fail "test_backward_override.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" @@ -56,6 +59,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_flex_attention.x NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_linear_mxfp8_attention.xml $TE_PATH/tests/pytorch/attention/test_linear_mxfp8_attention.py || test_fail "test_linear_mxfp8_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml $TE_PATH/tests/pytorch/attention/test_cu_seqlens_cache.py || test_fail "test_cu_seqlens_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" export NVTE_TEST_CHECKPOINT_ARTIFACT_PATH=$TE_PATH/artifacts/tests/pytorch/test_checkpoint if [ ! -d "$NVTE_TEST_CHECKPOINT_ARTIFACT_PATH" ]; then diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index 50a51353d1..c59aa9af6d 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -48,6 +48,7 @@ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2.xml $TE_ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_comm_gemm_overlap.xml $TE_PATH/tests/pytorch/distributed/test_comm_gemm_overlap.py || test_fail "test_comm_gemm_overlap.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops_with_userbuffers.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py || test_fail "test_fusible_ops_with_userbuffers.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cp_utils.xml $TE_PATH/tests/pytorch/attention/test_cp_utils.py || test_fail "test_cp_utils.py" +python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml $TE_PATH/tests/pytorch/attention/test_cu_seqlens_cache.py || test_fail "test_cu_seqlens_cache.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cast_master_weights_to_fp8.xml $TE_PATH/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py || test_fail "test_cast_master_weights_to_fp8.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_newton_schulz.xml $TE_PATH/tests/pytorch/distributed/test_newton_schulz.py || test_fail "test_newton_schulz.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_ep.xml $TE_PATH/tests/pytorch/distributed/test_ep.py || test_fail "test_ep.py" diff --git a/tests/pytorch/test_nvfp4_fsdp2_hooks.py b/tests/pytorch/nvfp4/test_nvfp4_fsdp2_hooks.py similarity index 100% rename from tests/pytorch/test_nvfp4_fsdp2_hooks.py rename to tests/pytorch/nvfp4/test_nvfp4_fsdp2_hooks.py