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

miles/backends/fsdp_utils/actor.py

@@ -1,5 +1,7 @@
 import functools
 import logging
+import os
+import time
 from argparse import Namespace
 from collections import defaultdict
 from contextlib import nullcontext
@@ -107,8 +109,13 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
 
         self.train_pipeline_config = load_function(args.train_pipeline_config_path)()
         self.model_backend = load_function(args.model_backend_path)(self.train_pipeline_config)
+        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")
         raw_models, self.scheduler = self.model_backend.load_models_and_scheduler(
-            args, master_dtype=self._master_dtype
+            args, master_dtype=self._frozen_dtype or self._master_dtype
         )
 
         self.models: dict[str, torch.nn.Module] = {}
@@ -119,6 +126,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()
 
@@ -135,6 +148,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,
             )
             self.models[component] = model
         # Force a sync to ensure sharding is complete and old memory is freed.
@@ -324,7 +338,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
@@ -371,6 +385,25 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
             parallel_state=self.parallel_state,
         )
 
+        # ------------- Recompute old log-probs trainer-side (impl-consistent PPO ratio) -------------
+        if self.args.diffusion_recompute_old_log_prob:
+            with timer("recompute_old_log_prob"), torch.no_grad():
+                throwaway: dict[str, list[torch.Tensor]] = defaultdict(list)
+                for lo in range(0, num_pairs, micro_bs):
+                    self._forward_train_pair_batch(
+                        train_pairs[lo : lo + micro_bs],
+                        use_cfg=use_cfg,
+                        guidance_scale=guidance_scale,
+                        true_cfg_scale=true_cfg_scale,
+                        clip_range=clip_range,
+                        noise_level=noise_level,
+                        num_train_timesteps=num_train_timesteps,
+                        log_stats=throwaway,
+                        device=device,
+                        pad_to_len=None,
+                        write_old_log_prob=True,
+                    )
+
         # ------------- Forward / Backward -------------
         with timer("actor_train"):
             for microbatch_ranges in microbatch_schedule:
@@ -447,8 +480,12 @@ def _forward_train_pair_batch(
         device: torch.device,
         kl_beta: float = 0.0,
         pad_to_len: int | None = None,
-    ) -> torch.Tensor:
-        """One DiT forward + PPO loss over ``len(batch)`` train pairs. Returns sum of per-pair losses."""
+        write_old_log_prob: bool = False,
+    ) -> torch.Tensor | None:
+        """One DiT forward + PPO loss over ``len(batch)`` train pairs. Returns sum of per-pair losses.
+
+        With ``write_old_log_prob`` the computed log-prob is written into each pair's
+        ``log_prob_old`` and no loss is returned (caller runs it under ``no_grad``)."""
         forward_dtype = self._forward_dtype
         train_pipeline_config = self.train_pipeline_config
         bsz = len(batch)
@@ -546,7 +583,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()
@@ -575,6 +615,11 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor:
             noise_level=noise_level,
         )
 
+        if write_old_log_prob:
+            for i, pair in enumerate(batch):
+                pair["log_prob_old"] = log_prob_new_microbatch[i].detach().float().cpu()
+            return None
+
         log_prob_new = log_prob_new_microbatch  # (bsz,) -- sde_step_with_logprob means over non-batch dims
         log_prob_old = log_prob_old_microbatch  # (bsz,)
         ratio = torch.exp(log_prob_new - log_prob_old)  # (bsz,)
@@ -597,10 +642,15 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor:
                     prev_sample=next_latents_microbatch.float(),
                     noise_level=noise_level,
                 )
+            # Low-sigma steps have tiny std -> 1/2std^2 makes the KL spring
+            # orders of magnitude stiffer there; the floor equalizes stiffness.
+            kl_std = std_dev_t_new
+            if self.args.diffusion_kl_std_floor > 0:
+                kl_std = torch.clamp(kl_std, min=self.args.diffusion_kl_std_floor)
             kl_per_pair = ((prev_sample_mean_new - prev_sample_mean_ref) ** 2).mean(
                 dim=tuple(range(1, prev_sample_mean_new.ndim)),
                 keepdim=True,
-            ) / (2 * std_dev_t_new**2)
+            ) / (2 * kl_std**2)
             loss_sum = loss_sum + kl_beta * kl_per_pair.sum()
             kl_loss = kl_per_pair.mean()
 
@@ -625,6 +675,21 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor:
             # functions of it). Matches the legacy actor metric name.
             rollout_model_output = stack_train_pair_rollout_debug(batch, "rollout_step_model_output")
             if rollout_model_output is not None:
+                if os.environ.get("MILES_DIAG_DUMP"):
+                    torch.save(
+                        {
+                            "noise_pred": noise_pred_microbatch.detach().float().cpu(),
+                            "rollout_model_output": rollout_model_output.detach().float().cpu(),
+                            "timesteps": timesteps_microbatch.detach().cpu(),
+                            "prompts": [pair.get("prompt") for pair in batch],
+                            "latents": latents_microbatch.detach().cpu(),
+                            "pos_cond": pos_list,
+                            "neg_cond": neg_list,
+                            "guidance_scale": guidance_scale,
+                            "use_cfg": use_cfg,
+                        },
+                        f"/tmp/miles_diag_dump_{dist.get_rank()}_{time.time_ns()}.pt",
+                    )
                 mean_abs_diff = _append_rollout_train_abs_diff_stats(
                     log_stats,
                     "model_output",
@@ -716,7 +781,7 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) -
     return model
 
 
-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
@@ -736,6 +801,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/checkpoint.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import shutil
 import json
 import logging
 import time
@@ -185,6 +186,17 @@ def load(actor: Any) -> dict[str, Any] | None:
     elif hasattr(actor, "lr_scheduler"):
         logger.info(f"[FSDP] LR scheduler checkpoint not found at {lr_scheduler_dir}, skipping LR scheduler load.")
 
+    # Restored optimizer/scheduler state carries the checkpoint's lr; the CLI value wins.
+    # FSDPLRScheduler.step() recomputes group lr from max_lr, so that is the field that matters.
+    for group in actor.optimizer.param_groups:
+        group["lr"] = actor.args.lr
+        if "max_lr" in group:
+            group["max_lr"] = actor.args.lr
+    if hasattr(actor, "lr_scheduler"):
+        actor.lr_scheduler.max_lr = float(actor.args.lr)
+        actor.lr_scheduler.base_lrs = [actor.args.lr] * len(actor.lr_scheduler.base_lrs)
+        logger.info(f"[FSDP] Applied CLI lr {actor.args.lr} to restored optimizer/scheduler.")
+
     rng_state = None
     rng_path = checkpoint_dir / "rng.pt"
     if rng_path.exists():
@@ -289,4 +301,11 @@ def save(actor: Any, iteration: int) -> None:
         tracker_file.write_text(str(step_id))
         logger.info(f"[FSDP] Saved checkpoint to {checkpoint_dir}")
 
+        keep_last_n = getattr(actor.args, "keep_last_n_checkpoints", 0)
+        if keep_last_n and keep_last_n > 0:
+            snapshots = sorted(base_dir.glob("iter_*"))
+            for stale in snapshots[:-keep_last_n]:
+                shutil.rmtree(stale, ignore_errors=True)
+                logger.info(f"[FSDP] Removed stale checkpoint {stale} (keep_last_n={keep_last_n})")
+
     dist.barrier()