[WIP] [fix] Simplify GPU ID remapping for SGLang colocate mode

[da-niao-dan]

Summary

Fix SGLang CUDA graph capture failure in colocate mode by using local loop indices directly instead of physical GPU IDs from Ray placement groups. Removes the workaround that disabled custom all-reduce v2.

Symptom & Reproduction

• Symptom: SGLang v0.5.12 server fails to start in colocate rollout mode with error: Capture cuda graph failed: Runtime check failed at custom_all_reduce.cuh:37: CUDA error: invalid argument
• Reproduction:

  export CUDA_VISIBLE_DEVICES=1,2,3,4
  python train.py \
    --hf-checkpoint /root/models/Qwen3-30B-A3B \
    --rollout-num-gpus-per-engine 4 \
    --tensor-model-parallel-size 2 \
    --actor-num-gpus-per-node 4 \
    --colocate \
    --num-rollout 1

• Linked issue: SGLang custom all-reduce v2 CUDA graph capture fails under Miles colocate rollout #1176

Root Cause

The original code passed physical GPU IDs from Ray placement groups (e.g., 1,2,3,4) to SGLang engines and then converted them to local IDs via _to_local_gpu_id(). The conversion was based on CUDA_VISIBLE_DEVICES in the parent Ray actor's context, which might differ from what the spawned SGLang server process sees, causing device index mismatch.

History: Why the Remapping Was Added

Before Jan 5, 2026:

• base_gpu_id was computed by get_base_gpu_id(args, rank) returning logical indices (0, 1, 2, ...)
• No CUDA_VISIBLE_DEVICES handling

Jan 5, 2026 (commit c77bdcf):

• Added _to_local_gpu_id() to handle CUDA_VISIBLE_DEVICES=4,5,6,7 style GPU selection
• Purpose: convert physical GPU 4 → local GPU 0, etc.

Later (around Jan 3, 2026, commit 9dfd339):

• Added reordered_gpu_ids from placement groups
• Started passing physical GPU IDs: base_gpu_id = int(reordered_gpu_ids[gpu_index])
• Then _to_local_gpu_id() would convert to local ID

Why it worked before v0.5.12:

• v0.5.10: CustomAllReduceV2 disabled by default → device mapping lenient
• v0.5.12: CustomAllReduceV2 enabled by default → strict device mapping, bug exposed

The Fix

With RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1, the parent Ray actor and spawned SGLang server see the same CUDA_VISIBLE_DEVICES. Therefore, we don't need _to_local_gpu_id() conversion - we can use local loop indices directly.

Code flow:

User: CUDA_VISIBLE_DEVICES=4,5,6,7
  ↓
Ray: Creates placement group with 4 GPUs
  ↓
Engine 0 (i=0): base_gpu_id = 0*4 = 0 → Uses local GPUs 0,1,2,3 ✓
Engine 1 (i=1): base_gpu_id = 1*4 = 4 → Uses local GPUs 4,5,6,7 ✓

Changes:

• rollout.py: Use base_gpu_id = i * num_gpu_per_engine (local loop index)
• sglang_engine.py: Remove _to_local_gpu_id() function, keep get_base_gpu_id() as fallback
• Remove the workaround that disabled SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2

How Users Specify Which GPUs to Use

GPU selection happens at the Ray level via CUDA_VISIBLE_DEVICES:

# Use specific GPUs (4,5,6,7 instead of 0,1,2,3)
export CUDA_VISIBLE_DEVICES=4,5,6,7
python train.py --colocate ...

Status

• Status: WIP - awaiting CI test results
• Local testing: No access to multi-GPU hardware for reproduction
• Relying on: CI tests for verification

Verification

CI tests should verify:

• ✅ SGLang server starts successfully
• ✅ "Custom allreduce v2 initialized" appears in logs
• ✅ "Capture cuda graph end" appears in logs
• ✅ No "CUDA error: invalid argument"
• ✅ Training step completes

[gemini-code-assist]

Code Review

This pull request simplifies GPU device mapping for SGLang engines by leveraging Ray's environment isolation, specifically by hardcoding the base GPU ID to 0 and removing the manual local ID conversion logic. Feedback indicates that the fallback mechanism in sglang_engine.py may still return global GPU indices when base_gpu_id is missing, which could lead to out-of-bounds errors now that the conversion utility has been removed.