[samsja]

Summary

Adds MultiNodeDeploymentConfig.orchestrator_on_inference (default False) to choose where the orchestrator launches:

• False (default, unchanged): orchestrator runs on trainer rank 0.
• True: orchestrator runs on the last inference node (a decode node — tail of the last replica, never a router/prefill head).

Why

On large multi-node runs the trainer nodes can be host-RAM-saturated (e.g. 8 ranks of CPU-offloaded optimizer state ≈ 800 GiB). Co-locating the orchestrator there means its rollout buffer + env workers compete for the same host RAM and can exhaust it (we hit a trainer SIGBUS from exactly this). Inference nodes keep weights/KV on GPU and typically have ample free host RAM, so moving the orchestrator there removes the contention. This makes it a per-deployment choice rather than hardcoding either placement.

Implementation

The orchestrator launch is now a single top-level block in multi_node_rl.sbatch.j2 (after the role if/else, so the existing wait -n teardown still covers it), gated on a flag-selected SLURM_PROCID:

{% if orchestrator_on_inference %}ORCH_PROCID=$((NUM_INFER_NODES - 1)){% else %}ORCH_PROCID=$NUM_INFER_NODES{% endif %}
if [ "$SLURM_PROCID" -eq "$ORCH_PROCID" ]; then ... fi

Both placements share one code path — no duplicated launch logic. NUM_INFER_NODES maps to trainer rank 0 (since TRAIN_NODE_RANK = SLURM_PROCID - NUM_INFER_NODES), NUM_INFER_NODES-1 is the last inference node. No effect when num_infer_nodes == 0.

Validation

Rendered the template both ways and bash -n'd the output — both valid; orchestrator launches exactly once in each, on the expected node.

---

Note

Medium Risk
Changes where a critical RL control process runs in multi-node SLURM jobs; default placement is unchanged, but enabling the flag shifts rollout/env RAM load onto an inference node that also runs vLLM.

Overview
Adds deployment.orchestrator_on_inference (default False) on multi-node RL configs so SLURM jobs can place the orchestrator on the last inference node instead of trainer rank 0, to avoid host-RAM contention with heavy trainer CPU offload.

The multi_node_rl.sbatch.j2 launch is refactored: orchestrator startup is no longer inside the trainer-only branch on rank 0, but a single shared block after the infer/train role split that picks ORCH_PROCID as either NUM_INFER_NODES (unchanged default) or NUM_INFER_NODES - 1 when the flag is set—still gated on num_infer_nodes > 0 and still covered by wait -n teardown.

write_slurm_script passes the flag into both standard and disaggregated template renders.

^{Reviewed by Cursor Bugbot for commit febd91c. Bugbot is set up for automated code reviews on this repo. Configure here.}