Allocate main_grad on the reduce-scatter stream

[wujingyue]

Root cause diagnosed by @Autumn1998 (Tong Liu) in #6137 (comment):

Partial buffer A
          │
          │ enqueue Reduce-Scatter (asynchronous)
          ▼
    Flat buffer B

    self.main_grad = B
          │
          └── A loses its last reference and is released by the allocator
                        │
                        ▼
                 A's storage is reused
                        │
                        ▼
           Reduce-Scatter is still reading A → corrupted gradients

FsdpParameterGroup allocated main_grad in __init__, which runs on whatever stream fully_shard() is called from — normally the default stream. reduce_partial_gradients then rebinds it from inside a reduce_scatter_stream context, dropping the old buffer while that stream is still reading it. The caching allocator binds a block to its allocation stream, so the dropped block became immediately reusable by default-stream allocations.

Background on the allocator's stream semantics and this free-while-in-use hazard: FSDP & CUDACachingAllocator: an outsider newb perspective. That write-up frames record_stream() as the mechanism for keeping a block alive across streams; this PR instead keeps allocation and free on a single stream, so the block only ever becomes reusable by allocations already ordered behind the reduce-scatter, and record_stream() is not needed.

The result is silently wrong gradients, not a crash: 12%–125% off per parameter, varying run to run, present at the end of the first backward. It surfaced in #5956 as flaky test_hsdp_losses_match_baseline failures.

Fix: materialize main_grad lazily from a property, so the first access — inside reduce_partial_gradients — allocates it on the same stream that later frees it. An assertion pins that invariant. Same allocate-and-release-on-one-stream rule FsdpModule._reshard_parameter_groups already documents for unsharded parameter storage.

Verified on 8×A100: full mfsdp_v2 bucket goes from 5/5 failing to 0/4 under a timing perturbation that makes the reuse collide reliably.

Also verified against #5956, the PR where this surfaced. That branch is flaky on its own, so no perturbation is needed — this is the real-world rate. Same 8×A100 node, arms alternated, full mfsdp_v2 bucket:

	failures
#5956 as-is	7/10
#5956 + this fix	0/10

The base-arm failures are the same signature as the original report — test_hsdp_losses_match_baseline in varying combinations of [1-True] (10), [3-True] (4), [1-False] (3), [3-False] (3).

[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.