From 6a83fc6682710215724a8e4cd0ad3e26eea6305c Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Sat, 8 Aug 2026 10:55:45 +0000 Subject: [PATCH] =?UTF-8?q?feat(cosmos3):=20Cosmos3-Nano=20GRPO=20support?= =?UTF-8?q?=20=E2=80=94=20train=20pipeline=20config=20+=20T2I=20recipe?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rebased onto the loss-hub/input-dtype-policy refactor as one commit (the pre-refactor history is preserved on backup/feat-cosmos3-pre-rebase; most intermediate commits were add-then-drop pairs already split to other PRs). - Cosmos3TrainPipelineConfig: packed single-sample forward over the MoT DiT (UND/GEN towers), token-level conditioning (text_ids/text_mask/fps on CondKwargs), GEN-only training with UND frozen, timestep-scale handled in-model. - Boundary dtypes moved to the refactor's input_dtype_policy: latents=default, cond=passthrough (the packed forward casts its own inputs; mRoPE position ids ~15000 would be scrambled by a bf16 boundary cast), timestep=fp32 (karras grid is non-integer; bf16 rounds 993.25 -> 992). Replaces the pre-refactor fsdp_cast_forward_inputs / cast_timesteps_to_forward_dtype flags. - use_cfg: guidance <= 1 runs single-branch, matching sglang do_cfg (now in loss_hub.flow_grpo). - --fsdp-frozen-params-dtype: frozen base stays at checkpoint-native dtype; only LoRA params get fp32 master copies. - Engine: multi-GPU engines pin the full contiguous CVD slice; num_gpus owns the rollout allocation with tp/sp from --sglang-* passthrough. - Weight sync: send one copy of the all-gathered payload (size-1 contract; the engine shards internally). Co-authored-by: Cursor --- miles/backends/fsdp_utils/actor.py | 11 +- miles/backends/fsdp_utils/configs/cosmos3.py | 189 ++++++++++++++++++ .../diffusion_update_weight_utils.py | 5 +- .../backends/fsdp_utils/loss_hub/flow_grpo.py | 2 +- .../sglang_diffusion_engine.py | 19 +- miles/utils/arguments.py | 12 ++ miles/utils/diffusion_rollout_response.py | 3 + miles/utils/types.py | 4 + ...ffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh | 128 ++++++++++++ 9 files changed, 365 insertions(+), 8 deletions(-) create mode 100644 miles/backends/fsdp_utils/configs/cosmos3.py create mode 100755 scripts/run-diffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index dd637f0f..79d84b38 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -97,6 +97,9 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty self._master_dtype = _resolve_dtype(args.fsdp_master_dtype) self._forward_dtype = _resolve_dtype(args.diffusion_forward_dtype) + self._frozen_dtype = _resolve_dtype(args.fsdp_frozen_params_dtype) if args.fsdp_frozen_params_dtype else None + if self._frozen_dtype is not None and not args.use_lora: + raise ValueError("--fsdp-frozen-params-dtype requires --use-lora") from miles.utils.misc import load_function @@ -117,7 +120,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty model = self.model_backend.load_component( component, args, - master_dtype=self._master_dtype, + master_dtype=self._frozen_dtype or self._master_dtype, materialize_weights=materialize_weights, ) if args.fsdp_attention_backend is not None: @@ -131,6 +134,12 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty if args.use_lora: model = apply_lora(model, args, self.train_pipeline_config) + if self._frozen_dtype is not None: + # Base stays at the (checkpoint-native) frozen dtype; only + # LoRA params get a true master copy for the optimizer. + for pname, p in model.named_parameters(): + if "lora_" in pname: + p.data = p.data.to(self._master_dtype) model.train() diff --git a/miles/backends/fsdp_utils/configs/cosmos3.py b/miles/backends/fsdp_utils/configs/cosmos3.py new file mode 100644 index 00000000..52793088 --- /dev/null +++ b/miles/backends/fsdp_utils/configs/cosmos3.py @@ -0,0 +1,189 @@ +"""Cosmos3 training pipeline config.""" + +from __future__ import annotations + +import math + +import torch +from miles.utils.types import CondKwargs + +from .train_pipeline_config import TrainPipelineConfig, register_train_pipeline_config + +# Cosmos3 reuses the Wan2.2 VAE (4x temporal compression). +_VAE_TEMPORAL_FACTOR = 4 + +# GEN-tower parameter name fragments (diffusers Cosmos3OmniTransformer layout). +# Everything else — UND tower, lm_head, unused sound/action heads — stays frozen. +_GEN_PARAM_FRAGMENTS = ( + ".add_q_proj.", + ".add_k_proj.", + ".add_v_proj.", + ".to_add_out.", + ".norm_added_q.", + ".norm_added_k.", + ".mlp_moe_gen.", + ".input_layernorm_moe_gen.", + ".post_attention_layernorm_moe_gen.", + ".norm_moe_gen.", + ".proj_in.", + ".proj_out.", + ".time_embedder.", +) + + +def _is_gen_param(name: str) -> bool: + dotted = f".{name}" + return any(fragment in dotted for fragment in _GEN_PARAM_FRAGMENTS) + + +@register_train_pipeline_config("cosmos3") +class Cosmos3TrainPipelineConfig(TrainPipelineConfig): + hf_ckpt_name_patterns = ("cosmos3", "cosmos-3") + # The transformer scales raw timesteps by config.timestep_scale internally. + needs_timestep_scaling = False + # Timesteps stay fp32: the karras flow grid is non-integer and sgl-d + # conditions on exact fp32 values (bf16 rounds 993.25 -> 992). Conds pass + # through — the packed forward casts its own inputs (mRoPE position ids + # sit at ~15000, where bf16 spacing is 128; a boundary cast scrambles + # rotary phases). + input_dtype_policy = {"latents": "default", "cond": None, "timestep": "fp32"} + lora_target_modules = ["add_q_proj", "add_k_proj", "add_v_proj", "to_add_out"] + + @classmethod + def validate_args(cls, args) -> None: + if getattr(args, "fsdp_cfg_batching", False): + raise ValueError("Cosmos3's packed forward is single-sample; disable --fsdp-cfg-batching.") + if list(args.update_weight_target_modules) != ["transformer"]: + raise ValueError("Cosmos3 requires --update-weight-target-module transformer.") + + def prepare_cond_kwargs(self, cond: CondKwargs | None, device: torch.device) -> dict: + if cond is None or cond.text_ids is None: + return {} + return { + "text_ids": cond.text_ids.to(device), + "text_mask": cond.text_mask.to(device), + "fps": cond.fps, + } + + def collate_cond_for_sample_batch( + self, + per_sample_cond_kwargs: list[dict], + device: torch.device, + pad_to_len: int | None = None, + ) -> dict: + # Packed single-sample forward: keep per-sample kwargs; no batched tensor to build. + return {"per_sample": per_sample_cond_kwargs} + + def compute_noise_pred( + self, + *, + model: torch.nn.Module, + latents_input: torch.Tensor, + timesteps_input: torch.Tensor, + pos_cond: dict | None, + neg_cond: dict | None, + joint_cond: dict | None, + use_cfg: bool, + cfg_batching: bool, + guidance_scale: float, + true_cfg_scale: float | None, + ) -> torch.Tensor: + assert not cfg_batching, "Cosmos3 packed forward is single-sample; cfg_batching unsupported" + config = model.config + preds = [] + for i, pos in enumerate(pos_cond["per_sample"]): + latent = latents_input[i : i + 1] + timestep = float(timesteps_input[i]) + pred = self._packed_forward(model, latent, timestep, pos, config) + if use_cfg: + pred_neg = self._packed_forward(model, latent, timestep, neg_cond["per_sample"][i], config) + pred = self.cfg_combine(pred, pred_neg, guidance_scale, true_cfg_scale=true_cfg_scale) + preds.append(pred) + return torch.stack(preds, dim=0) + + def _packed_forward( + self, + model: torch.nn.Module, + latent: torch.Tensor, + timestep: float, + cond: dict, + config, + ) -> torch.Tensor: + """One (text, video) joint-sequence forward; mirrors the packing in + diffusers' Cosmos3OmniDiffusersPipeline denoising loop (T2V/T2I: all + frames noisy, no conditioning frames).""" + from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import ( + get_3d_mrope_ids_text_tokens, + get_3d_mrope_ids_vae_tokens, + ) + + device = latent.device + und_len = int(cond["text_mask"].sum().item()) + input_ids = cond["text_ids"].reshape(-1)[:und_len] + + text_mrope_ids, next_offset = get_3d_mrope_ids_text_tokens( + num_tokens=und_len, + temporal_offset=0, + use_float_positions=config.enable_fps_modulation, + ) + vision_offset = next_offset + config.unified_3d_mrope_temporal_modality_margin + + p = config.latent_patch_size + _, _, latent_t, latent_h, latent_w = latent.shape + patch_h = math.ceil(latent_h / p) + patch_w = math.ceil(latent_w / p) + num_vision_tokens = latent_t * patch_h * patch_w + + vision_mrope_ids, _ = get_3d_mrope_ids_vae_tokens( + grid_t=latent_t, + grid_h=patch_h, + grid_w=patch_w, + temporal_offset=vision_offset, + reset_spatial_indices=config.unified_3d_mrope_reset_spatial_ids, + fps=cond["fps"] if config.enable_fps_modulation else None, + base_fps=float(config.base_fps), + temporal_compression_factor=_VAE_TEMPORAL_FACTOR, + ) + + sequence_length = und_len + num_vision_tokens + vision_sequence_indexes = torch.arange(und_len, sequence_length, dtype=torch.long, device=device) + preds_vision, _, _ = model( + input_ids=input_ids, + text_indexes=torch.arange(und_len, dtype=torch.long, device=device), + position_ids=torch.cat([text_mrope_ids, vision_mrope_ids], dim=1).to(device), + und_len=und_len, + sequence_length=sequence_length, + vision_tokens=[latent], + vision_token_shapes=[(latent_t, patch_h, patch_w)], + vision_sequence_indexes=vision_sequence_indexes, + vision_mse_loss_indexes=vision_sequence_indexes, + vision_timesteps=torch.full((num_vision_tokens,), timestep, device=device, dtype=torch.float32), + vision_noisy_frame_indexes=[torch.arange(latent_t, dtype=torch.long, device=device)], + ) + return preds_vision[0] + + def cfg_combine( + self, + noise_pred_pos: torch.Tensor, + noise_pred_neg: torch.Tensor, + guidance_scale: float, + true_cfg_scale: float | None = None, + ) -> torch.Tensor: + scale = true_cfg_scale if true_cfg_scale is not None else guidance_scale + return noise_pred_neg + scale * (noise_pred_pos - noise_pred_neg) + + def postprocess_model_after_materialize(self, model: torch.nn.Module) -> None: + # The UND tower sits inside the training forward graph (paired + # attention), so gradients reach it unless explicitly frozen. + for name, param in model.named_parameters(): + if "lora_" not in name and not _is_gen_param(name): + param.requires_grad_(False) + + # sglang-d casts the fp32 timestep sinusoid to the MLP weight dtype + # before linear_1; diffusers feeds it through as-is, which crashes on + # the fp32/bf16 mismatch under FSDP mixed precision. + def _cast_to_weight_dtype(module, args): + dtype = module.linear_1.weight.dtype + return tuple(a.to(dtype) if torch.is_tensor(a) else a for a in args) + + model.time_embedder.register_forward_pre_hook(_cast_to_weight_dtype) diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index 9142ce0f..43d6ed23 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -289,9 +289,12 @@ def update_bucket_weights( # TODO: here we assume all ranks have the same number of dtypes. num_dtypes = len(gathered_serialized_batches[0]) assert num_dtypes > 0 + # Each rank's payload is the full (all-gathered) weights, so the + # copies are identical. Send exactly one; the engine shards or + # replicates internally per its own parallelism (size-1 contract). for i in range(num_dtypes): kwargs = { - "serialized_named_tensors": [tensors[i] for tensors in gathered_serialized_batches], + "serialized_named_tensors": [gathered_serialized_batches[0][i]], "load_format": "flattened_bucket", "target_modules": [target_module], "weight_version": str(weight_version), diff --git a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py index 39a69d18..6b35467c 100644 --- a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py +++ b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py @@ -42,7 +42,7 @@ def prepare_flow_grpo_batch( guidance_scale = args.diffusion_guidance_scale true_cfg_scale = args.diffusion_true_cfg_scale cfg_scale = true_cfg_scale if true_cfg_scale is not None else guidance_scale - use_cfg = cfg_scale > 0 + use_cfg = cfg_scale > 1.0 # matches sglang do_cfg: guidance<=1 runs single-branch if len(ctx.models) == 1: component_name, model = next(iter(ctx.models.items())) diff --git a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py index 3de4dd1a..096e09d6 100644 --- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py +++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py @@ -178,12 +178,18 @@ def _pin_to_assigned_gpu(self): os.environ["CUDA_VISIBLE_DEVICES"] = str(self.base_gpu_id) 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] + span = getattr(self.args, "rollout_num_gpus_per_engine", 1) or 1 + if span > 1: + # Multi-GPU engine: Ray's fractional allocation narrows CVD to one + # device, so rebuild the contiguous physical slice from base_gpu_id. + pinned = ",".join(str(self.base_gpu_id + j) for j in range(span)) + else: + local_idx = _to_local_gpu_id(self.base_gpu_id) + pinned = visible[local_idx] os.environ["CUDA_VISIBLE_DEVICES"] = pinned logger.info( f"Engine rank={self.rank}: pinned CUDA_VISIBLE_DEVICES={pinned} " - f"(base_gpu_id={self.base_gpu_id}, local_idx={local_idx})" + f"(base_gpu_id={self.base_gpu_id}, span={span})" ) def _make_request(self, endpoint: str, payload: dict | None = None): @@ -319,8 +325,11 @@ 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, + # parallel — the engine owns all GPUs of its rollout allocation; how it + # splits them (ulysses/ring/tp) comes from --sglang-* passthrough. + "num_gpus": args.rollout_num_gpus_per_engine, + "tp_size": getattr(args, "sglang_tp_size", None) or 1, + # Sequence-parallel degree (None = disabled, SGL-D decides internally). "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/miles/utils/arguments.py b/miles/utils/arguments.py index 18ba848a..a15c3648 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -176,6 +176,18 @@ def add_train_arguments(parser): "--diffusion-forward-dtype." ), ) + parser.add_argument( + "--fsdp-frozen-params-dtype", + type=str, + default=None, + choices=["bf16", "fp16"], + help=( + "Storage dtype for frozen (non-LoRA) parameters. The model is " + "loaded at this dtype and only LoRA parameters are upcast to " + "--fsdp-master-dtype, so a large frozen base does not pay the " + "fp32 master cost. Requires --use-lora." + ), + ) parser.add_argument( "--fsdp-reduce-dtype", type=str, diff --git a/miles/utils/diffusion_rollout_response.py b/miles/utils/diffusion_rollout_response.py index 8a1727f6..13c29440 100644 --- a/miles/utils/diffusion_rollout_response.py +++ b/miles/utils/diffusion_rollout_response.py @@ -120,6 +120,9 @@ def _parse_cond_kwargs( data.get("audio_encoder_attention_mask"), deserialize_func=deserialize_func ), pooled_projections=_parse_tensor_or_list(data.get("pooled_projections"), deserialize_func=deserialize_func), + text_ids=deserialize_func(data.get("text_ids")), + text_mask=deserialize_func(data.get("text_mask")), + fps=data.get("fps"), ) diff --git a/miles/utils/types.py b/miles/utils/types.py index 0864a1c6..e1d2ad03 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -33,6 +33,10 @@ class CondKwargs: encoder_attention_mask: list[torch.Tensor] | None = None audio_encoder_attention_mask: list[torch.Tensor] | None = None pooled_projections: list[torch.Tensor] | None = None + # Cosmos3: token-level conditioning (no separate text encoder). + text_ids: torch.Tensor | None = None + text_mask: torch.Tensor | None = None + fps: float | None = None @dataclass diff --git a/scripts/run-diffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh b/scripts/run-diffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh new file mode 100755 index 00000000..dc796fd1 --- /dev/null +++ b/scripts/run-diffusion-grpo-cosmos3-pickscore-t2i-5gpu.sh @@ -0,0 +1,128 @@ +#!/usr/bin/env bash +# 4-GPU train + 1-GPU pickscore reward, Cosmos3-Nano T2I GRPO: +# pretrained = nvidia/Cosmos3-Nano (16B MoT: 8B UND tower frozen, 8B GEN tower +# trained via LoRA), resolution=832x480, single frame (T2I), +# num_steps=16, eval_steps=35, guidance=1.0 (CFG-free training; checkpoints +# still transfer to CFG deployment — merged LoRA sampled at g=4 beats the +# base model at g=4), Flow-SDE noise_level=0.7, +# KL beta=1e-3, global reward std, per-prompt mean. +# train: lr=3e-4, adam_beta2=0.95, weight_decay=1e-4, clip_range=1e-3, +# clip_grad=2e-3, mixed precision (master fp32 / forward bf16). +# LoRA: r=64, alpha=128, init=gaussian. +# +# SDE schedule: epoch_global_window draws a 2-step window per rollout from +# --diffusion-sde-candidate-steps 4-15. The Cosmos3 checkpoint ships a Karras +# flow-sigma grid whose head steps 1-3 sit at sigma>0.96 with |dt|<0.02 and +# train nothing; step numbers are NOT transferable across sigma-grid +# families — re-derive candidates from |dt| when changing model/grid. +# +# Ratio/stability choices: +# --diffusion-recompute-old-log-prob: the trainer recomputes old log-probs at +# rollout ingestion so the PPO ratio is implementation-self-consistent +# (rollout fa kernels vs train SDPA would otherwise leak into the ratio). +# --adam-beta2 0.95 + --clip-grad 2e-3: absorb Adam-preconditioner spikes +# after quiet stretches (single-step policy jumps that the PPO loss clip +# cannot stop). +# +# Per rollout: 48 prompts × 16 samples = 768 items. +# num_steps_per_rollout=2 → 384 items/optim step ÷ 4 train gpus = 96 items/rank. +# --diffusion-microgroup-size 1: the Cosmos3 transformer is a packed-sequence +# single-sample interface; one request cannot batch multiple outputs. +# +# 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 +# RL rollout scores raw samples; skip the serving-side guardrail models. +export SGLANG_DISABLE_COSMOS3_GUARDRAILS=1 +RUN_NAME="diffusion_grpo_cosmos3_pickscore_t2i_5gpu_$(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 + +"${PYTHON_BIN}" -u "${ROOT_DIR}/train_diffusion.py" \ + --train-backend fsdp \ + --rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout \ + --hf-checkpoint nvidia/Cosmos3-Nano \ + --diffusion-model nvidia/Cosmos3-Nano \ + --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 1 \ + --micro-batch-size 1 \ + --actor-num-gpus-per-node 4 \ + --rollout-num-gpus 4 \ + --rollout-num-gpus-per-engine 1 \ + --num-gpus-per-node 5 \ + --colocate \ + --use-lora \ + --lora-rank 64 \ + --lora-alpha 128 \ + --diffusion-init-lora-weight gaussian \ + --lr 3e-4 \ + --adam-beta2 0.95 \ + --clip-grad 2e-3 \ + --diffusion-clip-range 1e-3 \ + --weight-decay 1e-4 \ + --use-miles-router \ + --sglang-server-concurrency 8 \ + --sglang-attention-backend fa \ + --update-weight-buffer-size 2147483648 \ + --update-weight-target-module transformer \ + --diffusion-reward pickscore:1.0 \ + --advantage-estimator grpo \ + --globalize-reward-std \ + --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 16 \ + --diffusion-eval-num-steps 35 \ + --diffusion-output-num-frames 1 \ + --diffusion-guidance-scale 1.0 \ + --diffusion-noise-level 0.7 \ + --diffusion-height 480 \ + --diffusion-width 832 \ + --diffusion-step-strategy-path miles.rollout.step_strategy_hub.epoch_global_window \ + --diffusion-sde-window-size 2 \ + --diffusion-sde-candidate-steps 4,5,6,7,8,9,10,11,12,13,14,15 \ + --diffusion-recompute-old-log-prob \ + --diffusion-kl-beta 1e-3 \ + --diffusion-debug-mode \ + --save "${SAVE_DIR}" \ + --save-interval 5 \ + --eval-prompt-data pickscore_test "${DATASETS_DIR}/flowgrpo_pickscore/test.jsonl" \ + --eval-interval 30 \ + --skip-eval-before-train \ + "${WANDB_ARGS[@]}"