RFC: GPU Time-Sharing Support for NeMo RL Async GRPO

[taoluo]

Goal

Pack more RL training jobs into the same GPU cluster. In our ROLL integration, RLix achieved ~3x rollout throughput gain on SWE-agent RL training by time-sharing GPUs across pipelines instead of exclusive allocation.

Summary

RLix is a Ray-based multi-pipeline GPU orchestration system. Instead of each researcher waiting for an exclusive GPU allocation, RLix lets multiple RL training pipelines share the same GPUs. More experiments run concurrently, less time spent waiting for resources. The scheduler dynamically assigns GPU time across pipelines based on priority and demand. Currently RLix supports ROLL as the training framework; this RFC proposes adding NeMo RL as a second supported framework.

We propose changes to NeMo RL that enable its async GRPO training to participate in RLix multi-pipeline scheduling. The key capabilities needed:

• Scheduler-driven shard resize — the scheduler controls when individual inference DP shards sleep/wake, rather than the training loop doing all-or-nothing sleep/wake
• CPU weight caching + selective sync — after training, weights are cached on CPU and synced only to the specific shards that need updating. In a multi-pipeline setting, we cannot afford to broadcast weights to all vLLM workers across all pipelines — each pipeline must sync independently to its own shards without blocking others
• Progress reporting — pipelines report generation demand so RLix can make informed allocation decisions

One mechanism this enables is partial overlap, where inference GPUs are a superset of training GPUs — overlap GPUs alternate between roles via vLLM sleep(level=2) / wake_up() while non-overlap GPUs continue inference.

Scope: vLLM inference + Megatron training + async GRPO only. All changes are gated behind RLIX_CONTROL_PLANE=rlix env var — standalone NeMo RL behavior is unaffected.

---

Review Asks

1. API shape — Are the proposed new methods on VLLMGeneration, RayWorkerGroup, and MegatronPolicyWorker reasonable?
2. Lifecycle hooks — Is the RLixHooks protocol (4 methods: before/after_training called from async_grpo_train(), report/clear_progress called from AsyncTrajectoryCollector) an acceptable integration pattern?

---

Design Overview

Why CPU weight cache + selective sync (not native refit)

In a multi-pipeline setting, each pipeline's training and inference workers share GPU time with other pipelines. Two problems make NeMo RL's native refit_policy_generation() insufficient:

1. Full broadcast causes OOM. refit_policy_generation() broadcasts weights to all inference workers via NCCL collective, requiring the full model staged on GPU while receivers already have their weights loaded. With multiple pipelines sharing GPUs, this is not affordable — each pipeline must sync only to its own shards that actually need updating, not all workers.

2. GPU VRAM conflict. Refit requires training weights on GPU when sending (CUDA IPC handles or NCCL broadcast). But when inference workers need to wake up on the same GPUs, both cannot coexist → OOM.

Solution: CPU weight cache + selective sync. After each training step, snapshot weights to CPU, release the GPU. Sync only to the specific DP shards that need it, in chunks (staged CPU → GPU, freed after send) to control peak VRAM. Other shards and other pipelines are unaffected. The routing layer reuses NeMo RL's existing ZMQ IPC and NCCL broadcast transport — no modifications to transport code.

Abort-drain-sleep for shard preemption

When RLix reclaims GPUs from inference, we cannot simply call sleep() on workers with active requests — vLLM would crash accessing offloaded GPU memory. The safe sequence:

1. Mark preempted — block new request dispatch to these shards
2. Abort — worker calls engine.abort() for all running requests internally (no request ID plumbing needed from the generation layer)
3. Drain — poll existing vLLM engine metric vllm:num_requests_running until 0 (already tracked per worker at vllm_worker_async.py:241). Drain timeout with force-sleep fallback in case the metric never reaches 0.
4. Sleep — engine confirmed idle, GPU memory safely offloaded

Callers with in-flight requests on preempted shards receive errors, classified as ShardPreemptedError (via _preempted_shards flag check) and automatically re-dispatched to the next active shard.

Training cycle under RLix

DO_TIME_SHARING = os.environ.get("RLIX_CONTROL_PLANE") == "rlix"

# After each training step:
if DO_TIME_SHARING:
    build_cpu_weight_cache(step)    # snapshot weights to CPU
    offload_training_gpu()          # free GPU VRAM
    destroy_nccl_groups()           # free NCCL communicator buffers
    hooks.after_training(step)      # notify RLix → triggers expand + selective sync
else:
    refit_policy_generation(...)    # existing standalone path (unchanged)

---

Proposed NeMo RL Changes

Ownership overview

Location	Lines	Notes
NeMo RL repo (detailed below)	~490	All RLix-gated or additive
RLix repo (pipeline adapter, model update service, config bridge)	~720	Does not affect NeMo RL

