fix(trainer): stage CPU-offloaded weights through CUDA in NCCL broadcast#3045
fix(trainer): stage CPU-offloaded weights through CUDA in NCCL broadcast#3045mikasenghaas wants to merge 4 commits into
Conversation
|
Superseded by the removal PR — team decided to remove fsdp_cpu_offload entirely rather than fix it. |
|
Reopening: #3049 (fsdp_cpu_offload removal) is being reverted in #3052 after 2-node long-context measurements, so this staging fix is needed again — it's the only fast weight-broadcast path with FSDP CPU offload (filesystem broadcast full-gathers CPU-resident fp32 weights over Gloo and effectively hangs; measured on a Nemotron Super 131k run). |
|
Second fix pushed (a6bdfe8): the original staging happened after |
|
Status: the per-tensor CUDA staging (a6bdfe8) is insufficient — |
With fsdp_cpu_offload the gathered state-dict tensors live on CPU, and broadcast_state_dict passed them to the NCCL communicator directly: "AssertionError: this nccl communicator is created to work on cuda:0, but the input tensor is on cpu". Move the per-layer, per-dtype concatenated buffer to CUDA before broadcasting. Needed to combine fsdp_cpu_offload (the memory profile that fits 196k on 2 nodes) with NCCL weight broadcast in RL. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
_resolve_dtensors called full_tensor() on CPU DTensors, so the gather ran over Gloo — hours for 100B+ fp32 params (observed: Nemotron Super 131k run stalled ~1h at the v0 broadcast until the NCCL watchdog killed it). Move each local shard to CUDA before full_tensor() so the gather runs over NCCL, one tensor at a time; the downstream concatenated-buffer staging then no-ops. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
DTensor.to("cuda") moves the local shard but collectives still route by the
original backend selection, so full_tensor() on CPU-offloaded params gathered
over Gloo regardless (observed: ranks stuck hours in the v0 broadcast).
Reconstruct the DTensor via from_local with the CUDA-staged shard on the same
(CUDA) device mesh so the gather dispatches to NCCL.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Per-tensor synchronous H2D with CPU-side bf16 casts made the v0 broadcast latency-dominated (minutes). FSDP CPU offload pins the shards, so stage them with non_blocking H2D and cast to bf16 on GPU after the move; the NCCL gather serializes correctly behind the async copies on the same stream. Rebased onto main (includes the uint16 routed-experts wire format). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
92c891c to
bb8a777
Compare
|
Deprioritized on our side: even with async pinned staging (bb8a777, unvalidated — test run cancelled), FSDP CPU offload implies re-staging the full CPU-resident weights every step's broadcast, a recurring per-step tax the team judged not worth it for RL-with-inference. The branch is rebased onto main with the mesh-correct NCCL gather + async staging if anyone wants to pick it up; trainer-only/SFT users of fsdp_cpu_offload are unaffected (no broadcast). |
1 participant
Follow-up to #3042 (found during bring-up of the first run to combine the two settings; the commit just missed the merge).
With
fsdp_cpu_offloadthe gathered state-dict tensors live on CPU, andbroadcast_state_dictpassed them to the NCCL communicator directly:Pre-existing: RL never combined
fsdp_cpu_offload(params on CPU) with NCCL weight broadcast before — muon (the usual RL optimizer) rejects FSDP CPU offload, and it's the memory profile that makes 196k-context NemotronH runs fit on 2 nodes (adamw). Stage the per-layer, per-dtype concatenated buffer through CUDA before broadcasting (+4 lines; a transient per-layer GPU buffer, ~3 GB max for a 120B/88-layer model).Verified: the 196k cp=8 Nemotron Super RL run that hit this assert proceeds past weight broadcast with the fix (run in flight — will confirm rollouts/step 1 here).
🤖 Generated with Claude Code