Skip to content

[PyTorch][CI] Connect test files that are not wired into qa, fix the broken one #3287

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

Merged
pggPL merged 2 commits into from
Jul 31, 2026
Merged

[PyTorch][CI] Connect test files that are not wired into qa, fix the broken one #3287

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2387f4d
[PyTorch] Fix the selective activation checkpointing test
pggPL Jul 30, 2026
da1cadc
[CI] Connect orphaned pytorch test files to qa
pggPL Jul 30, 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
4 changes: 4 additions & 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,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"
Expand All @@ -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
Expand Down
1 change: 1 addition & 0 deletions qa/L1_pytorch_distributed_unittest/test.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
51 changes: 42 additions & 9 deletions tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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 [
Expand All @@ -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})"
)
0 tests/pytorch/test_nvfp4_fsdp2_hooks.py → ...s/pytorch/nvfp4/test_nvfp4_fsdp2_hooks.py
View file
Open in desktop
File renamed without changes.
Loading