diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py
index e00837e5..31263367 100644
--- a/miles/backends/fsdp_utils/actor.py
+++ b/miles/backends/fsdp_utils/actor.py
@@ -96,6 +96,11 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
         if args.deterministic_mode:
             # flash-attn is opaque to torch's determinism flag; backends patch their own dispatch.
             self.model_backend.enable_deterministic_attention(args.fsdp_attention_backend)
+        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")
         self.scheduler = self.model_backend.load_scheduler(args)
         rank = dist.get_rank()
 
@@ -103,12 +108,20 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
         for component in args.update_weight_target_modules:
             # per raw component (wan2.2 has two transformers), before LoRA/FSDP wrap
             with self._init_weight_context():
-                model = self.model_backend.load_component(component, args, master_dtype=self._master_dtype)
+                model = self.model_backend.load_component(
+                    component, args, master_dtype=self._frozen_dtype or self._master_dtype
+                )
             if args.fsdp_attention_backend is not None:
                 self.model_backend.set_attention_backend(model, args.fsdp_attention_backend)
 
             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()
 
@@ -125,6 +138,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
                 cpu_offload=self.args.fsdp_cpu_offload,
                 args=self.args,
                 no_split_modules=self.model_backend.fsdp_no_split_modules(model),
+                cast_forward_inputs=self.train_pipeline_config.fsdp_cast_forward_inputs,
             )
             load_sharded_model(model, full_state, cpu_offload=self.args.fsdp_cpu_offload)
             del full_state
@@ -334,7 +348,7 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
         guidance_scale = self.args.diffusion_guidance_scale
         true_cfg_scale = self.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
 
         # ------------- Loss / SDE Parameters -------------
         clip_range = self.args.diffusion_clip_range
@@ -592,7 +606,10 @@ def _stack(key):
         # leaves fp32 inputs, which would run first matmul at higher precision
         # than rollout → systematic noise_pred drift.
         latents_input = latents_microbatch.to(forward_dtype)
-        timesteps_input = timesteps_for_model.to(forward_dtype)
+        if train_pipeline_config.cast_timesteps_to_forward_dtype:
+            timesteps_input = timesteps_for_model.to(forward_dtype)
+        else:
+            timesteps_input = timesteps_for_model
 
         def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor:
             adapter_ctx = model.disable_adapter() if disable_adapter else nullcontext()
@@ -808,7 +825,7 @@ def sync_model_dtypes(model: torch.nn.Module) -> None:
             tensor.data = tensor.data.to(dtype)
 
 
-def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None):
+def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None, cast_forward_inputs=True):
     from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard
 
     offload_policy = CPUOffloadPolicy() if cpu_offload else None
@@ -828,6 +845,7 @@ def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules
         "mp_policy": MixedPrecisionPolicy(
             param_dtype=param_dtype,
             reduce_dtype=reduce_dtype,
+            cast_forward_inputs=cast_forward_inputs,
         ),
         "offload_policy": offload_policy,
         "mesh": mesh,
diff --git a/miles/backends/fsdp_utils/configs/cosmos3.py b/miles/backends/fsdp_utils/configs/cosmos3.py
new file mode 100644
index 00000000..7cefb818
--- /dev/null
+++ b/miles/backends/fsdp_utils/configs/cosmos3.py
@@ -0,0 +1,188 @@
+"""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
+    # The karras flow grid has non-integer timesteps; keep them fp32.
+    cast_timesteps_to_forward_dtype = False
+    # mRoPE position ids sit at ~15000 (temporal modality margin); the FSDP
+    # boundary cast to bf16 (spacing 128 there) scrambles rotary phases.
+    fsdp_cast_forward_inputs = False
+    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/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py
index a553735c..011ba9a7 100644
--- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py
+++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py
@@ -76,6 +76,15 @@ class TrainPipelineConfig(abc.ABC):
 
     lora_target_modules: list[str] = ["to_q", "to_k", "to_v", "to_out.0"]
     needs_timestep_scaling: bool = True
+    # Rollout engines condition on exact fp32 timesteps; families with
+    # non-integer timestep grids set this False to avoid bf16 rounding
+    # (e.g. 993.25 -> 992) drifting the train-side forward.
+    cast_timesteps_to_forward_dtype: bool = True
+    # FSDP2 casts float32 forward inputs to param_dtype at module boundaries.
+    # Families whose forward takes large-magnitude float inputs (e.g. Cosmos3
+    # mRoPE position ids ~15000, where bf16 spacing is 128) must disable this
+    # and cast inputs themselves.
+    fsdp_cast_forward_inputs: bool = True
     optimizer_state_allowed_missing: list[str] = []
     # Case-insensitive substrings matched against the checkpoint name (--diffusion-model).
     hf_ckpt_name_patterns: tuple[str, ...] = ()
diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py
index 18dcbf61..bdda4371 100644
--- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py
+++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py
@@ -248,9 +248,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/sglang_diffusion_utils/sglang_diffusion_engine.py b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py
index 2ab79a14..7951d50c 100644
--- a/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py
+++ b/miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py
@@ -171,12 +171,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):
@@ -310,8 +316,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 2ac47fbe..d9723cd0 100644
--- a/miles/utils/arguments.py
+++ b/miles/utils/arguments.py
@@ -175,6 +175,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 d133447e..ed6d24a7 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[@]}"