feat(rollout): add TP/EP/PP parallel strategies on SGLang rollout side

[NancyFyong]

Summary

Add tensor-parallel (TP) distribution on the SGLang rollout engine in colocate mode, with EP sharding inside SGLang and PP rank layout reserved for future work. Larger rollout models (e.g. Qwen3.5-MoE 35B) need >1 GPU per engine; before this change UniRL could only do dp_size = world_size, one GPU per engine.

The train side stays flat (FSDP). Each SGLang engine now spans tp_size GPUs: tp_rank==0 owns the server, tp_rank>0 workers are no-op shells that participate in training but hold no SGLang process. EP is forwarded to ServerArgs; PP>1 populates RankInfo but weight sync is fail-closed (NotImplementedError).

Key changes:

1. Rank layout (handle.py): _build_rank_infos lays out (dp, pp, tp) with TP rank fastest. Handle.__init__ injects tp_rank/tp_size/tp_device_ids/pp_rank/pp_size/ep_size into the engine constructor. sp_size>1 combined with tp/pp>1 is rejected; tp=pp=ep=1 reproduces the flat layout.

2. SGLang engine (sglang/{config,engine}.py): tp_rank>0 constructs as a no-op shell (no server, no ports). All rollout verbs early-return on tp_rank>0. tp_rank==0 with tp_size>1 passes base_gpu_id via runtime_overrides instead of overriding CUDA_VISIBLE_DEVICES globally.

3. Weight sync — all four handlers guard on tp_rank:

   • NCCLWeightSync: rank_offset = i*tp + 1, world_size = num_engines*tp + 1; PP>1 raises.
   • TensorWeightSync: ships [payload]*tp_size on tp_rank==0; tp_rank>0 drains the FSDP all-gather in lockstep.
   • IPCWeightSync: receiver started once per stage; tp_rank>0 shells drop at the collect filter.
   • LocalLoraWeightSync/RemoteLoraWeightSync: extract on all ranks, push only from tp_rank==0/rank-0.

4. TP>1 colocate crash fix (sglang/backends/{native,http}.py): SGLang Engine spawn temporarily exposes all TP GPUs via CUDA_VISIBLE_DEVICES then restores Ray's CVD. LoRA tensor serialization switched from ForkingPickler to plain pickle to avoid fd handle races across TP subprocesses.

Related Issue

N/A

Test Plan

• SKIP=no-commit-to-branch pre-commit run --all-files --show-diff-on-failure
• pytest scripts/tests/rollout/ -k on_cpu — CPU unit tests (rank layout, shell construction, weight sync tp_rank guards)
• python scripts/check_recipe_targets.py — _target_ paths resolve
• GPU E2E (manual, tiny-random/qwen3.5-moe):
  • 2 GPUs, tp=2: boot + generate + sleep/wake — EXIT=0
  • 2 GPUs, tp=2 weight-sync round-trip — tokens returned post-sync
  • 8 GPUs, tp=2 ep=2 (4 engines × 2 GPU): full GRPO loop — reward/loss logged, EXIT=0
• Not run; reason: PP>1 weight sync is fail-closed; TP=4/EP=4 blocked by tiny model's num_heads=2 (cannot divide tp_size=4)

Compatibility / Risk

• Config additions (backward-compatible): SGLangEngineConfig gains pp_size/ep_size/dp_size/enable_expert_parallel, all default to 1/None. tp_size=1 path is asserted identical to pre-change layout.
• Weight sync: NCCLWeightSync.connect now takes tp_size; pp_size>1 raises. All handlers' tp_size=1 path is unchanged.
• Risk area: tp_rank>0 workers must construct without booting SGLang or reserving ports. A regression would spawn duplicate multi-GPU servers and OOM. Covered by test_tp_engine_shell_on_cpu.py and E2E boot test.

Reviewer Notes

• AI assistance: implementation and tests drafted with CodeBuddy assistance; every changed line reviewed by the submitter before push.
• Duplicate-work check: no overlap found. PR feat(veomni-ep): expert parallelism for the VeOmni FSDP2 training backend #158 (feat/veomni-backend-ep-pr) is training-backend EP (VeOmni FSDP2), distinct from this PR's rollout-side (SGLang) EP.
• Limitations: PP>1 per-stage weight sync routing is fail-closed; PP is exercised at rank-layout / dispatch-filter level only. SGLang E2E tests force-exit pytest on success (SGLang's in-process shutdown fires SIGQUIT at siblings); long-term fix is to run the engine inside a Ray actor.
• Suggested review order: handle.py (rank layout + kwargs injection) → weight_sync/*.sync (tp_rank guards) → sglang/engine.py (tp_rank>0 no-op shell)

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated tests, docs, and configs where needed, or explained why not.