Push train setup with 49 frame validation

fastvideo/train/methods/rl/common/validation.py

@@ -21,6 +21,7 @@ class RLValidationConfig:
     fps: int = 1
     seed: int = 42
     data_path: str | None = None
+    num_latent_t: int | None = None
     sampling: dict[str, Any] | None = None
 
     @classmethod
@@ -36,6 +37,11 @@ class RLValidationConfig:
         max_samples = raw.get("max_samples", None)
         if max_samples is not None:
             max_samples = max(0, int(max_samples))
+        num_latent_t = raw.get("num_latent_t", None)
+        if num_latent_t is not None:
+            num_latent_t = int(num_latent_t)
+            if num_latent_t <= 0:
+                raise ValueError("method.validation.num_latent_t must be positive when set")
         return cls(
             every_steps=max(0, int(raw.get("every_steps", 0) or 0)),
             num_steps=max(1, int(raw.get("num_steps", 40) or 40)),
@@ -46,6 +52,7 @@ class RLValidationConfig:
             fps=max(1, int(raw.get("fps", 1) or 1)),
             seed=int(raw.get("seed", 42) or 42),
             data_path=(None if data_path in (None, "") else str(data_path)),
+            num_latent_t=num_latent_t,
             sampling=(dict(sampling) if sampling is not None else None),
         )
 

fastvideo/train/methods/rl/diffusion_nft.py

@@ -373,11 +373,7 @@ class DiffusionNFTMethod(TrainingMethod):
             batch_items = items[start:start + config.batch_size]
             raw_batch = self._collate_validation_rows([item[2] for item in batch_items])
             prompts = self._extract_prompts(raw_batch)
-            batch = self.student.prepare_batch(
-                raw_batch,
-                generator=prepare_generator,
-                latents_source="zeros",
-            )
+            batch = self._prepare_validation_batch(raw_batch, prepare_generator)
             sampling_result = self._validation_sampler.sample(
                 self.student,
                 batch,
@@ -429,6 +425,18 @@ class DiffusionNFTMethod(TrainingMethod):
         self._log_progress(f"[DiffusionNFT] validation step {iteration}: finished")
         return metrics
 
+    def _prepare_validation_batch(
+        self,
+        raw_batch: dict[str, Any],
+        generator: torch.Generator,
+    ) -> TrainingBatch:
+        return self.student.prepare_batch(
+            raw_batch,
+            generator=generator,
+            latents_source="zeros",
+            num_latent_t=self._validation_config.num_latent_t,
+        )
+
     def _get_validation_items(self) -> list[tuple[int, bool, dict[str, Any]]]:
         if self._validation_items is not None:
             return self._validation_items

fastvideo/train/models/base.py

@@ -126,8 +126,13 @@ class ModelBase(ABC):
         *,
         generator: torch.Generator,
         latents_source: Literal["data", "zeros"] = "data",
+        num_latent_t: int | None = None,
     ) -> TrainingBatch:
-        """Convert a dataloader batch into forward primitives."""
+        """Convert a dataloader batch into forward primitives.
+
+        ``num_latent_t`` may override the configured temporal length when
+        creating zero latents for sampling-only paths such as validation.
+        """
 
     @abstractmethod
     def add_noise(

fastvideo/train/models/cosmos/cosmos.py

@@ -81,6 +81,7 @@ class CosmosModel(WanModel):
         *,
         generator: torch.Generator,
         latents_source: Literal["data", "zeros"] = "data",
+        num_latent_t: int | None = None,
     ) -> TrainingBatch:
         """Same flow as Wan, but uses Cosmos VAE
         normalisation."""
@@ -97,6 +98,9 @@ class CosmosModel(WanModel):
         infos = raw_batch.get("info_list")
 
         if latents_source == "zeros":
