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/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})" + ) 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