fix(fsdp): stage weight-sync buckets in a persistent buffer to stop CUDA-IPC pinning

[zhihengy]

Problem

DiffusionUpdateWeightFromTensor.update_bucket_weights serializes each bucket's flattened tensor with MultiprocessingSerializer, which exports a CUDA-IPC handle for the underlying storage. The rollout engine imports the handle and — like any long-lived torch consumer — keeps the mapping open instead of closing it per request (open/close is an expensive driver call, so torch caches imports). A storage whose handle is still held by a peer can never be reclaimed on the producer side: torch.cuda.ipc_collect() only sweeps storages whose import refcount already dropped to zero. We had del + ipc_collect on the trainer side and gc.collect + ipc_collect + empty_cache on the engine side; none of it releases a single byte while the peer holds the handle.

Since every bucket cat-ed a fresh flattened tensor, every bucket exported a new storage, and each one stayed pinned in the trainer process. A full sync accumulates one model's worth of trainer memory per rank: ~32GB for a 16B model (silently eating a quarter of an H200), ~128GB for a 64B model — the trainer OOMs partway through its first sync. Three consecutive OOMs during the 64B bring-up landed at exactly the same allocation (67.5GB ≈ 34 buckets), which is what pointed at the sync path.

Fix

Keep one persistent staging buffer per parameter dtype. Each bucket is copied into the staging buffer and the serialized payload references a view of it. The whole run therefore exports exactly one storage per dtype; the engine's cached import maps that same storage forever, and the transient cat result — never exported — is freed normally each bucket. Overwrite safety across buckets comes from the existing ray.get on update_weights_from_tensor: the engine has finished reading before the next bucket is staged.

Two details that fell out of testing:

• Key by the group's parameter dtype, not flat.dtype. FlattenedTensorBucket flattens every dtype group to a byte view, so flat.dtype is always uint8 and all groups of a mixed-dtype bucket would collapse onto one buffer. Groups are serialized before any payload is sent, so the second group's copy overwrites the first payload's bytes while the engine may still read them. Worse, visibility of writes made after serialization is undefined across processes (the IPC event is recorded at serialize time) — in testing the engine read NaN weights on some buckets and stale-but-clean bytes on others. Keying by parameter dtype gives concurrent payloads disjoint buffers. All-fp32/all-bf16 state dicts produce one group per bucket and never hit this, but configs that load a frozen base in bf16 do: diffusers' _keep_in_fp32_modules keeps e.g. time_embedder in fp32 (Cosmos3OmniTransformer, WanTransformer3DModel pins five module names), so the bucket holding those params carries two dtype groups.
• Preallocate each buffer to its bound instead of growing. Replacing an already-exported buffer orphans the old storage behind the engine's cached handle — a bounded replay of the original leak. The first allocation is sized to a provable upper bound of any bucket's payload for that dtype: min(state-dict per-dtype total, max(--update-weight-buffer-size, largest param)). The bound was checked against a brute-force replay of the bucketing loop on mixed-dtype state dicts.

Also logs trainer memory_allocated every 16 buckets ([wsync-mem]) so a regression here shows up as a rising line instead of a mystery OOM.

Benchmarks

A/B harness driving the real update_bucket_weights (this branch vs its parent) over ray against a consumer that replicates the sgl-d WeightsUpdater flattened-bucket path, including its end-of-call gc.collect/ipc_collect/empty_cache. 12 buckets x 2GB bf16, single H200, payload checksums verified:

version	engine holds refs to last N payloads	trainer memory	steady ms/bucket
parent	N=4	+2GB/bucket until plateau at 8GB (= N x bucket)	57.4
parent	N=unbounded	+2GB/bucket, linear, no plateau (+22GB over 12)	59.5
this PR	N=4	flat 2GB	60.2
this PR	N=unbounded	flat 2GB	57.7

The leak is exactly (engine-side reference depth) x (bucket size); the staged version is flat regardless of engine behavior because every held payload references the same storage. Latency is unchanged within noise — the staging copy is a 2GB device-to-device copy, ~0.7ms/bucket in an isolated microbench of the transfer path (same microbench puts the alternative fix, forcing the engine to close and reopen the handle every bucket, at 50-80ms/bucket).

Additional runs for the two details above:

• Alternating 1GB/2GB buckets (growth pressure): lazy growth orphans the replaced 1GB buffer permanently; with the preallocated bound the run is flat with a single 2GB buffer and zero orphans.
• Mixed bf16+fp32 buckets: with the shared uint8 buffer the engine's first-bucket weights came back with ~260k NaNs (and healed nondeterministically on later buckets — it's a race); with per-dtype keying, zero NaNs on every bucket, buffers land at 2GB (bf16) + 0.25GB (fp32).

In the real 64B run the fix took a full 128GB sync from OOM to a flat 3.9-6.8GB [wsync-mem] line, 25.5s per sync.

Notes

Split out of the cosmos3 branch: the bug affects any model large enough that a full sync's volume doesn't fit in spare GPU memory, and colocated small models merely hide it.