From 4deb1511bdefff0b792c340c00963112c0e9a5d8 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 15 Jul 2026 22:55:04 +0000 Subject: [PATCH 1/4] fix(trainer): stage CPU-offloaded weights through CUDA in NCCL broadcast 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 --- src/prime_rl/trainer/rl/broadcast/nccl.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/prime_rl/trainer/rl/broadcast/nccl.py b/src/prime_rl/trainer/rl/broadcast/nccl.py index 75abf42195..8704396e90 100644 --- a/src/prime_rl/trainer/rl/broadcast/nccl.py +++ b/src/prime_rl/trainer/rl/broadcast/nccl.py @@ -61,6 +61,10 @@ def broadcast_state_dict(state_dict: dict[str, Tensor], communicator: PyNcclComm # Flatten all tensors and concatenate flat_tensors = [value.flatten() for _, value in items] concatenated = torch.cat(flat_tensors) + if not concatenated.is_cuda: + # fsdp_cpu_offload keeps parameters on CPU; NCCL requires device tensors, + # so stage the (per-layer, per-dtype) buffer through CUDA transiently. + concatenated = concatenated.cuda() communicator.broadcast(concatenated, src=0) del concatenated # Clean up individual tensors From 479cc9a1a5d459e981a59e651705437066026418 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 16 Jul 2026 10:39:21 +0000 Subject: [PATCH 2/4] fix: gather CPU-offloaded DTensor shards on CUDA MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _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 --- src/prime_rl/trainer/rl/broadcast/nccl.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/prime_rl/trainer/rl/broadcast/nccl.py b/src/prime_rl/trainer/rl/broadcast/nccl.py index 8704396e90..9d74b5dd2a 100644 --- a/src/prime_rl/trainer/rl/broadcast/nccl.py +++ b/src/prime_rl/trainer/rl/broadcast/nccl.py @@ -168,7 +168,13 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None: def _resolve_dtensors(self, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: for key, value in list(state_dict.items()): if isinstance(value, DTensor): - state_dict[key] = cast(DTensor, value.to(self.dtype)).full_tensor() + value = cast(DTensor, value.to(self.dtype)) + if value.device.type != "cuda": + # CPU-offloaded shards would full_tensor() over Gloo (hours for + # 100B+ params); move the local shard to CUDA so the gather runs + # over NCCL, one tensor at a time. + value = cast(DTensor, value.to("cuda")) + state_dict[key] = value.full_tensor() return state_dict From 3f07ecc66e0e45f566635a6155c0ff5e30f98c05 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 16 Jul 2026 20:43:12 +0000 Subject: [PATCH 3/4] fix: rebuild CPU-offloaded shards as CUDA DTensors before gather 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 --- src/prime_rl/trainer/rl/broadcast/nccl.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/prime_rl/trainer/rl/broadcast/nccl.py b/src/prime_rl/trainer/rl/broadcast/nccl.py index 9d74b5dd2a..19631aaf3c 100644 --- a/src/prime_rl/trainer/rl/broadcast/nccl.py +++ b/src/prime_rl/trainer/rl/broadcast/nccl.py @@ -169,11 +169,19 @@ def _resolve_dtensors(self, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: for key, value in list(state_dict.items()): if isinstance(value, DTensor): value = cast(DTensor, value.to(self.dtype)) - if value.device.type != "cuda": - # CPU-offloaded shards would full_tensor() over Gloo (hours for - # 100B+ params); move the local shard to CUDA so the gather runs - # over NCCL, one tensor at a time. - value = cast(DTensor, value.to("cuda")) + local = value.to_local() + if local.device.type != "cuda": + # FSDP CPU offload keeps the local shard on CPU while the device + # mesh is CUDA; full_tensor() would then gather over Gloo (hours + # for 100B+ params). Rebuild the DTensor from a CUDA-staged local + # shard so the gather runs over NCCL, one tensor at a time. + # (DTensor.to("cuda") does not migrate the collective backend.) + value = DTensor.from_local( + local.to("cuda", non_blocking=False), + device_mesh=value.device_mesh, + placements=value.placements, + run_check=False, + ) state_dict[key] = value.full_tensor() return state_dict From bb8a7773fda243c0b4b0e2708850f2a1f3056764 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 16 Jul 2026 21:56:39 +0000 Subject: [PATCH 4/4] fix: async pinned staging + GPU-side cast for offloaded shards 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 --- deps/renderers | 2 +- deps/verifiers | 2 +- src/prime_rl/trainer/rl/broadcast/nccl.py | 19 ++++++++----------- 3 files changed, 10 insertions(+), 13 deletions(-) diff --git a/deps/renderers b/deps/renderers index 201f88516c..fa04acd916 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 201f88516c012bf9f3e6204283f0fcc85c9b40ef +Subproject commit fa04acd916ce749c302eec905ecf3962bc438429 diff --git a/deps/verifiers b/deps/verifiers index 19fdd9cae1..1cfdd39cc2 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 19fdd9cae10ad26480ac74673267cfc046b8b299 +Subproject commit 1cfdd39cc234d53906a90322587ebc04d4239255 diff --git a/src/prime_rl/trainer/rl/broadcast/nccl.py b/src/prime_rl/trainer/rl/broadcast/nccl.py index 19631aaf3c..9ae2bf34a3 100644 --- a/src/prime_rl/trainer/rl/broadcast/nccl.py +++ b/src/prime_rl/trainer/rl/broadcast/nccl.py @@ -168,21 +168,18 @@ def broadcast_weights(self, model: nn.Module, step: int) -> None: def _resolve_dtensors(self, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: for key, value in list(state_dict.items()): if isinstance(value, DTensor): - value = cast(DTensor, value.to(self.dtype)) local = value.to_local() if local.device.type != "cuda": - # FSDP CPU offload keeps the local shard on CPU while the device - # mesh is CUDA; full_tensor() would then gather over Gloo (hours - # for 100B+ params). Rebuild the DTensor from a CUDA-staged local - # shard so the gather runs over NCCL, one tensor at a time. - # (DTensor.to("cuda") does not migrate the collective backend.) + # FSDP CPU offload keeps local shards on CPU (pinned); gathering + # them there would run over Gloo (hours for 100B+ params). Stage + # the shard to CUDA asynchronously and rebuild the DTensor on the + # same (CUDA) mesh so full_tensor() gathers over NCCL. The bf16 + # cast happens on GPU, after the copy. + local = local.to("cuda", non_blocking=True) value = DTensor.from_local( - local.to("cuda", non_blocking=False), - device_mesh=value.device_mesh, - placements=value.placements, - run_check=False, + local, device_mesh=value.device_mesh, placements=value.placements, run_check=False ) - state_dict[key] = value.full_tensor() + state_dict[key] = cast(DTensor, value.to(self.dtype)).full_tensor() return state_dict