diff --git a/.agents/skills/add-reward-model/SKILL.md b/.agents/skills/add-reward-model/SKILL.md new file mode 100644 index 000000000..221b14f22 --- /dev/null +++ b/.agents/skills/add-reward-model/SKILL.md @@ -0,0 +1,32 @@ +--- +name: add-reward-model +description: Use when adding reusable reward models under fastvideo/train/methods/rl/rewards for RLHF or online RL training. +--- + +# Add Reward Model + +Use for reward models consumed by RL methods. + +## Placement + +- Put reusable reward code under `fastvideo/train/methods/rl/rewards/`. +- Expose public builders from `fastvideo/train/methods/rl/rewards/__init__.py`. +- Keep method-specific aggregation or advantage logic out of reward classes. + +## Media Inputs + +- Reward callables receive decoded media tensors. +- Accept single-frame tensors as `[B, C, H, W]` and multi-frame tensors as `[B, C, T, H, W]` when practical. +- Frame selection is reward-specific. Frame scorers such as PickScore and CLIPScore should explicitly select frame `0`; temporal rewards should inspect whichever frames they need. +- Return one scalar reward per prompt/sample. + +## Attribution + +- If code is ported or closely adapted from another repo, add a short comment or docstring naming the source file/function. +- Preserve SPDX headers used by FastVideo files. + +## Tests + +- Unit-test tensor layout handling without loading large reward checkpoints. +- Allow fake scorer injection for multi-reward tests. +- Test weighted reward aggregation and metric keys. diff --git a/.agents/skills/add-rl-method/SKILL.md b/.agents/skills/add-rl-method/SKILL.md new file mode 100644 index 000000000..5dcaa359e --- /dev/null +++ b/.agents/skills/add-rl-method/SKILL.md @@ -0,0 +1,38 @@ +--- +name: add-rl-method +description: Use when adding or modifying an RL/RLHF method under fastvideo/train/methods/rl, including DiffusionNFT-like methods. +--- + +# Add RL Method + +Use for new RL methods in the modular `fastvideo/train` stack. + +## Required Shape + +- Add the method under `fastvideo/train/methods/rl/`. +- Subclass `TrainingMethod`. +- Keep model-family logic in `ModelBase` wrappers. +- Decode generated latents through `ModelBase.decode_latents`; add that hook to the new model wrapper instead of decoding inside the RL method. +- Use `fastvideo/train/methods/rl/common/sampling.py` for generation unless the method has a documented reason to avoid sampling. +- Use `fastvideo/train/methods/rl/common/prompt_sampling.py` for reusable grouped prompt sampling patterns such as DiffusionNFT K-repeat. +- Use `fastvideo/train/methods/rl/rewards/` for reward models. + +## Optimization + +- Return `manages_optimization() == True` only when the method must own a nonstandard outer/inner loop. +- If using managed optimization, implement `managed_train_step(data_stream, iteration)`. +- Existing trainer callbacks, checkpointing, tracking, and validation should still work. + +## Config + +- Put method knobs under `method`. +- Put sampler knobs under `method.sampling`. +- Do not put scheduler or trajectory policy into model configs. +- Do not split a diffusers-style scheduler from its built-in `step()` solver in YAML; use `trajectory` only for higher-level ODE vs re-noise behavior. +- Avoid fixed timestep lists in examples unless reproducing a known baseline; prefer scheduler-generated defaults. + +## Tests + +- Add fake-model tests for sampler/method behavior. +- Add config parse tests for the public YAML. +- Confirm existing train methods stay on the default Trainer path. diff --git a/.agents/skills/index.jsonl b/.agents/skills/index.jsonl index 9d72f4333..83e7767ed 100644 --- a/.agents/skills/index.jsonl +++ b/.agents/skills/index.jsonl @@ -8,3 +8,6 @@ {"name": "decompose-pipeline-pr", "description": "Decompose an oversized FastVideo pipeline PR into a stack of independently-reviewable PRs. Tiers the diff by blast radius (invisible / dead code / cross-cutting infra / activation), produces a branch graph and worktree bootstrap, drafts the AGENTS.md manifest, flags missing tests on cross-cutting infra changes, and extracts lessons from the PR body. Worked example: PR #1280 daVinci-MagiHuman (9.8k LOC) decomposed into 10 stacked PRs.", "path": "decompose-pipeline-pr/SKILL.md", "status": "tested", "trust": "medium"} {"name": "reseed-performance-baseline", "description": "Re-seed the HF performance-tracking baseline for an intentional runtime, dependency, or environment-caused benchmark shift. Use when performance CI fails because metrics such as latency, throughput, component time, or peak memory changed for an accepted reason and the rolling median baseline must be advanced by replicating one reviewed shifted source result into three success=true records, or five records when explicitly requested", "path": "reseed-performance-baseline/SKILL.md", "status": "draft", "trust": "low"} {"name": "add-model", "description": "Add a new model (or variant) to FastVideo: DiT + configs + pipeline + presets + registry + tests. Walks through FastVideo's single stage-based pipeline architecture with exact file paths and registration hooks.", "path": "add-model/SKILL.md", "status": "draft", "trust": "low"} +{"name": "rlhf-training-abstractions", "description": "Use when changing FastVideo RLHF/RL training infrastructure, especially sampler, reward, scheduler trajectory, or method boundaries under fastvideo/train.", "path": "rlhf-training-abstractions/SKILL.md", "status": "draft", "trust": "low"} +{"name": "add-rl-method", "description": "Use when adding or modifying an RL/RLHF method under fastvideo/train/methods/rl, including DiffusionNFT-like methods.", "path": "add-rl-method/SKILL.md", "status": "draft", "trust": "low"} +{"name": "add-reward-model", "description": "Use when adding reusable reward models under fastvideo/train/methods/rl/rewards for RLHF or online RL training.", "path": "add-reward-model/SKILL.md", "status": "draft", "trust": "low"} diff --git a/.agents/skills/rlhf-training-abstractions/SKILL.md b/.agents/skills/rlhf-training-abstractions/SKILL.md new file mode 100644 index 000000000..e6fcf778e --- /dev/null +++ b/.agents/skills/rlhf-training-abstractions/SKILL.md @@ -0,0 +1,41 @@ +--- +name: rlhf-training-abstractions +description: Use when changing FastVideo RLHF/RL training infrastructure, especially sampler, reward, scheduler trajectory, or method boundaries under fastvideo/train. +--- + +# RLHF Training Abstractions + +Use this skill before editing RLHF-style training code in `fastvideo/train`. + +## Boundaries + +- RL methods live under `fastvideo/train/methods/rl/` and own algorithm logic: reward collection, advantage computation, policy loss, KL/reference terms, and optimizer cadence. +- Rewards live under `fastvideo/train/methods/rl/rewards/` and must be reusable across RL methods. +- RL methods pass decoded media to rewards; each reward decides whether to use the first frame, sampled frames, or the full video. +- Sampling lives under `fastvideo/train/methods/rl/common/` and must use `ModelBase` primitives plus scheduler math, not model-family inference pipelines. +- Model wrappers under `fastvideo/train/models/` own model-specific forward details. +- Model wrappers also own model-specific latent decoding via `ModelBase.decode_latents`; RL methods should not reach into VAE normalization internals. +- Shared RL helpers such as K-repeat prompt sampling belong under `fastvideo/train/methods/rl/common/` when they are reusable across RL methods. + +## Anti-Patterns + +- Do not bind RL methods to inference pipeline classes such as `WanDMDPipeline`. +- Do not hardcode timestep lists in a method when the scheduler can generate them. +- Do not put reward-model code inside one RL method. +- Do not make existing non-RL methods use method-managed optimization unless explicitly requested. + +## Sampling Policy + +- Prefer YAML-configured `method.sampling` with `scheduler`, `trajectory`, `num_steps`, `timesteps`, and `sigmas`. +- Treat diffusers-style scheduler classes as owning both the timestep schedule and their `step()` update rule; avoid a separate `solver` field unless a new sampler truly implements solver math outside the scheduler object. +- Missing `timesteps` means “ask the scheduler”; explicit `timesteps` or `sigmas` are overrides. +- ODE-style trajectories should not re-noise between denoising steps. +- SDE/re-noise behavior must be explicit in config. + +## Validation + +- Run focused local tests for sampler config and Trainer opt-in behavior. +- Verify existing train methods still report `manages_optimization() == False`. +- Keep fixed-prompt validation helpers in `fastvideo/train/methods/rl/common/validation.py` so new RL methods can reuse sharding and captions. +- Test distributed prompt grouping helpers separately from heavyweight model loading. +- Run `pre-commit run --files `; respect configured excludes. diff --git a/examples/train/configs/distribution_matching/wan/dmd2_text_only_validation.json b/examples/train/configs/distribution_matching/wan/dmd2_text_only_validation.json new file mode 100644 index 000000000..3553f02ff --- /dev/null +++ b/examples/train/configs/distribution_matching/wan/dmd2_text_only_validation.json @@ -0,0 +1,52 @@ +{ + "data": [ + { + "caption": "gold tip pyramid in the night, extremely detailed , rain, stars" + }, + { + "caption": "a fruit stacking in the shape of a dog, stock image, shutterstock" + }, + { + "caption": "Landscape, By Lee madgwick, by Luis Royo, by Louise nevelson" + }, + { + "caption": "A colorful poster that says \"philo is a weird\"" + }, + { + "caption": "Danish male with blue eyes, realistic, viking" + }, + { + "caption": "futuristic, cityscape, flying cars, neon lights, towering skyscrapers, glowing purple sky." + }, + { + "caption": "a crow with cameras for eyes, sitting on a mans shoulder, anime, studio ghibli, fantasy, fairytale, sketch, digital art, watercolor, dnd, rustic, professional photograph, medieval, hd, 4k" + }, + { + "caption": "a background image mixing the matrix and AI" + }, + { + "caption": "Golden sunset, a bright orange and yellow sky is visible, lit up by the setting sun, the horizon is a mix of bright colors and deep shadows" + }, + { + "caption": "Grim reaper playing an electric guitar" + }, + { + "caption": "an epic view of a demonic Rose-ringed parakeet cyborg inside an ironmaiden robot,wearing a noble robe,large view,a surrealist painting, aralan bean and Philippe Druillet,hiromu arakawa,volumetric lighting,detailed shadows" + }, + { + "caption": "Ben Shapiro as the cover of ministry's filth pig album, but covered in milk" + }, + { + "caption": "king charles spaniel with , ethereal, midjourney style lighting and shadows, insanely detailed, 8k, photorealistic" + }, + { + "caption": "A website for a party resort service" + }, + { + "caption": "full shot of a steampunk horse" + }, + { + "caption": "60s psycedelic spiritual jazz album art" + } + ] +} diff --git a/examples/train/configs/example.yaml b/examples/train/configs/example.yaml index c01546bcb..b174f0c4c 100644 --- a/examples/train/configs/example.yaml +++ b/examples/train/configs/example.yaml @@ -87,6 +87,7 @@ training: # --- training.data [TYPED] -> DataConfig --- data: data_path: data/my_dataset # default: "" + preprocessed_data_type: t2v # default: "t2v" ("text_only" for simulate-only DMD text prompts) train_batch_size: 1 # default: 1 dataloader_num_workers: 4 # default: 0 training_cfg_rate: 0.1 # default: 0.0 diff --git a/examples/train/configs/rl/wan/diffusion_nft_pick_clip.yaml b/examples/train/configs/rl/wan/diffusion_nft_pick_clip.yaml new file mode 100644 index 000000000..0fdf95ab6 --- /dev/null +++ b/examples/train/configs/rl/wan/diffusion_nft_pick_clip.yaml @@ -0,0 +1,116 @@ +# DiffusionNFT multi-reward single-frame RL: Wan 2.1 T2V 1.3B on text-only PickScore prompts. +# +# Single-frame RL is represented as a one-latent-frame Wan run: +# num_latent_t: 1 +# num_frames: 1 +# +# The method trains the full transformer (no LoRA) and keeps an old-policy +# transformer plus a frozen reference transformer, matching the non-LoRA +# DiffusionNFT loss path. + +models: + student: + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: true + old: + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: false + disable_custom_init_weights: true + reference: + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: false + disable_custom_init_weights: true + +method: + _target_: fastvideo.train.methods.rl.diffusion_nft.DiffusionNFTMethod + reward_fn: + pickscore: 1.0 + clipscore: 1.0 + + sampling: + num_steps: 25 + scheduler: flow_match_euler + trajectory: ode + flow_shift: inherit + + validation: + every_steps: 10 + num_steps: 40 + num_prompts: 16 + batch_size: 16 + log_samples: true + seed: 42 + # Null reuses training.data.data_path. Override this with a held-out + # preprocessed parquet path when one is available. + data_path: + + # DiffusionNFT sd3_multi_reward on 4 GPUs resolves to per-GPU sample batch + # size 6, 48 sample batches per outer epoch, and grad accumulation 48. + sample_train_batch_size: 6 + train_batch_size: 6 + num_batches_per_epoch: 48 + num_video_per_prompt: 24 + num_inner_epochs: 1 + timestep_fraction: 0.99 + + beta: 0.1 + kl_beta: 0.0001 + decay_type: 1 + adv_mode: all + adv_clip_max: 5 + max_grad_norm: 1.0 + ema: + enabled: true + decay: 0.9 + update_after_step: 0 + validation: true + terminal_progress: true + +training: + distributed: + num_gpus: 4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + + data: + data_path: data/pickscore_text_only_preprocessed + preprocessed_data_type: text_only + dataloader_num_workers: 0 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 42 + num_latent_t: 1 + num_height: 448 + num_width: 832 + num_frames: 1 + + optimizer: + learning_rate: 3.0e-5 + betas: [0.9, 0.999] + weight_decay: 0.0001 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 100000 + gradient_accumulation_steps: 48 + + checkpoint: + output_dir: outputs/wan2.1_diffusion_nft_pick_clip + training_state_checkpointing_steps: 30 + checkpoints_total_limit: 3 + + tracker: + project_name: diffusion_nft_wan + run_name: wan2.1_diffusion_nft_pick_clip + + model: + enable_gradient_checkpointing_type: full + +pipeline: + flow_shift: 8 diff --git a/fastvideo/tests/train/callbacks/test_validation.py b/fastvideo/tests/train/callbacks/test_validation.py index 42a2d0888..4b2af70e2 100644 --- a/fastvideo/tests/train/callbacks/test_validation.py +++ b/fastvideo/tests/train/callbacks/test_validation.py @@ -67,6 +67,8 @@ class TestConstructor: assert cb.num_frames is None assert cb.sampling_timesteps is None assert cb.output_dir is None + assert cb.offload_training_state is False + assert cb.unload_pipeline_after_validation is False # Lazy fields not yet populated. assert cb._pipeline is None assert cb._sampling_param is None @@ -83,12 +85,16 @@ class TestConstructor: guidance_scale="4.5", # type: ignore[arg-type] num_frames="77", # type: ignore[arg-type] sampling_timesteps=["1000", "500"], + offload_training_state="1", # type: ignore[arg-type] + unload_pipeline_after_validation="false", # type: ignore[arg-type] ) assert cb.every_steps == 50 assert cb.sampling_steps == [20, 40] assert cb.guidance_scale == 4.5 assert cb.num_frames == 77 assert cb.sampling_timesteps == [1000, 500] + assert cb.offload_training_state is True + assert cb.unload_pipeline_after_validation is False def test_pipeline_kwargs_collected(self) -> None: cb = ValidationCallback( diff --git a/fastvideo/train/callbacks/validation.py b/fastvideo/train/callbacks/validation.py index 7c9752394..2b0ef702f 100644 --- a/fastvideo/train/callbacks/validation.py +++ b/fastvideo/train/callbacks/validation.py @@ -68,6 +68,8 @@ class ValidationCallback(Callback): num_frames: int | None = None, output_dir: str | None = None, sampling_timesteps: list[int] | None = None, + offload_training_state: bool = False, + unload_pipeline_after_validation: bool = False, **pipeline_kwargs: Any, ) -> None: self.pipeline_target = str(pipeline_target) @@ -78,6 +80,8 @@ class ValidationCallback(Callback): self.num_frames = (int(num_frames) if num_frames is not None else None) self.output_dir = (str(output_dir) if output_dir is not None else None) self.sampling_timesteps = ([int(s) for s in sampling_timesteps] if sampling_timesteps is not None else None) + self.offload_training_state = self._coerce_bool(offload_training_state) + self.unload_pipeline_after_validation = self._coerce_bool(unload_pipeline_after_validation) self.pipeline_kwargs = dict(pipeline_kwargs) # Set after on_train_start. @@ -88,6 +92,12 @@ class ValidationCallback(Callback): self.validation_random_generator: (torch.Generator | None) = None self.seed: int = 0 + @staticmethod + def _coerce_bool(value: Any) -> bool: + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on"} + return bool(value) + # ---------------------------------------------------------- # Callback hooks # ---------------------------------------------------------- @@ -140,16 +150,183 @@ class ValidationCallback(Callback): ) -> None: transformer = method.student.transformer - # Look for an EMA callback to temporarily swap - # EMA weights during validation. - ema_cb = self._find_ema_callback() - ctx = ema_cb.ema_context(transformer) if ema_cb is not None else contextlib.nullcontext(transformer) - with ctx as t: - self._run_validation_inner( + try: + with self._validation_memory_context( + method, + validation_transformer=transformer, + ): + # Look for an EMA callback to temporarily swap + # EMA weights during validation. + ema_cb = self._find_ema_callback() + ctx = ema_cb.ema_context(transformer) if ema_cb is not None else contextlib.nullcontext(transformer) + with ctx as t: + self._run_validation_inner( + method, + step, + t, + ) + finally: + if self.unload_pipeline_after_validation: + self._clear_pipeline_cache() + + @contextlib.contextmanager + def _validation_memory_context( + self, + method: TrainingMethod, + *, + validation_transformer: torch.nn.Module, + ): + if not self.offload_training_state: + yield + return + + optimizer_tensor_records: list[tuple[Any, Any, torch.device]] = [] + module_records: list[tuple[str, torch.nn.Module, torch.device]] = [] + try: + self._offload_optimizer_states_to_cpu( method, - step, - t, + optimizer_tensor_records, ) + self._offload_inactive_role_modules_to_cpu( + method, + validation_transformer=validation_transformer, + module_records=module_records, + ) + self._empty_cuda_cache() + yield + finally: + self._restore_inactive_role_modules(module_records) + self._restore_optimizer_states(optimizer_tensor_records) + self._empty_cuda_cache() + + def _offload_optimizer_states_to_cpu( + self, + method: TrainingMethod, + records: list[tuple[Any, Any, torch.device]], + ) -> None: + optimizers = getattr(method, "_optimizer_dict", {}) + if not optimizers: + return + moved = 0 + for optimizer in optimizers.values(): + state = getattr(optimizer, "state", None) + if not isinstance(state, dict): + continue + for param_state in state.values(): + moved += self._offload_tensor_container_to_cpu( + param_state, + records, + ) + if moved: + logger.info( + "Offloaded %d optimizer state tensors to CPU for validation.", + moved, + ) + + def _offload_tensor_container_to_cpu( + self, + obj: Any, + records: list[tuple[Any, Any, torch.device]], + ) -> int: + moved = 0 + if isinstance(obj, dict): + for key, value in list(obj.items()): + if torch.is_tensor(value) and value.device.type == "cuda": + records.append((obj, key, value.device)) + obj[key] = value.detach().cpu() + moved += 1 + else: + moved += self._offload_tensor_container_to_cpu(value, records) + return moved + if isinstance(obj, list): + for idx, value in enumerate(list(obj)): + if torch.is_tensor(value) and value.device.type == "cuda": + records.append((obj, idx, value.device)) + obj[idx] = value.detach().cpu() + moved += 1 + else: + moved += self._offload_tensor_container_to_cpu(value, records) + return moved + + def _restore_optimizer_states( + self, + records: list[tuple[Any, Any, torch.device]], + ) -> None: + for container, key, device in reversed(records): + value = container[key] + if torch.is_tensor(value): + container[key] = value.to(device=device) + if records: + logger.info( + "Restored %d optimizer state tensors after validation.", + len(records), + ) + + def _offload_inactive_role_modules_to_cpu( + self, + method: TrainingMethod, + *, + validation_transformer: torch.nn.Module, + module_records: list[tuple[str, torch.nn.Module, torch.device]], + ) -> None: + role_models = getattr(method, "_role_models", {}) + if not isinstance(role_models, dict): + return + + for role, model in role_models.items(): + module = getattr(model, "transformer", None) + if not isinstance(module, torch.nn.Module): + continue + if module is validation_transformer: + continue + device = self._first_cuda_tensor_device(module) + if device is None: + continue + try: + module.to("cpu") + except Exception as exc: + logger.warning( + "Could not offload role %r transformer to CPU before validation: %s", + role, + exc, + ) + continue + module_records.append((str(role), module, device)) + logger.info( + "Offloaded role %r transformer from %s to CPU for validation.", + role, + device, + ) + + def _restore_inactive_role_modules( + self, + module_records: list[tuple[str, torch.nn.Module, torch.device]], + ) -> None: + for role, module, device in reversed(module_records): + module.to(device) + logger.info( + "Restored role %r transformer to %s after validation.", + role, + device, + ) + + @staticmethod + def _first_cuda_tensor_device(module: torch.nn.Module) -> torch.device | None: + for tensor in list(module.parameters(recurse=True)) + list(module.buffers(recurse=True)): + device = getattr(tensor, "device", None) + if isinstance(device, torch.device) and device.type == "cuda": + return device + return None + + def _clear_pipeline_cache(self) -> None: + self._pipeline = None + self._pipeline_key = None + self._empty_cuda_cache() + + @staticmethod + def _empty_cuda_cache() -> None: + if torch.cuda.is_available(): + torch.cuda.empty_cache() def _find_ema_callback(self) -> Any | None: """Find the EMA callback in the callback dict.""" @@ -293,6 +470,7 @@ class ValidationCallback(Callback): } if flow_shift is not None: kwargs["flow_shift"] = float(flow_shift) + kwargs.update(self.pipeline_kwargs) self._pipeline = PipelineCls.from_pretrained( tc.model_path, diff --git a/fastvideo/train/methods/base.py b/fastvideo/train/methods/base.py index ca295b9b2..840f7dee0 100644 --- a/fastvideo/train/methods/base.py +++ b/fastvideo/train/methods/base.py @@ -197,6 +197,46 @@ class TrainingMethod(torch.nn.Module, ABC): # -- Shared hooks (override in subclasses as needed) -- + def manages_optimization(self) -> bool: + """Whether the method owns backward/optimizer stepping internally. + + Most methods return loss tensors and let :class:`Trainer` handle + gradient accumulation, callbacks, optimizer stepping, and scheduler + stepping. RL-style methods such as DiffusionNFT need to preserve their + own sample-then-inner-train loop, so they can provide a specific + ``managed_train_step``. + """ + return False + + def managed_train_step( + self, + data_stream: Any, + iteration: int, + ) -> tuple[ + dict[str, torch.Tensor], + dict[str, Any], + dict[str, LogScalar], + ]: + """Run one method-managed step. + + Subclasses that return ``True`` from :meth:`manages_optimization` + should override this. The fallback consumes one dataloader batch and + delegates to ``single_train_step`` so tests can exercise the hook with + tiny fake methods. + """ + return self.single_train_step(next(data_stream), iteration) + + def on_validation_begin(self, iteration: int = 0) -> dict[str, LogScalar]: + """Run method-owned validation, if any. + + Pipeline-style validation should remain in callbacks. Methods that + intentionally avoid inference pipelines, such as RL methods with their + own sampler/reward loop, can override this hook and return metrics for + the trainer to log at ``iteration``. + """ + del iteration + return {} + def get_grad_clip_targets( self, iteration: int, diff --git a/fastvideo/train/methods/distribution_matching/dmd2.py b/fastvideo/train/methods/distribution_matching/dmd2.py index 6bfbb44ad..275fb01ed 100644 --- a/fastvideo/train/methods/distribution_matching/dmd2.py +++ b/fastvideo/train/methods/distribution_matching/dmd2.py @@ -55,6 +55,8 @@ class DMD2Method(TrainingMethod): raise ValueError("DMD2Method requires critic to be trainable") self._cfg_uncond = self._parse_cfg_uncond() self._rollout_mode = self._parse_rollout_mode() + self._validate_preprocessed_data_type() + self._configure_student_negative_conditioning() self._denoising_step_list: torch.Tensor | None = (None) # Initialize preprocessors on student. @@ -206,6 +208,13 @@ class DMD2Method(TrainingMethod): return targets def _parse_rollout_mode(self, ) -> Literal["simulate", "data_latent"]: + """Parse how DMD2 obtains the latent point used for rollout. + + ``simulate`` starts from fresh noise and lets the student create an + artificial latent trajectory, so it can run with text-only data. + ``data_latent`` starts from preprocessed VAE latents and perturbs them + at a sampled denoising timestep. + """ raw = self.method_config.get("rollout_mode", None) if raw is None: raise ValueError("method_config.rollout_mode must be set " @@ -223,6 +232,34 @@ class DMD2Method(TrainingMethod): "{simulate, data_latent}, got " f"{raw!r}") + def _validate_preprocessed_data_type(self) -> None: + data_type = str(getattr( + self.training_config.data, + "preprocessed_data_type", + "t2v", + )).strip().lower() + if data_type == "text_only" and self._rollout_mode != "simulate": + raise ValueError("training.data.preprocessed_data_type='text_only' " + "requires method.rollout_mode='simulate'; " + "data_latent rollout requires vae_latent data.") + + def _uses_negative_prompt_conditioning(self) -> bool: + if self._cfg_uncond is None: + return True + text_policy = self._cfg_uncond.get("text", None) + if text_policy is None: + return True + return str(text_policy).strip().lower() == "negative_prompt" + + def _configure_student_negative_conditioning(self) -> None: + setter = getattr( + self.student, + "set_requires_negative_conditioning", + None, + ) + if setter is not None: + setter(self._uses_negative_prompt_conditioning()) + def _parse_cfg_uncond(self, ) -> dict[str, Any] | None: raw = self.method_config.get("cfg_uncond", None) if raw is None: diff --git a/fastvideo/train/methods/rl/__init__.py b/fastvideo/train/methods/rl/__init__.py new file mode 100644 index 000000000..0cf2dda3e --- /dev/null +++ b/fastvideo/train/methods/rl/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +"""RL training methods.""" + +from fastvideo.train.methods.rl.diffusion_nft import DiffusionNFTMethod + +__all__ = ["DiffusionNFTMethod"] diff --git a/fastvideo/train/methods/rl/common/__init__.py b/fastvideo/train/methods/rl/common/__init__.py new file mode 100644 index 000000000..d4a57f915 --- /dev/null +++ b/fastvideo/train/methods/rl/common/__init__.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reusable RL training primitives.""" + +from fastvideo.train.methods.rl.common.sampling import ( + DiffusionSampler, + SamplingConfig, + SamplingResult, +) +from fastvideo.train.methods.rl.common.prompt_sampling import ( + KRepeatSample, + distributed_k_repeat_indices, +) +from fastvideo.train.methods.rl.common.validation import ( + RLValidationConfig, + media_to_video_array, + validation_caption, + validation_shard_indices, +) + +__all__ = [ + "DiffusionSampler", + "KRepeatSample", + "RLValidationConfig", + "SamplingConfig", + "SamplingResult", + "distributed_k_repeat_indices", + "media_to_video_array", + "validation_caption", + "validation_shard_indices", +] diff --git a/fastvideo/train/methods/rl/common/prompt_sampling.py b/fastvideo/train/methods/rl/common/prompt_sampling.py new file mode 100644 index 000000000..aff1bda55 --- /dev/null +++ b/fastvideo/train/methods/rl/common/prompt_sampling.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Prompt-row sampling helpers for online RL methods. + +This module chooses and repeats dataset prompt rows across ranks for RL training +batches. Here, "sampling" means selection, not generator sampling. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True, slots=True) +class KRepeatSample: + """Local prompt indices for one distributed K-repeat sampling batch.""" + + local_indices: list[int] + unique_prompt_count: int + + +def distributed_k_repeat_indices( + *, + dataset_length: int, + batch_size: int, + repeats_per_prompt: int, + world_size: int, + rank: int, + seed: int, +) -> KRepeatSample: + """Mirror DiffusionNFT's distributed K-repeat prompt sampler. + + Adapted from DiffusionNFT's + ``scripts/train_nft_sd3.py::DistributedKRepeatSampler``. + """ + dataset_length = int(dataset_length) + batch_size = int(batch_size) + repeats_per_prompt = int(repeats_per_prompt) + world_size = int(world_size) + rank = int(rank) + if dataset_length <= 0: + raise ValueError("dataset_length must be positive") + if batch_size <= 0: + raise ValueError("batch_size must be positive") + if repeats_per_prompt <= 0: + raise ValueError("repeats_per_prompt must be positive") + if world_size <= 0: + raise ValueError("world_size must be positive") + if rank < 0 or rank >= world_size: + raise ValueError(f"rank must be in [0, {world_size}), got {rank}") + + total_samples = world_size * batch_size + if total_samples % repeats_per_prompt != 0: + raise ValueError("world_size * batch_size must be divisible by repeats_per_prompt " + f"({world_size} * {batch_size} vs {repeats_per_prompt})") + unique_prompt_count = total_samples // repeats_per_prompt + if unique_prompt_count > dataset_length: + raise ValueError("K-repeat sampling needs at least as many rows as unique prompts " + f"per sampling batch ({dataset_length} < {unique_prompt_count})") + + generator = torch.Generator() + generator.manual_seed(int(seed)) + indices = torch.randperm(dataset_length, generator=generator)[:unique_prompt_count].tolist() + repeated_indices = [idx for idx in indices for _ in range(repeats_per_prompt)] + shuffled_order = torch.randperm(len(repeated_indices), generator=generator).tolist() + shuffled_samples = [int(repeated_indices[idx]) for idx in shuffled_order] + + start = rank * batch_size + end = start + batch_size + return KRepeatSample( + local_indices=shuffled_samples[start:end], + unique_prompt_count=unique_prompt_count, + ) diff --git a/fastvideo/train/methods/rl/common/sampling.py b/fastvideo/train/methods/rl/common/sampling.py new file mode 100644 index 000000000..c0d7be50a --- /dev/null +++ b/fastvideo/train/methods/rl/common/sampling.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Configurable diffusion samplers for RL training methods.""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass +from typing import Any, Literal + +import torch + +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, ) +from fastvideo.pipelines import TrainingBatch +from fastvideo.train.models.base import ModelBase + +SchedulerName = Literal["flow_match_euler", "model_default"] +TrajectoryName = Literal["ode", "sde_reflow"] + + +@dataclass(slots=True) +class SamplingConfig: + """YAML-backed sampling knobs shared by RL methods.""" + + num_steps: int = 25 + scheduler: SchedulerName = "model_default" + trajectory: TrajectoryName = "ode" + flow_shift: float | None = None + timesteps: list[float] | None = None + sigmas: list[float] | None = None + + @classmethod + def from_mapping(cls, raw: dict[str, Any] | None) -> SamplingConfig: + if raw is None: + return cls() + if not isinstance(raw, dict): + raise ValueError(f"method.sampling must be a mapping, got {type(raw).__name__}") + supported_keys = { + "flow_shift", + "num_steps", + "scheduler", + "sigmas", + "timesteps", + "trajectory", + } + unsupported_keys = sorted(set(raw) - supported_keys) + if unsupported_keys: + raise ValueError(f"Unsupported method.sampling key(s): {unsupported_keys}. " + f"Supported keys: {sorted(supported_keys)}") + scheduler = str(raw.get("scheduler", "model_default") or "model_default").strip().lower() + if scheduler not in {"flow_match_euler", "model_default"}: + raise ValueError("method.sampling.scheduler must be one of " + "{flow_match_euler, model_default}, got " + f"{raw.get('scheduler')!r}") + trajectory = str(raw.get("trajectory", "ode") or "ode").strip().lower() + if trajectory not in {"ode", "sde_reflow"}: + raise ValueError("method.sampling.trajectory must be one of " + "{ode, sde_reflow}, got " + f"{raw.get('trajectory')!r}") + timesteps = raw.get("timesteps") + sigmas = raw.get("sigmas") + if timesteps is not None: + if not isinstance(timesteps, list) or not timesteps: + raise ValueError("method.sampling.timesteps must be a non-empty list when set") + timesteps = [float(t) for t in timesteps] + if sigmas is not None: + if not isinstance(sigmas, list) or not sigmas: + raise ValueError("method.sampling.sigmas must be a non-empty list when set") + sigmas = [float(s) for s in sigmas] + if timesteps is not None and sigmas is not None and len(timesteps) != len(sigmas): + raise ValueError("method.sampling.timesteps and method.sampling.sigmas must have the same length") + num_steps = int(raw.get("num_steps", 25) or 25) + if num_steps <= 0: + raise ValueError("method.sampling.num_steps must be positive") + return cls( + num_steps=num_steps, + scheduler=scheduler, # type: ignore[arg-type] + trajectory=trajectory, # type: ignore[arg-type] + flow_shift=(None if raw.get("flow_shift", None) in (None, "inherit") else float(raw["flow_shift"])), + timesteps=timesteps, + sigmas=sigmas, + ) + + +@dataclass(slots=True) +class SamplingResult: + latents: torch.Tensor + timesteps: torch.Tensor + sigmas: torch.Tensor + + +class DiffusionSampler: + """Thin model/scheduler sampler used by RL methods. + + This intentionally does not call FastVideo's full inference pipelines. + RL training needs a reusable sampling primitive that works with + ``ModelBase`` wrappers and scheduler math without binding a method to + model-family pipeline classes such as ``WanDMDPipeline``. + """ + + def __init__(self, config: SamplingConfig) -> None: + self.config = config + + @torch.no_grad() + def sample( + self, + model: ModelBase, + batch: TrainingBatch, + *, + generator: torch.Generator | None, + ) -> SamplingResult: + latents = batch.latents + if latents is None: + raise RuntimeError("TrainingBatch.latents is required for RL sampling") + current = torch.randn( + latents.shape, + device=latents.device, + dtype=latents.dtype, + generator=generator, + ) + + scheduler = self._prepare_scheduler(model, current.device) + timesteps = scheduler.timesteps.to(device=current.device) + sigmas = scheduler.sigmas.to(device=current.device) + + original_timesteps = batch.timesteps + try: + if self.config.trajectory == "ode": + pred_clean = current + for timestep in timesteps: + model_timestep = self._model_timestep(timestep, current) + batch.timesteps = model_timestep + pred_noise = model.predict_noise( + current, + model_timestep, + batch, + conditional=True, + attn_kind="dense", + ) + current = scheduler.step( + pred_noise.flatten(0, 1), + timestep, + current.flatten(0, 1), + return_dict=False, + )[0].unflatten(0, pred_noise.shape[:2]) + pred_clean = current + return SamplingResult(latents=pred_clean, timesteps=timesteps, sigmas=sigmas) + + return SamplingResult( + latents=self._sample_sde_reflow( + model, + batch, + current, + timesteps, + generator=generator, + ), + timesteps=timesteps, + sigmas=sigmas, + ) + finally: + batch.timesteps = original_timesteps + + def _prepare_scheduler( + self, + model: ModelBase, + device: torch.device, + ) -> Any: + if self.config.scheduler == "flow_match_euler": + shift = self.config.flow_shift + if shift is None: + shift = float(getattr(model.noise_scheduler, "shift", 1.0)) + scheduler = FlowMatchEulerDiscreteScheduler(shift=float(shift)) + else: + scheduler = copy.deepcopy(model.noise_scheduler) + kwargs: dict[str, Any] = {"device": device} + if self.config.timesteps is not None: + kwargs["timesteps"] = self.config.timesteps + kwargs["num_inference_steps"] = len(self.config.timesteps) + if self.config.sigmas is not None: + kwargs["sigmas"] = self.config.sigmas + kwargs["num_inference_steps"] = len(self.config.sigmas) + if "num_inference_steps" not in kwargs: + kwargs["num_inference_steps"] = self.config.num_steps + scheduler.set_timesteps(**kwargs) + return scheduler + + def _sample_sde_reflow( + self, + model: ModelBase, + batch: TrainingBatch, + current: torch.Tensor, + timesteps: torch.Tensor, + *, + generator: torch.Generator | None, + ) -> torch.Tensor: + pred_clean = current + for step_idx, timestep in enumerate(timesteps): + timestep_tensor = self._model_timestep(timestep, current) + batch.timesteps = timestep_tensor + pred_clean = model.predict_x0( + current, + timestep_tensor, + batch, + conditional=True, + attn_kind="dense", + ) + if step_idx < len(timesteps) - 1: + next_timestep = timesteps[step_idx + 1].reshape(1).to(device=current.device) + noise = torch.randn( + pred_clean.shape, + device=pred_clean.device, + dtype=pred_clean.dtype, + generator=generator, + ) + current = model.add_noise(pred_clean, noise, next_timestep) + return pred_clean + + @staticmethod + def _model_timestep( + timestep: torch.Tensor, + current: torch.Tensor, + ) -> torch.Tensor: + return timestep.reshape(1).to(device=current.device).expand(current.shape[0]).contiguous() diff --git a/fastvideo/train/methods/rl/common/validation.py b/fastvideo/train/methods/rl/common/validation.py new file mode 100644 index 000000000..1cf42a682 --- /dev/null +++ b/fastvideo/train/methods/rl/common/validation.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared validation helpers for RL training methods.""" + +from __future__ import annotations + +from dataclasses import dataclass +import math +from typing import Any + +import torch + + +@dataclass(slots=True) +class RLValidationConfig: + every_steps: int = 0 + num_steps: int = 40 # Reference DiffusionNFT sampling num steps for best visual quality + num_prompts: int = 16 + batch_size: int = 16 + log_samples: bool = True + seed: int = 42 + data_path: str | None = None + sampling: dict[str, Any] | None = None + + @classmethod + def from_mapping(cls, raw: dict[str, Any] | None) -> RLValidationConfig: + if raw is None: + return cls() + if not isinstance(raw, dict): + raise ValueError(f"method.validation must be a mapping, got {type(raw).__name__}") + data_path = raw.get("data_path", None) + sampling = raw.get("sampling", None) + if sampling is not None and not isinstance(sampling, dict): + raise ValueError(f"method.validation.sampling must be a mapping, got {type(sampling).__name__}") + 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)), + num_prompts=max(1, int(raw.get("num_prompts", 16) or 16)), + batch_size=max(1, int(raw.get("batch_size", 16) or 16)), + log_samples=bool(raw.get("log_samples", True)), + seed=int(raw.get("seed", 42) or 42), + data_path=(None if data_path in (None, "") else str(data_path)), + sampling=(dict(sampling) if sampling is not None else None), + ) + + +def validation_shard_indices( + num_prompts: int, + *, + rank: int, + world_size: int, +) -> list[tuple[int, bool]]: + """Return fixed validation prompt indices for one distributed rank.""" + num_prompts = max(1, int(num_prompts)) + world_size = max(1, int(world_size)) + per_rank = int(math.ceil(num_prompts / world_size)) + padded_total = per_rank * world_size + return [((idx % num_prompts), idx < num_prompts) for idx in range(rank, padded_total, world_size)] + + +def validation_caption( + prompt: str, + rewards: dict[str, float], +) -> str: + reward_parts = [f"{key}: {float(rewards[key]):.4f}" for key in sorted(rewards)] + return f"{' | '.join(reward_parts)} | {prompt[:1000]}" + + +def media_to_video_array(media: torch.Tensor) -> Any: + """Convert decoded media to a tracker video array. + + Accepts ``[C, T, H, W]`` tensors. ``[C, H, W]`` tensors are treated as + ``T=1`` media. Output follows the existing tracker convention used + elsewhere in FastVideo: ``[T, C, H, W]`` uint8. + """ + if media.ndim == 3: + media = media.unsqueeze(1) + if media.ndim != 4: + raise ValueError("media must have shape [C, T, H, W] or [C, H, W], " + f"got {tuple(media.shape)}") + video = (media.detach().float().clamp(0, 1) * 255).round().to(torch.uint8) + return video.permute(1, 0, 2, 3).contiguous().cpu().numpy() diff --git a/fastvideo/train/methods/rl/diffusion_nft.py b/fastvideo/train/methods/rl/diffusion_nft.py new file mode 100644 index 000000000..fb9488c0b --- /dev/null +++ b/fastvideo/train/methods/rl/diffusion_nft.py @@ -0,0 +1,1075 @@ +# SPDX-License-Identifier: Apache-2.0 +"""DiffusionNFT multi-reward policy optimization method.""" + +from __future__ import annotations + +import contextlib +from collections import defaultdict +from collections.abc import Iterator +from typing import Any + +import torch +import torch.distributed as dist +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + set_model_state_dict, +) +from torch.distributed.checkpoint.stateful import Stateful +from tqdm.auto import tqdm + +from fastvideo.dataset.parquet_dataset_map_style import ( + get_parquet_files_and_length, + read_row_from_parquet_file, +) +from fastvideo.dataset.utils import collate_rows_from_parquet_schema +from fastvideo.logger import init_logger +from fastvideo.pipelines import TrainingBatch +from fastvideo.train.methods.base import LogScalar, TrainingMethod +from fastvideo.train.models.base import ModelBase +from fastvideo.train.methods.rl.rewards import build_multi_reward_scorer +from fastvideo.train.methods.rl.common import ( + DiffusionSampler, + RLValidationConfig, + SamplingConfig, + distributed_k_repeat_indices, + media_to_video_array, + validation_caption, + validation_shard_indices, +) +from fastvideo.train.utils.config import ( + get_optional_float, + get_optional_int, + parse_betas, +) +from fastvideo.train.utils.optimizer import build_optimizer_and_scheduler +from fastvideo.training.training_utils import ( + EMA_FSDP, + clip_grad_norm_while_handling_failing_dtensor_cases, +) + +logger = init_logger(__name__) + + +class _DiffusionNFTEMAState: + """DCP state wrapper for method-owned DiffusionNFT EMA.""" + + def __init__(self, method: DiffusionNFTMethod) -> None: + self._method = method + + def state_dict(self) -> dict[str, Any]: + ema = self._method._student_ema + return { + "shadow": ema.state_dict() if ema is not None else {}, + "update_count": torch.tensor(self._method._ema_update_count, dtype=torch.long), + } + + def load_state_dict( + self, + state_dict: dict[str, Any], + ) -> None: + ema = self._method._student_ema + shadow = state_dict.get("shadow", {}) + if ema is not None and isinstance(shadow, dict): + ema.load_state_dict(shadow) + update_count = state_dict.get("update_count", 0) + if torch.is_tensor(update_count): + update_count = int(update_count.item()) + self._method._ema_update_count = int(update_count) + + +class _FullModelState(Stateful): + """DCP wrapper that saves frozen model parameters too. + + The shared ``ModelWrapper`` intentionally filters to ``requires_grad`` + parameters. DiffusionNFT's old policy is frozen but must be restored on + resume, so it needs full model state. + """ + + def __init__(self, model: torch.nn.Module) -> None: + self.model = model + + def state_dict(self) -> dict[str, Any]: + return get_model_state_dict(self.model) # type: ignore[no-any-return] + + def load_state_dict( + self, + state_dict: dict[str, Any], + ) -> None: + set_model_state_dict( + self.model, + model_state_dict=state_dict, + options=StateDictOptions(strict=False), + ) + + +class DiffusionNFTMethod(TrainingMethod): + """DiffusionNFT-style RL for diffusion models. + + This method owns the algorithm's sample-then-inner-train loop. One + ``Trainer`` step corresponds to one DiffusionNFT outer epoch. + """ + + def __init__( + self, + *, + cfg: Any, + role_models: dict[str, ModelBase], + ) -> None: + super().__init__(cfg=cfg, role_models=role_models) + if "old" not in role_models: + raise ValueError("DiffusionNFTMethod requires role 'old'") + if "reference" not in role_models: + raise ValueError("DiffusionNFTMethod requires role 'reference'") + if not self.student._trainable: + raise ValueError("DiffusionNFTMethod requires a trainable student") + + self.old = role_models["old"] + self.reference = role_models["reference"] + self.student.init_preprocessors(self.training_config) + + self._sampling_config = self._parse_sampling_config() + self._sampler = DiffusionSampler(self._sampling_config) + self._validation_config = RLValidationConfig.from_mapping(self.method_config.get("validation")) + self._validation_sampling_config = self._parse_validation_sampling_config() + self._validation_sampler = DiffusionSampler(self._validation_sampling_config) + self._validation_items: list[tuple[int, bool, dict[str, Any]]] | None = None + self._sample_steps = int(self._sampling_config.num_steps) + self._sample_train_batch_size = self._read_int( + "sample_train_batch_size", + max(1, int(self.training_config.data.train_batch_size or 1)), + ) + self._train_batch_size = self._read_int("train_batch_size", self._sample_train_batch_size) + self._num_batches_per_epoch = self._read_int("num_batches_per_epoch", 48) + self._num_inner_epochs = self._read_int("num_inner_epochs", 1) + self._num_video_per_prompt = self._read_int("num_video_per_prompt", 24) + self._adv_clip_max = self._read_float("adv_clip_max", 5.0) + self._timestep_fraction = self._read_float("timestep_fraction", 0.99) + self._kl_beta = self._read_float("kl_beta", 0.0001) + self._nft_beta = self._read_float("beta", 0.1) + self._max_grad_norm = self._read_float("max_grad_norm", 1.0) + self._decay_type = self._read_int("decay_type", 1) + self._adv_mode = str(self.method_config.get("adv_mode", "all") or "all").strip().lower() + self._terminal_progress = bool(self.method_config.get("terminal_progress", True)) + ema_config = self._parse_ema_config() + self._ema_enabled = bool(ema_config["enabled"]) + self._ema_decay = float(ema_config["decay"]) + self._ema_update_after_step = int(ema_config["update_after_step"]) + self._validation_use_ema = bool(ema_config["validation"]) + self._student_ema: EMA_FSDP | None = None + self._ema_update_count = 0 + self._trained_prompt_hashes: set[int] = set() + if self._adv_mode not in {"all", "positive_only", "negative_only", "one_only", "binary"}: + raise ValueError("method.adv_mode must be one of " + "{all, positive_only, negative_only, one_only, binary}") + + reward_fn = self.method_config.get("reward_fn", None) + if not isinstance(reward_fn, dict) or not reward_fn: + raise ValueError("method.reward_fn must be a non-empty mapping, " + "for example {pickscore: 1.0, clipscore: 1.0}") + self._reward_fn_config = {str(k): float(v) for k, v in reward_fn.items()} + unsupported = sorted(set(self._reward_fn_config) - {"pickscore", "clipscore"}) + if unsupported: + raise ValueError(f"Unsupported DiffusionNFT reward(s): {unsupported}. " + "Only pickscore and clipscore are currently ported.") + + self._reward_scorer: Any | None = None + self._init_optimizer_and_scheduler() + + @property + def _optimizer_dict(self, ) -> dict[str, torch.optim.Optimizer]: + return {"student": self._student_optimizer} + + @property + def _lr_scheduler_dict(self) -> dict[str, Any]: + return {"student": self._student_lr_scheduler} + + def manages_optimization(self) -> bool: + return True + + def single_train_step( + self, + batch: dict[str, Any], + iteration: int, + ) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, LogScalar]]: + del batch + raise RuntimeError("DiffusionNFTMethod uses managed_train_step()") + + def managed_train_step( + self, + data_stream: Iterator[dict[str, Any]], + iteration: int, + ) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, LogScalar]]: + self._log_progress(f"[DiffusionNFT] outer step {iteration}: start sampling " + f"{self._num_batches_per_epoch} batches") + sample_items = self._sample_epoch(data_stream, iteration) + self._log_progress(f"[DiffusionNFT] outer step {iteration}: scoring rewards") + rewards = self._score_samples(sample_items) + self._log_progress(f"[DiffusionNFT] outer step {iteration}: computing advantages") + advantages = self._compute_advantages(sample_items, rewards) + self._log_progress(f"[DiffusionNFT] outer step {iteration}: start inner training") + loss_map, metrics = self._inner_train(sample_items, advantages, iteration) + self._update_old_model(iteration) + metrics.update(self._reward_metrics(rewards)) + metrics.update(self._reward_diagnostic_metrics(sample_items, rewards)) + metrics["nft/num_sampled"] = float(sum(item["latents_clean"].shape[0] for item in sample_items)) + return loss_map, {}, metrics + + def on_validation_begin(self, iteration: int = 0) -> dict[str, LogScalar]: + return self._maybe_run_validation(iteration) + + def get_optimizers( + self, + iteration: int, + ) -> list[torch.optim.Optimizer]: + del iteration + return [self._student_optimizer] + + def get_lr_schedulers( + self, + iteration: int, + ) -> list[Any]: + del iteration + return [self._student_lr_scheduler] + + def checkpoint_state(self) -> dict[str, Any]: + states = super().checkpoint_state() + states["roles.old.transformer"] = _FullModelState(self.old.transformer) + if self._ema_enabled: + states["diffusion_nft.ema"] = _DiffusionNFTEMAState(self) + return states + + def seed_optimizer_state_for_resume(self) -> None: + super().seed_optimizer_state_for_resume() + + def on_train_start(self) -> None: + super().on_train_start() + self._sync_old_from_student() + if self._ema_enabled: + self._student_ema = EMA_FSDP( + self.student.transformer, + decay=self._ema_decay, + mode="local_shard", + ) + self._log_progress(f"[DiffusionNFT] EMA enabled (decay={self._ema_decay})") + self._reward_scorer = build_multi_reward_scorer( + self._reward_fn_config, + device=self.student.device, + ) + + def _init_optimizer_and_scheduler(self) -> None: + lr = float(self.training_config.optimizer.learning_rate) + betas = self.training_config.optimizer.betas + betas_raw = self.method_config.get("betas", None) + if betas_raw is not None: + betas = parse_betas(betas_raw, where="method.betas") + params = [p for p in self.student.transformer.parameters() if p.requires_grad] + self._student_optimizer, self._student_lr_scheduler = build_optimizer_and_scheduler( + params=params, + optimizer_config=self.training_config.optimizer, + loop_config=self.training_config.loop, + learning_rate=lr, + betas=betas, + scheduler_name=str(self.training_config.optimizer.lr_scheduler), + ) + + def _sample_epoch( + self, + data_stream: Iterator[dict[str, Any]], + iteration: int, + ) -> list[dict[str, Any]]: + self.student.transformer.eval() + self.old.transformer.eval() + sample_items: list[dict[str, Any]] = [] + with torch.no_grad(): + for batch_idx in tqdm( + range(self._num_batches_per_epoch), + desc=f"DiffusionNFT step {iteration}: sampling", + position=1, + leave=False, + disable=not self._show_terminal_progress(), + ): + raw_batch = self._sample_prompt_batch(data_stream, iteration, batch_idx) + prompts = self._extract_prompts(raw_batch) + batch = self.student.prepare_batch( + raw_batch, + generator=self.cuda_generator, + latents_source="zeros", + ) + sampling_result = self._sampler.sample( + self.old, + batch, + generator=self.cuda_generator, + ) + latents_clean = sampling_result.latents + train_timesteps = sampling_result.timesteps[:self._num_train_timesteps()] + media = self.student.decode_latents(latents_clean) + sample_items.append({ + "encoder_hidden_states": batch.encoder_hidden_states.detach(), + "encoder_attention_mask": batch.encoder_attention_mask.detach(), + "latents_clean": latents_clean.detach(), + "timesteps": train_timesteps.detach().unsqueeze(0).repeat(latents_clean.shape[0], 1), + "media": media.detach().cpu(), + "prompts": prompts, + }) + return sample_items + + def _sample_prompt_batch( + self, + data_stream: Iterator[dict[str, Any]], + iteration: int, + batch_idx: int, + ) -> dict[str, Any]: + dataset = getattr(getattr(self.student, "dataloader", None), "dataset", None) + if dataset is None or not all( + hasattr(dataset, attr) for attr in ("parquet_files", "lengths", "parquet_schema")): + return self._repeat_first_prompt(next(data_stream), self._sample_train_batch_size) + + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 + total_rows = int(sum(dataset.lengths)) + sample = distributed_k_repeat_indices( + dataset_length=total_rows, + batch_size=self._sample_train_batch_size, + repeats_per_prompt=self._num_video_per_prompt, + world_size=world_size, + rank=rank, + seed=int(self.training_config.data.seed) + (int(iteration) - 1) * self._num_batches_per_epoch + batch_idx, + ) + rows = [] + for prompt_idx in sample.local_indices: + row = read_row_from_parquet_file(dataset.parquet_files, prompt_idx, dataset.lengths) + row["_sample_index"] = prompt_idx + rows.append(row) + return self._collate_rows(rows) + + def _score_samples(self, sample_items: list[dict[str, Any]]) -> dict[str, torch.Tensor]: + if self._reward_scorer is None: + raise RuntimeError("Reward scorer has not been initialized") + media = torch.cat([item["media"] for item in sample_items], dim=0) + prompts = [prompt for item in sample_items for prompt in item["prompts"]] + reward_dict = self._reward_scorer(media, prompts) + device = self.student.device + return {k: v.to(device=device, dtype=torch.float32) for k, v in reward_dict.items()} + + @torch.no_grad() + def _maybe_run_validation(self, iteration: int) -> dict[str, LogScalar]: + config = self._validation_config + if config.every_steps <= 0 or iteration % config.every_steps != 0: + return {} + if self._reward_scorer is None: + raise RuntimeError("Reward scorer has not been initialized") + + use_ema = self._validation_use_ema and self._student_ema is not None + suffix = " with EMA" if use_ema else "" + self._log_progress(f"[DiffusionNFT] validation step {iteration}: start " + f"{config.num_prompts} prompts, {config.num_steps} sampling steps{suffix}") + with self._ema_context(): + return self._run_validation(iteration) + + @torch.no_grad() + def _run_validation(self, iteration: int) -> dict[str, LogScalar]: + config = self._validation_config + if self._reward_scorer is None: + raise RuntimeError("Reward scorer has not been initialized") + self.student.transformer.eval() + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + prepare_generator = torch.Generator(device=self.student.device).manual_seed(config.seed + 100_000 + rank) + validation_generator = torch.Generator(device=self.student.device).manual_seed(config.seed + rank) + + local_rewards: dict[str, list[torch.Tensor]] = defaultdict(list) + local_masks: list[torch.Tensor] = [] + local_logs: list[dict[str, Any]] = [] + items = self._get_validation_items() + + total_batches = max(1, (len(items) + config.batch_size - 1) // config.batch_size) + for start in tqdm( + range(0, len(items), config.batch_size), + total=total_batches, + desc=f"DiffusionNFT validation {iteration}", + position=1, + leave=False, + disable=not self._show_terminal_progress(), + ): + 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", + ) + sampling_result = self._validation_sampler.sample( + self.student, + batch, + generator=validation_generator, + ) + media = self.student.decode_latents(sampling_result.latents).detach().cpu() + + # DiffusionNFT's original ``train_nft_sd3.py::eval_fn`` scores + # validation samples and puts reward values in sample captions. + rewards = self._reward_scorer(media, prompts) + valid_mask = torch.tensor([item[1] for item in batch_items], + device=self.student.device, + dtype=torch.float32) + local_masks.append(valid_mask) + for key, value in rewards.items(): + local_rewards[key].append(value.to(device=self.student.device, dtype=torch.float32)) + + if config.log_samples: + for sample_idx, (global_idx, valid, _) in enumerate(batch_items): + if not valid: + continue + local_logs.append({ + "index": int(global_idx), + "prompt": prompts[sample_idx], + "media": media[sample_idx], + "rewards": { + key: float(value[sample_idx].detach().float().cpu()) + for key, value in rewards.items() + }, + }) + + if not local_rewards or not local_masks: + return {} + + local_mask = torch.cat(local_masks, dim=0) + gathered_mask = self._gather_tensor(local_mask).bool() + metrics: dict[str, LogScalar] = {} + for key, chunks in local_rewards.items(): + values = torch.cat(chunks, dim=0) + gathered_values = self._gather_tensor(values.detach().float()) + valid_values = gathered_values[gathered_mask] + if valid_values.numel() > 0: + metrics[f"validation/reward/{key}"] = valid_values.mean() + metrics["validation/num_prompts"] = gathered_mask.float().sum() + + if config.log_samples: + self._log_progress(f"[DiffusionNFT] validation step {iteration}: logging samples") + self._log_validation_samples(local_logs, iteration) + self._log_progress(f"[DiffusionNFT] validation step {iteration}: finished") + return metrics + + def _get_validation_items(self) -> list[tuple[int, bool, dict[str, Any]]]: + if self._validation_items is not None: + return self._validation_items + + dataset = getattr(getattr(self.student, "dataloader", None), "dataset", None) + if dataset is None: + raise RuntimeError("DiffusionNFT validation requires the student dataloader dataset") + + data_path = self._validation_config.data_path or getattr(dataset, "path", None) + if not data_path: + data_path = self.training_config.data.data_path + if self._validation_config.data_path is None or data_path == getattr(dataset, "path", None): + parquet_files = list(dataset.parquet_files) + lengths = list(dataset.lengths) + else: + parquet_files, lengths = get_parquet_files_and_length(data_path) + parquet_files = list(parquet_files) + lengths = list(lengths) + + total_rows = int(sum(lengths)) + if total_rows <= 0: + raise RuntimeError(f"Validation data_path {data_path!r} has no rows") + num_prompts = min(self._validation_config.num_prompts, total_rows) + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 + items: list[tuple[int, bool, dict[str, Any]]] = [] + for prompt_idx, valid in validation_shard_indices(num_prompts, rank=rank, world_size=world_size): + row = read_row_from_parquet_file(parquet_files, prompt_idx, lengths) + row["_sample_index"] = prompt_idx + items.append((prompt_idx, valid, row)) + self._validation_items = items + return items + + def _collate_validation_rows(self, rows: list[dict[str, Any]]) -> dict[str, Any]: + return self._collate_rows(rows) + + def _collate_rows(self, rows: list[dict[str, Any]]) -> dict[str, Any]: + dataset = getattr(getattr(self.student, "dataloader", None), "dataset", None) + if dataset is None or not hasattr(dataset, "parquet_schema"): + raise RuntimeError("DiffusionNFT requires a parquet-backed student dataset") + return collate_rows_from_parquet_schema( + rows, + dataset.parquet_schema, + int(getattr(dataset, "text_padding_length", 512)), + cfg_rate=0.0, + seed=self._validation_config.seed, + ) + + def _log_validation_samples( + self, + local_logs: list[dict[str, Any]], + iteration: int, + ) -> None: + logs = self._gather_objects(local_logs) + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + if rank != 0 or not logs: + return + tracker = getattr(self, "tracker", None) + if tracker is None: + return + + artifacts = [] + for item in sorted(logs, key=lambda x: int(x["index"])): + artifact = tracker.video( + media_to_video_array(item["media"]), + caption=validation_caption(str(item["prompt"]), item["rewards"]), + fps=1, + ) + if artifact is not None: + artifacts.append(artifact) + if artifacts: + tracker.log_artifacts({"validation/videos": artifacts}, step=iteration) + + def _compute_advantages( + self, + sample_items: list[dict[str, Any]], + rewards: dict[str, torch.Tensor], + ) -> torch.Tensor: + avg = rewards["avg"].detach().float() + local_prompts = [prompt for item in sample_items for prompt in item["prompts"]] + gathered_rewards = self._gather_tensor(avg) + gathered_prompts = self._gather_prompts(local_prompts) + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + local_count = int(avg.shape[0]) + if len(gathered_prompts) != int(gathered_rewards.shape[0]): + raise RuntimeError("Gathered prompt count does not match gathered rewards") + + global_advantages = torch.empty_like(gathered_rewards) + prompt_to_indices: dict[str, list[int]] = defaultdict(list) + for idx, prompt in enumerate(gathered_prompts): + prompt_to_indices[prompt].append(idx) + for indices in prompt_to_indices.values(): + index_tensor = torch.tensor(indices, device=gathered_rewards.device, dtype=torch.long) + group_rewards = gathered_rewards[index_tensor] + group_std = group_rewards.std(unbiased=False) + global_advantages[index_tensor] = (group_rewards - group_rewards.mean()) / (group_std + 1e-4) + + start = rank * local_count + end = start + local_count + advantages = global_advantages[start:end].to(avg.device) + num_train_timesteps = self._num_train_timesteps() + return advantages.unsqueeze(1).repeat(1, num_train_timesteps) + + def _inner_train( + self, + sample_items: list[dict[str, Any]], + advantages: torch.Tensor, + iteration: int, + ) -> tuple[dict[str, torch.Tensor], dict[str, LogScalar]]: + self.student.transformer.train() + self.old.transformer.eval() + self.reference.transformer.eval() + self._student_optimizer.zero_grad(set_to_none=True) + + samples = self._collate_samples(sample_items) + total_samples = int(samples["latents_clean"].shape[0]) + num_train_timesteps = self._num_train_timesteps() + if total_samples != int(advantages.shape[0]): + raise RuntimeError("advantages and samples have mismatched batch sizes") + + effective_grad_accum = max(1, int(self.training_config.loop.gradient_accumulation_steps or 1)) + effective_grad_accum *= max(1, num_train_timesteps) + current_accum = 0 + optimizer_steps = 0 + loss_terms: dict[str, list[torch.Tensor]] = defaultdict(list) + num_batches = max(1, total_samples // max(1, self._train_batch_size)) + training_batch_size = max(1, total_samples // num_batches) + total_train_batches = self._num_inner_epochs * num_batches + + with tqdm( + total=total_train_batches, + desc=f"DiffusionNFT step {iteration}: training", + position=1, + leave=False, + disable=not self._show_terminal_progress(), + ) as progress: + for _ in range(self._num_inner_epochs): + perm = torch.randperm( + total_samples, + device=self.student.device, + generator=self.cuda_generator, + ) + shuffled = {k: v[perm] for k, v in samples.items()} + shuffled_adv = advantages[perm] + perms_time = torch.stack([ + torch.randperm(num_train_timesteps, device=self.student.device, generator=self.cuda_generator) + for _ in range(total_samples) + ]) + shuffled["timesteps"] = shuffled["timesteps"][ + torch.arange(total_samples, device=self.student.device)[:, None], + perms_time, + ] + shuffled_adv = shuffled_adv[ + torch.arange(total_samples, device=self.student.device)[:, None], + perms_time, + ] + + for batch_idx in range(num_batches): + start = batch_idx * training_batch_size + end = total_samples if batch_idx == num_batches - 1 else (batch_idx + 1) * training_batch_size + train_sample = {k: v[start:end] for k, v in shuffled.items()} + train_adv = shuffled_adv[start:end] + for timestep_idx in range(num_train_timesteps): + losses, backward_ctx = self._training_timestep_loss( + train_sample, + train_adv[:, timestep_idx], + timestep_idx, + ) + self.student.backward( + losses["total_loss"], + backward_ctx, + grad_accum_rounds=effective_grad_accum, + ) + current_accum += 1 + for key, value in losses.items(): + loss_terms[key].append(value.detach()) + + if current_accum % effective_grad_accum == 0: + self._clip_student_grads() + self._student_optimizer.step() + self._student_lr_scheduler.step() + self._update_ema() + self._student_optimizer.zero_grad(set_to_none=True) + optimizer_steps += 1 + progress.update(1) + + if current_accum % effective_grad_accum != 0: + self._clip_student_grads() + self._student_optimizer.step() + self._student_lr_scheduler.step() + self._update_ema() + self._student_optimizer.zero_grad(set_to_none=True) + optimizer_steps += 1 + + self._log_progress(f"[DiffusionNFT] outer step {iteration}: finished inner training " + f"({current_accum} micro-steps, {optimizer_steps} optimizer steps)") + + reduced_local = { + key: torch.stack(values).mean() if values else torch.zeros((), device=self.student.device) + for key, values in loss_terms.items() + } + reduced = {key: self._mean_scalar_across_ranks(value) for key, value in reduced_local.items()} + reduced.setdefault("total_loss", torch.zeros((), device=self.student.device)) + metrics: dict[str, LogScalar] = { + "nft/iteration": float(iteration), + "nft/num_inner_epochs": float(self._num_inner_epochs), + "nft/inner_micro_steps": float(current_accum), + "nft/optimizer_steps": float(optimizer_steps), + "ema/update_count": float(self._ema_update_count), + } + return reduced, metrics + + @contextlib.contextmanager + def _ema_context(self) -> Iterator[None]: + if self._validation_use_ema and self._student_ema is not None: + with self._student_ema.apply_to_model(self.student.transformer): + yield + return + yield + + def _update_ema(self) -> None: + if not self._ema_enabled or self._student_ema is None: + return + if self._ema_update_count >= self._ema_update_after_step: + self._student_ema.update(self.student.transformer) + self._ema_update_count += 1 + + def _log_progress(self, message: str) -> None: + if not self._terminal_progress: + return + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + if rank == 0: + logger.info(message) + + def _show_terminal_progress(self) -> bool: + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + return bool(self._terminal_progress and rank == 0) + + def _training_timestep_loss( + self, + sample: dict[str, torch.Tensor], + advantages: torch.Tensor, + timestep_idx: int, + ) -> tuple[dict[str, torch.Tensor], tuple[torch.Tensor, Any]]: + # Ported from DiffusionNFT's ``scripts/train_nft_sd3.py`` inner + # timestep loss, adapted to WanModel's predict_noise interface. + x0 = sample["latents_clean"] + timestep = sample["timesteps"][:, timestep_idx].to(device=x0.device) + t = timestep.float() / float(self.student.num_train_timesteps) + t_expanded = t.view(-1, *([1] * (x0.ndim - 1))) + noise = torch.randn( + x0.shape, + device=x0.device, + dtype=x0.dtype, + generator=self.cuda_generator, + ) + xt = ((1 - t_expanded) * x0 + t_expanded * noise).to(dtype=x0.dtype) + batch = self._make_training_batch(sample, timestep) + + with torch.no_grad(): + old_prediction = self.old.predict_noise( + xt, + timestep, + batch, + conditional=True, + attn_kind="dense", + ).detach() + ref_forward_prediction = self.reference.predict_noise( + xt, + timestep, + batch, + conditional=True, + attn_kind="dense", + ).detach() + + forward_prediction = self.student.predict_noise( + xt, + timestep, + batch, + conditional=True, + attn_kind="dense", + ) + + advantages_clip = torch.clamp(advantages, -self._adv_clip_max, self._adv_clip_max) + if self._adv_mode == "positive_only": + advantages_clip = torch.clamp(advantages_clip, 0, self._adv_clip_max) + elif self._adv_mode == "negative_only": + advantages_clip = torch.clamp(advantages_clip, -self._adv_clip_max, 0) + elif self._adv_mode == "one_only": + advantages_clip = torch.where(advantages_clip > 0, torch.ones_like(advantages_clip), + torch.zeros_like(advantages_clip)) + elif self._adv_mode == "binary": + advantages_clip = torch.sign(advantages_clip) + + normalized_advantages_clip = (advantages_clip / self._adv_clip_max) / 2.0 + 0.5 + r = torch.clamp(normalized_advantages_clip, 0, 1) + + positive_prediction = self._nft_beta * forward_prediction + (1 - self._nft_beta) * old_prediction.detach() + implicit_negative_prediction = ((1.0 + self._nft_beta) * old_prediction.detach() - + self._nft_beta * forward_prediction) + + x0_prediction = xt - t_expanded * positive_prediction + with torch.no_grad(): + weight_factor = torch.abs(x0_prediction.double() - x0.double()) + weight_factor = weight_factor.mean(dim=tuple(range(1, x0.ndim)), keepdim=True).clip(min=0.00001) + positive_loss = ((x0_prediction - x0)**2 / weight_factor).mean(dim=tuple(range(1, x0.ndim))) + + negative_x0_prediction = xt - t_expanded * implicit_negative_prediction + with torch.no_grad(): + negative_weight_factor = torch.abs(negative_x0_prediction.double() - x0.double()) + negative_weight_factor = negative_weight_factor.mean(dim=tuple(range(1, x0.ndim)), + keepdim=True).clip(min=0.00001) + negative_loss = ((negative_x0_prediction - x0)**2 / negative_weight_factor).mean(dim=tuple(range(1, x0.ndim))) + + ori_policy_loss = r * positive_loss / self._nft_beta + (1.0 - r) * negative_loss / self._nft_beta + policy_loss = (ori_policy_loss * self._adv_clip_max).mean() + + kl_div_loss = ((forward_prediction - ref_forward_prediction)**2).mean(dim=tuple(range(1, x0.ndim))).mean() + total_loss = policy_loss + self._kl_beta * kl_div_loss + losses = { + "total_loss": total_loss, + "policy_loss": policy_loss, + "unweighted_policy_loss": ori_policy_loss.mean(), + "kl_div_loss": kl_div_loss, + "old_deviate": ((forward_prediction - old_prediction)**2).mean(), + "old_kl_div": ((old_prediction - ref_forward_prediction)**2).mean(), + "x0_norm": torch.mean(x0**2), + } + return losses, (batch.timesteps, batch.attn_metadata) + + def _repeat_first_prompt( + self, + raw_batch: dict[str, Any], + repeat: int, + ) -> dict[str, Any]: + out: dict[str, Any] = {} + for key, value in raw_batch.items(): + if torch.is_tensor(value) and value.shape[0] > 0: + out[key] = value[:1].repeat((repeat, ) + (1, ) * (value.ndim - 1)) + elif key == "info_list" and isinstance(value, list) and value: + out[key] = [dict(value[0]) for _ in range(repeat)] + elif isinstance(value, list) and value: + out[key] = [value[0] for _ in range(repeat)] + else: + out[key] = value + return out + + @staticmethod + def _extract_prompts(raw_batch: dict[str, Any]) -> list[str]: + infos = raw_batch.get("info_list") + if isinstance(infos, list) and infos: + prompts: list[str] = [] + for info in infos: + if isinstance(info, dict): + prompts.append(str(info.get("prompt") or info.get("caption") or "")) + else: + prompts.append("") + return prompts + captions = raw_batch.get("caption_text") + if isinstance(captions, list): + return [str(c) for c in captions] + raise ValueError("Could not find prompts in batch info_list or caption_text") + + def _collate_samples(self, sample_items: list[dict[str, Any]]) -> dict[str, torch.Tensor]: + keys = ["encoder_hidden_states", "encoder_attention_mask", "latents_clean", "timesteps"] + return {key: torch.cat([item[key] for item in sample_items], dim=0).to(self.student.device) for key in keys} + + @staticmethod + def _make_training_batch( + sample: dict[str, torch.Tensor], + timestep: torch.Tensor, + ) -> TrainingBatch: + batch = TrainingBatch() + batch.encoder_hidden_states = sample["encoder_hidden_states"] + batch.encoder_attention_mask = sample["encoder_attention_mask"] + batch.conditional_dict = { + "encoder_hidden_states": batch.encoder_hidden_states, + "encoder_attention_mask": batch.encoder_attention_mask, + } + batch.timesteps = timestep + batch.raw_latent_shape = tuple(sample["latents_clean"].shape) + return batch + + def _num_train_timesteps(self) -> int: + schedule_len = self._sample_steps + if self._sampling_config.timesteps is not None: + schedule_len = len(self._sampling_config.timesteps) + elif self._sampling_config.sigmas is not None: + schedule_len = len(self._sampling_config.sigmas) + return max(1, min(schedule_len, int(schedule_len * self._timestep_fraction))) + + def _clip_student_grads(self) -> None: + if self._max_grad_norm <= 0.0: + return + clip_grad_norm_while_handling_failing_dtensor_cases( + [p for p in self.student.transformer.parameters()], + self._max_grad_norm, + foreach=None, + ) + + def _sync_old_from_student(self) -> None: + with torch.no_grad(): + for src, tgt in zip(self.student.transformer.parameters(), self.old.transformer.parameters(), strict=True): + tgt.data.copy_(src.detach().data) + + def _update_old_model(self, iteration: int) -> None: + decay = self._return_decay(iteration, self._decay_type) + with torch.no_grad(): + for src, tgt in zip(self.student.transformer.parameters(), self.old.transformer.parameters(), strict=True): + tgt.data.copy_(tgt.detach().data * decay + src.detach().data * (1.0 - decay)) + + @staticmethod + def _return_decay(step: int, decay_type: int) -> float: + # Ported from DiffusionNFT's ``scripts/train_nft_sd3.py::return_decay``. + if decay_type == 0: + flat, uprate, uphold = 0, 0.0, 0.0 + elif decay_type == 1: + flat, uprate, uphold = 0, 0.001, 0.5 + elif decay_type == 2: + flat, uprate, uphold = 75, 0.0075, 0.999 + else: + raise ValueError(f"Unsupported decay_type: {decay_type}") + if step < flat: + return 0.0 + return min((step - flat) * uprate, uphold) + + def _reward_metrics(self, rewards: dict[str, torch.Tensor]) -> dict[str, LogScalar]: + metrics: dict[str, LogScalar] = {} + for key, value in rewards.items(): + gathered = self._gather_tensor(value.detach().float()) + metrics[f"reward/{key}"] = gathered.mean() + return metrics + + def _reward_diagnostic_metrics( + self, + sample_items: list[dict[str, Any]], + rewards: dict[str, torch.Tensor], + ) -> dict[str, LogScalar]: + avg = rewards["avg"].detach().float() + local_prompts = [prompt for item in sample_items for prompt in item["prompts"]] + gathered_rewards = self._gather_tensor(avg) + gathered_prompts = self._gather_prompts(local_prompts) + if len(gathered_prompts) != int(gathered_rewards.shape[0]): + raise RuntimeError("Gathered prompt count does not match gathered rewards") + + prompt_to_indices: dict[str, list[int]] = defaultdict(list) + for idx, prompt in enumerate(gathered_prompts): + prompt_to_indices[prompt].append(idx) + self._trained_prompt_hashes.add(hash(prompt)) + if not prompt_to_indices: + return {} + + group_sizes = [] + group_stds = [] + grouped_rewards = [] + for indices in prompt_to_indices.values(): + index_tensor = torch.tensor(indices, device=gathered_rewards.device, dtype=torch.long) + group_rewards = gathered_rewards[index_tensor].float() + grouped_rewards.append(group_rewards) + group_sizes.append(float(group_rewards.numel())) + group_stds.append(group_rewards.std(unbiased=False)) + + std_tensor = torch.stack(group_stds) + metrics: dict[str, LogScalar] = { + "group_size": sum(group_sizes) / max(1, len(group_sizes)), + "trained_prompt_num": float(len(self._trained_prompt_hashes)), + "zero_std_ratio": (std_tensor == 0).float().mean(), + "reward_std_mean": std_tensor.mean(), + } + for top_percentage in (100, 75, 50, 25, 10): + metrics[f"mean_reward_{top_percentage}"] = self._mean_of_top_group_rewards( + grouped_rewards, + top_percentage, + ) + return metrics + + @staticmethod + def _mean_of_top_group_rewards( + grouped_rewards: list[torch.Tensor], + top_percentage: int, + ) -> torch.Tensor: + per_prompt_means = [] + for rewards in grouped_rewards: + if rewards.numel() == 0: + continue + if top_percentage == 100: + per_prompt_means.append(rewards.mean()) + continue + threshold = torch.quantile( + rewards.float(), + (100 - float(top_percentage)) / 100.0, + ) + top_rewards = rewards[rewards >= threshold] + if top_rewards.numel() > 0: + per_prompt_means.append(top_rewards.mean()) + if not per_prompt_means: + device = grouped_rewards[0].device if grouped_rewards else torch.device("cpu") + return torch.zeros((), device=device) + return torch.stack(per_prompt_means).mean() + + @staticmethod + def _gather_tensor(tensor: torch.Tensor) -> torch.Tensor: + if not dist.is_available() or not dist.is_initialized(): + return tensor.detach() + gathered = [torch.zeros_like(tensor) for _ in range(dist.get_world_size())] + dist.all_gather(gathered, tensor.contiguous()) + return torch.cat(gathered, dim=0).detach() + + @staticmethod + def _mean_scalar_across_ranks(value: torch.Tensor) -> torch.Tensor: + if not dist.is_available() or not dist.is_initialized(): + return value.detach() + reduced = value.detach().clone() + dist.all_reduce(reduced, op=dist.ReduceOp.AVG) + return reduced + + @staticmethod + def _gather_prompts(prompts: list[str]) -> list[str]: + if not dist.is_available() or not dist.is_initialized(): + return list(prompts) + gathered: list[list[str] | None] = [None for _ in range(dist.get_world_size())] + dist.all_gather_object(gathered, list(prompts)) + flat: list[str] = [] + for rank_prompts in gathered: + if rank_prompts is None: + continue + flat.extend(rank_prompts) + return flat + + @staticmethod + def _gather_objects(items: list[Any]) -> list[Any]: + if not dist.is_available() or not dist.is_initialized(): + return list(items) + gathered: list[list[Any] | None] = [None for _ in range(dist.get_world_size())] + dist.all_gather_object(gathered, list(items)) + flat: list[Any] = [] + for rank_items in gathered: + if rank_items is None: + continue + flat.extend(rank_items) + return flat + + def _read_int(self, key: str, default: int) -> int: + value = get_optional_int(self.method_config, key, where=f"method.{key}") + return int(default if value is None else value) + + def _read_float(self, key: str, default: float) -> float: + value = get_optional_float(self.method_config, key, where=f"method.{key}") + return float(default if value is None else value) + + def _parse_sampling_config(self) -> SamplingConfig: + raw = self.method_config.get("sampling", None) + if isinstance(raw, dict): + raw = dict(raw) + if raw.get("flow_shift", None) in (None, "inherit"): + pipeline_config = getattr(self.training_config, "pipeline_config", None) + flow_shift = getattr(pipeline_config, "flow_shift", None) + if flow_shift is not None: + raw["flow_shift"] = float(flow_shift) + return SamplingConfig.from_mapping(raw) + + def _parse_validation_sampling_config(self) -> SamplingConfig: + raw = dict(self.method_config.get("sampling", {}) or {}) + validation_sampling = self._validation_config.sampling + if validation_sampling is not None: + raw.update(validation_sampling) + if validation_sampling is None or "num_steps" not in validation_sampling: + raw["num_steps"] = self._validation_config.num_steps + if raw.get("flow_shift") in (None, "inherit"): + pipeline_config = getattr(self.training_config, "pipeline_config", None) + flow_shift = getattr(pipeline_config, "flow_shift", None) + if flow_shift is not None: + raw["flow_shift"] = float(flow_shift) + return SamplingConfig.from_mapping(raw) + + def _parse_ema_config(self) -> dict[str, bool | float | int]: + raw = self.method_config.get("ema", None) + config: dict[str, bool | float | int] = { + "enabled": True, + "decay": 0.9, + "update_after_step": 0, + "validation": True, + } + if raw is None: + return config + if isinstance(raw, bool | str | int): + config["enabled"] = self._coerce_bool(raw, where="method.ema") + return config + if not isinstance(raw, dict): + raise ValueError("method.ema must be a bool or mapping") + allowed = set(config) + unknown = sorted(set(raw) - allowed) + if unknown: + raise ValueError(f"Unsupported method.ema key(s): {unknown}") + if "enabled" in raw: + config["enabled"] = self._coerce_bool(raw["enabled"], where="method.ema.enabled") + if "decay" in raw: + decay = float(raw["decay"]) + if not 0.0 <= decay <= 1.0: + raise ValueError("method.ema.decay must be in [0, 1]") + config["decay"] = decay + if "update_after_step" in raw: + update_after_step = int(raw["update_after_step"]) + if update_after_step < 0: + raise ValueError("method.ema.update_after_step must be >= 0") + config["update_after_step"] = update_after_step + if "validation" in raw: + config["validation"] = self._coerce_bool(raw["validation"], where="method.ema.validation") + return config + + @staticmethod + def _coerce_bool(value: Any, *, where: str) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, int): + return bool(value) + if isinstance(value, str): + lowered = value.strip().lower() + if lowered in {"1", "true", "yes", "on"}: + return True + if lowered in {"0", "false", "no", "off"}: + return False + raise ValueError(f"{where} must be a boolean") diff --git a/fastvideo/train/methods/rl/rewards/__init__.py b/fastvideo/train/methods/rl/rewards/__init__.py new file mode 100644 index 000000000..f986a8978 --- /dev/null +++ b/fastvideo/train/methods/rl/rewards/__init__.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reusable reward models for training methods.""" + +from fastvideo.train.methods.rl.rewards.frame_rewards import ( + ClipScoreScorer, + PickScoreScorer, +) +from fastvideo.train.methods.rl.rewards.media import ( + MultiRewardScorer, + RewardScorer, + select_first_frame, +) + + +def build_multi_reward_scorer( + reward_weights, + *, + device="cuda", + scorers: dict[str, RewardScorer] | None = None, +) -> MultiRewardScorer: + available: dict[str, RewardScorer] = dict(scorers or {}) + if not available: + available = { + "pickscore": PickScoreScorer(device=device), + "clipscore": ClipScoreScorer(device=device), + } + return MultiRewardScorer(reward_weights, scorers=available) + + +__all__ = [ + "ClipScoreScorer", + "MultiRewardScorer", + "PickScoreScorer", + "RewardScorer", + "build_multi_reward_scorer", + "select_first_frame", +] diff --git a/fastvideo/train/methods/rl/rewards/frame_rewards.py b/fastvideo/train/methods/rl/rewards/frame_rewards.py new file mode 100644 index 000000000..6c041e85e --- /dev/null +++ b/fastvideo/train/methods/rl/rewards/frame_rewards.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Frame-based reward scorers used by RL training methods.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +from PIL import Image +import torch + +from fastvideo.train.methods.rl.rewards.media import select_first_frame + + +class PickScoreScorer(torch.nn.Module): + """PickScore reward, matching DiffusionNFT normalization. + + Ported from DiffusionNFT's ``flow_grpo/pickscore_scorer.py``. + """ + + def __init__( + self, + *, + device: torch.device | str = "cuda", + dtype: torch.dtype = torch.float32, + ) -> None: + super().__init__() + from transformers import AutoModel, AutoProcessor + + processor_path = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K" + model_path = "yuvalkirstain/PickScore_v1" + self.device = torch.device(device) + self.dtype = dtype + self.processor = AutoProcessor.from_pretrained(processor_path) + self.model = AutoModel.from_pretrained(model_path).eval().to(self.device) + self.model = self.model.to(dtype=dtype) + + @torch.no_grad() + def forward( + self, + media: torch.Tensor, + prompts: Sequence[str], + ) -> torch.Tensor: + frame_tensor = select_first_frame(media) + frame_np = (frame_tensor.detach().float().clamp(0, 1) * 255).round() + frame_np = frame_np.to(torch.uint8).cpu().numpy().transpose(0, 2, 3, 1) + pil_frames = [Image.fromarray(frame) for frame in frame_np] + + frame_inputs = self.processor( + images=pil_frames, + padding=True, + truncation=True, + max_length=77, + return_tensors="pt", + ) + frame_inputs = {k: v.to(device=self.device) for k, v in frame_inputs.items()} + + text_inputs = self.processor( + text=list(prompts), + padding=True, + truncation=True, + max_length=77, + return_tensors="pt", + ) + text_inputs = {k: v.to(device=self.device) for k, v in text_inputs.items()} + + text_embs = self.model.get_text_features(**text_inputs) + text_embs = text_embs / text_embs.norm(p=2, dim=-1, keepdim=True) + + frame_embs = self.model.get_image_features(**frame_inputs) + frame_embs = frame_embs / frame_embs.norm(p=2, dim=-1, keepdim=True) + + scores = self.model.logit_scale.exp() * (text_embs @ frame_embs.T) + return scores.diag().float() / 26.0 + + +class ClipScoreScorer(torch.nn.Module): + """CLIPScore reward, matching DiffusionNFT normalization. + + Ported from DiffusionNFT's ``flow_grpo/clip_scorer.py``. + """ + + def __init__( + self, + *, + device: torch.device | str = "cuda", + ) -> None: + super().__init__() + import torch.nn as nn + import torchvision.transforms as T + from transformers import CLIPModel, CLIPProcessor + + def get_size(size: Any) -> Any: + if isinstance(size, int): + return (size, size) + if isinstance(size, Mapping) and "height" in size and "width" in size: + return (size["height"], size["width"]) + if isinstance(size, Mapping) and "shortest_edge" in size: + return size["shortest_edge"] + raise ValueError(f"Invalid processor size: {size!r}") + + def get_frame_transform(processor: Any) -> torch.nn.Module: + config = processor.to_dict() + resize = T.Resize(get_size(config.get("size"))) if config.get("do_resize") else nn.Identity() + crop = T.CenterCrop(get_size(config.get("crop_size"))) if config.get("do_center_crop") else nn.Identity() + normalize = (T.Normalize(mean=processor.image_mean, std=processor.image_std) + if config.get("do_normalize") else nn.Identity()) + return T.Compose([resize, crop, normalize]) + + self.device = torch.device(device) + self.model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(self.device).eval() + self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14") + self.transform = get_frame_transform(self.processor.image_processor) + + @torch.no_grad() + def forward( + self, + media: torch.Tensor, + prompts: Sequence[str], + ) -> torch.Tensor: + frame_tensor = select_first_frame(media).detach().float().clamp(0, 1) + texts = self.processor( + text=list(prompts), + padding="max_length", + truncation=True, + return_tensors="pt", + ).to(self.device) + pixels = self.transform(frame_tensor).to(device=self.device, dtype=frame_tensor.dtype) + outputs = self.model(pixel_values=pixels, **texts) + return outputs.logits_per_image.diagonal().float() / 100.0 diff --git a/fastvideo/train/methods/rl/rewards/media.py b/fastvideo/train/methods/rl/rewards/media.py new file mode 100644 index 000000000..1f9287144 --- /dev/null +++ b/fastvideo/train/methods/rl/rewards/media.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Generic media reward composition utilities.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence + +import torch + +RewardScorer = Callable[[torch.Tensor, Sequence[str]], torch.Tensor] + + +def select_first_frame(media: torch.Tensor) -> torch.Tensor: + """Return first-frame media as ``[B, C, H, W]``. + + This is a helper for reward models that are intrinsically frame-based + (for example PickScore and CLIPScore). Video-aware rewards should inspect + the full ``[B, C, T, H, W]`` tensor themselves. + """ + if not torch.is_tensor(media): + raise TypeError(f"media must be a torch.Tensor, got {type(media).__name__}") + if media.ndim == 5: + return media[:, :, 0] + if media.ndim == 4: + return media + raise ValueError("media must have shape [B, C, H, W] or [B, C, T, H, W], " + f"got {tuple(media.shape)}") + + +class MultiRewardScorer: + """Weighted sum of reusable media reward scorers. + + Mirrors DiffusionNFT's ``flow_grpo/rewards.py::multi_score`` behavior, + while leaving frame selection to each concrete reward. + """ + + def __init__( + self, + reward_weights: Mapping[str, float], + *, + scorers: Mapping[str, RewardScorer], + ) -> None: + self.reward_weights = {str(k): float(v) for k, v in reward_weights.items()} + if not self.reward_weights: + raise ValueError("reward_weights must contain at least one reward") + + self.scorers = dict(scorers) + unsupported = sorted(set(self.reward_weights) - set(self.scorers)) + if unsupported: + raise ValueError(f"Unsupported reward(s): {unsupported}. " + f"Available rewards: {sorted(self.scorers)}") + + @torch.no_grad() + def __call__( + self, + media: torch.Tensor, + prompts: Sequence[str], + ) -> dict[str, torch.Tensor]: + prompt_count = len(prompts) + if media.shape[0] != prompt_count: + raise ValueError(f"media batch size ({media.shape[0]}) must match prompt count ({prompt_count})") + total: torch.Tensor | None = None + details: dict[str, torch.Tensor] = {} + for name, weight in self.reward_weights.items(): + scores = self.scorers[name](media, prompts).detach().float() + if scores.ndim != 1 or int(scores.shape[0]) != prompt_count: + raise ValueError(f"Reward {name!r} must return shape [{prompt_count}], got {tuple(scores.shape)}") + details[name] = scores + weighted = scores * float(weight) + total = weighted if total is None else total.to(weighted.device) + weighted + assert total is not None + details["avg"] = total + return details diff --git a/fastvideo/train/models/base.py b/fastvideo/train/models/base.py index b8e7dc74b..ffffb31ba 100644 --- a/fastvideo/train/models/base.py +++ b/fastvideo/train/models/base.py @@ -91,6 +91,17 @@ class ModelBase(ABC): def on_train_start(self) -> None: # noqa: B027 """Called once before the training loop begins.""" + def decode_latents( + self, + latents_b_t_c_h_w: torch.Tensor, + ) -> torch.Tensor: + """Decode ``[B, T, C, H, W]`` latents to ``[B, C, T, H, W]`` media. + + RL reward methods call this hook instead of reaching into + model-specific VAE normalization details. + """ + raise NotImplementedError(f"{type(self).__name__} does not implement decode_latents()") + # ------------------------------------------------------------------ # Timestep helpers # ------------------------------------------------------------------ diff --git a/fastvideo/train/models/wan/wan.py b/fastvideo/train/models/wan/wan.py index 6e38dd81f..bc745f5c1 100644 --- a/fastvideo/train/models/wan/wan.py +++ b/fastvideo/train/models/wan/wan.py @@ -101,6 +101,7 @@ class WanModel(ModelBase): self.negative_prompt_embeds: (torch.Tensor | None) = None self.negative_prompt_attention_mask: (torch.Tensor | None) = None + self._requires_negative_conditioning = True # Timestep mechanics. self.timestep_shift: float = float(flow_shift) @@ -160,17 +161,31 @@ class WanModel(ModelBase): self._init_timestep_mechanics() from fastvideo.dataset.dataloader.schema import ( - pyarrow_schema_t2v, ) + pyarrow_schema_t2v, + pyarrow_schema_text_only, + ) from fastvideo.train.utils.dataloader import ( build_parquet_t2v_train_dataloader, ) + preprocessed_data_type = str(getattr( + training_config.data, + "preprocessed_data_type", + "t2v", + )).strip().lower() + parquet_schema = pyarrow_schema_t2v + if preprocessed_data_type == "text_only": + parquet_schema = pyarrow_schema_text_only + elif preprocessed_data_type != "t2v": + raise ValueError("Unsupported Wan preprocessed_data_type: " + f"{preprocessed_data_type!r}") + text_len = ( training_config.pipeline_config.text_encoder_configs[ # type: ignore[union-attr] 0].arch_config.text_len) self.dataloader = build_parquet_t2v_train_dataloader( training_config.data, text_len=int(text_len), - parquet_schema=pyarrow_schema_t2v, + parquet_schema=parquet_schema, ) self.start_step = 0 @@ -178,6 +193,9 @@ class WanModel(ModelBase): def num_train_timesteps(self) -> int: return int(self.num_train_timestep) + def set_requires_negative_conditioning(self, requires: bool) -> None: + self._requires_negative_conditioning = bool(requires) + def shift_and_clamp_timestep(self, timestep: torch.Tensor) -> torch.Tensor: timestep = shift_timestep( timestep, @@ -187,7 +205,25 @@ class WanModel(ModelBase): return timestep.clamp(self.min_timestep, self.max_timestep) def on_train_start(self) -> None: - self.ensure_negative_conditioning() + if self._requires_negative_conditioning: + self.ensure_negative_conditioning() + + @torch.no_grad() + def decode_latents( + self, + latents_b_t_c_h_w: torch.Tensor, + ) -> torch.Tensor: + if self.vae is None: + raise RuntimeError("Wan VAE is not initialized") + latents = latents_b_t_c_h_w.permute(0, 2, 1, 3, 4).float() + if bool(getattr(self.vae, "handles_latent_denorm", False)): + denorm = latents + else: + mean = torch.tensor(self.vae.latents_mean, device=latents.device, dtype=latents.dtype).view(1, -1, 1, 1, 1) + std = torch.tensor(self.vae.latents_std, device=latents.device, dtype=latents.dtype).view(1, -1, 1, 1, 1) + denorm = latents * std + mean + media = self.vae.to(latents.device).decode(denorm) + return (media / 2 + 0.5).clamp(0, 1) # ------------------------------------------------------------------ # Runtime primitives @@ -200,7 +236,8 @@ class WanModel(ModelBase): generator: torch.Generator, latents_source: Literal["data", "zeros"] = "data", ) -> TrainingBatch: - self.ensure_negative_conditioning() + if self._requires_negative_conditioning: + self.ensure_negative_conditioning() assert self.training_config is not None tc = self.training_config @@ -285,7 +322,7 @@ class WanModel(ModelBase): attn_kind: Literal["dense", "vsa"] = "dense", ) -> torch.Tensor: device_type = self.device.type - dtype = noisy_latents.dtype + dtype = self._get_training_dtype() if conditional: text_dict = batch.conditional_dict if text_dict is None: @@ -301,6 +338,11 @@ class WanModel(ModelBase): else: raise ValueError(f"Unknown attn_kind: {attn_kind!r}") + if noisy_latents.is_floating_point(): + noisy_latents = noisy_latents.to(dtype=dtype) + + # Keep Wan training autocast tied to the model's training dtype, not + # to caller-created intermediates that may accidentally be fp32. with torch.autocast(device_type, dtype=dtype), set_forward_context( current_timestep=batch.timesteps, attn_metadata=attn_metadata, diff --git a/fastvideo/train/models/wan/wan_causal.py b/fastvideo/train/models/wan/wan_causal.py index 639642d85..25c161de9 100644 --- a/fastvideo/train/models/wan/wan_causal.py +++ b/fastvideo/train/models/wan/wan_causal.py @@ -128,7 +128,9 @@ class WanCausalModel(WanModel, CausalModelBase): } device_type = self.device.type - dtype = noisy_latents.dtype + dtype = self._get_training_dtype() + if noisy_latents.is_floating_point(): + noisy_latents = noisy_latents.to(dtype=dtype) if conditional: text_dict = batch.conditional_dict diff --git a/fastvideo/train/trainer.py b/fastvideo/train/trainer.py index b77ad5341..be6054f77 100644 --- a/fastvideo/train/trainer.py +++ b/fastvideo/train/trainer.py @@ -12,7 +12,7 @@ from tqdm.auto import tqdm from fastvideo.distributed import get_sp_group, get_world_group from fastvideo.train.callbacks.callback import CallbackDict -from fastvideo.train.methods.base import TrainingMethod +from fastvideo.train.methods.base import LogScalar, TrainingMethod from fastvideo.train.utils.tracking import build_tracker if TYPE_CHECKING: @@ -82,6 +82,22 @@ class Trainer: batch = next(data_iter) yield batch + def _run_method_validation( + self, + method: TrainingMethod, + iteration: int, + ) -> None: + hook = getattr(method, "on_validation_begin", None) + if hook is None: + return + validation_metrics: dict[str, LogScalar] = hook(iteration) + validation_metrics = { + k: float(_coerce_log_scalar(v, where=(f"method.on_validation_begin().metrics[{k!r}]"))) + for k, v in validation_metrics.items() + } + if self.global_rank == 0 and validation_metrics: + self.tracker.log(validation_metrics, iteration) + def run( self, method: TrainingMethod, @@ -115,6 +131,7 @@ class Trainer: method, iteration=start_step, ) + self._run_method_validation(method, start_step) method.optimizers_zero_grad(start_step) data_stream = self._iter_dataloader(dataloader) @@ -130,6 +147,8 @@ class Trainer: desc="Steps", disable=self.local_rank > 0, ) + # Allow method-specific optimization flow (e.g. DiffusionNFT). + method_manages_optimization = bool(method.manages_optimization()) for step in progress: t0 = time.perf_counter() @@ -137,47 +156,69 @@ class Trainer: # to CPU once per step right before logging. loss_sums: dict[str, float | torch.Tensor] = {} metric_sums: dict[str, float | torch.Tensor] = {} - for accum_iter in range(grad_accum): - batch = next(data_stream) - loss_map, outputs, step_metrics = (method.single_train_step( - batch, + if method_manages_optimization: + loss_map, outputs, step_metrics = method.managed_train_step( + data_stream, step, - )) - - method.backward( - loss_map, - outputs, - grad_accum_rounds=grad_accum, ) - for k, v in loss_map.items(): if isinstance(v, torch.Tensor): - prev = loss_sums.get(k, 0.0) - loss_sums[k] = prev + v.detach() + loss_sums[k] = v.detach() for k, v in step_metrics.items(): if k in loss_sums: raise ValueError(f"Metric key {k!r} collides " "with loss key. Use a " "different name (e.g. prefix " "with 'train/').") - prev = metric_sums.get(k, 0.0) - metric_sums[k] = (prev + _coerce_log_scalar( + metric_sums[k] = _coerce_log_scalar( v, - where=("method.single_train_step()" + where=("method.managed_train_step()" f".metrics[{k!r}]"), + ) + else: + for accum_iter in range(grad_accum): + batch = next(data_stream) + loss_map, outputs, step_metrics = (method.single_train_step( + batch, + step, )) - self.callbacks.on_before_optimizer_step( - method, - iteration=step, - ) - method.optimizers_schedulers_step(step) - method.optimizers_zero_grad(step) + method.backward( + loss_map, + outputs, + grad_accum_rounds=grad_accum, + ) + + for k, v in loss_map.items(): + if isinstance(v, torch.Tensor): + prev = loss_sums.get(k, 0.0) + loss_sums[k] = prev + v.detach() + for k, v in step_metrics.items(): + if k in loss_sums: + raise ValueError(f"Metric key {k!r} collides " + "with loss key. Use a " + "different name (e.g. prefix " + "with 'train/').") + prev = metric_sums.get(k, 0.0) + metric_sums[k] = (prev + _coerce_log_scalar( + v, + where=("method.single_train_step()" + f".metrics[{k!r}]"), + )) + + if not method_manages_optimization: + self.callbacks.on_before_optimizer_step( + method, + iteration=step, + ) + method.optimizers_schedulers_step(step) + method.optimizers_zero_grad(step) # Single CPU sync point: materialise GPU tensors # to float right before logging. - metrics = {k: float(v) / grad_accum for k, v in loss_sums.items()} - metrics.update({k: float(v) / grad_accum for k, v in metric_sums.items()}) + divisor = 1 if method_manages_optimization else grad_accum + metrics = {k: float(v) / divisor for k, v in loss_sums.items()} + metrics.update({k: float(v) / divisor for k, v in metric_sums.items()}) metrics["step_time_sec"] = (time.perf_counter() - t0) metrics["vsa_sparsity"] = float(tc.vsa_sparsity) if self.global_rank == 0 and metrics: @@ -196,6 +237,7 @@ class Trainer: method, iteration=step, ) + self._run_method_validation(method, step) self.callbacks.on_validation_end( method, iteration=step, diff --git a/fastvideo/train/utils/config.py b/fastvideo/train/utils/config.py index 41b250009..932a2f5f1 100644 --- a/fastvideo/train/utils/config.py +++ b/fastvideo/train/utils/config.py @@ -343,6 +343,12 @@ def _build_training_config( if init_from is not None: model_path = str(init_from) + preprocessed_data_type = str(da.get("preprocessed_data_type", "t2v") or "t2v").strip().lower() + if preprocessed_data_type not in {"t2v", "text_only"}: + raise ValueError("training.data.preprocessed_data_type must be one of " + "{'t2v', 'text_only'}, got " + f"{preprocessed_data_type!r}") + return TrainingConfig( distributed=DistributedConfig( num_gpus=num_gpus, @@ -354,6 +360,7 @@ def _build_training_config( ), data=DataConfig( data_path=str(da.get("data_path", "") or ""), + preprocessed_data_type=preprocessed_data_type, train_batch_size=int(da.get("train_batch_size", 1) or 1), dataloader_num_workers=int(da.get("dataloader_num_workers", 0) or 0), training_cfg_rate=float(da.get("training_cfg_rate", 0.0) or 0.0), diff --git a/fastvideo/train/utils/training_config.py b/fastvideo/train/utils/training_config.py index 8dd0f440b..c0a98ffa2 100644 --- a/fastvideo/train/utils/training_config.py +++ b/fastvideo/train/utils/training_config.py @@ -23,6 +23,7 @@ class DistributedConfig: @dataclass(slots=True) class DataConfig: data_path: str = "" + preprocessed_data_type: str = "t2v" train_batch_size: int = 1 dataloader_num_workers: int = 0 training_cfg_rate: float = 0.0 diff --git a/fastvideo/training/trackers.py b/fastvideo/training/trackers.py index f1ccfaa42..62ec9d5ed 100644 --- a/fastvideo/training/trackers.py +++ b/fastvideo/training/trackers.py @@ -112,7 +112,7 @@ class BaseTracker: self._timed_metrics = {} def log_artifacts(self, artifacts: dict[str, Any], step: int) -> None: - """Log artifacts such as videos or images. + """Log tracker artifacts such as sampled media. By default this is treated the same as :meth:`log`. """ diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index 2cc1a0a87..7eb3847fa 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -1694,7 +1694,7 @@ class EMA_FSDP: if p_local.numel() == 0: # Nothing to swap on this rank for this param continue - self.saved[name] = p_local.clone().to(device=p_local.device, dtype=p_local.dtype) + self.saved[name] = p_local.clone().to("cpu") if name in self.ema.shadow: ema_cpu = self.ema.shadow[name] if ema_cpu.numel() != p_local.numel(): @@ -1714,7 +1714,7 @@ class EMA_FSDP: saved_local = self.saved[name] if saved_local.numel() != p_local.numel(): continue - p_local.copy_(saved_local) + p_local.copy_(saved_local.to(dtype=p_local.dtype, device=p_local.device)) self.saved.clear() return False diff --git a/scripts/preprocess/preprocess_train_text_only_dmd.sh b/scripts/preprocess/preprocess_train_text_only_dmd.sh new file mode 100644 index 000000000..f8212ad42 --- /dev/null +++ b/scripts/preprocess/preprocess_train_text_only_dmd.sh @@ -0,0 +1,180 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Build a text-only Parquet dataset from a one-prompt-per-line file. The script +# shards prompts across GPU_NUM single-GPU torchrun workers, runs +# v1_preprocess.py with --preprocess_task text_only, and writes prompt +# embeddings/captions under OUTPUT_DIR for DMD2/DiffusionNFT text-only runs. + +INPUT_FILE="${1:-train.txt}" +OUTPUT_DIR="${OUTPUT_DIR:-data/train_text_only_dmd_preprocessed}" +MODEL_PATH="${MODEL_PATH:-Wan-AI/Wan2.1-T2V-1.3B-Diffusers}" +GPU_NUM="${GPU_NUM:-2}" +BATCH_SIZE="${BATCH_SIZE:-1}" +SAMPLES_PER_FILE="${SAMPLES_PER_FILE:-8}" +FLUSH_FREQUENCY="${FLUSH_FREQUENCY:-8}" +TEXT_MAX_LENGTH="${TEXT_MAX_LENGTH:-512}" +CONDA_ROOT="${CONDA_ROOT:-/root/miniconda3}" +CONDA_ENV="${CONDA_ENV:-fastvideo}" +MIN_FREE_GPU_MB="${MIN_FREE_GPU_MB:-22000}" + +if [[ ! -f "$INPUT_FILE" ]]; then + echo "Input text file not found: $INPUT_FILE" >&2 + exit 1 +fi + +if [[ ! -f "$CONDA_ROOT/etc/profile.d/conda.sh" ]]; then + echo "Conda activation script not found under $CONDA_ROOT" >&2 + exit 1 +fi + +# shellcheck source=/dev/null +source "$CONDA_ROOT/etc/profile.d/conda.sh" +conda activate "$CONDA_ENV" + +if [[ "${HF_HUB_ENABLE_HF_TRANSFER:-0}" == "1" ]]; then + if ! python -c "import hf_transfer" >/dev/null 2>&1; then + echo "HF_HUB_ENABLE_HF_TRANSFER=1 but hf_transfer is not installed; disabling fast transfer." + export HF_HUB_ENABLE_HF_TRANSFER=0 + fi +else + export HF_HUB_ENABLE_HF_TRANSFER=0 +fi + +visible_gpus=$(nvidia-smi --query-gpu=index --format=csv,noheader 2>/dev/null | wc -l | tr -d ' ') +if [[ "$visible_gpus" -lt "$GPU_NUM" ]]; then + echo "Expected at least $GPU_NUM GPUs, found $visible_gpus" >&2 + exit 1 +fi + +for gpu_id in $(seq 0 $((GPU_NUM - 1))); do + free_mb=$(nvidia-smi --id="$gpu_id" --query-gpu=memory.free --format=csv,noheader,nounits | tr -d ' ') + if [[ "$free_mb" -lt "$MIN_FREE_GPU_MB" ]]; then + echo "GPU $gpu_id has only ${free_mb} MiB free; text-only Wan preprocessing needs" \ + "about ${MIN_FREE_GPU_MB} MiB." >&2 + echo "Free the GPU or lower MIN_FREE_GPU_MB if you know this run will fit." >&2 + exit 1 + fi +done + +MARKER="$OUTPUT_DIR/.fastvideo_text_only_dmd_output" +if [[ -d "$OUTPUT_DIR" && ! -f "$MARKER" ]]; then + echo "Refusing to overwrite existing non-script output directory: $OUTPUT_DIR" >&2 + echo "Set OUTPUT_DIR to a new path or remove the directory manually." >&2 + exit 1 +fi + +rm -rf "$OUTPUT_DIR" +mkdir -p "$OUTPUT_DIR" +touch "$MARKER" + +echo "Text-only DMD preprocessing config:" +echo " input: $INPUT_FILE" +echo " output: $OUTPUT_DIR" +echo " model: $MODEL_PATH" +echo " gpus: $GPU_NUM" +echo " batch size per GPU: $BATCH_SIZE" +echo " text max length: $TEXT_MAX_LENGTH" +echo " samples per parquet file: $SAMPLES_PER_FILE" +echo " flush frequency: $FLUSH_FREQUENCY" + +SHARD_DIR="$OUTPUT_DIR/_text_shards" +mkdir -p "$SHARD_DIR" + +python - "$INPUT_FILE" "$SHARD_DIR" "$GPU_NUM" <<'PY' +from pathlib import Path +import sys + +input_path = Path(sys.argv[1]) +shard_dir = Path(sys.argv[2]) +num_shards = int(sys.argv[3]) + +prompts = [line.rstrip("\n") for line in input_path.read_text(encoding="utf-8").splitlines() if line.strip()] +if not prompts: + raise SystemExit(f"No non-empty prompts found in {input_path}") + +for shard_idx in range(num_shards): + shard_prompts = prompts[shard_idx::num_shards] + shard_path = shard_dir / f"train_text_shard_{shard_idx}.txt" + shard_path.write_text("\n".join(shard_prompts) + "\n", encoding="utf-8") + print(f"Wrote {len(shard_prompts)} prompts to {shard_path}") +PY + +run_preprocess_worker() { + local gpu_id="$1" + local shard_file="$2" + local shard_output="$3" + local log_file="$4" + local master_port="$5" + local -a cmd=( + torchrun + --nnodes=1 + --nproc_per_node=1 + --master_port "$master_port" + fastvideo/pipelines/preprocess/v1_preprocess.py + --model_path "$MODEL_PATH" + --data_merge_path "$shard_file" + --preprocess_video_batch_size "$BATCH_SIZE" + --seed 42 + --max_height 448 + --max_width 832 + --num_frames 77 + --dataloader_num_workers 0 + --output_dir "$shard_output" + --train_fps 16 + --samples_per_file "$SAMPLES_PER_FILE" + --flush_frequency "$FLUSH_FREQUENCY" + --text_max_length "$TEXT_MAX_LENGTH" + --video_length_tolerance_range 5 + --preprocess_task text_only + ) + + { + echo "[gpu${gpu_id}] log file: $log_file" + echo "[gpu${gpu_id}] command: CUDA_VISIBLE_DEVICES=${gpu_id} ${cmd[*]}" + } | tee "$log_file" + + CUDA_VISIBLE_DEVICES="$gpu_id" "${cmd[@]}" 2>&1 \ + | sed -u "s/^/[gpu${gpu_id}] /" \ + | tee -a "$log_file" + local status=${PIPESTATUS[0]} + if [[ "$status" -ne 0 ]]; then + echo "[gpu${gpu_id}] preprocessing failed with exit code $status" | tee -a "$log_file" + fi + return "$status" +} + +pids=() +for gpu_id in $(seq 0 $((GPU_NUM - 1))); do + shard_file="$SHARD_DIR/train_text_shard_${gpu_id}.txt" + shard_output="$OUTPUT_DIR/shard_${gpu_id}" + mkdir -p "$shard_output" + log_file="$OUTPUT_DIR/preprocess_gpu_${gpu_id}.log" + + echo "Launching text-only preprocessing on GPU ${gpu_id}: ${shard_file}" + run_preprocess_worker "$gpu_id" "$shard_file" "$shard_output" "$log_file" "$((29610 + gpu_id))" & + pids+=("$!") +done + +failed=0 +for pid in "${pids[@]}"; do + if ! wait "$pid"; then + failed=1 + fi +done + +if [[ "$failed" -ne 0 ]]; then + echo "One or more preprocessing workers failed. Check logs under $OUTPUT_DIR." >&2 + exit 1 +fi + +num_parquet=$(find "$OUTPUT_DIR" -name '*.parquet' | wc -l | tr -d ' ') +if [[ "$num_parquet" -eq 0 ]]; then + echo "No parquet files were produced under $OUTPUT_DIR" >&2 + exit 1 +fi + +echo "Text-only preprocessing complete." +echo "Parquet files: $num_parquet" +echo "Use this training data_path:" +echo "$OUTPUT_DIR" diff --git a/scripts/train/train_diffusion_nft_wan_pick_clip.sh b/scripts/train/train_diffusion_nft_wan_pick_clip.sh new file mode 100755 index 000000000..40df5e66a --- /dev/null +++ b/scripts/train/train_diffusion_nft_wan_pick_clip.sh @@ -0,0 +1,118 @@ +#!/usr/bin/env bash +set -euo pipefail + +CONFIG="${CONFIG:-examples/train/configs/rl/wan/diffusion_nft_pick_clip.yaml}" +DATA_PATH="${DATA_PATH:-data/pickscore_text_only_preprocessed}" +OUTPUT_DIR="${OUTPUT_DIR:-outputs/wan2.1_diffusion_nft_pick_clip}" +NUM_GPUS="${NUM_GPUS:-4}" +NNODES="${NNODES:-1}" +NODE_RANK="${NODE_RANK:-0}" +MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}" +MASTER_PORT="${MASTER_PORT:-29531}" +SP_SIZE="${SP_SIZE:-1}" +TP_SIZE="${TP_SIZE:-1}" +HSDP_REPLICATE_DIM="${HSDP_REPLICATE_DIM:-1}" +HSDP_SHARD_DIM="${HSDP_SHARD_DIM:-$NUM_GPUS}" +DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-0}" +NUM_FRAMES="${NUM_FRAMES:-1}" +NUM_LATENT_T="${NUM_LATENT_T:-1}" +PROJECT_NAME="${PROJECT_NAME:-diffusion_nft_wan}" +RUN_NAME="${RUN_NAME:-wan2.1_diffusion_nft_pick_clip}" +CONDA_ROOT="${CONDA_ROOT:-/root/miniconda3}" +CONDA_ENV="${CONDA_ENV:-fastvideo}" +LOG_DIR="${LOG_DIR:-logs/train}" + +if [[ ! -f "$CONFIG" ]]; then + echo "Training config not found: $CONFIG" >&2 + exit 1 +fi + +if [[ ! -d "$DATA_PATH" ]]; then + echo "Preprocessed dataset directory not found: $DATA_PATH" >&2 + echo "Run preprocessing first, for example:" >&2 + echo " GPU_NUM=4 BATCH_SIZE=1 OUTPUT_DIR=$DATA_PATH \\" >&2 + echo " bash scripts/preprocess/preprocess_train_text_only_dmd.sh DiffusionNFT/dataset/pickscore/train.txt" >&2 + exit 1 +fi + +num_parquet=$(find "$DATA_PATH" -name '*.parquet' | wc -l | tr -d ' ') +if [[ "$num_parquet" -eq 0 ]]; then + echo "No parquet files found under $DATA_PATH" >&2 + echo "Run preprocessing first, for example:" >&2 + echo " GPU_NUM=4 BATCH_SIZE=1 OUTPUT_DIR=$DATA_PATH \\" >&2 + echo " bash scripts/preprocess/preprocess_train_text_only_dmd.sh DiffusionNFT/dataset/pickscore/train.txt" >&2 + exit 1 +fi + +if [[ ! -f "$CONDA_ROOT/etc/profile.d/conda.sh" ]]; then + echo "Conda activation script not found under $CONDA_ROOT" >&2 + exit 1 +fi + +# shellcheck source=/dev/null +source "$CONDA_ROOT/etc/profile.d/conda.sh" +conda activate "$CONDA_ENV" + +if [[ "${HF_HUB_ENABLE_HF_TRANSFER:-0}" == "1" ]]; then + if ! python -c "import hf_transfer" >/dev/null 2>&1; then + echo "HF_HUB_ENABLE_HF_TRANSFER=1 but hf_transfer is not installed; disabling fast transfer." + export HF_HUB_ENABLE_HF_TRANSFER=0 + fi +else + export HF_HUB_ENABLE_HF_TRANSFER=0 +fi + +export TOKENIZERS_PARALLELISM="${TOKENIZERS_PARALLELISM:-false}" +export WANDB_MODE="${WANDB_MODE:-online}" +export WANDB_API_KEY="${WANDB_API_KEY:-}" +export WANDB_BASE_URL="${WANDB_BASE_URL:-https://api.wandb.ai}" +export FASTVIDEO_ATTENTION_BACKEND="${FASTVIDEO_ATTENTION_BACKEND:-FLASH_ATTN}" +export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-/tmp/triton_cache_diffusion_nft_wan}" +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" + +mkdir -p "$LOG_DIR" "$OUTPUT_DIR" +timestamp="$(date +%Y%m%d_%H%M%S)" +log_file="$LOG_DIR/diffusion_nft_wan_pick_clip_${timestamp}.log" + +cmd=( + torchrun + --nnodes "$NNODES" + --node_rank "$NODE_RANK" + --nproc_per_node "$NUM_GPUS" + --master_addr "$MASTER_ADDR" + --master_port "$MASTER_PORT" + -m fastvideo.train.entrypoint.train + --config "$CONFIG" + --training.data.data_path "$DATA_PATH" + --training.data.preprocessed_data_type text_only + --training.data.dataloader_num_workers "$DATALOADER_NUM_WORKERS" + --training.data.num_frames "$NUM_FRAMES" + --training.data.num_latent_t "$NUM_LATENT_T" + --training.distributed.num_gpus "$NUM_GPUS" + --training.distributed.sp_size "$SP_SIZE" + --training.distributed.tp_size "$TP_SIZE" + --training.distributed.hsdp_replicate_dim "$HSDP_REPLICATE_DIM" + --training.distributed.hsdp_shard_dim "$HSDP_SHARD_DIM" + --training.checkpoint.output_dir "$OUTPUT_DIR" + --training.tracker.project_name "$PROJECT_NAME" + --training.tracker.run_name "$RUN_NAME" +) + +echo "DiffusionNFT Wan single-frame RL training config:" +echo " config: $CONFIG" +echo " data path: $DATA_PATH" +echo " parquet files: $num_parquet" +echo " output dir: $OUTPUT_DIR" +echo " frames / latent T: $NUM_FRAMES / $NUM_LATENT_T" +echo " rewards: pickscore + clipscore" +echo " learning rate: 3e-5" +echo " GPUs: $NUM_GPUS" +echo " SP/TP: $SP_SIZE/$TP_SIZE" +echo " HSDP replicate/shard: $HSDP_REPLICATE_DIM/$HSDP_SHARD_DIM" +echo " W&B mode: $WANDB_MODE" +echo " log file: $log_file" +echo "Command:" +printf ' %q' "${cmd[@]}" "$@" +echo + +"${cmd[@]}" "$@" 2>&1 | tee "$log_file" diff --git a/scripts/train/train_dmd2_t2v_text_only.sh b/scripts/train/train_dmd2_t2v_text_only.sh new file mode 100755 index 000000000..6bf1cb0be --- /dev/null +++ b/scripts/train/train_dmd2_t2v_text_only.sh @@ -0,0 +1,170 @@ +#!/usr/bin/env bash +set -euo pipefail + +CONFIG="${CONFIG:-examples/train/configs/distribution_matching/wan/dmd2_t2v.yaml}" +DATA_PATH="${DATA_PATH:-data/train_text_only_dmd_preprocessed}" +OUTPUT_DIR="${OUTPUT_DIR:-outputs/wan2.1_dmd2_text_only}" +NUM_GPUS="${NUM_GPUS:-2}" +NNODES="${NNODES:-1}" +NODE_RANK="${NODE_RANK:-0}" +MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}" +MASTER_PORT="${MASTER_PORT:-29521}" +SP_SIZE="${SP_SIZE:-1}" +TP_SIZE="${TP_SIZE:-1}" +HSDP_REPLICATE_DIM="${HSDP_REPLICATE_DIM:-1}" +HSDP_SHARD_DIM="${HSDP_SHARD_DIM:-$NUM_GPUS}" +DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-0}" +NUM_FRAMES="${NUM_FRAMES:-1}" +NUM_LATENT_T="${NUM_LATENT_T:-1}" +VALIDATION_NUM_FRAMES="${VALIDATION_NUM_FRAMES:-$NUM_FRAMES}" +VALIDATION_PROMPT_FILE="${VALIDATION_PROMPT_FILE:-}" +VALIDATION_FILE="${VALIDATION_FILE:-examples/train/configs/distribution_matching/wan/dmd2_text_only_validation.json}" +VALIDATION_OFFLOAD_TRAINING_STATE="${VALIDATION_OFFLOAD_TRAINING_STATE:-true}" +VALIDATION_UNLOAD_PIPELINE_AFTER="${VALIDATION_UNLOAD_PIPELINE_AFTER:-true}" +CFG_UNCOND_TEXT="${CFG_UNCOND_TEXT:-zero}" +CFG_UNCOND_ON_MISSING="${CFG_UNCOND_ON_MISSING:-ignore}" +PROJECT_NAME="${PROJECT_NAME:-distillation_wan_text_only}" +RUN_NAME="${RUN_NAME:-wan2.1_dmd2_text_only}" +CONDA_ROOT="${CONDA_ROOT:-/root/miniconda3}" +CONDA_ENV="${CONDA_ENV:-fastvideo}" +LOG_DIR="${LOG_DIR:-logs/train}" + +if [[ ! -f "$CONFIG" ]]; then + echo "Training config not found: $CONFIG" >&2 + exit 1 +fi + +if [[ ! -d "$DATA_PATH" ]]; then + echo "Preprocessed dataset directory not found: $DATA_PATH" >&2 + echo "Run scripts/preprocess/preprocess_train_text_only_dmd.sh first." >&2 + exit 1 +fi + +num_parquet=$(find "$DATA_PATH" -name '*.parquet' | wc -l | tr -d ' ') +if [[ "$num_parquet" -eq 0 ]]; then + echo "No parquet files found under $DATA_PATH" >&2 + echo "Run scripts/preprocess/preprocess_train_text_only_dmd.sh first." >&2 + exit 1 +fi + +if [[ ! -f "$CONDA_ROOT/etc/profile.d/conda.sh" ]]; then + echo "Conda activation script not found under $CONDA_ROOT" >&2 + exit 1 +fi + +# shellcheck source=/dev/null +source "$CONDA_ROOT/etc/profile.d/conda.sh" +conda activate "$CONDA_ENV" + +if [[ "${HF_HUB_ENABLE_HF_TRANSFER:-0}" == "1" ]]; then + if ! python -c "import hf_transfer" >/dev/null 2>&1; then + echo "HF_HUB_ENABLE_HF_TRANSFER=1 but hf_transfer is not installed; disabling fast transfer." + export HF_HUB_ENABLE_HF_TRANSFER=0 + fi +else + export HF_HUB_ENABLE_HF_TRANSFER=0 +fi + +export TOKENIZERS_PARALLELISM="${TOKENIZERS_PARALLELISM:-false}" +export WANDB_MODE="${WANDB_MODE:-offline}" +export WANDB_API_KEY="${WANDB_API_KEY:-}" +export WANDB_BASE_URL="${WANDB_BASE_URL:-https://api.wandb.ai}" +export FASTVIDEO_ATTENTION_BACKEND="${FASTVIDEO_ATTENTION_BACKEND:-FLASH_ATTN}" +export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-/tmp/triton_cache_dmd2_text_only}" +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" + +mkdir -p "$LOG_DIR" "$OUTPUT_DIR" + +if [[ -n "$VALIDATION_PROMPT_FILE" ]]; then + if [[ ! -f "$VALIDATION_PROMPT_FILE" ]]; then + echo "Validation prompt file not found: $VALIDATION_PROMPT_FILE" >&2 + echo "Set VALIDATION_PROMPT_FILE to a text file, or leave it empty and set VALIDATION_FILE=." >&2 + exit 1 + fi + python - "$VALIDATION_PROMPT_FILE" "$VALIDATION_FILE" <<'PY' +import json +import os +import sys + +prompt_file, validation_file = sys.argv[1:3] +with open(prompt_file, encoding="utf-8") as f: + prompts = [line.strip() for line in f if line.strip()] + +if not prompts: + raise SystemExit(f"No validation prompts found in {prompt_file}") + +validation_dir = os.path.dirname(os.path.abspath(validation_file)) +if validation_dir: + os.makedirs(validation_dir, exist_ok=True) + +with open(validation_file, "w", encoding="utf-8") as f: + json.dump( + {"data": [{"caption": prompt} for prompt in prompts]}, + f, + indent=2, + ensure_ascii=False, + ) + f.write("\n") + +print(f"Wrote {len(prompts)} validation prompts to {validation_file}") +PY +elif [[ ! -f "$VALIDATION_FILE" ]]; then + echo "Validation dataset file not found: $VALIDATION_FILE" >&2 + exit 1 +fi + +timestamp="$(date +%Y%m%d_%H%M%S)" +log_file="$LOG_DIR/dmd2_t2v_text_only_${timestamp}.log" + +cmd=( + torchrun + --nnodes "$NNODES" + --node_rank "$NODE_RANK" + --nproc_per_node "$NUM_GPUS" + --master_addr "$MASTER_ADDR" + --master_port "$MASTER_PORT" + -m fastvideo.train.entrypoint.train + --config "$CONFIG" + --training.data.data_path "$DATA_PATH" + --training.data.preprocessed_data_type text_only + --training.data.dataloader_num_workers "$DATALOADER_NUM_WORKERS" + --training.data.num_frames "$NUM_FRAMES" + --training.data.num_latent_t "$NUM_LATENT_T" + --callbacks.validation.dataset_file "$VALIDATION_FILE" + --callbacks.validation.num_frames "$VALIDATION_NUM_FRAMES" + --callbacks.validation.offload_training_state "$VALIDATION_OFFLOAD_TRAINING_STATE" + --callbacks.validation.unload_pipeline_after_validation "$VALIDATION_UNLOAD_PIPELINE_AFTER" + --method.cfg_uncond.text "$CFG_UNCOND_TEXT" + --method.cfg_uncond.on_missing "$CFG_UNCOND_ON_MISSING" + --training.distributed.num_gpus "$NUM_GPUS" + --training.distributed.sp_size "$SP_SIZE" + --training.distributed.tp_size "$TP_SIZE" + --training.distributed.hsdp_replicate_dim "$HSDP_REPLICATE_DIM" + --training.distributed.hsdp_shard_dim "$HSDP_SHARD_DIM" + --training.checkpoint.output_dir "$OUTPUT_DIR" + --training.tracker.project_name "$PROJECT_NAME" + --training.tracker.run_name "$RUN_NAME" +) + +echo "DMD2 T2V text-only training config:" +echo " config: $CONFIG" +echo " data path: $DATA_PATH" +echo " parquet files: $num_parquet" +echo " output dir: $OUTPUT_DIR" +echo " frames / latent T: $NUM_FRAMES / $NUM_LATENT_T" +echo " validation prompt file: ${VALIDATION_PROMPT_FILE:-}" +echo " validation dataset: $VALIDATION_FILE" +echo " validation frames: $VALIDATION_NUM_FRAMES" +echo " validation offload training state: $VALIDATION_OFFLOAD_TRAINING_STATE" +echo " validation unload pipeline after: $VALIDATION_UNLOAD_PIPELINE_AFTER" +echo " GPUs: $NUM_GPUS" +echo " SP/TP: $SP_SIZE/$TP_SIZE" +echo " HSDP replicate/shard: $HSDP_REPLICATE_DIM/$HSDP_SHARD_DIM" +echo " CFG uncond text/on_missing: $CFG_UNCOND_TEXT/$CFG_UNCOND_ON_MISSING" +echo " W&B mode: $WANDB_MODE" +echo " log file: $log_file" +echo "Command:" +printf ' %q' "${cmd[@]}" "$@" +echo + +"${cmd[@]}" "$@" 2>&1 | tee "$log_file" diff --git a/tests/local_tests/test_diffusion_nft_method_utils.py b/tests/local_tests/test_diffusion_nft_method_utils.py new file mode 100644 index 000000000..2221c99ae --- /dev/null +++ b/tests/local_tests/test_diffusion_nft_method_utils.py @@ -0,0 +1,99 @@ +from types import SimpleNamespace + +import torch + +from fastvideo.train.methods.rl.diffusion_nft import DiffusionNFTMethod + + +class _FakeEMA: + + def __init__(self): + self.updates = 0 + + def update(self, module): + del module + self.updates += 1 + + +def test_reward_diagnostic_metrics_match_per_prompt_groups(): + method = object.__new__(DiffusionNFTMethod) + method._trained_prompt_hashes = set() + + sample_items = [{ + "prompts": ["a", "a"], + }, { + "prompts": ["b", "b"], + }] + rewards = {"avg": torch.tensor([1.0, 3.0, 2.0, 6.0])} + + metrics = method._reward_diagnostic_metrics(sample_items, rewards) + + assert metrics["group_size"] == 2.0 + assert metrics["trained_prompt_num"] == 2.0 + assert torch.isclose(metrics["zero_std_ratio"], torch.tensor(0.0)) + assert torch.isclose(metrics["reward_std_mean"], torch.tensor(1.5)) + assert torch.isclose(metrics["mean_reward_100"], torch.tensor(3.0)) + assert torch.isclose(metrics["mean_reward_50"], torch.tensor(4.5)) + + method._reward_diagnostic_metrics(sample_items, rewards) + assert len(method._trained_prompt_hashes) == 2 + + +def test_update_ema_honors_update_after_step(): + method = object.__new__(DiffusionNFTMethod) + method._ema_enabled = True + method._student_ema = _FakeEMA() + method._ema_update_count = 0 + method._ema_update_after_step = 1 + method.student = SimpleNamespace(transformer=object()) + + method._update_ema() + assert method._student_ema.updates == 0 + assert method._ema_update_count == 1 + + method._update_ema() + assert method._student_ema.updates == 1 + assert method._ema_update_count == 2 + + +def test_num_train_timesteps_uses_explicit_schedule_length(): + method = object.__new__(DiffusionNFTMethod) + method._sample_steps = 25 + method._timestep_fraction = 0.5 + method._sampling_config = SimpleNamespace( + timesteps=[900, 800, 700, 600, 500, 400, 300, 200, 100, 10], + sigmas=None, + ) + + assert method._num_train_timesteps() == 5 + + +def test_checkpoint_state_saves_frozen_old_policy_weights(): + student = torch.nn.Linear(2, 2) + old = torch.nn.Linear(2, 2) + for param in old.parameters(): + param.requires_grad_(False) + + method = object.__new__(DiffusionNFTMethod) + method._role_models = { + "student": SimpleNamespace( + transformer=student, + _trainable=True, + ), + "old": SimpleNamespace( + transformer=old, + _trainable=False, + ), + } + method.student = method._role_models["student"] + method.old = method._role_models["old"] + method._student_optimizer = None + method._student_lr_scheduler = None + method._ema_enabled = False + + states = method.checkpoint_state() + old_state = states["roles.old.transformer"].state_dict() + + assert "weight" in old_state + assert "bias" in old_state + assert torch.equal(old_state["weight"], old.weight) diff --git a/tests/local_tests/test_train_rewards.py b/tests/local_tests/test_train_rewards.py new file mode 100644 index 000000000..8a25a3bfe --- /dev/null +++ b/tests/local_tests/test_train_rewards.py @@ -0,0 +1,57 @@ +import torch +import pytest + +from fastvideo.train.methods.rl.rewards import MultiRewardScorer, select_first_frame + + +def test_select_first_frame_for_video_tensor(): + video = torch.arange(2 * 3 * 4 * 5 * 6).reshape(2, 3, 4, 5, 6) + + frame = select_first_frame(video) + + assert frame.shape == (2, 3, 5, 6) + torch.testing.assert_close(frame, video[:, :, 0]) + + +def test_select_first_frame_keeps_frame_tensor(): + frame = torch.randn(2, 3, 5, 6) + + selected = select_first_frame(frame) + + assert selected is frame + + +def test_multi_reward_weighted_sum_with_injected_scorers(): + def pickscore(media, prompts): + assert media.shape == (2, 3, 4, 5, 6) + assert prompts == ["a", "b"] + return torch.tensor([1.0, 2.0]) + + def clipscore(media, prompts): + assert media.shape == (2, 3, 4, 5, 6) + assert prompts == ["a", "b"] + return torch.tensor([0.5, 1.5]) + + scorer = MultiRewardScorer( + {"pickscore": 2.0, "clipscore": 3.0}, + scorers={ + "pickscore": pickscore, + "clipscore": clipscore, + }, + ) + + scores = scorer(torch.zeros(2, 3, 4, 5, 6), ["a", "b"]) + + torch.testing.assert_close(scores["pickscore"], torch.tensor([1.0, 2.0])) + torch.testing.assert_close(scores["clipscore"], torch.tensor([0.5, 1.5])) + torch.testing.assert_close(scores["avg"], torch.tensor([3.5, 8.5])) + + +def test_multi_reward_validates_score_shape(): + scorer = MultiRewardScorer( + {"pickscore": 1.0}, + scorers={"pickscore": lambda media, prompts: torch.tensor([[1.0], [2.0]])}, + ) + + with pytest.raises(ValueError, match="must return shape"): + scorer(torch.zeros(2, 3, 4, 5, 6), ["a", "b"]) diff --git a/tests/local_tests/test_train_rl_sampling.py b/tests/local_tests/test_train_rl_sampling.py new file mode 100644 index 000000000..c015a451c --- /dev/null +++ b/tests/local_tests/test_train_rl_sampling.py @@ -0,0 +1,238 @@ +import torch +import pytest + +from fastvideo.pipelines import TrainingBatch +from fastvideo.train.methods.rl.common import ( + DiffusionSampler, + SamplingConfig, + distributed_k_repeat_indices, + media_to_video_array, + validation_caption, + validation_shard_indices, +) +from fastvideo.train.utils.config import load_run_config + + +class _FakeScheduler: + + def __init__(self): + self.num_train_timesteps = 1000 + self.set_timesteps_calls = [] + self.timesteps = torch.empty(0) + self.sigmas = torch.empty(0) + self.step_calls = 0 + + def set_timesteps(self, num_inference_steps=None, device=None, timesteps=None, sigmas=None): + self.set_timesteps_calls.append({ + "num_inference_steps": num_inference_steps, + "timesteps": timesteps, + "sigmas": sigmas, + }) + if timesteps is not None: + self.timesteps = torch.tensor(timesteps, device=device, dtype=torch.float32) + else: + self.timesteps = torch.linspace(1000, 0, int(num_inference_steps), device=device) + if sigmas is not None: + self.sigmas = torch.tensor(sigmas, device=device, dtype=torch.float32) + else: + self.sigmas = torch.cat([self.timesteps / 1000.0, torch.zeros(1, device=device)]) + + def step(self, model_output, timestep, sample, return_dict=False): + del timestep + self.step_calls += 1 + prev = sample + model_output + return (prev, ) if not return_dict else {"prev_sample": prev} + + +class _FakeModel: + + def __init__(self): + self.noise_scheduler = _FakeScheduler() + self.add_noise_calls = 0 + self.timestep_shapes = [] + + def predict_noise(self, noisy_latents, timestep, batch, *, conditional, attn_kind): + del conditional, attn_kind + self.timestep_shapes.append(tuple(timestep.shape)) + assert batch.timesteps is timestep + return torch.zeros_like(noisy_latents) + + def predict_x0(self, noisy_latents, timestep, batch, *, conditional, attn_kind): + del conditional, attn_kind + self.timestep_shapes.append(tuple(timestep.shape)) + assert batch.timesteps is timestep + return noisy_latents + + def add_noise(self, clean_latents, noise, timestep): + del timestep + self.add_noise_calls += 1 + return clean_latents + noise + + +def _batch(): + batch = TrainingBatch() + batch.latents = torch.zeros(2, 1, 3, 4, 4) + return batch + + +def test_sampler_preserves_latent_dtype(): + model = _FakeModel() + sampler = DiffusionSampler(SamplingConfig(num_steps=4)) + batch = _batch() + batch.latents = batch.latents.to(torch.bfloat16) + + result = sampler.sample(model, batch, generator=torch.Generator().manual_seed(0)) + + assert result.latents.dtype is torch.bfloat16 + + +def test_sampler_uses_scheduler_generated_timesteps_by_default(): + model = _FakeModel() + sampler = DiffusionSampler(SamplingConfig(num_steps=4)) + + result = sampler.sample(model, _batch(), generator=torch.Generator().manual_seed(0)) + + assert result.timesteps.tolist() == [1000.0, 666.6666259765625, 333.3333435058594, 0.0] + + +def test_sampler_honors_explicit_timestep_override(): + model = _FakeModel() + sampler = DiffusionSampler(SamplingConfig(num_steps=3, timesteps=[900, 300, 10])) + + result = sampler.sample(model, _batch(), generator=torch.Generator().manual_seed(0)) + + assert result.timesteps.tolist() == [900.0, 300.0, 10.0] + + +def test_sampler_honors_explicit_timesteps_without_matching_num_steps(): + model = _FakeModel() + sampler = DiffusionSampler(SamplingConfig(timesteps=[900, 300, 10])) + + result = sampler.sample(model, _batch(), generator=torch.Generator().manual_seed(0)) + + assert result.timesteps.tolist() == [900.0, 300.0, 10.0] + assert model.noise_scheduler.set_timesteps_calls == [] + + +def test_sampling_config_rejects_unknown_keys(): + with pytest.raises(ValueError, match="Unsupported method.sampling key"): + SamplingConfig.from_mapping({"solver": "dpm2"}) + + +def test_sampler_restores_original_batch_timestep_after_sampling(): + model = _FakeModel() + sampler = DiffusionSampler(SamplingConfig(num_steps=2)) + batch = _batch() + original_timesteps = torch.tensor([123.0]) + batch.timesteps = original_timesteps + + sampler.sample(batch=batch, model=model, generator=torch.Generator().manual_seed(0)) + + assert batch.timesteps is original_timesteps + + +def test_euler_sampler_does_not_renoise_between_steps(): + model = _FakeModel() + sampler = DiffusionSampler(SamplingConfig(num_steps=4)) + + sampler.sample(model, _batch(), generator=torch.Generator().manual_seed(0)) + + assert model.add_noise_calls == 0 + assert model.timestep_shapes == [(2,), (2,), (2,), (2,)] + + +def test_sde_reflow_sampler_renoises_between_steps(): + model = _FakeModel() + sampler = DiffusionSampler(SamplingConfig(num_steps=4, trajectory="sde_reflow")) + + sampler.sample(model, _batch(), generator=torch.Generator().manual_seed(0)) + + assert model.add_noise_calls == 3 + + +def test_diffusion_nft_config_uses_rl_sampler_not_dmd_pipeline(): + config_path = "examples/train/configs/rl/wan/diffusion_nft_pick_clip.yaml" + + cfg = load_run_config(config_path) + raw_text = open(config_path, encoding="utf-8").read() + + assert cfg.method["_target_"] == "fastvideo.train.methods.rl.diffusion_nft.DiffusionNFTMethod" + assert cfg.training.optimizer.learning_rate == 3.0e-5 + assert cfg.training.data.num_latent_t == 1 + assert cfg.training.data.num_frames == 1 + assert "sampling_timesteps" not in raw_text + assert "WanDMDPipeline" not in raw_text + assert "solver" not in cfg.method["sampling"] + assert cfg.method["sampling"]["scheduler"] == "flow_match_euler" + assert cfg.method["sampling"]["trajectory"] == "ode" + assert cfg.method["sampling"]["flow_shift"] == "inherit" + assert "deterministic" not in cfg.method["sampling"] + assert "noise_level" not in cfg.method["sampling"] + assert cfg.method["validation"]["every_steps"] == 10 + assert cfg.method["validation"]["num_steps"] == 40 + assert cfg.method["validation"]["num_prompts"] == 16 + assert cfg.method["validation"]["log_samples"] is True + + +def test_validation_shard_indices_are_stable_and_padded(): + rank0 = validation_shard_indices(5, rank=0, world_size=2) + rank1 = validation_shard_indices(5, rank=1, world_size=2) + + assert rank0 == [(0, True), (2, True), (4, True)] + assert rank1 == [(1, True), (3, True), (0, False)] + + +def test_distributed_k_repeat_indices_repeats_prompts_globally(): + rank0 = distributed_k_repeat_indices( + dataset_length=100, + batch_size=6, + repeats_per_prompt=24, + world_size=4, + rank=0, + seed=123, + ) + all_indices = [] + for rank in range(4): + sample = distributed_k_repeat_indices( + dataset_length=100, + batch_size=6, + repeats_per_prompt=24, + world_size=4, + rank=rank, + seed=123, + ) + all_indices.extend(sample.local_indices) + + assert rank0.unique_prompt_count == 1 + assert len(all_indices) == 24 + assert len(set(all_indices)) == 1 + + +def test_validation_caption_puts_rewards_first(): + caption = validation_caption( + "a small blue cube", + { + "avg": 0.75, + "pickscore": 0.5, + }, + ) + + assert caption.startswith("avg: 0.7500 | pickscore: 0.5000 | ") + assert caption.endswith("a small blue cube") + + +def test_media_to_video_array_treats_frame_as_single_frame_video(): + frame = torch.ones(3, 4, 5) + + video = media_to_video_array(frame) + + assert video.shape == (1, 3, 4, 5) + assert video.dtype.name == "uint8" + + +def test_media_to_video_array_preserves_video_frames(): + media = torch.ones(3, 2, 4, 5) + + video = media_to_video_array(media) + + assert video.shape == (2, 3, 4, 5) diff --git a/tests/local_tests/test_trainer_managed_optimization.py b/tests/local_tests/test_trainer_managed_optimization.py new file mode 100644 index 000000000..d659af97e --- /dev/null +++ b/tests/local_tests/test_trainer_managed_optimization.py @@ -0,0 +1,157 @@ +from dataclasses import dataclass, field + +import torch + +from fastvideo.train.trainer import Trainer +from fastvideo.train.utils.training_config import ( + CheckpointConfig, + DistributedConfig, + ModelTrainingConfig, + OptimizerConfig, + TrackerConfig, + TrainingConfig, + TrainingLoopConfig, +) + + +class _Tracker: + + def __init__(self): + self.logged = [] + self.finished = False + + def log(self, metrics, step): + self.logged.append((step, metrics)) + + def finish(self): + self.finished = True + + +class _Callbacks: + + def __init__(self): + self.before_optimizer_steps = 0 + self.training_step_ends = 0 + + def on_train_start(self, method, iteration=0): + pass + + def on_before_optimizer_step(self, method, iteration=0): + self.before_optimizer_steps += 1 + + def on_training_step_end(self, method, metrics, iteration=0): + self.training_step_ends += 1 + + def on_validation_begin(self, method, iteration=0): + pass + + def on_validation_end(self, method, iteration=0): + pass + + def on_train_end(self, method, iteration=0): + pass + + +class _World: + rank = 0 + local_rank = 0 + + +class _Method(torch.nn.Module): + + def __init__(self): + super().__init__() + self.calls = 0 + self.backward_calls = 0 + self.optimizer_steps = 0 + self.tracker = None + + def set_tracker(self, tracker): + self.tracker = tracker + + def on_train_start(self): + pass + + def manages_optimization(self): + return True + + def managed_train_step(self, data_stream, iteration): + batch = next(data_stream) + self.calls += 1 + return ( + {"total_loss": torch.tensor(float(batch["x"]))}, + {}, + {"managed_metric": float(iteration)}, + ) + + def backward(self, *args, **kwargs): + self.backward_calls += 1 + + def optimizers_schedulers_step(self, iteration): + self.optimizer_steps += 1 + + def optimizers_zero_grad(self, iteration): + pass + + +class _MethodWithValidation(_Method): + + def __init__(self): + super().__init__() + self.validation_iterations = [] + + def on_validation_begin(self, iteration=0): + self.validation_iterations.append(iteration) + return {"validation/fake": float(iteration)} + + +def test_trainer_skips_default_optimizer_path_for_managed_methods(monkeypatch): + monkeypatch.setattr("fastvideo.train.trainer.get_world_group", lambda: _World()) + monkeypatch.setattr("fastvideo.train.trainer.get_sp_group", lambda: _World()) + monkeypatch.setattr("fastvideo.train.trainer.build_tracker", lambda *args, **kwargs: _Tracker()) + + cfg = TrainingConfig( + distributed=DistributedConfig(), + optimizer=OptimizerConfig(), + loop=TrainingLoopConfig(max_train_steps=1, gradient_accumulation_steps=3), + checkpoint=CheckpointConfig(), + tracker=TrackerConfig(trackers=[]), + model=ModelTrainingConfig(), + ) + trainer = Trainer(cfg) + trainer.callbacks = _Callbacks() + method = _Method() + dataloader = [{"x": 2}] + + trainer.run(method, dataloader=dataloader, max_steps=1) + + assert method.calls == 1 + assert method.backward_calls == 0 + assert method.optimizer_steps == 0 + assert trainer.callbacks.before_optimizer_steps == 0 + assert trainer.callbacks.training_step_ends == 1 + assert trainer.tracker.logged[0][1]["total_loss"] == 2.0 + + +def test_trainer_logs_method_validation_at_step_zero(monkeypatch): + monkeypatch.setattr("fastvideo.train.trainer.get_world_group", lambda: _World()) + monkeypatch.setattr("fastvideo.train.trainer.get_sp_group", lambda: _World()) + monkeypatch.setattr("fastvideo.train.trainer.build_tracker", lambda *args, **kwargs: _Tracker()) + + cfg = TrainingConfig( + distributed=DistributedConfig(), + optimizer=OptimizerConfig(), + loop=TrainingLoopConfig(max_train_steps=1, gradient_accumulation_steps=1), + checkpoint=CheckpointConfig(), + tracker=TrackerConfig(trackers=[]), + model=ModelTrainingConfig(), + ) + trainer = Trainer(cfg) + trainer.callbacks = _Callbacks() + method = _MethodWithValidation() + dataloader = [{"x": 2}] + + trainer.run(method, dataloader=dataloader, max_steps=1) + + assert method.validation_iterations == [0, 1] + assert trainer.tracker.logged[0] == (0, {"validation/fake": 0.0}) diff --git a/tests/local_tests/test_wan_training_dtype.py b/tests/local_tests/test_wan_training_dtype.py new file mode 100644 index 000000000..d20099172 --- /dev/null +++ b/tests/local_tests/test_wan_training_dtype.py @@ -0,0 +1,62 @@ +import torch + +from fastvideo.pipelines.pipeline_batch_info import TrainingBatch +from fastvideo.train.models.wan.wan import WanModel + + +class _CPUWanModel(WanModel): + + @property + def device(self) -> torch.device: + return torch.device("cpu") + + +class _AutocastProbe(torch.nn.Module): + + def __init__(self) -> None: + super().__init__() + self.autocast_enabled: bool | None = None + self.autocast_dtype: torch.dtype | None = None + + def forward( + self, + *, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, + timestep: torch.Tensor, + return_dict: bool, + ) -> torch.Tensor: + del encoder_hidden_states, encoder_attention_mask, timestep, return_dict + self.autocast_enabled = torch.is_autocast_enabled("cpu") + self.autocast_dtype = torch.get_autocast_dtype("cpu") + self.hidden_states_dtype = hidden_states.dtype + return hidden_states + + +def test_wan_predict_noise_uses_training_dtype_autocast_for_fp32_inputs(): + model = object.__new__(_CPUWanModel) + model.transformer = _AutocastProbe() + + batch = TrainingBatch() + batch.timesteps = torch.tensor([1], dtype=torch.long) + batch.conditional_dict = { + "encoder_hidden_states": torch.randn(1, 4, 8, dtype=torch.float32), + "encoder_attention_mask": torch.ones(1, 4, dtype=torch.float32), + } + + noisy_latents = torch.randn(1, 1, 2, 4, 4, dtype=torch.float32) + timestep = torch.tensor([1], dtype=torch.long) + + pred_noise = model.predict_noise( + noisy_latents, + timestep, + batch, + conditional=True, + ) + + assert pred_noise.shape == noisy_latents.shape + assert pred_noise.dtype is torch.bfloat16 + assert model.transformer.hidden_states_dtype is torch.bfloat16 + assert model.transformer.autocast_enabled is True + assert model.transformer.autocast_dtype is torch.bfloat16