feat(rollout): send per-GPU weight-sync payloads to multi-GPU engines

[Rockdu]

What

• Weight sync labels each CUDA-IPC payload with the physical GPU UUID it was exported from (payload_gpu_uuids), so every sgl-d worker imports the payload of the training rank on its own GPU.
• Rollout engines may span more than one GPU: num_gpus is the engine's GPU span, tp_size comes from --sglang-tp-size instead of being hardcoded to that span, and the engine pins the whole span instead of a single GPU.
• validate_args rejects --sglang-tp-size * --sglang-sp-degree != --rollout-num-gpus-per-engine.
• Adds a Wan2.2 recipe with rollout sequence parallelism (2 engines x 2 GPUs, sp_degree=2).

Why

--rollout-num-gpus-per-engine > 1 used to mean TP only, because tp_size was pinned to the span and num_gpus was never passed, so --sglang-sp-degree could not be satisfied. With the span split by tp_size * sp_degree, the same allocation serves TP, SP, or TP x SP rollout engines.

The payload labels remove an assumption in the weight-sync path: the gathered list is ordered by training rank, and a worker only ends up with its co-located rank's buffer if the engine's GPU span is ordered the same way. Labelling by GPU makes the pairing explicit, and a mismatch (engine span != the span miles gathered over) fails loudly instead of silently importing a peer GPU's memory.

Requires the matching sglang-d change (payload_gpu_uuids on UpdateWeightFromTensorReqInput); a build without it drops the field and the worker rejects the request instead of guessing.

Files

• miles/backends/fsdp_utils/diffusion_update_weight_utils.py — gather (gpu_uuid, payloads) per rank, send the labels with the request.
• miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py — forward payload_gpu_uuids; num_gpus/tp_size server args; pin the engine's whole GPU span.
• miles/backends/sglang_diffusion_utils/arguments.py — reject a span that tp_size * sp_degree does not fill.
• scripts/run-diffusion-grpo-wan22-pickscore-5gpu-rollout-sp2.sh — the 5-GPU Wan2.2 recipe plus --rollout-num-gpus-per-engine 2 --sglang-sp-degree 2.

Validation

On 8xH200 with the Wan2.2 recipe at rollout_num_gpus_per_engine=2, sp_degree=2: both engines started as num_gpus=2, tp_size=1, sp_degree=2, ulysses_degree=2 with CFG parallel auto-disabled, each pinned to its own GPU pair, 8 weight syncs returned 200 OK across both experts with zero skipped updates, and a temporary probe confirmed all four workers imported a buffer on their own physical GPU.

That run used an earlier variant of the sgl-d side (world-rank indexing); the UUID-labelled path here has only had a logic walkthrough. End-to-end (rollout samples -> reward -> first train step) is still unverified: the run was cut short, and the old sp branch it started from could not finish a rollout because it predates the msgpack rollout transport (af1fdc0).

Checklist

• pre-commit run passes on the touched files
• Added/updated tests for new behaviour — no tests added
• pytest -x is green — not run
• If launch flags changed, python3 train.py --help still parses — no new flags; only existing --sglang-* args are used
• If a public flag was added, it appears in the CLI reference docs — n/a
• If an example was added, it has a real walkthrough — the new recipe mirrors the existing 5-GPU one