From 753eed7def383cfebc754cd35d8a69f6c5cb5c8e Mon Sep 17 00:00:00 2001 From: Adam Lee Date: Tue, 14 Jul 2026 01:45:04 +0900 Subject: [PATCH] Push train setup with 49 frame validation --- .../train/methods/rl/common/validation.py | 7 +++++ fastvideo/train/methods/rl/diffusion_nft.py | 18 +++++++++---- fastvideo/train/models/base.py | 7 ++++- fastvideo/train/models/cosmos/cosmos.py | 8 +++++- fastvideo/train/models/hunyuan/hunyuan.py | 8 +++++- .../train/models/matrixgame2/matrixgame2.py | 20 +++++++++++--- fastvideo/train/models/wan/wan.py | 8 +++++- modal_train_diffusion_nft.py | 16 +++++++++++ .../test_diffusion_nft_method_utils.py | 27 +++++++++++++++++++ tests/local_tests/test_train_rl_sampling.py | 9 +++++-- 10 files changed, 114 insertions(+), 14 deletions(-) diff --git a/fastvideo/train/methods/rl/common/validation.py b/fastvideo/train/methods/rl/common/validation.py index 5c676e372..306e93ef5 100644 --- a/fastvideo/train/methods/rl/common/validation.py +++ b/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), ) diff --git a/fastvideo/train/methods/rl/diffusion_nft.py b/fastvideo/train/methods/rl/diffusion_nft.py index 3b3df149a..0830f956e 100644 --- a/fastvideo/train/methods/rl/diffusion_nft.py +++ b/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 diff --git a/fastvideo/train/models/base.py b/fastvideo/train/models/base.py index ffffb31ba..288e96336 100644 --- a/fastvideo/train/models/base.py +++ b/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( diff --git a/fastvideo/train/models/cosmos/cosmos.py b/fastvideo/train/models/cosmos/cosmos.py index 2bb73d425..1efec43a1 100644 --- a/fastvideo/train/models/cosmos/cosmos.py +++ b/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'") diff --git a/fastvideo/train/models/hunyuan/hunyuan.py b/fastvideo/train/models/hunyuan/hunyuan.py index d3beeee60..e17ed7eff 100644 --- a/fastvideo/train/models/hunyuan/hunyuan.py +++ b/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'") diff --git a/fastvideo/train/models/matrixgame2/matrixgame2.py b/fastvideo/train/models/matrixgame2/matrixgame2.py index 54f127f34..0a893e246 100644 --- a/fastvideo/train/models/matrixgame2/matrixgame2.py +++ b/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, diff --git a/fastvideo/train/models/wan/wan.py b/fastvideo/train/models/wan/wan.py index 316dafd1e..abccdae77 100644 --- a/fastvideo/train/models/wan/wan.py +++ b/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'") diff --git a/modal_train_diffusion_nft.py b/modal_train_diffusion_nft.py index ee114131d..8e825d5dd 100644 --- a/modal_train_diffusion_nft.py +++ b/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, diff --git a/tests/local_tests/test_diffusion_nft_method_utils.py b/tests/local_tests/test_diffusion_nft_method_utils.py index ea1c6c251..87d7f4c22 100644 --- a/tests/local_tests/test_diffusion_nft_method_utils.py +++ b/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 diff --git a/tests/local_tests/test_train_rl_sampling.py b/tests/local_tests/test_train_rl_sampling.py index bee223800..56d5793d9 100644 --- a/tests/local_tests/test_train_rl_sampling.py +++ b/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():