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

miles/backends/fsdp_utils/diffusion_update_weight_utils.py

@@ -31,6 +31,8 @@
     compute_weights_checksum = None
     _checksum_import_error = _e
 
+from miles.ray.utils import get_physical_gpu_id
+
 
 logger = logging.getLogger(__name__)
 
@@ -233,24 +235,27 @@ def update_bucket_weights(
             serialized_tensors.append(MultiprocessingSerializer.serialize(flattened_tensor_data, output_str=True))
 
         if self._ipc_gather_src == dist.get_rank():
-            gathered_serialized_batches = [None for _ in range(dist.get_world_size(self._ipc_gather_group))]
+            gathered_batches = [None for _ in range(dist.get_world_size(self._ipc_gather_group))]
         else:
-            gathered_serialized_batches = None
+            gathered_batches = None
 
         dist.gather_object(
-            obj=serialized_tensors,
-            object_gather_list=gathered_serialized_batches,
+            obj=(get_physical_gpu_id(), serialized_tensors),
+            object_gather_list=gathered_batches,
             dst=self._ipc_gather_src,
             group=self._ipc_gather_group,
         )
 
         if dist.get_rank() == self._ipc_gather_src:
+            payload_gpu_uuids = [gpu_uuid for gpu_uuid, _ in gathered_batches]
+            gathered_serialized_batches = [tensors for _, tensors in gathered_batches]
             # TODO: here we assume all ranks have the same number of dtypes.
             num_dtypes = len(gathered_serialized_batches[0])
             assert num_dtypes > 0
             for i in range(num_dtypes):
                 kwargs = {
                     "serialized_named_tensors": [tensors[i] for tensors in gathered_serialized_batches],
+                    "payload_gpu_uuids": payload_gpu_uuids,
                     "load_format": "flattened_bucket",
                     "target_modules": [target_module],
                     "weight_version": str(weight_version),

miles/backends/sglang_diffusion_utils/arguments.py

@@ -93,6 +93,13 @@ def new_add_argument_wrapper(*name_or_flags, **kwargs):
 
 
 def validate_args(args):
+    if args.sglang_tp_size is not None and args.sglang_sp_degree is not None:
+        if args.sglang_tp_size * args.sglang_sp_degree != args.rollout_num_gpus_per_engine:
+            raise ValueError(
+                f"--sglang-tp-size ({args.sglang_tp_size}) * --sglang-sp-degree ({args.sglang_sp_degree}) "
+                f"must equal --rollout-num-gpus-per-engine ({args.rollout_num_gpus_per_engine})"
+            )
+
     # `sglang_dp_size` is used by rollout port allocation; default to 1 if
     # SGL-D's ServerArgs didn't register the CLI arg in this build.
     if not hasattr(args, "sglang_dp_size"):

miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py

@@ -165,14 +165,15 @@ def _init_normal(self, server_args_dict):
     def _pin_to_assigned_gpu(self):
         if self.base_gpu_id is None:
             return
+        span = self.args.rollout_num_gpus_per_engine
         cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "")
         if not cvd:
-            # No ambient device list (full-node multi-node ray start): pin to the physical GPU.
-            os.environ["CUDA_VISIBLE_DEVICES"] = str(self.base_gpu_id)
+            # No ambient device list (full-node multi-node ray start): pin to the physical GPUs.
+            os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(self.base_gpu_id + i) for i in range(span))
             return
         visible = [x.strip() for x in cvd.split(",") if x.strip()]
         local_idx = _to_local_gpu_id(self.base_gpu_id)
