From b5067b86c79f1059881607c7e1ca9417e4962a57 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Wed, 29 Jul 2026 15:12:59 +0200 Subject: [PATCH 1/6] [PyTorch] Fix selective activation checkpointing test and add it to L0 tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py has never been run by any script in qa/, and it fails on all 16 parametrizations on current 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. The memory assertion used a hardcoded 6x ratio. The measured ratio is a structural constant of 5.715-5.719, independent of seq_len (128..65536) and hidden size (128..2048), so the threshold was simply unreachable. Changes: - Assert on the memory actually freed by recompute (fc1_out + act_out per layer, derived from the config) instead of a magic ratio. - Check outputs and gradients before the memory check, so a numerical regression cannot be masked by a memory/perf failure. - Drop the bare `ln_bwd_time < sln_bwd_time` assertion. The margin is as low as 13% on an idle GPU, which makes it a CI flake. Timings and memory ratios are reported via record_property instead. - Skip parametrizations that do not fit in device memory. large/65536 and huge/65536 need >32 GiB for the non-checkpointed model alone and OOM on 48 GiB cards. - Run the file in qa/L0_pytorch_unittest. Verified on RTX 5880 Ada: 12 passed, 4 skipped on memory. Signed-off-by: Pawel Gadzinski --- qa/L0_pytorch_unittest/test.sh | 1 + .../test_selective_activation_checkpoint.py | 81 ++++++++++++++----- 2 files changed, 62 insertions(+), 20 deletions(-) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 5d767ba4d1..6901d3c293 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -45,6 +45,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_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" diff --git a/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py b/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py index 306d0627f5..02923f3673 100644 --- a/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py +++ b/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py @@ -137,12 +137,49 @@ 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, per layer.""" + return cfg._layers * 2 * seq_size * cfg._ffn_hidden_size * itemsize + + +GRAD_KEYS = [ + "layer_norm_weight", + "layer_norm_bias", + "fc1_weight", + "fc1_bias", + "fc2_weight", + "fc2_bias", +] + + @pytest.mark.parametrize("size", config.keys()) @pytest.mark.parametrize("seq_size", seq_sizes) -def test_selective_activation_checkpoint(size, seq_size): +def test_selective_activation_checkpoint(size, seq_size, record_property): - ln_model, sln_model = config[size].build() - data = torch.randn((seq_size, config[size]._hidden_size), device=device) + cfg = config[size] + itemsize = torch.empty((), dtype=torch.get_default_dtype()).element_size() + no_ckpt_bytes = _no_checkpoint_activation_bytes(cfg, 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 = cfg.build() + data = torch.randn((seq_size, cfg._hidden_size), device=device) _warmup(ln_model, data) ln_fwd_out, ln_fwd_time, ln_fwd_mem = _run_fwd(ln_model, data) @@ -152,24 +189,28 @@ 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 [ - "layer_norm_weight", - "layer_norm_bias", - "fc1_weight", - "fc1_bias", - "fc2_weight", - "fc2_bias", - ]: + for key in GRAD_KEYS: 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(cfg, 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})" + ) + + # Checkpointing trades backward time for memory, but wall-clock is too noisy in + # CI to assert on - report it instead. + record_property("ln_fwd_ms", round(ln_fwd_time, 3)) + record_property("sln_fwd_ms", round(sln_fwd_time, 3)) + record_property("ln_bwd_ms", round(ln_bwd_time, 3)) + record_property("sln_bwd_ms", round(sln_bwd_time, 3)) + record_property("fwd_mem_ratio", round(ln_fwd_mem / sln_fwd_mem, 3)) + record_property("bwd_mem_ratio", round(ln_bwd_mem / sln_bwd_mem, 3)) From c5a1d80a3ff78dd6b5bbf78871d9cf23c4aec317 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 30 Jul 2026 07:49:54 +0200 Subject: [PATCH 2/6] [CI] Run four orphaned pytorch test files in L0 These files are not referenced by any script in qa/ and have therefore never run in CI. Confirmed with `git log -S --all -- qa/`: none of them was ever added and later removed, and none of the PRs that introduced them touched qa/. - test_qk_norm.py, test_float8_current_scaling_exact.py and attention/test_cu_seqlens_cache.py get an entry in L0_pytorch_unittest. - test_nvfp4_fsdp2_hooks.py moves into tests/pytorch/nvfp4/, which L0 already runs as a whole directory. No qa/ change needed, and the file is now covered by the same rule as the other NVFP4 tests. All four are single-GPU and self-skip on unsupported hardware: test_float8_current_scaling_exact.py guards 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. Measured on RTX 5880 Ada: 45 passed, 5 passed, 1 passed + 1 skipped, and 16 skipped respectively - about 13 s in total. tests/pytorch/test_fused_router_perf.py is deliberately left out; it is gated behind TE_RUN_PERF_TESTS and needs a separate decision. Signed-off-by: Pawel Gadzinski --- qa/L0_pytorch_unittest/test.sh | 3 +++ tests/pytorch/{ => nvfp4}/test_nvfp4_fsdp2_hooks.py | 0 2 files changed, 3 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 6901d3c293..9dfe95346c 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -40,8 +40,10 @@ 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" @@ -57,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/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 From 367621d648f82152c602aef842d8e059eb191d24 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 30 Jul 2026 17:05:16 +0200 Subject: [PATCH 3/6] [CI] Also run test_cu_seqlens_cache.py in L1 The file has two tests. test_cu_seqlens_cache_isolated_across_devices_for_forward needs two CUDA devices and therefore always skips in L0, which is where the file was just wired in. That test is the actual regression guard for #2728 - the cu_seqlens cache key not being scoped by device - so leaving it permanently skipped defeats the purpose of connecting the file at all. L1 is the only suite that guarantees more than one GPU. It is a plain pytest run, no torchrun, matching how attention/test_cp_utils.py is invoked there. Cost is about 2 s. Signed-off-by: Pawel Gadzinski --- qa/L1_pytorch_distributed_unittest/test.sh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index 50a51353d1..1f7d44a1a6 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -48,6 +48,8 @@ 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" +# Also run in L0, but the cross-device case there always skips on a single GPU. +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" From 09c8847da1af29f5d519b0b220d6cb5259dc7719 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 30 Jul 2026 17:06:17 +0200 Subject: [PATCH 4/6] [CI] Drop redundant comment in L1 test.sh Signed-off-by: Pawel Gadzinski --- qa/L1_pytorch_distributed_unittest/test.sh | 1 - 1 file changed, 1 deletion(-) diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index 1f7d44a1a6..c59aa9af6d 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -48,7 +48,6 @@ 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" -# Also run in L0, but the cross-device case there always skips on a single GPU. 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" From c46e1694795a7b45ccdcc3c43f5885640715ef83 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 30 Jul 2026 17:08:49 +0200 Subject: [PATCH 5/6] [PyTorch] Drop record_property from the SAC test, fix the saving threshold Two follow-ups on the selective activation checkpointing test. record_property was added here to keep reporting the backward timings after the flaky `ln_bwd_time < sln_bwd_time` assertion was removed. It was not in the original test, and it emits a PytestWarning on every run because it is not compatible with the default xunit2 junit family. Remove it; the timings were never asserted on and nothing consumes them. The memory threshold divided by the full layer count, but the checkpointed peak still holds the transient of one layer, so recompute only saves the ffn-sized activations of the remaining layers. With `layers`, the ratio of measured to expected saving is (L-1)/L * (1 + h/2f), which happens to clear 0.95 at L=12 (1.031) but would fail at L=4 (0.844) - the threshold silently encoded the shape of the test models. With `layers - 1` the ratio is 1 + h/2f, always above 1 regardless of layer count, sequence length and hidden size, so 0.95 is what it was meant to be: slack for allocator noise. Verified on RTX 5880 Ada: 12 passed, 4 skipped, no PytestWarning. Signed-off-by: Pawel Gadzinski --- .../test_selective_activation_checkpoint.py | 20 ++++++++----------- 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py b/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py index 02923f3673..e5889f51b7 100644 --- a/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py +++ b/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py @@ -148,8 +148,13 @@ def _no_checkpoint_activation_bytes(cfg, seq_size, itemsize): def _recomputed_activation_bytes(cfg, seq_size, itemsize): - """Activations checkpointing must free: fc1_out and act_out, per layer.""" - return cfg._layers * 2 * seq_size * cfg._ffn_hidden_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 GRAD_KEYS = [ @@ -164,7 +169,7 @@ def _recomputed_activation_bytes(cfg, seq_size, itemsize): @pytest.mark.parametrize("size", config.keys()) @pytest.mark.parametrize("seq_size", seq_sizes) -def test_selective_activation_checkpoint(size, seq_size, record_property): +def test_selective_activation_checkpoint(size, seq_size): cfg = config[size] itemsize = torch.empty((), dtype=torch.get_default_dtype()).element_size() @@ -205,12 +210,3 @@ def test_selective_activation_checkpoint(size, seq_size, record_property): f" {saving} B, expected at least {0.95 * expected_saving} B (ln_fwd_mem={ln_fwd_mem}," f" sln_fwd_mem={sln_fwd_mem})" ) - - # Checkpointing trades backward time for memory, but wall-clock is too noisy in - # CI to assert on - report it instead. - record_property("ln_fwd_ms", round(ln_fwd_time, 3)) - record_property("sln_fwd_ms", round(sln_fwd_time, 3)) - record_property("ln_bwd_ms", round(ln_bwd_time, 3)) - record_property("sln_bwd_ms", round(sln_bwd_time, 3)) - record_property("fwd_mem_ratio", round(ln_fwd_mem / sln_fwd_mem, 3)) - record_property("bwd_mem_ratio", round(ln_bwd_mem / sln_bwd_mem, 3)) From 88ab87670bc3d36cf1c5a8e5de95b624e9b2e52a Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 30 Jul 2026 17:13:25 +0200 Subject: [PATCH 6/6] [PyTorch] Drop incidental churn from the SAC test change Restore the inline gradient key list and the config[size] accesses that the previous commit had refactored for no reason. They were unrelated to the fix and only made the diff harder to read. No behaviour change: 12 passed, 4 skipped on RTX 5880 Ada, same as before. Signed-off-by: Pawel Gadzinski --- .../test_selective_activation_checkpoint.py | 28 ++++++++----------- 1 file changed, 12 insertions(+), 16 deletions(-) diff --git a/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py b/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py index e5889f51b7..34aaf32ec9 100644 --- a/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py +++ b/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py @@ -157,23 +157,12 @@ def _recomputed_activation_bytes(cfg, seq_size, itemsize): return (cfg._layers - 1) * 2 * seq_size * cfg._ffn_hidden_size * itemsize -GRAD_KEYS = [ - "layer_norm_weight", - "layer_norm_bias", - "fc1_weight", - "fc1_bias", - "fc2_weight", - "fc2_bias", -] - - @pytest.mark.parametrize("size", config.keys()) @pytest.mark.parametrize("seq_size", seq_sizes) def test_selective_activation_checkpoint(size, seq_size): - cfg = config[size] itemsize = torch.empty((), dtype=torch.get_default_dtype()).element_size() - no_ckpt_bytes = _no_checkpoint_activation_bytes(cfg, seq_size, itemsize) + 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) @@ -183,8 +172,8 @@ def test_selective_activation_checkpoint(size, seq_size): f" {free_bytes / 2**30:.1f} GiB available" ) - ln_model, sln_model = cfg.build() - data = torch.randn((seq_size, cfg._hidden_size), device=device) + ln_model, sln_model = config[size].build() + data = torch.randn((seq_size, config[size]._hidden_size), device=device) _warmup(ln_model, data) ln_fwd_out, ln_fwd_time, ln_fwd_mem = _run_fwd(ln_model, data) @@ -198,12 +187,19 @@ def test_selective_activation_checkpoint(size, seq_size): # 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 GRAD_KEYS: + for key in [ + "layer_norm_weight", + "layer_norm_bias", + "fc1_weight", + "fc1_bias", + "fc2_weight", + "fc2_bias", + ]: 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(cfg, seq_size, itemsize) + 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"