Separate mFSDP v2 unit tests

[wujingyue]

Summary

Split the Megatron-FSDP unit tests into explicit v1 and v2 buckets:

• tests/unit_tests/distributed/mfsdp_v1/**/*.py
• tests/unit_tests/distributed/mfsdp_v2/**/*.py

The v2 / experimental tests now live in their own mfsdp_v2 bucket, with a duplicated local conftest.py and a dedicated H100 unit-test recipe entry. The remaining v1 tests now live in mfsdp_v1, with the H100 recipe and docs references updated to match.

Moved v2 tests:

• test_cuda_graph.py
• test_dbuffer.py
• test_fully_shard.py (renamed from test_experimental_fully_shard.py because the folder now carries the v2 context)
• test_symmetric_memory.py

Renamed v1 bucket files:

• conftest.py
• test_mcore_fully_sharded_data_parallel.py
• test_mcore_tensor_parallelism_detect.py
• test_mfsdp_fully_shard.py
• test_mfsdp_uneven_dtensor.py
• utils.py

Why

The biggest motivation is CI signal: the v1 bucket is much slower and much noisier than the v2 tests. In the final split run, the v1 bucket took 14m19s wall time and reported 574 warnings, while the standalone v2 bucket took 3m28s wall time and reported 43 warnings.

The previous mixed Megatron-FSDP CI bucket combined two different kinds of coverage:

• v1 mFSDP tests (test_mfsdp_fully_shard.py, test_mfsdp_uneven_dtensor.py, test_mcore_fully_sharded_data_parallel.py, etc.)
• v2 / experimental FSDP tests (test_dbuffer.py, test_experimental_fully_shard.py, test_cuda_graph.py, test_symmetric_memory.py)

Keeping v2 coverage in the same bucket makes warning growth and runtime changes harder to attribute. One old combined-bucket sample attributed a 145.20s teardown to v2 test_symmetric_memory.py, but that test had run in the same distributed pytest invocation as the v1 files. The split-bucket CI result below shows that attribution was not representative of standalone v2 runtime.

Renaming the remaining v1 bucket avoids a layout where the v1 path is implicit while v2 is explicit. Both buckets now carry the version in the path.

Reference split-bucket CI data from run 28675683046:

Bucket	Job	Wall time	Main pytest result	Slowest duration
tests/unit_tests/distributed/mfsdp_v2/**/*.py - latest	85055242029	3m28s	36 passed, 43 warnings in 12.66s	2.86s teardown tests/unit_tests/distributed/mfsdp_v2/test_symmetric_memory.py::test_fully_shard_symmetric_memory_matches_default_and_profiles_nccl[3]
tests/unit_tests/distributed/mfsdp_v1/**/*.py - latest	85055242017	14m19s	251 passed, 121 skipped, 21 deselected, 6 xfailed, 574 warnings in 656.01s (0:10:56)	147.32s teardown tests/unit_tests/distributed/mfsdp_v1/test_mfsdp_uneven_dtensor.py::test_split_dtensor_zero_local_shard

Current v2 warning state:

• The v2 bucket now passes CUDA devices explicitly to dist.barrier(...) in both function-scoped and session-scoped cleanup, removing the repeated PyTorch NCCL barrier-warning noise.
• CI run 28675683046, job 85055242029, validated the final cleanup with 36 passed, 43 warnings in 12.66s.
• The current 43 warnings are 34 import/collection baseline warnings, 8 PyTorch module-backward-hook warnings from FSDP training-style v2 tests, and 1 PyTorch profiler cycle warning from test_symmetric_memory.py. No torch.distributed barrier warning remains in the final v2 bucket run.

Splitting the bucket does not fix the existing v1 warnings or teardown behavior. It makes ownership and regression signal explicit: v1 warnings/teardown remain in the mfsdp_v1 bucket, while v2 warnings/teardown show up in the mfsdp_v2 bucket with separate CI timing.

This PR only separates the test buckets, cleans the barrier warning introduced by the new v2 fixture boundary, and renames the v1 bucket for clarity. It does not change test logic.

Validation

• uv run --no-sync python -m pytest --collect-only -q tests/unit_tests/distributed/mfsdp_v2
  • collected 36 tests
• uv run --no-sync pytest --collect-only -q tests/unit_tests/distributed/mfsdp_v1
  • collected 399 tests
• python tests/unit_tests/find_test_cases.py 'tests/unit_tests/**/*.py' h100 | rg 'mfsdp_v1|mfsdp_v2'
  • confirmed the catch-all H100 unit-test bucket ignores both dedicated FSDP buckets
• python tests/unit_tests/find_test_cases.py 'tests/unit_tests/distributed/mfsdp_v1/**/*.py' h100 && python tests/unit_tests/find_test_cases.py 'tests/unit_tests/distributed/mfsdp_v2/**/*.py' h100
  • confirmed neither dedicated bucket ignores files from itself
• GitHub Actions run 28675683046
  • tests/unit_tests/distributed/mfsdp_v2/**/*.py - latest: success after session-scoped barrier cleanup, with 43 warnings
  • tests/unit_tests/distributed/mfsdp_v1/**/*.py - latest: success

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[wujingyue]

/ok to test f3cd82f

[wujingyue]

/ok to test 86c7085

[wujingyue]

/ok to test 1e4e6c3

[wujingyue]

/ok to test 9b8be36

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28920782182

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28926630617