Add opt-in CPU-staging transport for colocated weight sync (ROLL_WEIGHT_SYNC_USE_CPU)

[Donec-x]

Add opt-in CPU-staging transport for colocated weight sync (ROLL_WEIGHT_SYNC_USE_CPU)

Summary

This PR adds an opt-in CPU-staging transport for serialize_named_weights, gated by the environment variable ROLL_WEIGHT_SYNC_USE_CPU=1. The default zero-copy CUDA IPC path is completely unchanged: when the env var is unset (the default), control flow and serialization behavior are byte-for-byte identical to current main. When the opt-in is enabled, the flattened weight bucket is staged to host memory as a contiguous numpy array and serialized via ForkingPickler, producing a portable payload that does not depend on CUDA IPC at deserialization time.

Motivation

In seccomp-managed containers (e.g. managed AutoDL instances), the container runtime blocks pidfd_getfd, which the CUDA IPC handshake uses to transfer file descriptors between processes. As a result, deserializing a CUDA IPC payload fails (CUDA error: invalid argument / driver errors) even for colocated trainer and inference workers — there is no way to receive the IPC handle across the process boundary. Issue #484 (Megatron + LoRA on accelerator environments hitting a CUDA driver error during weight sync) appears to be independent evidence of the same class of failure.

Our case: FSDP2 + vLLM colocated training of Qwen2.5-3B (full-param, ~6.2 GB per model update) on managed AutoDL containers (3x RTX 4090). With the default transport, weight sync crashes at the first model_update due to the blocked pidfd_getfd. With ROLL_WEIGHT_SYNC_USE_CPU=1, training becomes possible at all on such environments; the measured cost is a mean of 103.1 s per model_update (n=26), roughly 17% of step wall time — a trade-off we gladly pay where the zero-copy path simply cannot work.

Implementation

• roll/utils/send_recv_utils.py:
  • New opt-in branch in serialize_named_weights, inserted after _bucket_named_tensors and before the existing is_cuda/FSDP2-CPUOffload check:

    if os.getenv("ROLL_WEIGHT_SYNC_USE_CPU", "0") == "1":
        bucket = bucket.detach().to("cpu").contiguous().numpy()
    elif not getattr(bucket, "is_cuda", False):
        bucket = bucket.to(current_platform.device_type).contiguous()

    The original if becomes elif, so the CPU-staged numpy payload never enters the .to(current_platform.device_type) branch.
  • monkey_patch_torch_reductions() is now called only when the bucket is still a CUDA tensor; the numpy payload is pickled through the standard ForkingPickler reduction and needs no monkey patch.
  • named_tensors_from_bucket accepts a numpy bucket and converts it via torch.from_numpy for backwards compatibility of the receiver side.
• No changes to any sglang-path code, configs, or other files.

Validation

• 100+ steps of production colocated RL training (FSDP2 + vLLM, 3x RTX 4090, Qwen2.5-3B full-param) with ROLL_WEIGHT_SYNC_USE_CPU=1, including suspend/resume cycles.
• New CPU-only unit tests: tests/utils/test_send_recv_cpu_staging.py (both the numpy staging path and the unchanged default CPU path).

Fixes #485

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}