[cosmos3] Cosmos3-Nano GRPO support: train pipeline config, VideoAlign reward, T2V recipe

miles/backends/fsdp_utils/configs/cosmos3.py

@@ -0,0 +1,177 @@
+"""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 param name fragments; everything else (UND tower, lm_head, unused 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")
+    # Timesteps stay fp32 (bf16 rounds the karras grid); conds pass through, the packed forward casts its own inputs.
+    input_dtype_policy = {"latents": "default", "cond": None, "timestep": "fp32"}
+    # The packed forward is single-sample; never batch the CFG branches.
+    cfg_batching = False
+    lora_target_modules = ["add_q_proj", "add_k_proj", "add_v_proj", "to_add_out"]
+
+    @classmethod
+    def validate_args(cls, args) -> None:
+        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 packed (text, video) forward mirroring the diffusers denoising loop (T2V/T2I: all frames noisy)."""
+        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 graph, so it must be explicitly frozen.
+        for name, param in model.named_parameters():
+            if "lora_" not in name and not _is_gen_param(name):
+                param.requires_grad_(False)
+
+        # Cast the timestep sinusoid to the MLP weight dtype before linear_1, exactly as sglang-d does.
+        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)

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()))

miles/backends/fsdp_utils/models/diffusers/cosmos3/parallel_plan.py

@@ -0,0 +1,4 @@
+from miles.backends.fsdp_utils.models.parallel_plan import FSDPParallelPlan
+
+
+FSDP_PARALLEL_PLAN = FSDPParallelPlan()

miles/utils/diffusion_rollout_response.py

@@ -107,6 +107,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"),
     )
 
 

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

scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py