Supporting API additions (~40 lines, 3 files)

Small additive changes, not RLix-gated:

• vllm_worker.py (+5): Parameterize hardcoded sleep(level=1) → sleep(level=self._sleep_level) from config. Level 2 offloads weights + KV cache (needed for training to fit).
• vllm_worker_async.py (+15): Same sleep level parameterization. New abort_all_requests() method — gets running request IDs from the engine internally and aborts them. New is_idle() -> bool — checks existing vllm:num_requests_running metric (already scraped at line 241). No request ID plumbing needed from callers.
• worker_groups.py (+20): New run_on_dp_shard_leaders(dp_ranks, method, ...) — executes a method on the DP-leader worker of each specified rank. Subset variant of existing run_all_workers_single_data(..., run_rank_0_only_axes=...).

vllm_generation.py — Partial sleep/wake + routing + preemption (+150 lines)

New methods on VLLMGeneration: sleep_partial(dp_ranks), wake_up_partial(dp_ranks), activate_dp_ranks(), mark_dp_ranks_inactive().

New state: _active_dp_ranks: Set[int] (canonical routing set), _preempted_shards: Set[int] (abort window — used for error classification).

Modified _async_generate_base(): round-robin skips sleeping shards, blocks when all shards sleeping (does not raise), converts errors on preempted shards to ShardPreemptedError with automatic re-dispatch retry. No per-request tracking needed — drain uses engine-level idle check.

RLix-gated: Partially. New methods only called from RLix. Routing skip only activates when _active_dp_ranks != all_ranks (never happens in standalone).

megatron_policy_worker.py — CPU buffer cache (+60 lines)

New build_cpu_buffer_cache(step): all TP/PP/CP/EP ranks participate in collective gather, only cache owner (pp0/dp0/tp0/cp0) stores the complete model as CPU tensor buffers. Reuses existing gather_all_hf_weights with EP-aware gather.

RLix-gated: Yes.

nccl_offload.py (new) — Megatron NCCL destroy/re-init (+90 lines)

destroy_megatron_nccl_groups(): collects all NCCL process groups from megatron.core.parallel_state, filters to NCCL backend (excluding Gloo), deduplicates handles, and destroys each via torch.distributed.destroy_process_group(). reinit_megatron_nccl_groups(): re-initializes with saved parallel config. Manual approach because destroy_model_parallel() is not designed for repeated use in long-lived workers.

Why: Long-lived training actors keep NCCL communicator buffers on GPU. Without explicit cleanup, inference wake_up() hits OOM.

RLix-gated: Yes. Repeated destroy/re-init cycles in long-lived workers need validation for VRAM leaks and correctness.

grpo.py — Training loop RLix branches (+60 lines)

DO_TIME_SHARING flag adds if/else branches in async_grpo_train() (see "Training cycle under RLix" above). Also skips ray.shutdown() (RLix manages Ray lifecycle) and prepare_for_generation() / finish_generation() (RLix drives sleep/wake).

RLix-gated: Yes — all branches behind DO_TIME_SHARING.

async_utils.py — Progress reporting (+60 lines)

AsyncTrajectoryCollector accepts optional rlix_hooks (defaults to NoOpRLixHooks()). After each trajectory push, reports generation demand as a point-in-time snapshot with 2% granularity. ReplayBuffer gains valid_count(current_weight_version, max_age_steps) -> int for progress calculation.

Why: RLix uses generation demand progress for planning when to reclaim GPUs for training.

RLix-gated: Partially. Hooks are no-ops in standalone. valid_count() is a pure query method.

rlix_hooks.py (new) — Hook protocol (+30 lines)

4-method protocol (before_training, after_training, report_progress, clear_progress) + NoOpRLixHooks default. Standalone file because both NeMo RL and RLix code import it.

---

Non-Goals

• Standalone behavior — All changes behind DO_TIME_SHARING or NoOpRLixHooks. Standalone async_grpo_train() is identical to current code.
• Synchronous grpo_train() — Partial overlap provides no value for sync training.
• NeMo-Gym — HTTP-based shard preemption requires a different mechanism (503 middleware on NeMo RL's own FastAPI app). Planned as future work.
• Existing transports — ZMQ IPC and NCCL broadcast used as-is. No modifications.
• Existing refit_policy_generation() — Unchanged. Skipped in RLix mode, works normally in standalone.
• SGLang, DTensor/FSDP2, Multi-LoRA, DPO/SFT — Out of scope.

---

Validation gates (2-GPU environment)

Gate	What it validates
1: Single NeMo RL pipeline	Partial overlap e2e: sleep/wake, routing, NCCL lifecycle, selective sync
2: Two NeMo RL pipelines	Multi-pipeline GPU time-sharing
3: NeMo RL + ROLL mixed	Cross-framework scheduling