diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index 18dcbf61..9a369bc9 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/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), diff --git a/miles/backends/sglang_diffusion_utils/arguments.py b/miles/backends/sglang_diffusion_utils/arguments.py index 4baac11e..e15b8365 100644 --- a/miles/backends/sglang_diffusion_utils/arguments.py +++ b/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"): diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 2ab79a14..90c54cce 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/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. diff --git a/scripts/run-diffusion-grpo-wan22-pickscore-5gpu-rollout-sp2.sh b/scripts/run-diffusion-grpo-wan22-pickscore-5gpu-rollout-sp2.sh new file mode 100755 index 00000000..fce09d7e --- /dev/null +++ b/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[@]}"