MFSDP v2: overlap DP-inner and DP-outer communication

[wujingyue]

Follow-up to #5616: explore overlapping DP-inner and DP-outer communication in MFSDP v2.

The two communications use different network domains, and early experience suggests that overlapping them could reduce latency significantly. The implementation should follow the runtime scheduling model in megatron/core/distributed/fsdp/src/docs/runtime_schedule.md.

We should first measure the benefit and use those results to guide the implementation and scheduling design.

[Achyuthan-S]

@wujingyue Thanks — I'll take this. Agreed on measure-first: I'll quantify the DP-outer reduction's current cost on the critical path and the latency a separate DP-outer stream recovers, on an HSDP/HFSDP config following runtime_schedule.md, and post numbers here before proposing the implementation. I'll prototype the separate stream just enough to get a real before/after, then let the results shape the production design.

[Achyuthan-S]

Hey @wujingyue ,
Measured the DP-inner/DP-outer collective-overlap benefit (measure-first). Setup: 2×2 DP mesh, distinct NCCL communicators per axis (mesh.get_group), DP-inner reduce-scatter vs DP-outer all-reduce on separate CUDA streams, CUDA-event timed.

Single-node, 4×A100-80GB / 4×H100-80GB:

size	serial	overlapped	savings
64 MiB	1.65 / 1.07 ms	1.05 / 0.73 ms	37% / 31%
256 MiB	6.14 / 3.83 ms	3.92 / 2.56 ms	36% / 33%
1 GiB	23.3 / 14.4 ms	14.8 / 9.8 ms	37% / 32%

Overlapped ≈ max(RS, AR) — the reduce-scatter hides fully behind the all-reduce, recovering ~⅓ of combined collective time, consistent across sizes and both GPUs. In the actual HSDP path the DP-outer all-reduce is fully serialized behind DP-inner today (0/N overlap in the profiler, both on the reduce-scatter stream), so that third is entirely on the critical path.

Caveats: single-node (both axes on NVLink) — they overlap cleanly here with little contention; a multi-node run (DP-outer cross-node on IB) is the natural follow-up for the cross-domain number. And the microbench runs both collectives every iteration, whereas in training DP-outer is last-microbatch only, so the realized per-step win needs the actual separate-stream implementation to pin down.

Given the collectives clearly overlap on distinct communicators+streams and are serialized today, I'd like to proceed to the separate DP-outer stream implementation. Does that sequencing work, or would you prefer a multi-node measurement first?

[wujingyue]

I'd like to proceed to the separate DP-outer stream implementation.

Let's do that. Then, I can measure the impact on some larger, more realistic settings.