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 75abf42195..9ae2bf34a3 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 @@ -164,6 +168,17 @@ 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): + local = value.to_local() + if local.device.type != "cuda": + # 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, device_mesh=value.device_mesh, placements=value.placements, run_check=False + ) state_dict[key] = cast(DTensor, value.to(self.dtype)).full_tensor() return state_dict