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

miles/backends/fsdp_utils/actor.py

@@ -96,19 +96,32 @@ 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()
 
         self.models: dict[str, torch.nn.Module] = {}
         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,

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)

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, ...] = ()

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

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.

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,

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