-        pinned = visible[local_idx]
+        pinned = ",".join(visible[local_idx : local_idx + span])
         os.environ["CUDA_VISIBLE_DEVICES"] = pinned
         logger.info(
             f"Engine rank={self.rank}: pinned CUDA_VISIBLE_DEVICES={pinned} "
@@ -226,6 +227,7 @@ def health_generate(self, timeout: float = 5.0) -> bool:
     def update_weights_from_tensor(
         self,
         serialized_named_tensors: list[str],
+        payload_gpu_uuids: list[str],
         load_format: str | None = None,
         target_modules: list[str] | None = None,
         weight_version: str | None = None,
@@ -241,6 +243,7 @@ def update_weights_from_tensor(
         """
         payload = {
             "serialized_named_tensors": serialized_named_tensors,
+            "payload_gpu_uuids": payload_gpu_uuids,
             "load_format": load_format,
         }
         if target_modules is not None:
@@ -310,8 +313,9 @@ def _compute_server_args(args, host, port, nccl_port):
         "nccl_port": nccl_port,
         # Distinct per engine so concurrent settle_port() probes don't race on the default.
         "master_port": nccl_port + 10000 if nccl_port is not None else None,
-        # Must match rollout allocation, not user CLI.
-        "tp_size": args.rollout_num_gpus_per_engine,
+        # SGL-D runs one worker per num_gpus; tp_size * sp_degree splits that span.
+        "num_gpus": args.rollout_num_gpus_per_engine,
+        "tp_size": args.sglang_tp_size,
         "sp_degree": args.sglang_sp_degree,
         "enable_cfg_parallel": args.sglang_enable_cfg_parallel,
         # Skip warmup to avoid timeout during RL rollouts.

scripts/run-diffusion-grpo-wan22-pickscore-5gpu-rollout-sp2.sh

@@ -0,0 +1,130 @@
+#!/usr/bin/env bash
+# run-diffusion-grpo-wan22-pickscore-5gpu.sh with rollout sequence parallelism:
+# 2 sgl-d engines x 2 GPUs, sp_degree=2.
+# 4-GPU train + 1-GPU pickscore reward, Wan2.2-T2V-A14B dual-expert 5-frame video GRPO:
+#   pretrained = Wan-AI/Wan2.2-T2V-A14B-Diffusers, resolution=480, num_frames=5,
+#   num_steps=10, eval_steps=28, flow_shift=3.0 (engine-launch override of the
+#   sgl-d serving default 12.0), guidance=4.0 (high-noise) / 3.0 (low-noise expert),
+#   Flow-SDE noise_level=0.9, beta=0 (no KL), per-prompt mean/std.
+#   train: lr=1e-4, adam_beta2=0.999, weight_decay=1e-4, clip_range=1e-4,
+#          mixed precision (master fp32 / forward bf16).
+#   LoRA: r=64, alpha=128, init=gaussian, self-attn + cross-attn + FFN of both DiTs.
+#
+# SDE schedule: epoch_global_random_choice draws ONE step per rollout (shared across the
+#   batch) from --diffusion-sde-candidate-steps 1,2,3. At flow_shift=3.0 the
+#   dual-expert boundary is t=875: steps 1,2 train "transformer" (high-noise),
+#   step 3 trains "transformer_2" (low-noise), so both experts get gradient
+#   stochastically and --update-weight-target-module syncs both.
+#
+# Per rollout: 48 prompts × 16 samples = 768 items.
+#   num_steps_per_rollout=2 → 384 items/optim step ÷ 4 train gpus = 96 items/rank.
+#   --micro-batch-size 2: the one-step-per-rollout schedule keeps every
+#   micro-batch phase-pure (one DiT, one CFG scale); mbs=4 OOMs on H200, 2 fits.
+#
+# NOTE: gradient checkpointing stays OFF. Wan2.2 under FSDP2 mixed precision hits
+#   torch.utils.checkpoint CheckpointError (fp32 RoPE freqs buffers; fix pending
+#   in a separate PR). If you OOM, lower --rollout-batch-size,
+#   --n-samples-per-prompt, or --diffusion-microgroup-size.
+#
+# Layout: first 4 GPUs in CUDA_VISIBLE_DEVICES = train+sgld colocate,
+# the 5th GPU = pickscore reward worker. Default: GPU 0,1,2,3 + GPU 4.
+
+set -euo pipefail
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4}"
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False
+RUN_NAME="diffusion_grpo_wan22_pickscore_5gpu_rollout_sp2_$(date +%Y%m%d_%H%M%S)"
+SAVE_DIR="${ROOT_DIR}/logs/${RUN_NAME}/ckpt"
+
+WANDB_ARGS=()
+if [[ -n "${WANDB_API_KEY:-}" ]]; then
+  WANDB_ARGS+=(
+    --use-wandb
+    --wandb-project miles-diffusion-grpo
+    --wandb-group "${RUN_NAME}"
+    --wandb-key "${WANDB_API_KEY}"
+    --diffusion-log-images 8
+    --diffusion-log-image-interval 10
+    --disable-wandb-random-suffix
+  )
+fi
+
+PYTHON_BIN="${PYTHON_BIN:-python}"
+
+DATASETS_DIR="/root/datasets/miles-diffusion-datasets"
+if [[ ! -f "${DATASETS_DIR}/flowgrpo_pickscore/train.jsonl" ]]; then
+  hf download --repo-type dataset rockdu/miles-diffusion-datasets \
+    --include "flowgrpo_pickscore/**" \
+    --local-dir "${DATASETS_DIR}"
+fi
+
+# Wan2.2 DiT LoRA targets: self-attn (attn1), cross-attn (attn2), and FFN.
+WAN_LORA_TARGET_MODULES=(
+  attn1.to_q attn1.to_k attn1.to_v attn1.to_out.0
+  attn2.to_q attn2.to_k attn2.to_v attn2.to_out.0
+  ffn.net.0.proj ffn.net.2
+)
+
+"${PYTHON_BIN}" -u "${ROOT_DIR}/train_diffusion.py" \
+  --train-backend fsdp \
+  --rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout \
+  --hf-checkpoint Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --diffusion-model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --prompt-data "${DATASETS_DIR}/flowgrpo_pickscore/train.jsonl" \
+  --input-key input \
+  --rollout-batch-size 48 \
+  --n-samples-per-prompt 16 \
+  --num-rollout 10000 \
+  --num-steps-per-rollout 2 \
+  --diffusion-microgroup-size 8 \
+  --micro-batch-size 2 \
+  --actor-num-gpus-per-node 4 \
+  --rollout-num-gpus 4 \
+  --rollout-num-gpus-per-engine 2 \
+  --sglang-sp-degree 2 \
+  --num-gpus-per-node 5 \
+  --colocate \
+  --use-lora \
+  --lora-ipc-weight-sync \
+  --lora-rank 64 \
+  --lora-alpha 128 \
+  --lora-target-modules "${WAN_LORA_TARGET_MODULES[@]}" \
+  --diffusion-init-lora-weight gaussian \
+  --lr 1e-4 \
+  --adam-beta2 0.999 \
+  --diffusion-clip-range 1e-4 \
+  --weight-decay 1e-4 \
+  --use-miles-router \
+  --sglang-server-concurrency 8 \
+  --update-weight-buffer-size 2147483648 \
+  --update-weight-target-module transformer,transformer_2 \
+  --diffusion-reward pickscore:1.0 \
+  --advantage-estimator grpo \
+  --rm-type pickscore \
+  --pickscore-num-workers 1 \
+  --pickscore-num-gpus-per-worker 1.0 \
+  --pickscore-batch-size 8 \
+  --pickscore-processor-path laion/CLIP-ViT-H-14-laion2B-s32B-b79K \
+  --pickscore-model-path yuvalkirstain/PickScore_v1 \
+  --fsdp-master-dtype fp32 \
+  --fsdp-reduce-dtype fp32 \
+  --diffusion-forward-dtype bf16 \
+  --diffusion-num-steps 10 \
+  --diffusion-eval-num-steps 28 \
+  --diffusion-output-num-frames 5 \
+  --diffusion-guidance-scale 4.0 \
+  --diffusion-guidance-scale-2 3.0 \
+  --diffusion-noise-level 0.9 \
+  --diffusion-height 480 \
+  --diffusion-width 480 \
+  --diffusion-flow-shift 3.0 \
+  --diffusion-step-strategy-path miles.rollout.step_strategy_hub.epoch_global_random_choice \
+  --diffusion-num-sde-steps 1 \
+  --diffusion-sde-candidate-steps 1,2,3 \
+  --diffusion-debug-mode \
+  --save "${SAVE_DIR}" \
+  --save-interval 10 \
+  --eval-prompt-data pickscore_test "${DATASETS_DIR}/flowgrpo_pickscore/test.jsonl" \
+  --eval-interval 30 \
+  --skip-eval-before-train \
+  "${WANDB_ARGS[@]}"