[albmalamd]

Summary:

Threaded PG performs the data movement for all ranks on rank 0's thread, hence on rank 0's current stream. The current stream is thread-local, so a collective issued from a side stream is unordered w.r.t. rank 0's copies: rank 0 can read a rank's input before its copy-in has run, and a rank can read its output before rank 0 has written it. FSDP2 hits this on every collective because

FSDPCommContext.lazy_init gives each rank private all-gather and reduce-scatter streams. The host-side ordering in Collective.join was already correct, which is why the race went unnoticed for so long: it resolves favourably on NVIDIA CI and loses on ROCm, where it surfaces as check_sharded_parity gradient mismatches and diverging losses in the FSDP2 and replicate tests (ROCM-28104).

A real backend is stream-ordered w.r.t. the stream that issued the collective, and Collective.join now provides the same guarantee using events: every non-root rank records an event on the stream it issued from, rank 0 waits on those events before copying and records one afterwards, and the other ranks wait on that. The path is a no-op for CPU tensors. test_collectives_issued_from_side_stream covers the same ordering without involving FSDP.

Reproduction:

The affected tests pass by luck on a fast enough GPU, so delay the producing side. In test_param_registration_after_forward, right before model(inp):

if self.rank != 0:
    device_sleep(device_type.type, int(50 * get_cycles_per_ms(device_type.type)))

FSDP2's side streams wait on the rank's compute stream, so the sleep holds back the collective's inputs while rank 0 runs ahead. Without this change the test fails with 9/36 elements mismatched; with it, it passes. test_to_float64_after_init reproduces the same way with the sleep at the top of the training loop.