feat(executors): add colocated dual-engine with vLLM sleep/wake#218
feat(executors): add colocated dual-engine with vLLM sleep/wake#218icenfly wants to merge 1 commit into
Conversation
Add colocated training architecture that shares ALL GPUs between vLLM inference and DeepSpeed training via vLLM sleep/wake mechanism (v0.8+). - ColocatedOrchestrator: phase-sequenced vLLM rollout + DeepSpeed train - Persistent training worker: launched once, signal-driven, pinned memory - InferenceEngineAdapter: unified vLLM/SGLang interface with sleep/wake - Unit tests and GPU integration tests - Setup guide with architecture docs and memory budget Part of RL-Align#127
📝 WalkthroughWalkthroughAdds an experimental single-node colocated workflow that alternates vLLM rollouts with persistent DeepSpeed LoRA training on shared GPUs. It includes inference adapters, filesystem-based orchestration, a distributed worker, tests, and setup documentation. ChangesColocated dual-engine workflow
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant ColocatedOrchestrator
participant VLLM
participant TrainingWorker
participant DeepSpeed
ColocatedOrchestrator->>VLLM: Generate rollout completions
VLLM-->>ColocatedOrchestrator: Return completions
ColocatedOrchestrator->>VLLM: Sleep and release GPU memory
ColocatedOrchestrator->>TrainingWorker: Signal training step
TrainingWorker->>DeepSpeed: Train on rollout completions
DeepSpeed-->>TrainingWorker: Updated parameters and metrics
TrainingWorker-->>ColocatedOrchestrator: Write completion signal and metrics
ColocatedOrchestrator->>VLLM: Wake for next rollout
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (6)
docs/guides/colocated_setup.md (2)
87-87: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd language specifier to architecture diagram code block.
Same MD040 warning as line 11. Use
```textfor the ASCII architecture diagram.♻️ Proposed fix
-``` +```text ┌──────────────────────────────────────────────────┐🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/guides/colocated_setup.md` at line 87, Update the architecture diagram code block in the documentation to use the text language specifier, changing its opening fence to ```text while preserving the diagram content.Source: Linters/SAST tools
11-20: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUnlabeled fenced code blocks trigger MD040 warnings. Both ASCII diagram blocks omit a language specifier. Add
```textto each.
docs/guides/colocated_setup.md#L11-L20: Change the opening```to```textfor the disaggregated/colocated diagram.docs/guides/colocated_setup.md#L87-L105: Change the opening```to```textfor the architecture diagram.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/guides/colocated_setup.md` around lines 11 - 20, Update both ASCII diagram fenced blocks in docs/guides/colocated_setup.md:11-20 and docs/guides/colocated_setup.md:87-105 by changing each opening fence to specify the text language. No other diagram content needs modification.Source: Linters/SAST tools
tests/test_colocated_dual_engine.py (1)
160-160: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd explicit GPU cache cleanup after
del llm.
del llmmay not immediately release GPU memory if reference cycles exist. Addingtorch.cuda.empty_cache()ensures the cache is cleared before subsequent tests run.♻️ Proposed cleanup improvement
del llm + torch.cuda.empty_cache()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_colocated_dual_engine.py` at line 160, After deleting llm in the affected test cleanup flow, explicitly call torch.cuda.empty_cache() so GPU cache memory is released before subsequent tests run; preserve the existing deletion behavior and use the already available torch symbol.rl_engine/executors/inference_adapter.py (1)
156-174: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNarrow the blind exception catch in
update_weights.Ruff flags
except Exception(BLE001). While catching all failures in a hot-reload path is reasonable, narrowing toExceptionsubclasses (e.g.,(RuntimeError, TypeError, AttributeError)) would avoid swallowingKeyboardInterrupt/SystemExitand make debugging easier.♻️ Suggested narrowing
except (RuntimeError, TypeError, AttributeError) as e: logger.warning("[vLLM] Weight reload failed: %s", e)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rl_engine/executors/inference_adapter.py` around lines 156 - 174, Update the exception handler in update_weights to catch only the expected operational failures from wake_up, engine/collective_rpc lookup, and reload_weights—such as RuntimeError, TypeError, and AttributeError—instead of using a broad Exception catch. Preserve the existing warning, cleanup via sleep(), and False return behavior.examples/_colocated_train_worker.py (1)
36-43: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
wait_for_fileis dead code — inline polling is used instead.The function is defined but never called. The main loop (lines 133–142) uses inline polling without a timeout, meaning workers can hang indefinitely if the orchestrator crashes. Either use
wait_for_file(which has a timeout) or remove it.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/_colocated_train_worker.py` around lines 36 - 43, Update the worker’s file-waiting flow to use the existing wait_for_file function, including its timeout behavior, instead of the duplicated inline polling in the main loop. Remove the inline polling logic and preserve the worker’s existing handling after the file is found; do not leave wait_for_file unused.examples/colocated_dual_engine.py (1)
84-96: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winOrchestrator bypasses
VLLMInferenceAdapterand duplicates vLLM setup.The PR introduces
VLLMInferenceAdapterwith a configurableInferenceEngineConfig, but the orchestrator hardcodes a separatevllm.LLMconstruction with different defaults (e.g.,gpu_memory_utilization=0.40vs adapter's0.45). This creates duplication and makes future changes (e.g., switching to SGLang, addingupdate_weightssupport) harder.Consider using
create_inference_adapter("vllm", ...)and delegating to the adapter'sinitialize/sleep/wake_up/generatemethods.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/colocated_dual_engine.py` around lines 84 - 96, Replace the direct vllm.LLM construction in the orchestrator with create_inference_adapter("vllm", ...) using the existing InferenceEngineConfig, then delegate initialization and inference lifecycle operations to the adapter’s initialize, sleep, wake_up, and generate methods. Remove duplicated vLLM-specific setup while preserving the orchestrator’s current model, parallelism, memory, and sequence-length configuration.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/guides/colocated_setup.md`:
- Line 26: Update the sleep level 2 description in the colocated setup guide to
state that it discards model weights and KV cache entirely and requires
reloading from checkpoint on wake. Keep the CPU offload description associated
with sleep level 1, matching the InferenceEngineAdapter protocol documentation.
- Around line 57-58: Update the benchmark command in the colocated setup guide
to reference the existing colocated benchmark entry point instead of the
nonexistent benchmarks/benchmark_colocated.py script, or add a wrapper at that
path if that is the established entry point. Ensure the documented command runs
successfully with the existing model and steps arguments.
- Line 73: Remove the `--ds-zero-stage` row from the colocated setup guide,
since `examples/_colocated_train_worker.py` hardcodes ZeRO stage 2 and
`examples/colocated_dual_engine.py` does not propagate a configurable value.
In `@examples/_colocated_train_worker.py`:
- Around line 186-209: Add a torch.distributed barrier after the GPU offload and
pinned-state restoration in the training step, before the rank-0 completion log
and train_done signal. Use the existing dist.barrier() synchronization flow so
rank 0 signals completion only after every rank has finished offloading and
clearing GPU memory.
- Around line 158-165: Update the label construction in the training flow around
tokenizer and labels so prompt-token positions are masked with -100, leaving
only completion-token positions as valid loss targets. Preserve the existing
padding mask and account for truncation when determining each completion’s token
span, using the tokenizer’s corresponding tokenization or offsets rather than
character lengths.
In `@examples/colocated_dual_engine.py`:
- Line 95: Make remote code execution opt-in across all affected sites: in
examples/colocated_dual_engine.py at lines 95-95, add a --trust-remote-code CLI
flag defaulting to false and pass it to vllm.LLM; in
examples/_colocated_train_worker.py at lines 69-74, add the same flag and pass
its value to both AutoTokenizer.from_pretrained and
AutoModelForCausalLM.from_pretrained; in
rl_engine/executors/inference_adapter.py at lines 85-85, change
InferenceEngineConfig.trust_remote_code to default to false while preserving
explicit opt-in by callers.
- Around line 279-288: Update cleanup() to catch subprocess.TimeoutExpired
around self.worker_proc.wait(timeout=30), then ensure del self.llm and
shutil.rmtree(self.work_dir, ignore_errors=True) still execute when the worker
hangs. Preserve the existing shutdown signaling and normal wait behavior, while
handling the timed-out worker so cleanup does not abort.
In `@rl_engine/executors/inference_adapter.py`:
- Around line 198-202: Update initialize so the ImportError handler binds the
caught exception and re-raises the custom ImportError explicitly from the
original exception, preserving the exception chain and traceback context.
---
Nitpick comments:
In `@docs/guides/colocated_setup.md`:
- Line 87: Update the architecture diagram code block in the documentation to
use the text language specifier, changing its opening fence to ```text while
preserving the diagram content.
- Around line 11-20: Update both ASCII diagram fenced blocks in
docs/guides/colocated_setup.md:11-20 and docs/guides/colocated_setup.md:87-105
by changing each opening fence to specify the text language. No other diagram
content needs modification.
In `@examples/_colocated_train_worker.py`:
- Around line 36-43: Update the worker’s file-waiting flow to use the existing
wait_for_file function, including its timeout behavior, instead of the
duplicated inline polling in the main loop. Remove the inline polling logic and
preserve the worker’s existing handling after the file is found; do not leave
wait_for_file unused.
In `@examples/colocated_dual_engine.py`:
- Around line 84-96: Replace the direct vllm.LLM construction in the
orchestrator with create_inference_adapter("vllm", ...) using the existing
InferenceEngineConfig, then delegate initialization and inference lifecycle
operations to the adapter’s initialize, sleep, wake_up, and generate methods.
Remove duplicated vLLM-specific setup while preserving the orchestrator’s
current model, parallelism, memory, and sequence-length configuration.
In `@rl_engine/executors/inference_adapter.py`:
- Around line 156-174: Update the exception handler in update_weights to catch
only the expected operational failures from wake_up, engine/collective_rpc
lookup, and reload_weights—such as RuntimeError, TypeError, and
AttributeError—instead of using a broad Exception catch. Preserve the existing
warning, cleanup via sleep(), and False return behavior.
In `@tests/test_colocated_dual_engine.py`:
- Line 160: After deleting llm in the affected test cleanup flow, explicitly
call torch.cuda.empty_cache() so GPU cache memory is released before subsequent
tests run; preserve the existing deletion behavior and use the already available
torch symbol.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: bd4a70b4-6769-4891-af2c-bbe2a3bb4c60
📒 Files selected for processing (5)
docs/guides/colocated_setup.mdexamples/_colocated_train_worker.pyexamples/colocated_dual_engine.pyrl_engine/executors/inference_adapter.pytests/test_colocated_dual_engine.py
|
|
||
| The key enabler is **vLLM's sleep mode** (v0.8+): | ||
|
|
||
| 1. `llm.sleep(level=2)` — offloads model weights to CPU, discards KV cache, frees ~90% GPU memory |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Sleep level 2 description contradicts the adapter protocol docstring.
The guide says sleep(level=2) "offloads model weights to CPU," but the InferenceEngineAdapter protocol docstring (in rl_engine/executors/inference_adapter.py) defines level 2 as "discard weights and KV cache entirely (faster wake from checkpoint)." The guide's description actually matches level 1. This could mislead users about wake behavior — level 2 requires a full reload from checkpoint, not a CPU-to-GPU transfer.
📝 Proposed fix
-1. `llm.sleep(level=2)` — offloads model weights to CPU, discards KV cache, frees ~90% GPU memory
+1. `llm.sleep(level=2)` — discards weights and KV cache entirely, frees ~90% GPU memory (faster wake from checkpoint)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 1. `llm.sleep(level=2)` — offloads model weights to CPU, discards KV cache, frees ~90% GPU memory | |
| 1. `llm.sleep(level=2)` — discards weights and KV cache entirely, frees ~90% GPU memory (faster wake from checkpoint) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/guides/colocated_setup.md` at line 26, Update the sleep level 2
description in the colocated setup guide to state that it discards model weights
and KV cache entirely and requires reloading from checkpoint on wake. Keep the
CPU offload description associated with sleep level 1, matching the
InferenceEngineAdapter protocol documentation.
| python benchmarks/benchmark_colocated.py \ | ||
| --model Qwen/Qwen3-0.6B --steps 10 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check if the benchmark script exists
fd benchmark_colocated.pyRepository: RL-Align/RL-Kernel
Length of output: 156
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the documentation file and any benchmark-related files
git ls-files 'docs/guides/colocated_setup.md' 'benchmarks/*' | sed -n '1,200p'
printf '\n---\n'
# Search for references to the benchmark script and nearby benchmark names
rg -n "benchmark_colocated\.py|colocated" docs benchmarks . -g '!**/.git/**' | sed -n '1,200p'Repository: RL-Align/RL-Kernel
Length of output: 4175
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Show the Quick Start section around the benchmark command
sed -n '40,65p' docs/guides/colocated_setup.md
printf '\n---\n'
# Check whether the example script mentions benchmarking or equivalent CLI options
sed -n '1,120p' examples/colocated_dual_engine.pyRepository: RL-Align/RL-Kernel
Length of output: 4718
Update the benchmark command to a real entry point
benchmarks/benchmark_colocated.py does not exist, so this step will fail as written. Point it at the existing colocated script or add the missing benchmark wrapper.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/guides/colocated_setup.md` around lines 57 - 58, Update the benchmark
command in the colocated setup guide to reference the existing colocated
benchmark entry point instead of the nonexistent
benchmarks/benchmark_colocated.py script, or add a wrapper at that path if that
is the established entry point. Ensure the documented command runs successfully
with the existing model and steps arguments.
| | `--steps` | 5 | Training iterations | | ||
| | `--vllm-sleep-level` | 2 | 1=offload weights, 2=discard everything | | ||
| | `--vllm-gpu-memory-utilization` | 0.45 | vLLM memory fraction (leave room for training) | | ||
| | `--ds-zero-stage` | 2 | DeepSpeed ZeRO stage | |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check if --ds-zero-stage is passed to the worker or handled elsewhere
rg -n 'ds.zero.stage\|ds_zero_stage\|zero_stage' examples/colocated_dual_engine.py examples/_colocated_train_worker.pyRepository: RL-Align/RL-Kernel
Length of output: 156
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Top-level files:\n'
git ls-files | sed -n '1,200p'
printf '\nRelevant files:\n'
git ls-files | rg '(^|/)(docs/guides/colocated_setup\.md|examples/colocated_dual_engine\.py|examples/_colocated_train_worker\.py|examples/)'
printf '\nOutline docs/guides/colocated_setup.md:\n'
wc -l docs/guides/colocated_setup.md
sed -n '1,140p' docs/guides/colocated_setup.md
printf '\nSearch for ds-zero-stage across repo:\n'
rg -n --hidden --glob '!**/.git/**' 'ds-zero-stage|ds_zero_stage|zero_stage' .Repository: RL-Align/RL-Kernel
Length of output: 10879
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'Locate colocated setup docs and example scripts:\n'
fd -a 'colocated_setup.md|colocated_dual_engine.py|_colocated_train_worker.py' . || true
printf '\nFind any argparse/config definitions mentioning deepSpeed or zero stage:\n'
rg -n --hidden --glob '!**/.git/**' '(DeepSpeed|deepspeed|ZeRO|zero stage|zero_stage|ds[-_]zero[-_]stage)' docs examples .github .Repository: RL-Align/RL-Kernel
Length of output: 18196
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf 'examples/colocated_dual_engine.py outline:\n'
ast-grep outline examples/colocated_dual_engine.py --view expanded || true
printf '\nKey sections from examples/colocated_dual_engine.py:\n'
rg -n -A6 -B6 '_start_workers|ds_zero_stage|zero_stage|DeepSpeedTrainingConfig|torchrun|colocated' examples/colocated_dual_engine.py
printf '\nKey sections from examples/_colocated_train_worker.py:\n'
rg -n -A6 -B6 'zero_stage|DeepSpeedTrainingConfig|deepspeed.initialize|zero_optimization' examples/_colocated_train_worker.py
printf '\nRelevant tests:\n'
rg -n -A4 -B4 'zero_stage|ds_zero_stage|colocated_dual_engine' tests/test_colocated_dual_engine.py tests/test_deepspeed_training_worker.pyRepository: RL-Align/RL-Kernel
Length of output: 11059
Remove or wire --ds-zero-stage
docs/guides/colocated_setup.md:73 advertises --ds-zero-stage as configurable, but examples/_colocated_train_worker.py hardcodes ZeRO stage 2 and examples/colocated_dual_engine.py never forwards a value. Either thread it through end-to-end or drop the row from the guide.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/guides/colocated_setup.md` at line 73, Remove the `--ds-zero-stage` row
from the colocated setup guide, since `examples/_colocated_train_worker.py`
hardcodes ZeRO stage 2 and `examples/colocated_dual_engine.py` does not
propagate a configurable value.
| full_texts = [c["prompt"] + c["completion"] for c in completions] | ||
| enc = tokenizer( | ||
| full_texts, truncation=True, max_length=args.max_len, padding=True, return_tensors="pt" | ||
| ) | ||
| input_ids = enc["input_ids"].to(device) | ||
| attention_mask = enc["attention_mask"].to(device) | ||
| labels = input_ids.clone() | ||
| labels[attention_mask == 0] = -100 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Labels include prompt tokens — loss is computed on prompts, not just completions.
labels = input_ids.clone() with only padding masked (labels[attention_mask == 0] = -100) means the model also trains on prompt tokens. In RL/fine-tuning, loss should be computed only on completion tokens to avoid wasting capacity learning the prompt distribution and potentially degrading quality.
💚 Proposed fix: mask prompt tokens in labels
full_texts = [c["prompt"] + c["completion"] for c in completions]
enc = tokenizer(
full_texts, truncation=True, max_length=args.max_len, padding=True, return_tensors="pt"
)
input_ids = enc["input_ids"].to(device)
attention_mask = enc["attention_mask"].to(device)
labels = input_ids.clone()
labels[attention_mask == 0] = -100
+ # Mask prompt tokens — only compute loss on completions
+ prompt_texts = [c["prompt"] for c in completions]
+ prompt_enc = tokenizer(
+ prompt_texts, truncation=True, max_length=args.max_len, padding=True, return_tensors="pt"
+ )
+ prompt_lens = prompt_enc["attention_mask"].sum(dim=1)
+ for i, plen in enumerate(prompt_lens):
+ labels[i, :plen] = -100
+
engine.train()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| full_texts = [c["prompt"] + c["completion"] for c in completions] | |
| enc = tokenizer( | |
| full_texts, truncation=True, max_length=args.max_len, padding=True, return_tensors="pt" | |
| ) | |
| input_ids = enc["input_ids"].to(device) | |
| attention_mask = enc["attention_mask"].to(device) | |
| labels = input_ids.clone() | |
| labels[attention_mask == 0] = -100 | |
| full_texts = [c["prompt"] + c["completion"] for c in completions] | |
| enc = tokenizer( | |
| full_texts, truncation=True, max_length=args.max_len, padding=True, return_tensors="pt" | |
| ) | |
| input_ids = enc["input_ids"].to(device) | |
| attention_mask = enc["attention_mask"].to(device) | |
| labels = input_ids.clone() | |
| labels[attention_mask == 0] = -100 | |
| # Mask prompt tokens — only compute loss on completions | |
| prompt_texts = [c["prompt"] for c in completions] | |
| prompt_enc = tokenizer( | |
| prompt_texts, truncation=True, max_length=args.max_len, padding=True, return_tensors="pt" | |
| ) | |
| prompt_lens = prompt_enc["attention_mask"].sum(dim=1) | |
| for i, plen in enumerate(prompt_lens): | |
| labels[i, :plen] = -100 | |
| engine.train() |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/_colocated_train_worker.py` around lines 158 - 165, Update the label
construction in the training flow around tokenizer and labels so prompt-token
positions are masked with -100, leaving only completion-token positions as valid
loss targets. Preserve the existing padding mask and account for truncation when
determining each completion’s token span, using the tokenizer’s corresponding
tokenization or offsets rather than character lengths.
| dist.barrier() | ||
|
|
||
| # --- Offload to pinned CPU, free GPU memory --- | ||
| t_offload = time.time() | ||
| for name, param in engine.module.named_parameters(): | ||
| pinned_state[name].copy_(param.data, non_blocking=True) | ||
| torch.cuda.synchronize() | ||
| engine.module.to("cpu") | ||
| # Restore pinned data as param.data for next load | ||
| for name, param in engine.module.named_parameters(): | ||
| param.data = pinned_state[name] | ||
| torch.cuda.empty_cache() | ||
| offload_ms = (time.time() - t_offload) * 1000 | ||
|
|
||
| if rank == 0: | ||
| print( | ||
| f"[Worker] Step {step}: loss={loss_val:.4f} " | ||
| f"load={load_ms:.0f}ms train={train_ms:.0f}ms offload={offload_ms:.0f}ms", | ||
| flush=True, | ||
| ) | ||
| # Signal done | ||
| done_path = os.path.join(args.work_dir, f"train_done_{step}.signal") | ||
| with open(done_path, "w") as f: | ||
| f.write(f"done step={step} loss={loss_val}") |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Missing dist.barrier() after GPU offload — race condition with vLLM wake-up.
After training, all ranks offload to CPU (lines 189–198), then rank 0 writes the train_done signal (lines 207–209) without a barrier. The orchestrator sees the done signal and calls self.llm.wake_up(), but other ranks may still hold GPU memory. This can cause vLLM wake-up to OOM or fail.
The existing dist.barrier() at line 186 only synchronizes before offloading — it does not guarantee all ranks have finished offloading.
🔒 Fix: add barrier after offload, before done signal
torch.cuda.empty_cache()
offload_ms = (time.time() - t_offload) * 1000
+ dist.barrier() # Ensure all ranks have offloaded before signaling done
+
if rank == 0:
print(
f"[Worker] Step {step}: loss={loss_val:.4f} "📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| dist.barrier() | |
| # --- Offload to pinned CPU, free GPU memory --- | |
| t_offload = time.time() | |
| for name, param in engine.module.named_parameters(): | |
| pinned_state[name].copy_(param.data, non_blocking=True) | |
| torch.cuda.synchronize() | |
| engine.module.to("cpu") | |
| # Restore pinned data as param.data for next load | |
| for name, param in engine.module.named_parameters(): | |
| param.data = pinned_state[name] | |
| torch.cuda.empty_cache() | |
| offload_ms = (time.time() - t_offload) * 1000 | |
| if rank == 0: | |
| print( | |
| f"[Worker] Step {step}: loss={loss_val:.4f} " | |
| f"load={load_ms:.0f}ms train={train_ms:.0f}ms offload={offload_ms:.0f}ms", | |
| flush=True, | |
| ) | |
| # Signal done | |
| done_path = os.path.join(args.work_dir, f"train_done_{step}.signal") | |
| with open(done_path, "w") as f: | |
| f.write(f"done step={step} loss={loss_val}") | |
| dist.barrier() | |
| # --- Offload to pinned CPU, free GPU memory --- | |
| t_offload = time.time() | |
| for name, param in engine.module.named_parameters(): | |
| pinned_state[name].copy_(param.data, non_blocking=True) | |
| torch.cuda.synchronize() | |
| engine.module.to("cpu") | |
| # Restore pinned data as param.data for next load | |
| for name, param in engine.module.named_parameters(): | |
| param.data = pinned_state[name] | |
| torch.cuda.empty_cache() | |
| offload_ms = (time.time() - t_offload) * 1000 | |
| dist.barrier() # Ensure all ranks have offloaded before signaling done | |
| if rank == 0: | |
| print( | |
| f"[Worker] Step {step}: loss={loss_val:.4f} " | |
| f"load={load_ms:.0f}ms train={train_ms:.0f}ms offload={offload_ms:.0f}ms", | |
| flush=True, | |
| ) | |
| # Signal done | |
| done_path = os.path.join(args.work_dir, f"train_done_{step}.signal") | |
| with open(done_path, "w") as f: | |
| f.write(f"done step={step} loss={loss_val}") |
🧰 Tools
🪛 ast-grep (0.44.1)
[warning] 207-207: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(done_path, "w")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').
(open-filename-from-request)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/_colocated_train_worker.py` around lines 186 - 209, Add a
torch.distributed barrier after the GPU offload and pinned-state restoration in
the training step, before the rank-0 completion log and train_done signal. Use
the existing dist.barrier() synchronization flow so rank 0 signals completion
only after every rank has finished offloading and clearing GPU memory.
| dtype="bfloat16", | ||
| max_model_len=a.max_prompt_len + a.max_completion_len, | ||
| enable_sleep_mode=True, | ||
| trust_remote_code=True, |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
trust_remote_code=True used across all three files — security posture gap (CWE-94).
All three files enable trust_remote_code=True by default, which executes arbitrary code from the HuggingFace model repository. While common in ML codebases, this should be opt-in rather than default, especially since the model path comes from CLI args and could point to an untrusted source.
examples/colocated_dual_engine.py#L95: Hardcodedtrust_remote_code=Trueinvllm.LLM(...). Make this configurable via a CLI arg (e.g.,--trust-remote-code), defaulting toFalse.examples/_colocated_train_worker.py#L69-L74: Hardcoded in bothAutoTokenizer.from_pretrainedandAutoModelForCausalLM.from_pretrained. Add a--trust-remote-codeCLI arg and pass it through.rl_engine/executors/inference_adapter.py#L85:InferenceEngineConfig.trust_remote_codedefaults toTrue. Consider defaulting toFalseand letting callers opt in explicitly.
📍 Affects 3 files
examples/colocated_dual_engine.py#L95-L95(this comment)examples/_colocated_train_worker.py#L69-L74rl_engine/executors/inference_adapter.py#L85-L85
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/colocated_dual_engine.py` at line 95, Make remote code execution
opt-in across all affected sites: in examples/colocated_dual_engine.py at lines
95-95, add a --trust-remote-code CLI flag defaulting to false and pass it to
vllm.LLM; in examples/_colocated_train_worker.py at lines 69-74, add the same
flag and pass its value to both AutoTokenizer.from_pretrained and
AutoModelForCausalLM.from_pretrained; in
rl_engine/executors/inference_adapter.py at lines 85-85, change
InferenceEngineConfig.trust_remote_code to default to false while preserving
explicit opt-in by callers.
| def cleanup(self): | ||
| # Shutdown workers | ||
| shutdown_path = os.path.join(self.work_dir, "shutdown.signal") | ||
| with open(shutdown_path, "w") as f: | ||
| f.write("shutdown") | ||
| if self.worker_proc and self.worker_proc.poll() is None: | ||
| self.worker_proc.wait(timeout=30) | ||
| if self.llm is not None: | ||
| del self.llm | ||
| shutil.rmtree(self.work_dir, ignore_errors=True) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
cleanup() doesn't handle TimeoutExpired — resource leak on worker hang.
self.worker_proc.wait(timeout=30) raises subprocess.TimeoutExpired if the worker doesn't exit in 30s. This skips del self.llm and shutil.rmtree, leaking GPU memory, temp files, and potentially orphaning the worker process.
🛡️ Proposed fix: wrap wait in try/except
def cleanup(self):
# Shutdown workers
shutdown_path = os.path.join(self.work_dir, "shutdown.signal")
with open(shutdown_path, "w") as f:
f.write("shutdown")
- if self.worker_proc and self.worker_proc.poll() is None:
- self.worker_proc.wait(timeout=30)
+ if self.worker_proc and self.worker_proc.poll() is None:
+ try:
+ self.worker_proc.wait(timeout=30)
+ except subprocess.TimeoutExpired:
+ self.worker_proc.kill()
+ self.worker_proc.wait()
if self.llm is not None:
del self.llm
shutil.rmtree(self.work_dir, ignore_errors=True)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def cleanup(self): | |
| # Shutdown workers | |
| shutdown_path = os.path.join(self.work_dir, "shutdown.signal") | |
| with open(shutdown_path, "w") as f: | |
| f.write("shutdown") | |
| if self.worker_proc and self.worker_proc.poll() is None: | |
| self.worker_proc.wait(timeout=30) | |
| if self.llm is not None: | |
| del self.llm | |
| shutil.rmtree(self.work_dir, ignore_errors=True) | |
| def cleanup(self): | |
| # Shutdown workers | |
| shutdown_path = os.path.join(self.work_dir, "shutdown.signal") | |
| with open(shutdown_path, "w") as f: | |
| f.write("shutdown") | |
| if self.worker_proc and self.worker_proc.poll() is None: | |
| try: | |
| self.worker_proc.wait(timeout=30) | |
| except subprocess.TimeoutExpired: | |
| self.worker_proc.kill() | |
| self.worker_proc.wait() | |
| if self.llm is not None: | |
| del self.llm | |
| shutil.rmtree(self.work_dir, ignore_errors=True) |
🧰 Tools
🪛 ast-grep (0.44.1)
[warning] 281-281: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(shutdown_path, "w")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').
(open-filename-from-request)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/colocated_dual_engine.py` around lines 279 - 288, Update cleanup()
to catch subprocess.TimeoutExpired around self.worker_proc.wait(timeout=30),
then ensure del self.llm and shutil.rmtree(self.work_dir, ignore_errors=True)
still execute when the worker hangs. Preserve the existing shutdown signaling
and normal wait behavior, while handling the timed-out worker so cleanup does
not abort.
| def initialize(self) -> None: | ||
| try: | ||
| importlib.import_module("sglang") | ||
| except ImportError: | ||
| raise ImportError("SGLang is not installed. Install with: pip install sglang") |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Preserve exception chain when re-raising ImportError.
Ruff B904: re-raising without from err loses the original traceback context.
🛡️ Proposed fix
except ImportError as err:
- raise ImportError("SGLang is not installed. Install with: pip install sglang")
+ raise ImportError("SGLang is not installed. Install with: pip install sglang") from err📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def initialize(self) -> None: | |
| try: | |
| importlib.import_module("sglang") | |
| except ImportError: | |
| raise ImportError("SGLang is not installed. Install with: pip install sglang") | |
| def initialize(self) -> None: | |
| try: | |
| importlib.import_module("sglang") | |
| except ImportError as err: | |
| raise ImportError("SGLang is not installed. Install with: pip install sglang") from err |
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 202-202: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rl_engine/executors/inference_adapter.py` around lines 198 - 202, Update
initialize so the ImportError handler binds the caught exception and re-raises
the custom ImportError explicitly from the original exception, preserving the
exception chain and traceback context.
|
Thank you for work, but The work on ws3 is not currently in the planning stage and requires discussion and approval before a corresponding PR can be submitted. |
|
@Flink-ddd Thank you for your notice, would be happy to contribute to some planned issues and some need-gpu-cis. Looking forward for further discussion on this topic. |
|
Sure, welcome to contribution for gpu-cis issue. |
2 participants
Summary
Part of #127. Lays groundwork for #129.
Background
While using RL-Kernel for a full SFT → GRPO pipeline on Qwen3 (8×A100-SXM4-80GB), I designed a 6-dimension reward function and improved the validation full-field match rate from 83.78% to 85.95% (+2.17%). During training I noticed a clear GPU utilization problem in the standard disaggregated setup (6 GPUs training + 2 GPUs vLLM): training GPUs sit idle during rollout, and inference GPUs sit idle during training — roughly 25% waste.
This PR implements a colocated alternative: all 8 GPUs are shared between vLLM inference and DeepSpeed training, switching via vLLM's sleep/wake mechanism (v0.8+). Sleep/wake overhead is only 3–8% of step time, far less than the 25% idle waste in disaggregated mode.
What Changed
Core orchestration:
examples/colocated_dual_engine.py—ColocatedOrchestrator: vLLM(TP=N) + persistent DeepSpeed workers on shared GPUs. Phase loop:wake → rollout → sleep → train → offload → repeat.examples/_colocated_train_worker.py— Persistent training worker launched once via torchrun, stays alive across all steps. Uses file-based signals for coordination and pinned CPU memory for fast GPU↔CPU weight transfer.Inference adapter:
rl_engine/executors/inference_adapter.py—InferenceEngineAdapterProtocol withsleep()/wake_up()/generate()/update_weights().VLLMInferenceAdapteris fully functional;SGLangInferenceAdapteris a placeholder (SGLang does not yet support sleep mode).Tests and docs:
tests/test_colocated_dual_engine.py— Unit tests (orchestrator init, adapter creation, Protocol conformance, WeightBridge integration) + GPU integration test (vLLM sleep/wake memory lifecycle).docs/guides/colocated_setup.md— Architecture overview, Quick Start, configuration, memory budget.Validation on 8×A100-SXM4-80GB
Training quality: Colocated matches Disaggregated
Both architectures were evaluated on the same validation set (968 samples, 6-dimension reward, Qwen3 SFT → GRPO).
Step-aligned comparison (step 100):
At the same training step, colocated is 0.62% behind on match rate but produces shorter, more parseable outputs. The gap is within noise for 968 samples.
Best checkpoint comparison:
Colocated converges slightly slower (best at step 500 vs 100) but reaches the same peak accuracy. The slower convergence is expected — colocated uses all 8 GPUs for both phases with time-sharing overhead, while disaggregated dedicates 6 GPUs exclusively to training.
GPU efficiency
Colocated benchmark details
gpu_memory_utilization=0.40.Known Limitations
vllm_gpu_memory_utilizationmust be tuned per model size.Testing
Summary by CodeRabbit
New Features
Documentation
Tests