@@ -0,0 +1,160 @@
+"""Cosmos3-Nano T2I GRPO with PickScore, fully colocated on 4 GPUs.
+
+pretrained = nvidia/Cosmos3-Nano (16B MoT: 8B UND tower frozen, 8B GEN tower
+trained via LoRA r64), 832x480 single frame, num_steps=16, eval_steps=35,
+guidance 4.0, Flow-SDE noise_level=0.7, no KL, per-prompt mean + global std.
+
+Layout: train, rollout and PickScore reward all share the same 4 GPUs
+(--colocate --colocate-reward, one PickScore worker per rollout engine).
+
+SDE schedule: epoch_global_random_choice draws 2 steps per epoch from
+candidates 8-11. The Cosmos3 checkpoint ships a Karras flow-sigma grid whose
+head steps 1-7 sit at sigma>0.96 with |dt|<0.02 and train nothing; steps 8-11
+are the true high-noise segment (sigma 0.94-0.80). Step numbers are NOT
+transferable across sigma-grid families - re-derive candidates from |dt| when
+changing model/grid.
+
+Pacing: lr 1e-4 x 1 optimizer step per rollout (the whole rollout is one
+batch). CFG amplifies per-step policy displacement, so training with
+guidance > 1 needs this slower pacing than a comparable CFG-free recipe.
+
+--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). With 1 step per
+rollout this makes every optimizer step exactly on-policy.
+
+Usage:
+    python3 scripts/run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py
+"""
+
+from dataclasses import dataclass
+
+import typer
+
+import miles.utils.external_utils.command_utils as U
+
+MODEL = "nvidia/Cosmos3-Nano"
+DATASET = "rockdu/miles-diffusion-datasets"
+DATASET_SUBSET = "flowgrpo_pickscore"
+WANDB_PROJECT = "miles-diffusion-grpo"
+
+
+@dataclass
+class ScriptArgs(U.ExecuteTrainConfig):
+    cuda_visible_devices: str = "0,1,2,3"
+    num_rollout: int = 10000
+    data_dir: str = "/root/datasets"
+    extra_args: str = ""
+
+
+def prepare(args: ScriptArgs) -> str:
+    local_dir = U.hf_download_dataset(DATASET, include=f"{DATASET_SUBSET}/**", data_dir=args.data_dir)
+    return f"{local_dir}/{DATASET_SUBSET}"
+
+
+def execute(args: ScriptArgs, data_dir: str) -> None:
+    run_name = f"diffusion_grpo_cosmos3_pickscore_t2i_4gpu_{U.create_run_id()}"
+
+    ckpt_args = f"--hf-checkpoint {MODEL} --save {args.output_dir}/{run_name}/ckpt --save-interval 10 "
+
+    rollout_args = (
+        "--rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout "
+        f"--prompt-data {data_dir}/train.jsonl "
+        "--input-key input "
+        "--rollout-batch-size 48 "
+        "--n-samples-per-prompt 16 "
+        f"--num-rollout {args.num_rollout} "
+        "--num-steps-per-rollout 1 "
+        # The Cosmos3 transformer is a packed-sequence single-sample interface;
+        # one request cannot batch multiple outputs.
+        "--rollout-microgroup-size 1 "
+        "--micro-batch-size 1 "
+    )
+
+    diffusion_args = (
+        "--diffusion-num-steps 16 "
+        "--diffusion-output-num-frames 1 "
+        "--diffusion-guidance-scale 4.0 "
+        "--diffusion-noise-level 0.7 "
+        "--diffusion-height 480 "
+        "--diffusion-width 832 "
+        "--diffusion-step-strategy-path miles.rollout.step_strategy_hub.epoch_global_random_choice "
+        "--diffusion-num-sde-steps 2 "
+        "--diffusion-sde-candidate-steps 8,9,10,11 "
+        "--diffusion-recompute-old-log-prob "
+    )
+
+    eval_args = (
+        f"--eval-prompt-data pickscore_test {data_dir}/test.jsonl "
+        "--eval-interval 30 "
+        "--diffusion-eval-num-steps 35 "
+        "--skip-eval-before-train "
+    )
+
+    grpo_args = "--advantage-estimator grpo --globalize-reward-std --diffusion-clip-range 1e-3 "
+
+    optimizer_args = "--lr 1e-4 --adam-beta2 0.999 --weight-decay 1e-4 "
+
+    # UND/GEN towers share layers and differ by parameter name (to_q vs
+    # add_q_proj, mlp vs mlp_moe_gen); LoRA targeting defaults to the GEN
+    # fragments in the cosmos3 train pipeline config.
+    lora_args = "--use-lora --lora-ipc-weight-sync --lora-rank 64 --lora-alpha 128 --lora-init-weights gaussian "
+
+    reward_args = (
+        "--rm-type pickscore "
+        "--colocate-reward "
+        "--pickscore-num-workers 4 "
+        "--pickscore-batch-size 8 "
+        "--pickscore-processor-path laion/CLIP-ViT-H-14-laion2B-s32B-b79K "
+        "--pickscore-model-path yuvalkirstain/PickScore_v1 "
+    )
+
+    wandb_args = U.get_default_wandb_args(
+        __file__, run_id=run_name, project=WANDB_PROJECT, wandb_log_num_images=8, wandb_log_image_interval=10
+    )
+
+    sglang_args = (
+        "--use-miles-router "
+        "--sglang-server-concurrency 8 "
+        "--update-weight-buffer-size 2147483648 "
+        "--update-weight-target-module transformer "
+    )
+
+    train_backend_args = (
+        "--train-backend fsdp --fsdp-master-dtype fp32 --fsdp-reduce-dtype fp32 --diffusion-forward-dtype bf16 "
+    )
+
+    misc_args = (
+        "--actor-num-gpus-per-node 4 "
+        "--rollout-num-gpus 4 "
+        "--rollout-num-gpus-per-engine 1 "
+        "--num-gpus-per-node 4 "
+        "--colocate "
+    )
+
+    debug_args = "--diffusion-debug-mode "
+
+    U.execute_train(
+        train_args=(
+            f"{ckpt_args} {rollout_args} {diffusion_args} {eval_args} {grpo_args} {optimizer_args} "
+            f"{lora_args} {reward_args} {wandb_args} {sglang_args} {train_backend_args} {misc_args} "
+            f"{debug_args} {args.extra_args}"
+        ),
+        num_gpus_per_node=4,
+        config=args,
+        extra_env_vars={
+            "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:False",
+            # RL rollout scores raw samples; skip the serving-side guardrail models.
+            "SGLANG_DISABLE_COSMOS3_GUARDRAILS": "1",
+        },
+    )
+
+
+@U.dataclass_cli
+def main(args: ScriptArgs) -> None:
+    data_dir = prepare(args)
+    execute(args, data_dir)
+
+
+if __name__ == "__main__":
+    typer.run(main)