[PyTorch][CI] Connect test files that are not wired into qa, fix the broken one - #3287
Open
pggPL wants to merge 2 commits into
Open
[PyTorch][CI] Connect test files that are not wired into qa, fix the broken one#3287pggPL wants to merge 2 commits into
pggPL wants to merge 2 commits into
Conversation
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 <pgadzinski@nvidia.com>
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<name> --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 NVIDIA#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 <pgadzinski@nvidia.com>
Contributor
Greptile SummaryThis PR connects previously uncollected PyTorch tests to CI and repairs the selective activation-checkpointing test.
Confidence Score: 5/5The PR appears safe to merge, with the newly wired tests correctly collected and no actionable regressions identified. The QA invocations match existing orchestration patterns, the multi-device cache test works through ordinary visible CUDA devices, the moved NVFP4 test has no location-sensitive dependencies, and the revised checkpointing assertions reflect the configured activation allocations. Important Files Changed
Reviews (1): Last reviewed commit: "[CI] Connect orphaned pytorch test files..." | Re-trigger Greptile |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Six pytorch test files are not referenced by any script in
qa/, so they have never run in CI.git log -S<name> --all -- qa/shows that none of them was ever added and later removed, and noneof the PRs that introduced them touched
qa/— they were simply never wired up.maintest_qk_norm.pytest_float8_current_scaling_exact.pyattention/test_cu_seqlens_cache.pytest_nvfp4_fsdp2_hooks.pytest_fused_router_perf.pyTE_RUN_PERF_TESTS)layernorm_mlp/test_selective_activation_checkpoint.pyThis PR connects four of them and fixes the one that turned out to be broken. All are single-GPU
and self-skip on unsupported hardware. Added L0 time is roughly 13 s.
test_fused_router_perf.pyis left out — it is gated behindTE_RUN_PERF_TESTSand deserves itsown decision.
Type of change
Changes
test_qk_norm.py,test_float8_current_scaling_exact.py,attention/test_cu_seqlens_cache.pyandlayernorm_mlp/test_selective_activation_checkpoint.py.test_nvfp4_fsdp2_hooks.pyintotests/pytorch/nvfp4/, which L0 already runs as a wholedirectory.
attention/test_cu_seqlens_cache.pyin L1, where its cross-device test can actually runinstead of being permanently skipped.
test_selective_activation_checkpoint.py. The feature is fine — outputs and all sixparameter gradients are bit-exact between the checkpointed and non-checkpointed path. The test
asserted a
6xforward memory ratio that the model shape in the test can never reach; it nowasserts on the memory that recompute actually frees. Also drops a flaky wall-clock assertion,
checks correctness before memory, and skips configurations that do not fit in device memory.
The commit messages carry the full derivation and measurements.
Checklist: