[Question] What's the plan for rollout `pp_size>1` on the SGLang engine?

[ruiling-smartbear]

Pipeline parallelism on the SGLang rollout engine is currently blocked in three places:

file	guard
unirl/rollout/engine/sglang/config.py:146	pp_size>1 is not supported yet — UniRL Handle would spawn one engine per pp_rank while SGLang spawns its own PP scheduler subprocesses, double-booking GPUs
unirl/distributed/weight_sync/full/nccl.py:73	NCCLWeightSync.connect: rollout pp_size>1 is not implemented
unirl/distributed/group/dispatch.py:178	rejected because auto-backward cannot propagate gradients across pipeline stages

The config comment names two directions as future work — per-stage rank_offset routing, and single-engine PP fan-out.

Reading sglang/engine.py, TP already resolves the same ownership question by delegating: every TP rank gets an engine instance, but only tp_rank=0 is live (engine.py:76-88 turns the rest into no-op shells), and that one instance receives the whole TP group via cuda_visible_devices=tp_visible_devices plus runtime_overrides["tp_size"] (engine.py:121-124, :153), so SGLang does the sharding internally.

Questions:

1. Is single-engine PP fan-out (same delegation pattern as TP, gated on pp_rank == 0 and tp_rank == 0) the intended direction, or is per-stage rank_offset routing preferred?
2. NCCLWeightSync.connect currently derives the broadcast group from num_rollout_gpus and tp_size. Under the delegation model the rollout side would still expose one live engine per PP group — is reworking that rank accounting the main blocker, or is there more?
3. The dispatch.py guard is about gradients crossing pipeline stages. Rollout engines don't backprop, so is that guard incidental to this path, or does it apply for a reason I'm missing?
4. Is anyone working on this already?

Happy to take a piece of it if the design direction is settled. Asking before writing anything, since the config comment suggests the approach is still open.

[CjhHa1]

Thanks for digging into this. The single-engine PP fan-out direction is the more appropriate fit for the current architecture. For each DP replica, only the pp_rank == 0 && tp_rank == 0 head should own a live SGLang engine, and SGLang should own the full TP x PP topology internally; the remaining UniRL ranks should be no-op shells. This keeps parallelism ownership in one place and extends the delegation model already used for TP.

The NCCL rank accounting is one blocker, but not the only one. The implementation will also need to:

• build and pass the full TP x PP device set to the engine, including placement constraints;
• update the live-engine/no-op predicate and lifecycle routing;
• change weight-sync targets, rank_offset, and world-size accounting for one engine per DP replica;
• adjust synchronous rollout collection, which currently collects from the pipeline last stage while the delegated engine head would live at pp_rank == 0; and
• add focused TP+PP construction, generation, and weight-sync coverage.

The dispatch.py guard is specific to auto-backward calls made inside enable_grad(), so it is incidental to ordinary rollout inference. It does not need to be removed just to support PP rollout, although the forward dispatch/collection semantics above still need attention.

There is no active implementation we are aware of. Please feel free to open a PR following the single-engine design, or a follow-up implementation issue if you would like to first spell out the device/rank layout and split the work into stages.

[ruiling-smartbear]

Thanks for the detailed reply — really helpful.

I'll go with the single-engine fan-out and, as you suggested, write up the device/rank layout in an issue first. I'll leave dispatch.py alone.

Happy to start on whichever piece is most useful to you; otherwise I'll link the issue here once it's up.

[ruiling-smartbear]

Following the single-engine PP fan-out direction confirmed here, the rollout-engine half is up for review in #502: for a replica with tp_size=T, pp_size=P, only the tp_rank == 0 && pp_rank == 0 head boots a live SGLang engine, SGLang owns the TP×PP topology internally, and the other ranks are no-op shells. It also unifies the owner predicate (the repo answered "who owns the engine?" four different ways, two of them wrong under PP), hands the head the whole replica's devices in the order SGLang's own compute_local_gpu_id expects, and fails closed on replicas that straddle a node.

Generation works — verified end to end on one 4×L4 node (pp_size=1, tp=1 pp=2, tp=2 pp=2, tp=1 pp=4), with CUDA graphs on, on 8×L4 with two engine heads (dp=2 tp=2 pp=2), and on the repository's own DAPO-Math data with math-verify scoring. Holding tp and the attention backend fixed, tp=1, pp=1 and tp=1, pp=2 produce the same result.

Weight sync is not in that PR, because all three existing paths are blocked — measured, not inferred, and each now fails closed instead of dying inside SGLang:

1. TensorWeightSync: SGLang reads the payload list by the stage-local tp_rank, so each pipeline stage reads the same slots, and a CUDA-IPC handle is only openable on the GPU that produced it → AcceleratorError: CUDA error: invalid argument on the pp>0 stage. CPU tensors are not an escape (the reduction patch indexes arg 6, absent for CPU args → IndexError).
2. NCCLWeightSync: rank = rank_offset + self.tp_rank is stage-local too, so both stages join as rank=1 — and the run still exits 0, i.e. a silent mis-sync.
3. CkptEngineIPCWeightSync: addressing is per device UUID, so it could support PP, but the server-side fan-out does not reach the pp>0 scheduler (its log shows no IPC activity; the pp=0 sender's handshake times out after 600 s).

Details and evidence are in #503. The short version of what we are missing upstream: weight
sync has to be shard-aware — it must know which parameters each inference rank owns, since a
pipeline stage holds only a subset of layers. None of the three paths implements that, and
SGLang has no test covering weight update with pp_size>1.

[ruiling-smartbear]

PR #502 follows the single-engine design. Against your list:

• Full TP x PP device set + placement — done: PP-major token order matching SGLang's own compute_local_gpu_id, and a replica that straddles a node is rejected.
• live-engine/no-op predicate + lifecycle routing — done: tp_rank == 0 && pp_rank == 0, 16 gated methods.
• Synchronous collection — done: a delegated engine now collects from the pp_rank == 0 head; non-delegated roles keep collecting at the last stage.
• dispatch.py guard — untouched, per your note.
• Weight-sync targets — done: tp_zero_workers -> engine_head_workers.
• rank_offset / world-size accounting — done: a delegated engine owns the replica's whole tp*pp rank set, so connect now offsets target i by i*tp*pp + 1 and takes world = num_targets*tp*pp + 1, and it rejects a num_rollout_gpus that disagrees instead of forming a group of the wrong size. Verified on GPU at dp=2, tp=1, pp=1: rank_offset=1, rank=1, world_size=3 and rank_offset=2, rank=2, world_size=3, rc=0 with the sync landing between rollouts. If the delegated head is instead meant to expose only its tp ranks to the trainer, say so and we will flip it back.
• TP+PP weight-sync coverage — not done; all three paths are blocked inside SGLang. Measured: TensorWeightSync reads the payload list at the stage-local tp_rank, so stage pp>0 opens another GPU's CUDA-IPC handle and dies with invalid argument; the update_weights_from_ipc fan-out never reaches the pp>0 scheduler. NCCLWeightSync derives its group rank from the same stage-local tp_rank, so an engine owning tp*pp ranks can present only tp distinct ones. Details and logs in SGLang rollout pp_size>1: weight sync needs a shard-aware path #503.

We fail closed on pp_size>1 rather than risk a silent mis-sync, so under PP #502 is generation-only.