+            resolved_num_latent_t = tc.data.num_latent_t if num_latent_t is None else int(num_latent_t)
+            if resolved_num_latent_t <= 0:
+                raise ValueError("num_latent_t must be positive when creating zero latents")
             batch_size = encoder_hidden_states.shape[0]
             vae_config = (
                 tc.pipeline_config.vae_config  # type: ignore[union-attr]
@@ -112,13 +116,15 @@ class CosmosModel(WanModel):
             latents = torch.zeros(
                 batch_size,
                 num_channels,
-                tc.data.num_latent_t,
+                resolved_num_latent_t,
                 latent_height,
                 latent_width,
                 device=device,
                 dtype=dtype,
             )
         elif latents_source == "data":
+            if num_latent_t is not None:
+                raise ValueError("num_latent_t can only override latents_source='zeros'")
             if "vae_latent" not in raw_batch:
                 raise ValueError("vae_latent not found in batch "
                                  "and latents_source='data'")

fastvideo/train/models/hunyuan/hunyuan.py

@@ -75,6 +75,7 @@ class HunyuanModel(WanModel):
         *,
         generator: torch.Generator,
         latents_source: Literal["data", "zeros"] = "data",
+        num_latent_t: int | None = None,
     ) -> TrainingBatch:
         """Same flow as Wan, but uses Hunyuan VAE normalisation."""
         self.ensure_negative_conditioning()
@@ -90,6 +91,9 @@ class HunyuanModel(WanModel):
         infos = raw_batch.get("info_list")
 
         if latents_source == "zeros":
+            resolved_num_latent_t = tc.data.num_latent_t if num_latent_t is None else int(num_latent_t)
+            if resolved_num_latent_t <= 0:
+                raise ValueError("num_latent_t must be positive when creating zero latents")
             batch_size = encoder_hidden_states.shape[0]
             vae_config = (
                 tc.pipeline_config.vae_config  # type: ignore[union-attr]
@@ -105,13 +109,15 @@ class HunyuanModel(WanModel):
             latents = torch.zeros(
                 batch_size,
                 num_channels,
-                tc.data.num_latent_t,
+                resolved_num_latent_t,
                 latent_height,
                 latent_width,
                 device=device,
                 dtype=dtype,
             )
         elif latents_source == "data":
+            if num_latent_t is not None:
+                raise ValueError("num_latent_t can only override latents_source='zeros'")
             if "vae_latent" not in raw_batch:
                 raise ValueError("vae_latent not found in batch "
                                  "and latents_source='data'")

fastvideo/train/models/matrixgame2/matrixgame2.py

@@ -54,6 +54,7 @@ class MatrixGame2Model(WanModel):
         *,
         generator: torch.Generator,
         latents_source: Literal["data", "zeros"] = "data",
+        num_latent_t: int | None = None,
     ) -> TrainingBatch:
         assert self.training_config is not None
         tc = self.training_config
@@ -65,8 +66,16 @@ class MatrixGame2Model(WanModel):
         batch_size = self._infer_batch_size(raw_batch)
 
         if latents_source == "zeros":
-            latents = self._make_zero_latents(batch_size=batch_size)
+            resolved_num_latent_t = tc.data.num_latent_t if num_latent_t is None else int(num_latent_t)
+            if resolved_num_latent_t <= 0:
+                raise ValueError("num_latent_t must be positive when creating zero latents")
+            latents = self._make_zero_latents(
+                batch_size=batch_size,
+                num_latent_t=resolved_num_latent_t,
+            )
         elif latents_source == "data":
+            if num_latent_t is not None:
+                raise ValueError("num_latent_t can only override latents_source='zeros'")
             latents = raw_batch["vae_latent"][:, :, :tc.data.num_latent_t]
             latents = latents.to(device=device, dtype=dtype)
         else:
@@ -427,7 +436,12 @@ class MatrixGame2Model(WanModel):
             return int(raw_batch["clip_feature"].shape[0])
         raise ValueError("Unable to infer batch size from Matrix-Game 2.0 batch")
 
-    def _make_zero_latents(self, *, batch_size: int) -> torch.Tensor:
+    def _make_zero_latents(
+        self,
+        *,
+        batch_size: int,
+        num_latent_t: int,
+    ) -> torch.Tensor:
         assert self.training_config is not None
         vae_config = self.training_config.pipeline_config.vae_config.arch_config  # type: ignore[union-attr]
         num_channels = vae_config.z_dim
@@ -437,7 +451,7 @@ class MatrixGame2Model(WanModel):
         return torch.zeros(
             batch_size,
             num_channels,
-            self.training_config.data.num_latent_t,
+            num_latent_t,
             latent_height,
             latent_width,
             device=self.device,

fastvideo/train/models/wan/wan.py

@@ -235,6 +235,7 @@ class WanModel(ModelBase):
         *,
         generator: torch.Generator,
         latents_source: Literal["data", "zeros"] = "data",
+        num_latent_t: int | None = None,
     ) -> TrainingBatch:
         if self._requires_negative_conditioning:
             self.ensure_negative_conditioning()
