Skip to content

fix(trainer): stage CPU-offloaded weights through CUDA in NCCL broadcast #3045

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mikasenghaas wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

fix(trainer): stage CPU-offloaded weights through CUDA in NCCL broadcast #3045

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4deb151
fix(trainer): stage CPU-offloaded weights through CUDA in NCCL broadcast
mikasenghaas Jul 15, 2026
479cc9a
fix: gather CPU-offloaded DTensor shards on CUDA
mikasenghaas Jul 16, 2026
3f07ecc
fix: rebuild CPU-offloaded shards as CUDA DTensors before gather
mikasenghaas Jul 16, 2026
bb8a777
fix: async pinned staging + GPU-side cast for offloaded shards
mikasenghaas Jul 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion deps/renderers
Open in desktop
Submodule renderers updated 2 files
+0 −1 renderers/base.py
+6 −17 renderers/client.py
2 changes: 1 addition & 1 deletion deps/verifiers
Open in desktop
Submodule verifiers updated 5 files
+1 −1 pyproject.toml
+192 −162 uv.lock
+0 −1 verifiers/types.py
+1 −3 verifiers/v1/graph.py
+4 −15 verifiers/v1/harnesses/default/program.py
15 changes: 15 additions & 0 deletions src/prime_rl/trainer/rl/broadcast/nccl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
Loading