Add forward all-gather overlap to experimental FSDP

[wujingyue]

Summary

• Split the all-gather overlap portion from Add FSDP stream context delayed release #5124 onto current main.
• Add an FsdpContext that lazily owns a CUDA all-gather stream for one FSDP subtree.
• Run parameter unshard/all-gather on that context stream while compute remains on the default stream.
• Delay release of unsharded parameter storage so child-unit forward all-gathers can overlap default-stream GEMMs under a shared root context.
• Add context-sharing and memory-bound tests for nested FSDP units, plus a profiler test that asserts the forward child pipeline gives num_children - 1 all-gather/compute overlaps.

Scope

• This PR intentionally covers forward all-gather overlap only.
• Backward all-gathers are present in the profiled iteration, but they do not overlap compute in this PR because gradient reduction is still launched synchronously in post_backward before autograd reaches the next module’s pre_backward all-gather.
• The delayed gradient-reduction follow-up PR will address backward overlap.
• Reduce-scatter overlap is not included here.

Testing

• python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/megatron_fsdp/test_context.py tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py --tb=short --disable-warnings -rN
• git diff --check

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[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