@@ -250,6 +251,9 @@ class WanModel(ModelBase):
         infos = raw_batch.get("info_list")
 
         if latents_source == "zeros":
+            resolved_num_latent_t = tc.data.num_latent_t if num_latent_t is None else int(num_latent_t)
+            if resolved_num_latent_t <= 0:
+                raise ValueError("num_latent_t must be positive when creating zero latents")
             batch_size = encoder_hidden_states.shape[0]
             vae_config = (
                 tc.pipeline_config.vae_config.arch_config  # type: ignore[union-attr]
@@ -261,13 +265,15 @@ class WanModel(ModelBase):
             latents = torch.zeros(
                 batch_size,
                 num_channels,
-                tc.data.num_latent_t,
+                resolved_num_latent_t,
                 latent_height,
                 latent_width,
                 device=device,
                 dtype=dtype,
             )
         elif latents_source == "data":
+            if num_latent_t is not None:
+                raise ValueError("num_latent_t can only override latents_source='zeros'")
             if "vae_latent" not in raw_batch:
                 raise ValueError("vae_latent not found in batch "
                                  "and latents_source='data'")

modal_train_diffusion_nft.py

@@ -46,6 +46,7 @@ DEFAULT_SAMPLE_GUIDANCE_SCALE = -1.0
 DEFAULT_VALIDATION_NUM_STEPS = 50
 DEFAULT_VALIDATION_NUM_PROMPTS = 16
 DEFAULT_VALIDATION_BATCH_SIZE = 4
+DEFAULT_VALIDATION_NUM_FRAMES = 0
 DEFAULT_NUM_FRAMES = 17
 DEFAULT_NUM_LATENT_T = 0
 DEFAULT_LOG_SAMPLE_MAX_VIDEOS = 16
@@ -157,6 +158,7 @@ def train(
     validation_num_steps: int = DEFAULT_VALIDATION_NUM_STEPS,
     validation_num_prompts: int = DEFAULT_VALIDATION_NUM_PROMPTS,
     validation_batch_size: int = DEFAULT_VALIDATION_BATCH_SIZE,
+    validation_num_frames: int = DEFAULT_VALIDATION_NUM_FRAMES,
     num_frames: int = DEFAULT_NUM_FRAMES,
     num_latent_t: int = DEFAULT_NUM_LATENT_T,
     log_sample_max_videos: int = DEFAULT_LOG_SAMPLE_MAX_VIDEOS,
@@ -187,6 +189,7 @@ def train(
     validation_num_steps = int(validation_num_steps)
     validation_num_prompts = int(validation_num_prompts)
     validation_batch_size = int(validation_batch_size)
+    validation_num_frames = int(validation_num_frames)
     if preprocess_batch_size <= 0:
         raise ValueError("--preprocess-batch-size must be positive")
     if learning_rate < 0.0 and learning_rate != DEFAULT_LEARNING_RATE:
@@ -206,6 +209,12 @@ def train(
         raise ValueError("--validation-num-prompts must be >= 0")
     if validation_batch_size < 0:
         raise ValueError("--validation-batch-size must be >= 0")
+    if validation_num_frames < 0:
+        raise ValueError("--validation-num-frames must be >= 0")
+    if validation_num_frames > 0 and (validation_num_frames - 1) % 4 != 0:
+        raise ValueError("Wan validation frame counts must satisfy "
+                         "num_frames = (num_latent_t - 1) * 4 + 1; "
+                         f"got {validation_num_frames}")
     if preprocess_num_gpus != 1:
         raise ValueError("FastVideo text preprocessing currently supports "
                          "--preprocess-num-gpus 1 only.")
@@ -311,6 +320,8 @@ def train(
     output_dir = summary["output_dir"]
     resolved_num_frames = int(summary["num_frames"])
     resolved_num_latent_t = int(summary["num_latent_t"])
+    resolved_validation_num_frames = validation_num_frames or resolved_num_frames
+    resolved_validation_num_latent_t = ((resolved_validation_num_frames - 1) // 4 + 1)
     cache_vol.commit()
     data_vol.commit()
 
@@ -331,6 +342,8 @@ def train(
         f"validation_num_steps={validation_num_steps if validation_num_steps > 0 else 'config'} "
         f"validation_num_prompts={validation_num_prompts if validation_num_prompts > 0 else 'config'} "
         f"validation_batch_size={validation_batch_size if validation_batch_size > 0 else 'config'} "
+        f"validation_num_frames={resolved_validation_num_frames} "
+        f"validation_num_latent_t={resolved_validation_num_latent_t} "
         f"num_frames={resolved_num_frames} "
         f"num_latent_t={resolved_num_latent_t} "
         f"log_sample_max_videos={log_sample_max_videos} "
@@ -382,6 +395,7 @@ def train(
         cmd.extend(["--method.validation.num_prompts", str(validation_num_prompts)])
     if validation_batch_size > 0:
         cmd.extend(["--method.validation.batch_size", str(validation_batch_size)])
+    cmd.extend(["--method.validation.num_latent_t", str(resolved_validation_num_latent_t)])
     if num_batches_per_epoch > 0:
         cmd.extend(["--method.num_batches_per_epoch", str(num_batches_per_epoch)])
 
@@ -414,6 +428,7 @@ def main(
     validation_num_steps: int = DEFAULT_VALIDATION_NUM_STEPS,
     validation_num_prompts: int = DEFAULT_VALIDATION_NUM_PROMPTS,
     validation_batch_size: int = DEFAULT_VALIDATION_BATCH_SIZE,
+    validation_num_frames: int = DEFAULT_VALIDATION_NUM_FRAMES,
     num_frames: int = DEFAULT_NUM_FRAMES,
     num_latent_t: int = DEFAULT_NUM_LATENT_T,
     log_sample_max_videos: int = DEFAULT_LOG_SAMPLE_MAX_VIDEOS,
@@ -439,6 +454,7 @@ def main(
         validation_num_steps=validation_num_steps,
         validation_num_prompts=validation_num_prompts,
         validation_batch_size=validation_batch_size,
+        validation_num_frames=validation_num_frames,
         num_frames=num_frames,
         num_latent_t=num_latent_t,
         log_sample_max_videos=log_sample_max_videos,

tests/local_tests/test_diffusion_nft_method_utils.py

@@ -33,6 +33,16 @@ class _FakeTracker:
         self.artifacts.append((artifacts, step))
 
 
+class _FakeValidationStudent:
+
+    def __init__(self):
+        self.kwargs = None
+
+    def prepare_batch(self, raw_batch, **kwargs):
+        self.kwargs = kwargs
+        return raw_batch
+
+
 def test_reward_diagnostic_metrics_match_per_prompt_groups():
     method = object.__new__(DiffusionNFTMethod)
     method._trained_prompt_hashes = set()
@@ -108,6 +118,23 @@ def test_log_validation_samples_caps_video_count_and_uses_configured_fps():
     assert tracker.artifacts == [({"validation/videos": ["video-1"]}, 10)]
 
 
+def test_prepare_validation_batch_uses_validation_only_temporal_length():
+    method = object.__new__(DiffusionNFTMethod)
+    method.student = _FakeValidationStudent()
+    method._validation_config = SimpleNamespace(num_latent_t=13)
+    generator = torch.Generator().manual_seed(0)
+    raw_batch = {"prompt": ["test"]}
+
+    result = method._prepare_validation_batch(raw_batch, generator)
+
+    assert result is raw_batch
+    assert method.student.kwargs == {
+        "generator": generator,
+        "latents_source": "zeros",
+        "num_latent_t": 13,
+    }
+
+
 def test_num_train_timesteps_uses_explicit_schedule_length():
     method = object.__new__(DiffusionNFTMethod)
     method._sample_steps = 25

tests/local_tests/test_train_rl_sampling.py

@@ -257,11 +257,16 @@ def test_validation_shard_indices_are_stable_and_padded():
     assert rank1 == [(1, True), (3, True), (0, False)]
 
 
-def test_rl_validation_config_parses_video_fps():
-    config = RLValidationConfig.from_mapping({"fps": 16, "max_samples": 2})
+def test_rl_validation_config_parses_video_shape_and_fps():
+    config = RLValidationConfig.from_mapping({
+        "fps": 16,
+        "max_samples": 2,
+        "num_latent_t": 13,
+    })
 
     assert config.fps == 16
     assert config.max_samples == 2
+    assert config.num_latent_t == 13
 
 
 def test_distributed_k_repeat_indices_repeats_prompts_globally():