|
| 1 | +# Interleaving RL Training with Batch Inference on a Shared GPU |
| 2 | + |
| 3 | +This guide time-slices **one verl fully-async RL training job** with a |
| 4 | +**stock vLLM batch-inference server** on the same GPU. The RL trainer has |
| 5 | +absolute priority; vLLM harvests the trainer's idle valleys (sample waits |
| 6 | +between training bursts — 40–70% of wall-clock for async RL workloads) to |
| 7 | +serve latency-tolerant batch traffic. Neither workload's code is modified: |
| 8 | +verl is integrated through its fully-async lifecycle hooks: the |
| 9 | +`timeslice-verl` package from this repo ships a FullyAsyncTrainer subclass |
| 10 | +registered under the trainer name `timeslice` |
| 11 | +(`async_training.trainer_name=timeslice`), and vLLM runs its normal |
| 12 | +OpenAI-compatible server wrapped by a ~100-line supervisor (shipped in |
| 13 | +`examples/shadow-vllm.yaml` — demo-quality, example-only). |
| 14 | + |
| 15 | +What makes the batch side cheap to yield: vLLM's native **sleep mode** moves |
| 16 | +weights + KV cache to host RAM in ~1–2 s and back in ~100 ms — no process |
| 17 | +restart, no model reload. The RL trainer side uses `cuda-checkpoint` |
| 18 | +(~10–35 s per swap for a 1.5 B trainer), which is amortized against training |
| 19 | +bursts that run for minutes. |
| 20 | + |
| 21 | +**Topology** (two 1-GPU nodes; H100-class): |
| 22 | + |
| 23 | +``` |
| 24 | + ROLLOUT NODE (dedicated) TRAINER NODE (SHARED, time-sliced) |
| 25 | + ┌──────────────────────────┐ ┌──────────────────────────────────┐ |
| 26 | + │ rl-rollout │ ray │ rl-head shadow-vllm │ |
| 27 | + │ RL rollout engine (vLLM, │◀───────▶│ RL trainer ⇄ vLLM server │ |
| 28 | + │ generates continuously) │ │ cuda-checkpoint sleep mode │ |
| 29 | + └──────────────────────────┘ │ C/R (~1-2 s swaps)│ |
| 30 | + └──────────────────────────────────┘ |
| 31 | + lock owner over time on the trainer node's GPU: |
| 32 | + vLLM ████████░T██████████░T█████████░T████████ (T = training burst) |
| 33 | + ▲ trainer queues → vLLM yields in ≤ poll(0.5s) + drain(~3s) + sleep(~1-2s) |
| 34 | +``` |
| 35 | + |
| 36 | +The RL job is a **2-node Ray cluster**: head pod `rl-head` (verl driver + |
| 37 | +FSDP trainer) on the shared trainer node, rollout pod `rl-rollout` (vLLM |
| 38 | +rollout engine) on a dedicated node. Placement is pinned with Ray custom |
| 39 | +resources (`trainer_node` / `rollout_node`) through verl's per-pool |
| 40 | +placement-group bundle resources (hydra override `ray_pg_extra_resources`). |
| 41 | + |
| 42 | +## 1. Components |
| 43 | + |
| 44 | +| Piece | What it is | Where it runs | |
| 45 | +|---|---|---| |
| 46 | +| TimeSlice Orchestrator | gRPC group-lock service | `timeslice-system` ns | |
| 47 | +| Snapshot Agent | DaemonSet; executes cuda-checkpoint C/R AND relays sleep/wake to registered apps (workload channel) | every GPU node | |
| 48 | +| RL job (`examples/rl-job.yaml`) | verl `fully_async_policy` code-RLVR training (Eurus-2 code split, rewards = live test execution), integrated via verl lifecycle hooks (trainer subclass registered as `timeslice`); head pod + rollout pod + headless Service | `default` ns | |
| 49 | +| Shadow vLLM (`examples/shadow-vllm.yaml`) | EXAMPLE-ONLY: stock `vllm serve` + supervisor: holds the lock only while nobody waits; registers sleep/wake callbacks with the agent | `default` ns, trainer node | |
| 50 | +| Load generator (`examples/load-generator.yaml`) | EXAMPLE-ONLY: continuous `/v1/completions` client + throughput log | `default` ns | |
| 51 | + |
| 52 | +How a handoff works, end to end: |
| 53 | + |
| 54 | +1. Trainer's sample batch becomes ready → the timeslice trainer's hook calls |
| 55 | + `acquire()`. |
| 56 | +2. Supervisor's 0.5 s poll sees `waiter_queue_depth > 0` → it **drains** |
| 57 | + (closes its readiness gate, so the pod leaves the Service endpoints and |
| 58 | + new batch requests fail fast; waits ~3 s for in-flight requests to |
| 59 | + finish — vLLM's `/sleep` does NOT drain the scheduler, and sleeping with |
| 60 | + a request mid-decode is a fatal CUDA error, vllm#28714) → calls |
| 61 | + `release()`. |
| 62 | +3. Orchestrator snapshots the vLLM job — the agent's per-job backend |
| 63 | + resolution (explicit config → live workload channel → pod annotation → |
| 64 | + default cuda) lands on its **registered workload channel** → the agent |
| 65 | + invokes the supervisor's callback → `POST /sleep` → HBM freed in ~1–2 s. |
| 66 | +4. Orchestrator restores the trainer (cuda-checkpoint, trainer node only) and |
| 67 | + grants the lock. Trainer runs its burst (minutes), then yields; the |
| 68 | + platform snapshots it and wakes vLLM the same way in reverse; the |
| 69 | + supervisor reopens its readiness gate after the wake. |
| 70 | +5. Batch requests sent while vLLM is draining or asleep fail fast at the |
| 71 | + Service (connection refused — the pod is NotReady); the demo client |
| 72 | + counts them (`errors_or_asleep`) and carries on. Production clients |
| 73 | + should queue and retry — see §8. |
| 74 | + |
| 75 | +## 2. Prerequisites |
| 76 | + |
| 77 | +- Kubernetes cluster (validated on GKE/COS) with **two free 1-GPU nodes** |
| 78 | + (H100-class recommended; reference shape GKE `a3-highgpu-1g`, 26 vCPU / |
| 79 | + 234 Gi each): one shared trainer node (RL head pod + shadow vLLM pod) and |
| 80 | + one dedicated RL rollout node. |
| 81 | +- Cluster-admin; `kubectl`, `helm` v3, `git`, `envsubst` on your workstation. |
| 82 | +- Pods can reach Docker Hub, GitHub, HuggingFace. No HF token needed (models |
| 83 | + are fetched by HF id per node). |
| 84 | +- Host RAM on the trainer node: the RL head pod requests 80 Gi (110 Gi |
| 85 | + limit; trainer checkpoint ≈ full GPU allocation) + 60 Gi for the shadow |
| 86 | + vLLM's sleep offload. |
| 87 | + |
| 88 | +**Version pins** (keep on first run): verl branch |
| 89 | +[`feat/fully-async-lifecycle-hooks`](https://github.com/aishukamal/verl/tree/feat/fully-async-lifecycle-hooks) |
| 90 | +of `github.com/aishukamal/verl` — upstream |
| 91 | +`983cb0f24443f87b3d161fad318445130a620b07` plus two feature commits |
| 92 | +(fully-async lifecycle hooks + trainer registry; per-pool PG bundle |
| 93 | +resources); temporary fork until the commits land upstream. Job image |
| 94 | +`verlai/verl:vllm020.dev2`; shadow vLLM image `vllm/vllm-openai:v0.9.2`; |
| 95 | +platform images: official `ghcr.io/llm-d-incubation/llm-d-rl-time-slicing/*` |
| 96 | +containing per-job backend resolution in the |
| 97 | +agent (config-less requests resolve: explicit config → live workload |
| 98 | +channel → `timeslice.io/backend` pod annotation → default cuda; the |
| 99 | +workload-channel step is required by this guide — see the release-pin |
| 100 | +note in §3); |
| 101 | +integration package `timeslice-verl` from this repo, |
| 102 | +`pkg/integrations/verl/` (ConfigMap install, see §5); trainer |
| 103 | +model `deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B`; batch model |
| 104 | +`Qwen/Qwen2.5-0.5B-Instruct`. |
| 105 | + |
| 106 | +## 3. Step 1 — Install the platform |
| 107 | + |
| 108 | +```bash |
| 109 | +git clone https://github.com/llm-d-incubation/llm-d-rl-time-slicing.git |
| 110 | +cd llm-d-rl-time-slicing |
| 111 | +helm dependency update ./deploy |
| 112 | +# If a previous `timeslice` release exists, `helm uninstall` it first — |
| 113 | +# helm install fails on an existing release name. NOTE: the chart owns the |
| 114 | +# timeslice-system namespace, so uninstall deletes it; wait for the |
| 115 | +# namespace to finish terminating (~30-60 s) or the reinstall fails with |
| 116 | +# "namespace is being terminated". |
| 117 | +helm install timeslice ./deploy -n timeslice-system --create-namespace |
| 118 | + |
| 119 | +kubectl -n timeslice-system get pods -o wide # orchestrator + agents Running |
| 120 | +``` |
| 121 | + |
| 122 | +> **Release pin:** the chart defaults pull the official |
| 123 | +> `ghcr.io/llm-d-incubation/llm-d-rl-time-slicing/*` images at `latest`, |
| 124 | +> which include per-job backend |
| 125 | +> resolution in the snapshot agent (#159 — required here: the workload |
| 126 | +> channel is how the agent drives vLLM's sleep/wake instead of |
| 127 | +> cuda-checkpointing it). Once the first tagged release containing it is |
| 128 | +> cut, pin it with `--set timesliceorchestrator.image.tag=<version> |
| 129 | +> --set snapshot-agent.image.tag=<version>`. |
| 130 | +
|
| 131 | +No snapshot-device filter is needed in this topology: every node has one GPU, |
| 132 | +and only the labeled pods on the trainer node (rl-head, shadow-vllm) are |
| 133 | +snapshot targets — the rollout pod is unlabeled on another node. |
| 134 | + |
| 135 | +## 4. Step 2 — Pick the nodes and label the trainer node |
| 136 | + |
| 137 | +```bash |
| 138 | +export TRAINER_NODE=<the shared 1-GPU node (RL trainer + shadow vLLM)> |
| 139 | +export ROLLOUT_NODE=<the RL rollout's dedicated 1-GPU node> |
| 140 | +
|
| 141 | +kubectl label node "$TRAINER_NODE" group.timeslice.io/trainers=true --overwrite |
| 142 | +``` |
| 143 | +
|
| 144 | +Only the trainer node gets the label (the rollout node must NOT carry it). |
| 145 | +Verify an agent runs on the trainer node and the orchestrator synced the |
| 146 | +group: |
| 147 | +
|
| 148 | +```bash |
| 149 | +AGENT_POD=$(kubectl -n timeslice-system get pods -l app.kubernetes.io/name=snapshot-agent \ |
| 150 | + --field-selector spec.nodeName=$TRAINER_NODE -o jsonpath='{.items[0].metadata.name}') |
| 151 | +echo "$AGENT_POD" |
| 152 | +kubectl -n timeslice-system logs deploy/timeslice-timesliceorchestrator --tail=50 | grep -i trainers | tail -3 |
| 153 | +``` |
| 154 | +
|
| 155 | +## 5. Step 3 — Launch (order matters) |
| 156 | +
|
| 157 | +First, publish the `timeslice-verl` integration package into the cluster as |
| 158 | +ConfigMaps (the RL job pods install it from these — no GitHub fetch): |
| 159 | +
|
| 160 | +```bash |
| 161 | +# Package source: the repo you cloned in §3 |
| 162 | +PKG=llm-d-rl-time-slicing/pkg/integrations/verl |
| 163 | +kubectl create configmap timeslice-verl-root --from-file=$PKG/pyproject.toml |
| 164 | +kubectl create configmap timeslice-verl-src --from-file=$PKG/timeslice_verl/ |
| 165 | +``` |
| 166 | +
|
| 167 | +ConfigMap keys cannot contain `/`, so the package tree is split into two maps |
| 168 | +(project root + module sources) and reassembled inside the pod; this also |
| 169 | +works in air-gapped clusters. |
| 170 | +
|
| 171 | +**verl source**: the RL pods clone the `feat/fully-async-lifecycle-hooks` |
| 172 | +branch of `github.com/aishukamal/verl` and install it with |
| 173 | +`pip install -e ".[gpu]"`. Once the commits land upstream, point the |
| 174 | +`VERL_REPO`/`VERL_REF` pod envs at mainline verl instead. |
| 175 | +
|
| 176 | +Then start the **shadow vLLM first** so it owns the GPU during the trainer's |
| 177 | +long CPU-side setup, then the RL job (both pods at once — the rollout pod's |
| 178 | +installs run in parallel with the head's), then the load: |
| 179 | + |
| 180 | +```bash |
| 181 | +export RUN_SECONDS=5400 # RL training budget (~90 min) |
| 182 | +
|
| 183 | +envsubst '${TRAINER_NODE}' < examples/shadow-vllm.yaml | kubectl apply -f - |
| 184 | +# wait for "[supervisor] vLLM is up" (first model download ~2-4 min): |
| 185 | +kubectl logs shadow-vllm --tail=20 |
| 186 | +
|
| 187 | +envsubst '${TRAINER_NODE} ${ROLLOUT_NODE} ${RUN_SECONDS}' < examples/rl-job.yaml | kubectl apply -f - |
| 188 | +kubectl apply -f examples/load-generator.yaml |
| 189 | +``` |
| 190 | + |
| 191 | +The RL job spends 20–45 min on setup (verl install on both pods, the head's |
| 192 | +wait for the rollout pod to join the ray cluster, model + dataset prep — |
| 193 | +**both pods prepare the dataset locally**: the `FullyAsyncRollouter` actor |
| 194 | +runs on the rollout pod and reads the parquet there; no shared volume, same |
| 195 | +seed ⇒ identical files) before it first requests the GPU — vLLM serves batch |
| 196 | +traffic the whole time. |
| 197 | +
|
| 198 | +## 6. Step 4 — Watch it work |
| 199 | +
|
| 200 | +```bash |
| 201 | +kubectl logs batch-load-generator --tail=6 |
| 202 | +# [12:01:05] last 5s: completed=41 errors_or_asleep=0 <- vLLM holds GPU |
| 203 | +# [12:04:35] last 5s: completed=0 errors_or_asleep=12 <- trainer burst |
| 204 | +# [12:08:10] last 5s: completed=38 errors_or_asleep=1 <- vLLM back |
| 205 | +
|
| 206 | +kubectl logs shadow-vllm --tail=10 |
| 207 | +# [supervisor] trainer is waiting - yielding GPU |
| 208 | +# [supervisor] drained in 2.82s |
| 209 | +# [supervisor] vLLM slept (HBM -> host RAM) in 0.97s |
| 210 | +# [supervisor] lock reacquired (waited 214380 ms, context_restored=True) |
| 211 | +# [supervisor] vLLM woke in 0.05s |
| 212 | +
|
| 213 | +# lock handoffs: the integration package logs every acquire/release |
| 214 | +# ([timeslice] ... ACQUIRE/RELEASE lines, forwarded into the head pod's log): |
| 215 | +kubectl logs rl-head --tail=2000 | grep -E "ACQUIRE|RELEASE" | tail -4 |
| 216 | +# [timeslice] job=rl-trainer ACQUIRE group=trainers waited=8210ms context_restored=True |
| 217 | +# [timeslice] job=rl-trainer RELEASE group=trainers pending_waiters=1 snapshot_deferred=False |
| 218 | +``` |
| 219 | + |
| 220 | +Healthy steady state: the load generator alternates between full-throughput |
| 221 | +windows (trainer idle) and error windows a few minutes long (trainer burst); |
| 222 | +supervisor drain ≈ 3 s, sleep ≤ 3 s, wake ≤ 1 s; trainer `waited=` on |
| 223 | +ACQUIRE ≈ 5–11 s (poll latency + drain + vLLM sleep + the trainer's own |
| 224 | +cuda-checkpoint context restore, ~5 s once training state is large — the |
| 225 | +trainer never queues behind batch work). |
| 226 | + |
| 227 | +Success criteria to check after ~3 training steps: |
| 228 | + |
| 229 | +- Trainer step time within ~10% of a solo run (compare `train.log` timing to |
| 230 | + a run without the shadow vLLM, or to the reference: median 570 s/step). |
| 231 | +- Batch throughput > 0 in every trainer-idle window. |
| 232 | +- No cuda-checkpoint operations targeting the rollout node in any agent log |
| 233 | + (the rollout node must never show cuda-checkpoint activity). |
| 234 | + |
| 235 | +## 7. Step 5 — Collect results |
| 236 | + |
| 237 | +```bash |
| 238 | +mkdir -p results |
| 239 | +kubectl cp rl-head:/workspace/results/train.log results/train.log || true |
| 240 | +kubectl logs batch-load-generator --timestamps > results/batch_throughput.log |
| 241 | +kubectl logs shadow-vllm --timestamps > results/supervisor.log |
| 242 | +``` |
| 243 | + |
| 244 | +`train.log` (trainer bursts, with the `[timeslice]` ACQUIRE/RELEASE |
| 245 | +lock-handoff lines) + `batch_throughput.log` (harvested inference) together |
| 246 | +give the shared-GPU timeline: every second is either a training burst, batch |
| 247 | +serving, or a swap. |
| 248 | + |
| 249 | +## 8. Troubleshooting |
| 250 | + |
| 251 | +| Symptom | Cause | Fix | |
| 252 | +|---|---|---| |
| 253 | +| vLLM never sleeps at handoff; trainer restore then fails with OOM on the shared GPU | Workload registration didn't reach the agent (check supervisor log for `workload registered`) | Verify `TIMESLICE_AGENT_ADDR` resolves to the node IP :9001 and agent logs show the registration; the supervisor's local wake fallback covers restores but sleep MUST go through the channel | |
| 254 | +| `POST /sleep` returns 404 | vLLM started without dev endpoints | `VLLM_SERVER_DEV_MODE=1` and `--enable-sleep-mode` are both required (set in the manifest) | |
| 255 | +| vLLM dies seconds after a sleep (`CUDA error: an illegal memory access` in EngineCore), then the group faults on the failed wake | `/sleep` was called with requests in flight — vLLM does not drain the scheduler on sleep (vllm#28714, also affects ≥0.10.x; requests sent to a sleeping engine crash it too, vllm#15483) | Keep the supervisor's readiness-gate drain (in the manifest): gate closed + `running+waiting==0` BEFORE `release()`. Never call `/sleep` on a serving engine directly | |
| 256 | +| Trainer placement group lands on the rollout node (or rollout work on the trainer node) | PG pinning not active — `--resources` missing on a `ray start` line, or the `ray_pg_extra_resources` hydra override is missing/typo'd, or the verl build lacks the feature | Verify both `ray start` lines carry `--resources`, keep the override in `examples/rl-job.yaml`'s `run_head.sh` intact, and use the pinned verl branch (§2); check `ray status` custom resources | |
| 257 | +| cuda-checkpoint activity on the rollout node's agent, or rollout throughput collapses during handoffs | The rollout pod/node got labeled into the group | The rollout node must never show cuda-checkpoint activity: keep `timeslice.io/*` labels off rl-rollout and the group label off `$ROLLOUT_NODE` | |
| 258 | +| `rl-head` log stuck at `Phase 2b: wait for the rollout pod to join`, then `FATAL: rollout pod did not join` | rl-rollout Pending/crashed, or headless Service `rl-head` missing | `kubectl get pod rl-rollout; kubectl logs rl-rollout --tail=50`; apply the whole rendered `examples/rl-job.yaml` | |
| 259 | +| Load generator: 100% errors even when trainer is idle | vLLM crashed (supervisor exits, pod restarts) or Service selector mismatch | `kubectl logs shadow-vllm --previous`; check `kubectl get endpoints shadow-vllm` | |
| 260 | +| RL head or shadow vLLM pod Pending | Trainer node CPU/RAM too small for 8 CPU + 80 Gi (head) alongside 6 CPU + 60 Gi (vLLM) | Use an a3-highgpu-1g-class node (26 vCPU / 234 Gi), or shrink requests | |
| 261 | + |
| 262 | +Demo-client caveat, worth stating to any customer: while vLLM is draining or |
| 263 | +asleep, new batch requests fail fast at the Service (the pod is NotReady). |
| 264 | +The demo load generator just counts those errors. A production batch |
| 265 | +front-end should be queue-based with retries — this composes naturally with |
| 266 | +[llm-d async processor](https://github.com/llm-d/llm-d-async) as the batch |
| 267 | +dispatcher, which is the planned productization path. |
| 268 | + |
| 269 | +## 9. Adapting |
| 270 | + |
| 271 | +- **Bigger batch model**: anything that fits in `VLLM_GPU_FRAC` of the shared |
| 272 | + GPU alongside zero trainer residency (they never co-reside). Sleep/wake |
| 273 | + time scales with weight size (~1 s per 15 GB to host RAM). |
| 274 | +- **Your RL workload**: swap model/dataset in `examples/rl-job.yaml` exactly as in the |
| 275 | + two-RL-jobs guide (§9 there); everything under "Required" comments stays. |
| 276 | +- **Yield latency vs. poll cost**: `WAITER_POLL_SECONDS` (default 0.5) is the |
| 277 | + worst-case extra wait the trainer sees. |
| 278 | +- **Trainer solo baseline** (for the ≤10% overhead check): run `examples/rl-job.yaml` |
| 279 | + alone — without the shadow vLLM the trainer acquires instantly every step |
| 280 | + and the platform defers snapshots (nothing contends), so it behaves like an |
| 281 | + unshared run. |
| 282 | + |
| 283 | +## 10. Teardown |
| 284 | + |
| 285 | +```bash |
| 286 | +kubectl delete pod shadow-vllm batch-load-generator --ignore-not-found |
| 287 | +kubectl delete pod rl-head rl-rollout --ignore-not-found |
| 288 | +kubectl delete service shadow-vllm rl-head --ignore-not-found |
| 289 | +kubectl delete configmap shadow-vllm-scripts rlbatch-trainer --ignore-not-found |
| 290 | +kubectl delete configmap timeslice-verl-root timeslice-verl-src --ignore-not-found |
| 291 | +helm uninstall timeslice -n timeslice-system |
| 292 | +kubectl label node "$TRAINER_NODE" group.timeslice.io/trainers- || true |
| 293 | +``` |
0 commit comments