From 86cf4ff0fd517e87440efd53910ac75c12c4971d Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sun, 12 Jul 2026 21:47:47 -0700 Subject: [PATCH] Increase Megatron-FSDP overlap test dim to 8192 for reliable overlap test_overlaps_all_gather_and_compute asserts that at least num_children - 1 forward all-gathers overlap GEMM compute. At dim = 4096 the GEMM kernels are short enough that kernel launch and dispatch latency dominate, so the independent all-gather and GEMM kernels do not reliably co-reside on the GPU. On a single DGX H100 node (8xH100) with default NCCL settings the test fails nearly every run, because any one of the eight ranks missing an overlap fails the whole test. Raising the GEMM size to dim = 8192 makes each kernel long enough to amortize launch latency, so the forward all-gathers reliably overlap the preceding child's GEMM. Verified on 8xH100 under default NCCL: 25/25 runs pass, versus 0/12 at dim = 4096. The flakiness does not surface in CI because tests/unit_tests/run_ci_test.sh exports NCCL_MAX_NCHANNELS=1 to reduce NCCL memory. That serializes the all-gather onto a single channel, producing a longer, lighter-weight all-gather kernel that happens to overlap the small dim = 4096 GEMM. Isolated on 8xH100: NCCL_MAX_NCHANNELS=1 alone passes 10/10, while default NCCL fails 0/12. The test was therefore implicitly depending on a memory-reduction environment variable unrelated to the feature under test. Increasing dim removes that hidden dependency so the test passes under both default and CI NCCL configurations. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jingyue Wu --- tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py index 229cd0bff4b..ff892ec74ee 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -319,7 +319,7 @@ def test_overlaps_all_gather_and_compute(distributed_setup): pytest.skip("This test requires at least 2 ranks.") mesh = init_device_mesh(device.type, (world_size,)) - dim = 4096 + dim = 8192 num_children = 4 dtype = torch.bfloat16 model = MultiChildModel(dim=dim, num_children=num_children).to(dtype=dtype)