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

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

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)

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

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/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.

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,

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"),
     )
 
 

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