diff --git a/AGENTS.md b/AGENTS.md index fcfc099242..adfcd12e84 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -54,3 +54,30 @@ This repository is agent-friendly. Before doing any work, read: If you are exploring a new procedure that has no existing SOP, document your progress in `.agents/exploration/` and flag it for review at the end of your session. + +## Context Efficiency + +### Subagent Discipline + +**Context-aware delegation:** + - Under ~50k context: prefer inline work for tasks under ~5 tool calls. + - Over ~50k context: prefer subagents for self-contained tasks, even simple ones — the per-call token tax on large contexts adds up fast. + +When using subagents, include output rules: "Final response under 2000 characters. List outcomes, not process." +Never call TaskOutput twice for the same subagent. If it times out, increase the timeout — don't re-read. + +### File Reading +Read files with purpose. Before reading a file, know what you're looking for. +Use Grep to locate relevant sections before reading entire large files. +Never re-read a file you've already read in this session. +For files over 500 lines, use offset/limit to read only the relevant section. + +### Responses +Don't echo back file contents you just read — the user can see them. +Don't narrate tool calls ("Let me read the file..." / "Now I'll edit..."). Just do it. +Keep explanations proportional to complexity. Simple changes need one sentence, not three paragraphs. + +**Tables — STRICT RULES (apply everywhere, always):** +- Markdown tables: use minimum separator (`|-|-|`). Never pad with repeated hyphens (`|---|---|`). +- NEVER use box-drawing / ASCII-art tables with characters like `┌`, `┬`, `─`, `│`, `└`, `┘`, `├`, `┤`, `┼`. These are completely banned. +- No exceptions. Not for "clarity", not for alignment, not for terminal output. \ No newline at end of file diff --git a/GenRL b/GenRL new file mode 160000 index 0000000000..82a867ed24 --- /dev/null +++ b/GenRL @@ -0,0 +1 @@ +Subproject commit 82a867ed24f6095a7d931ec76319aae20d9a01c4 diff --git a/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml b/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml new file mode 100644 index 0000000000..c102752298 --- /dev/null +++ b/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml @@ -0,0 +1,164 @@ +# GenRL / Video GRPO: Wan 2.1 T2V 1.3B with LongCat reweighting. +# +# Ported from GenRL/config/longcat.yaml. +# +# - Student: trainable full-parameter model by default. +# - LoRA is still available via models.student.use_lora=true. +# - Full fine-tuning with beta > 0 requires models.reference and much +# more memory; keep beta at 0.0 for the 4xH100 probe run. +# +# Usage: +# torchrun --nnodes=1 --nproc_per_node=4 \ +# -m fastvideo.train.entrypoint.train \ +# --config examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml + +models: + student: + _target_: fastvideo.train.models.wan.wan_genrl.GenRLWanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: true + # Set true for LoRA. LoRA can use method.beta with disable_adapter(). + use_lora: false + lora_r: 128 + lora_alpha: 64 + lora_init_weights: gaussian + lora_path: null + lora_target_modules: + - to_k + - to_out + - to_q + - to_v + - ffn.fc_in + - ffn.fc_out + enable_gradient_checkpointing_type: full + +method: + _target_: fastvideo.train.methods.rl.genrl.GenRLMethod + + # ---- Reward functions ---- + reward_fn: + hpsv3_general: 1.0 + hpsv3_percentile: 1.0 + videoalign_mq: 1.0 + videoalign_ta: 1.0 + reward_module: null + reward_on_gpu: true + + # ---- Data ---- + prompt_dataset_path: GenRL/datasets/filtered_prompts + prompt_fn: filtered_prompts + + # ---- Sampling ---- + sample_batch_size: 4 + eval_batch_size: 2 + # Sample multiple rollout microbatches, average their PPO losses, then + # apply one optimizer update. This reduces reward/advantage variance. + num_batches_per_epoch: 4 + accumulate_ppo_microbatches: true + eval_every_steps: 20 + eval_num_batches: 1 + eval_num_steps: 16 + eval_guidance_scale: 4.5 + num_inference_steps: 16 + guidance_scale: 4.5 + num_video_per_prompt: 4 + noise_level: 1.0 + sde_type: flow_sde + sde_window_size: 1 + sde_window_range: [0, 6] + diffusion_clip: true + diffusion_clip_value: 0.45 + kl_reward: 0 + same_latent: true + + # ---- Video dimensions ---- + height: 480 + width: 832 + num_frames: 81 + + # ---- PPO training ---- + train_batch_size: 4 + num_inner_epochs: 1 + clip_range: 1.0e-4 + adv_clip_max: 5.0 + # Official LoRA LongCat uses beta: 3.0e-4 with disable_adapter(). + # For full fine-tuning on 4 H100s, avoid a second frozen Wan copy. + beta: 0.0 + use_cfg: true + # Flash-GRPO-style temporal gradient rectification: avoid the large + # LongCat sigma/dt multiplier while debugging full-FT stability. + loss_reweighting: flash_tgr + loss_reweighting_clip: null + weight_advantages: true + # Match official GenRL PPO cadence. With sde_window_size: 1 this is + # equivalent to one optimizer step per sampled trajectory timestep. + optimizer_step_per_timestep: true + log_post_update_kl: true + max_grad_norm: 1.0 + seed: 42 + + # ---- Advantage computation ---- + per_prompt_stat_tracking: true + global_std: false + max_group_std: true + +training: + distributed: + + num_gpus: 4 + sp_size: 1 + tp_size: 1 + # Full fine-tuning needs FSDP/HSDP sharding across all 4 GPUs. + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + + data: + # Not used by GenRL (prompt dataloaders are in method config) + # but required by the config parser. + data_path: "" + train_batch_size: 1 + seed: 42 + num_height: 480 + num_width: 832 + num_frames: 81 + + optimizer: + # Full-parameter visual GRPO is much more sensitive than LoRA. + # DanceGRPO reports 5e-6 to 2e-5 as the practical range. + learning_rate: 1.0e-5 + betas: [0.9, 0.999] + weight_decay: 1.0e-4 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + # max_train_steps == num_epochs in GenRL terms. + max_train_steps: 100000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/genrl_longcat + training_state_checkpointing_steps: 100 + checkpoints_total_limit: 3 + + tracker: + project_name: VideoRL + # Leave blank so W&B auto-generates a unique display name per run. + run_name: "" + + model: + enable_gradient_checkpointing_type: full + +callbacks: + # Gradnorm call back Disabled; GenRLMethod clips internally. + ema: + decay: 0.9 + start_iter: 0 + update_interval: 8 + log_rl_samples: + every_steps: 1 + max_videos: 4 + fps: 16 + +pipeline: + flow_shift: 3.0 diff --git a/examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml b/examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml new file mode 100644 index 0000000000..1b1c67db91 --- /dev/null +++ b/examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml @@ -0,0 +1,113 @@ +# GenRL / Video GRPO: Wan 2.1 T2V 1.3B — OCR reward, full finetune, 4 GPUs. +# +# Usage: +# torchrun --nnodes=1 --nproc_per_node=4 \ +# -m fastvideo.train.entrypoint.train \ +# --config examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml + +models: + student: + _target_: fastvideo.train.models.wan.wan_genrl.GenRLWanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: true + enable_gradient_checkpointing_type: full + +method: + _target_: fastvideo.train.methods.rl.genrl.GenRLMethod + + # ---- Reward functions ---- + reward_fn: + video_ocr: 1.0 + reward_module: null + + # ---- Data ---- + prompt_dataset_path: GenRL/datasets/ocr + prompt_fn: general_ocr + + # ---- Sampling ---- + sample_batch_size: 4 + eval_batch_size: 2 + num_batches_per_epoch: 1 + num_inference_steps: 16 + guidance_scale: 4.5 + num_video_per_prompt: 4 + noise_level: 1.0 + sde_type: flow_sde + sde_window_size: 1 + sde_window_range: [0, 6] + diffusion_clip: true + diffusion_clip_value: 0.45 + kl_reward: 0 + same_latent: true + + # ---- Video dimensions ---- + height: 480 + width: 832 + num_frames: 81 + + # ---- PPO training ---- + train_batch_size: 4 + num_inner_epochs: 1 + clip_range: 1.0e-3 + adv_clip_max: 5.0 + # No frozen reference model is configured in this launch. + beta: 0.0 + use_cfg: true + loss_reweighting: longcat + weight_advantages: false + max_grad_norm: 1.0 + seed: 42 + + # ---- Advantage computation ---- + per_prompt_stat_tracking: true + global_std: false + max_group_std: true + +training: + distributed: + num_gpus: 4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + + data: + data_path: "" + train_batch_size: 1 + seed: 42 + num_height: 480 + num_width: 832 + num_frames: 81 + + optimizer: + learning_rate: 1.0e-4 + betas: [0.9, 0.999] + weight_decay: 1.0e-4 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 100000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/genrl_ocr + training_state_checkpointing_steps: 100 + checkpoints_total_limit: 3 + + tracker: + project_name: VideoRL + run_name: wan_2_1_t2v_1_3b_ocr + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + max_grad_norm: 0.0 # Disabled; GenRLMethod clips internally. + ema: + decay: 0.9 + start_iter: 0 + +pipeline: + flow_shift: 3.0 diff --git a/examples/train/run.sh b/examples/train/run.sh index 263a6713a3..84b21cc182 100755 --- a/examples/train/run.sh +++ b/examples/train/run.sh @@ -55,7 +55,7 @@ python -m torch.distributed.run \ --nproc_per_node "${NUM_GPUS}" \ --master_addr "${MASTER_ADDR}" \ --master_port "${MASTER_PORT}" \ - fastvideo/train/entrypoint/train.py \ + -m fastvideo.train.entrypoint.train \ --config "${CONFIG}" \ "$@" \ 2>&1 | tee "${LOG_FILE}" diff --git a/examples/train/run_slurm.sh b/examples/train/run_slurm.sh index 3d9f3439bd..5f37479971 100755 --- a/examples/train/run_slurm.sh +++ b/examples/train/run_slurm.sh @@ -107,7 +107,7 @@ srun torchrun \\ --node_rank \$SLURM_PROCID \\ --rdzv_backend=c10d \\ --rdzv_endpoint="\$MASTER_ADDR:\$MASTER_PORT" \\ - fastvideo/train/entrypoint/train.py \\ + -m fastvideo.train.entrypoint.train \\ --config ${CONFIG} \\ --training.distributed.num_gpus ${TOTAL_GPUS} \\ ${EXTRA_ARGS[*]:-} diff --git a/fastvideo/layers/lora/linear.py b/fastvideo/layers/lora/linear.py index d1368242f2..3ae0bfa78a 100644 --- a/fastvideo/layers/lora/linear.py +++ b/fastvideo/layers/lora/linear.py @@ -336,8 +336,24 @@ def forward(self, input_: torch.Tensor): output_parallel = self.base_layer.quant_method.apply( self.base_layer, input_parallel) - if self.set_lora: - output_parallel = self.apply_lora(output_parallel, input_parallel) + if not self.merged and not self.disable_lora: + lora_A = self.lora_A + lora_B = self.lora_B + assert lora_A is not None and lora_B is not None + if isinstance(lora_B, DTensor): + lora_B = lora_B.to_local() + lora_A = lora_A.to_local() + + lora_A_sliced = self.slice_lora_a_weights( + lora_A.to(input_parallel, non_blocking=True)) + lora_B_sliced = self.slice_lora_b_weights( + lora_B.to(output_parallel, non_blocking=True)) + delta = input_parallel @ lora_A_sliced.T @ lora_B_sliced.T + if self.lora_alpha != self.lora_rank: + delta = delta * ( + self.lora_alpha / self.lora_rank # type: ignore + ) + output_parallel = output_parallel + delta if self.base_layer.reduce_results and self.base_layer.tp_size > 1: output_ = tensor_model_parallel_all_reduce(output_parallel) diff --git a/fastvideo/train/callbacks/__init__.py b/fastvideo/train/callbacks/__init__.py index d57e048a98..3d8e73cc49 100644 --- a/fastvideo/train/callbacks/__init__.py +++ b/fastvideo/train/callbacks/__init__.py @@ -8,6 +8,8 @@ EMACallback, ) from fastvideo.train.callbacks.grad_clip import ( GradNormClipCallback, ) +from fastvideo.train.callbacks.log_rl_samples import ( + LogRLSamplesCallback, ) from fastvideo.train.callbacks.validation import ( ValidationCallback, ) @@ -16,5 +18,6 @@ "CallbackDict", "EMACallback", "GradNormClipCallback", + "LogRLSamplesCallback", "ValidationCallback", ] diff --git a/fastvideo/train/callbacks/callback.py b/fastvideo/train/callbacks/callback.py index 024a92860c..c228e7f345 100644 --- a/fastvideo/train/callbacks/callback.py +++ b/fastvideo/train/callbacks/callback.py @@ -24,6 +24,7 @@ "grad_clip": "fastvideo.train.callbacks.grad_clip.GradNormClipCallback", "validation": "fastvideo.train.callbacks.validation.ValidationCallback", "ema": "fastvideo.train.callbacks.ema.EMACallback", + "log_rl_samples": "fastvideo.train.callbacks.log_rl_samples.LogRLSamplesCallback", } @@ -58,6 +59,7 @@ def on_before_optimizer_step( self, method: TrainingMethod, iteration: int = 0, + outputs: dict[str, Any] | None = None, ) -> None: pass diff --git a/fastvideo/train/callbacks/ema.py b/fastvideo/train/callbacks/ema.py index 2bd9f01f09..8452a19959 100644 --- a/fastvideo/train/callbacks/ema.py +++ b/fastvideo/train/callbacks/ema.py @@ -47,9 +47,11 @@ def __init__( *, decay: float = 0.9999, start_iter: int = 0, + update_interval: int = 1, ) -> None: self._decay = float(decay) self._start_iter = int(start_iter) + self._update_interval = max(1, int(update_interval)) self._ema_started = False self.student_ema: EMA_FSDP | None = None @@ -78,9 +80,10 @@ def on_train_start( ) logger.info( "EMA callback enabled (decay=%s, " - "start_iter=%d).", + "start_iter=%d, update_interval=%d).", self._decay, self._start_iter, + self._update_interval, ) def on_training_step_end( @@ -94,6 +97,8 @@ def on_training_step_end( if iteration < self._start_iter: return + if (iteration - self._start_iter) % self._update_interval != 0: + return if not self._ema_started: logger.info( "Starting EMA updates at iteration %d " diff --git a/fastvideo/train/callbacks/grad_clip.py b/fastvideo/train/callbacks/grad_clip.py index e157610625..41e6ea544c 100644 --- a/fastvideo/train/callbacks/grad_clip.py +++ b/fastvideo/train/callbacks/grad_clip.py @@ -8,7 +8,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING from fastvideo.logger import init_logger from fastvideo.train.callbacks.callback import Callback @@ -41,6 +41,7 @@ def on_before_optimizer_step( self, method: TrainingMethod, iteration: int = 0, + outputs: dict[str, Any] | None = None, ) -> None: max_norm = self._max_grad_norm if max_norm <= 0.0: diff --git a/fastvideo/train/callbacks/log_rl_samples.py b/fastvideo/train/callbacks/log_rl_samples.py new file mode 100644 index 0000000000..0a86bdb7b8 --- /dev/null +++ b/fastvideo/train/callbacks/log_rl_samples.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Callback to log sampled RL videos to the tracker.""" + +from __future__ import annotations + +import contextlib +import os +import tempfile +from typing import Any, TYPE_CHECKING + +import imageio +import numpy as np +import torch + +from fastvideo.logger import init_logger +from fastvideo.train.callbacks.callback import Callback + +if TYPE_CHECKING: + from fastvideo.train.methods.base import TrainingMethod + +logger = init_logger(__name__) + + +class LogRLSamplesCallback(Callback): + """Log RL-sampled videos to the experiment tracker. + + Expects ``outputs`` to contain: + + - ``sample_videos``: uint8 tensor (B, 3, T, H, W). + - ``sample_prompts``: list of prompt strings. + + Configuration (YAML ``callbacks.log_rl_samples``): + + .. code-block:: yaml + + callbacks: + log_rl_samples: + every_steps: 5 + max_videos: 4 + fps: 16 + """ + + def __init__( + self, + *, + every_steps: int = 1, + max_videos: int = 4, + fps: int = 16, + ) -> None: + self._every_steps = int(every_steps) + self._max_videos = int(max_videos) + self._fps = int(fps) + + def on_before_optimizer_step( + self, + method: TrainingMethod, + iteration: int = 0, + outputs: dict[str, Any] | None = None, + ) -> None: + if outputs is None: + return + if ( + self._every_steps > 0 + and iteration % self._every_steps != 0 + ): + return + + videos = outputs.get("sample_videos") + prompts = outputs.get("sample_prompts") + if videos is None: + return + + tracker = getattr(method, "tracker", None) + if tracker is None: + return + + self._log_videos(tracker, videos, prompts, iteration) + + def _log_videos( + self, + tracker: Any, + videos: torch.Tensor, + prompts: list[str] | None, + step: int, + ) -> None: + n = min(len(videos), self._max_videos) + tmp_dir = tempfile.mkdtemp(prefix="rl_samples_") + video_logs = [] + + try: + for i in range(n): + # (3, T, H, W) uint8 -> (T, H, W, 3) numpy. + v = videos[i].permute(1, 2, 3, 0) + frames = v.numpy().astype(np.uint8) + fname = os.path.join( + tmp_dir, f"sample_{step}_{i}.mp4" + ) + imageio.mimsave(fname, frames, fps=self._fps) + + caption = ( + prompts[i] + if prompts and i < len(prompts) + else None + ) + art = tracker.video( + fname, caption=caption, fps=self._fps + ) + if art is not None: + video_logs.append(art) + + if video_logs: + tracker.log_artifacts( + {"rl_sample_videos": video_logs}, + step, + ) + logger.info( + "Logged %d RL sample videos at step %d", + len(video_logs), + step, + ) + finally: + # Clean up temp files. + for f in os.listdir(tmp_dir): + with contextlib.suppress(OSError): + os.remove(os.path.join(tmp_dir, f)) + with contextlib.suppress(OSError): + os.rmdir(tmp_dir) diff --git a/fastvideo/train/methods/__init__.py b/fastvideo/train/methods/__init__.py index 61fd6ef2e7..0a5a8f99ac 100644 --- a/fastvideo/train/methods/__init__.py +++ b/fastvideo/train/methods/__init__.py @@ -8,20 +8,5 @@ "FineTuneMethod", "SelfForcingMethod", "DiffusionForcingSFTMethod", + "GenRLMethod", ] - - -def __getattr__(name: str) -> object: - if name == "DMD2Method": - from fastvideo.train.methods.distribution_matching.dmd2 import DMD2Method - return DMD2Method - if name == "FineTuneMethod": - from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod - return FineTuneMethod - if name == "SelfForcingMethod": - from fastvideo.train.methods.distribution_matching.self_forcing import SelfForcingMethod - return SelfForcingMethod - if name == "DiffusionForcingSFTMethod": - from fastvideo.train.methods.fine_tuning.dfsft import DiffusionForcingSFTMethod - return DiffusionForcingSFTMethod - raise AttributeError(name) diff --git a/fastvideo/train/methods/rl/__init__.py b/fastvideo/train/methods/rl/__init__.py new file mode 100644 index 0000000000..109067383c --- /dev/null +++ b/fastvideo/train/methods/rl/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reinforcement learning methods for video generation.""" diff --git a/fastvideo/train/methods/rl/genrl.py b/fastvideo/train/methods/rl/genrl.py new file mode 100644 index 0000000000..824257d754 --- /dev/null +++ b/fastvideo/train/methods/rl/genrl.py @@ -0,0 +1,1475 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Video GRPO / PPO training method for diffusion models. + +One trainer step == one GenRL outer epoch: + 1. Sample videos and compute rewards. + 2. Compute advantages (per-prompt normalization). + 3. PPO training across inner epochs. + +The method handles backward / optimizer internally and +returns all stats. The trainer's outer loop just calls +``single_train_step`` once per step. +""" + +from __future__ import annotations + +import contextlib +import copy +import time +from collections import defaultdict +from concurrent import futures +from typing import Any + +import numpy as np +import torch +import torch.distributed as dist + +from fastvideo.distributed import get_world_group +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.train.methods.base import ( + LogScalar, + TrainingMethod, +) +from fastvideo.train.methods.rl.utils.advantages import ( + compute_advantages, +) +from fastvideo.train.methods.rl.utils.data import ( + build_prompt_dataloaders, +) +from fastvideo.train.methods.rl.utils.diffusion import ( + compute_log_prob, +) +from fastvideo.train.methods.rl.utils.embeddings import ( + compute_text_embeddings, +) +from fastvideo.train.methods.rl.utils.evaluation import ( + eval_once, +) +from fastvideo.train.methods.rl.utils.rewards import ( + clear_reward_models, + move_reward_models, + multi_score, + reward_models_on_device, +) +from fastvideo.train.methods.rl.utils.sampling import ( + sample_epoch, +) +from fastvideo.train.methods.rl.utils.stat_tracking import ( + PerPromptStatTracker, +) +from fastvideo.train.models.base import ModelBase +from fastvideo.train.utils.optimizer import ( + build_optimizer_and_scheduler, + clip_grad_norm_if_needed, +) + +logger = init_logger(__name__) + +ADVANTAGE_EPSILON = 1e-6 +SEED_EPOCH_STRIDE = 10_000 + + +def _gather_tensor( + tensor: torch.Tensor, + world_size: int, +) -> torch.Tensor: + """Gather tensor from all ranks and concatenate.""" + if world_size <= 1: + return tensor + gathered = [torch.empty_like(tensor) for _ in range(world_size)] + dist.all_gather(gathered, tensor) + return torch.cat(gathered, dim=0) + + +class GenRLMethod(TrainingMethod): + """Video GRPO / PPO method for diffusion models. + + Handles the full RL training loop internally + within ``single_train_step``. + """ + + def __init__( + self, + *, + cfg: Any, + role_models: dict[str, ModelBase], + ) -> None: + super().__init__(cfg=cfg, role_models=role_models) + + mc = self.method_config + tc = self.training_config + + # Optional reference model for KL. + self._reference = role_models.get("reference") + + # Parse RL config. + self._parse_config(mc) + self._validate_config() + + # Init student preprocessors (VAE, text encoder). + self.student.init_preprocessors(tc) + + # Scheduler copy for RL pipeline. + self._scheduler = copy.deepcopy( + self.student.noise_scheduler + ) + + # Reward functions. + self._reward_on_gpu = bool( + mc.get("reward_on_gpu", False) + ) + reward_init_device = ( + self.student.device + if self._reward_on_gpu + else torch.device("cpu") + ) + self._reward_fn = multi_score( + reward_init_device, + self._reward_cfg, + mc.get("reward_module"), + return_raw_scores=True, + ) + if self._reward_on_gpu: + move_reward_models( + self._reward_cfg, self.student.device + ) + + # Prompt dataloaders. + wg = get_world_group() + self._world_size = wg.world_size + self._rank = wg.rank + self._is_main = wg.rank == 0 + total_samples = self._world_size * self._sample_batch_size + if total_samples % self._num_video_per_prompt != 0: + raise ValueError( + "world_size * sample_batch_size must be divisible by " + "num_video_per_prompt for DistributedKRepeatSampler. Got " + f"{self._world_size} * {self._sample_batch_size} and " + f"{self._num_video_per_prompt}." + ) + + train_dl, test_dl, train_sampler = ( + build_prompt_dataloaders( + prompt_dataset_path=mc[ + "prompt_dataset_path" + ], + prompt_fn=mc.get( + "prompt_fn", "general_ocr" + ), + sample_batch_size=self._sample_batch_size, + eval_batch_size=self._eval_batch_size, + num_video_per_prompt=( + self._num_video_per_prompt + ), + num_processes=self._world_size, + process_index=self._rank, + seed=self._seed, + ) + ) + self._train_dataloader = train_dl + self._test_dataloader = test_dl + self._train_sampler = train_sampler + self._train_iter = iter(train_dl) + + # Stat trackers. + self._build_stat_trackers() + + # Negative prompt embeddings. + self._compute_negative_embeds() + + # Optimizer and scheduler. + self._init_optimizer() + + # Async reward executor. + self._executor = futures.ThreadPoolExecutor( + max_workers=8 + ) + + # Training timestep indices. + self._compute_train_timesteps() + + # ------------------------------------------------------------------ + # Config parsing + # ------------------------------------------------------------------ + + def _parse_config(self, mc: dict[str, Any]) -> None: + # Sampling. + self._sample_batch_size = int( + mc.get("sample_batch_size", 8) + ) + self._eval_batch_size = int( + mc.get("eval_batch_size", 2) + ) + self._num_batches_per_epoch = int( + mc.get("num_batches_per_epoch", 2) + ) + self._num_inference_steps = int( + mc.get("num_inference_steps", 20) + ) + self._guidance_scale = float( + mc.get("guidance_scale", 4.5) + ) + self._num_video_per_prompt = int( + mc.get("num_video_per_prompt", 4) + ) + self._noise_level = float( + mc.get("noise_level", 0.7) + ) + self._sde_type = str( + mc.get("sde_type", "flow_sde") + ) + self._sde_window_size = int( + mc.get("sde_window_size", 0) + ) + raw_range = mc.get("sde_window_range") + self._sde_window_range = ( + tuple(raw_range) if raw_range else None + ) + self._diffusion_clip = bool( + mc.get("diffusion_clip", False) + ) + self._diffusion_clip_value = float( + mc.get("diffusion_clip_value", 0.45) + ) + self._kl_reward = float( + mc.get("kl_reward", 0.0) + ) + self._same_latent = bool( + mc.get("same_latent", False) + ) + + # Training. + self._num_inner_epochs = int( + mc.get("num_inner_epochs", 1) + ) + self._clip_range = float( + mc.get("clip_range", 1e-3) + ) + self._adv_clip_max = float( + mc.get("adv_clip_max", 5.0) + ) + self._beta = float(mc.get("beta", 0.0)) + self._use_cfg = bool(mc.get("use_cfg", True)) + self._loss_reweighting = mc.get( + "loss_reweighting" + ) + self._loss_reweighting_clip = mc.get( + "loss_reweighting_clip" + ) + self._loss_reweighting_clip = ( + None + if self._loss_reweighting_clip is None + else float(self._loss_reweighting_clip) + ) + self._weight_advantages = bool( + mc.get("weight_advantages", False) + ) + self._max_grad_norm = float( + mc.get("max_grad_norm", 1.0) + ) + self._optimizer_step_per_timestep = bool( + mc.get("optimizer_step_per_timestep", True) + ) + self._accumulate_ppo_microbatches = bool( + mc.get("accumulate_ppo_microbatches", False) + ) + self._log_post_update_kl = bool( + mc.get("log_post_update_kl", True) + ) + self._train_batch_size = int( + mc.get("train_batch_size", 8) + ) + self._eval_every_steps = int( + mc.get("eval_every_steps", 0) + ) + self._eval_num_batches = int( + mc.get("eval_num_batches", 1) + ) + self._eval_guidance_scale = float( + mc.get("eval_guidance_scale", self._guidance_scale) + ) + self._eval_num_steps = int( + mc.get("eval_num_steps", self._num_inference_steps) + ) + + # Data / dimensions. + self._height = int(mc.get("height", 480)) + self._width = int(mc.get("width", 832)) + self._num_frames = int(mc.get("num_frames", 81)) + self._seed = int(mc.get("seed", 42)) + + # Per-prompt tracking. + self._per_prompt_stat_tracking = bool( + mc.get("per_prompt_stat_tracking", True) + ) + if self._num_video_per_prompt == 1: + self._per_prompt_stat_tracking = False + + # Reward config. + self._reward_cfg = dict(mc.get("reward_fn", {})) + + def _validate_config(self) -> None: + """Fail early on unsupported RL config combinations.""" + if not self._reward_cfg: + raise ValueError( + "method.reward_fn must contain at least one reward." + ) + + if self._beta > 0 and not self._has_reference_policy(): + raise ValueError( + "method.beta > 0 requires either a configured reference " + "model or a LoRA student with disable_adapter()." + ) + + if self._kl_reward > 0 and not self._has_reference_policy(): + raise ValueError( + "method.kl_reward > 0 requires either a configured reference " + "model or a LoRA student with disable_adapter()." + ) + + if self._loss_reweighting not in { + None, + "longcat", + "flash_tgr", + }: + raise ValueError( + "method.loss_reweighting must be one of null, " + "'longcat', or 'flash_tgr'." + ) + + if self._sde_window_range is not None: + if len(self._sde_window_range) != 2: + raise ValueError( + "method.sde_window_range must contain exactly two values." + ) + start, end = self._sde_window_range + if start < 0 or end <= start: + raise ValueError( + "method.sde_window_range must satisfy 0 <= start < end." + ) + if end > self._num_inference_steps: + raise ValueError( + "method.sde_window_range end cannot exceed " + "method.num_inference_steps." + ) + + # ------------------------------------------------------------------ + # Setup helpers + # ------------------------------------------------------------------ + + def _build_stat_trackers(self) -> None: + self._stat_tracker = None + self._reward_stat_trackers = None + self._kl_stat_tracker = None + + if self._per_prompt_stat_tracking: + self._stat_tracker = PerPromptStatTracker( + use_global_std=bool( + self.method_config.get( + "global_std", False + ) + ), + max_group_std=bool( + self.method_config.get( + "max_group_std", False + ) + ), + ) + + if ( + self._weight_advantages + and self._per_prompt_stat_tracking + ): + self._reward_stat_trackers = { + name: PerPromptStatTracker( + use_global_std=bool( + self.method_config.get( + "global_std", False + ) + ), + max_group_std=bool( + self.method_config.get( + "max_group_std", False + ) + ), + ) + for name in self._reward_cfg + } + if self._kl_reward > 0: + self._kl_stat_tracker = ( + PerPromptStatTracker( + use_global_std=bool( + self.method_config.get( + "global_std", False + ) + ), + max_group_std=bool( + self.method_config.get( + "max_group_std", False + ) + ), + ) + ) + + def _compute_negative_embeds(self) -> None: + device = self.student.device + neg = compute_text_embeddings( + [""], + self.student.text_encoder, + self.student.tokenizer, + max_sequence_length=512, + device=device, + ) + self._sample_neg_embeds = neg.repeat( + self._sample_batch_size, 1, 1 + ) + self._train_neg_embeds = neg.repeat( + self._train_batch_size, 1, 1 + ) + + def _init_optimizer(self) -> None: + tc = self.training_config + params = [ + p + for p in self.student.transformer.parameters() + if p.requires_grad + ] + if not params: + raise ValueError( + "GenRL student transformer has no trainable parameters. " + "For LoRA, check models.student.use_lora and " + "lora_target_modules. For full tuning, check " + "models.student.trainable." + ) + trainable_count = sum(p.numel() for p in params) + logger.info( + "GenRL trainable transformer parameters: %.2fM", + trainable_count / 1e6, + ) + self._transformer_params = params + ( + self._optimizer, + self._lr_scheduler, + ) = build_optimizer_and_scheduler( + params=params, + optimizer_config=tc.optimizer, + loop_config=tc.loop, + learning_rate=float( + tc.optimizer.learning_rate + ), + betas=tc.optimizer.betas, + scheduler_name=str(tc.optimizer.lr_scheduler), + ) + + def _compute_train_timesteps(self) -> None: + if self._sde_window_size > 0: + num_ts = self._sde_window_size + else: + num_ts = int( + self._num_inference_steps * 0.99 + ) + self._num_train_timesteps = num_ts + self._train_timesteps = list(range(num_ts)) + + # ------------------------------------------------------------------ + # TrainingMethod interface + # ------------------------------------------------------------------ + + def single_train_step( + self, + batch: dict[str, Any], + iteration: int, + ) -> tuple[ + dict[str, torch.Tensor], + dict[str, Any], + dict[str, LogScalar], + ]: + """Run one full RL epoch (sample -> train). + + The *batch* argument (from the dummy dataloader) + is ignored. + """ + epoch = iteration - 1 # 0-indexed + device = self.student.device + all_metrics: dict[str, LogScalar] = { + "epoch": float(epoch), + } + + # 1. Sample epoch. + torch.cuda.synchronize() + t_sample_start = time.perf_counter() + self.student.transformer.eval() + reward_ctx = ( + contextlib.nullcontext() + if self._reward_on_gpu + else reward_models_on_device( + self._reward_cfg, device + ) + ) + with reward_ctx: + samples, batch_videos, batch_prompts = ( + sample_epoch( + model=self.student, + scheduler=self._scheduler, + train_sampler=self._train_sampler, + train_iter=self._train_iter, + reward_fn=self._reward_fn, + sample_neg_prompt_embeds=( + self._sample_neg_embeds + ), + text_encoder=self.student.text_encoder, + tokenizer=self.student.tokenizer, + executor=self._executor, + epoch=epoch, + global_step=iteration, + sample_batch_size=( + self._sample_batch_size + ), + num_batches_per_epoch=( + self._num_batches_per_epoch + ), + num_inference_steps=( + self._num_inference_steps + ), + guidance_scale=self._guidance_scale, + height=self._height, + width=self._width, + num_frames=self._num_frames, + noise_level=self._noise_level, + sde_type=self._sde_type, + diffusion_clip=self._diffusion_clip, + diffusion_clip_value=( + self._diffusion_clip_value + ), + sde_window_size=( + self._sde_window_size + ), + sde_window_range=( + self._sde_window_range + ), + kl_reward=self._kl_reward, + same_latent=self._same_latent, + seed=self._seed, + device=device, + is_main_process=self._is_main, + ref_transformer=( + self._reference.transformer + if self._reference + else None + ), + lora_model=self._get_lora_ref_transformer(), + tracker=self.tracker, + async_reward_scoring=( + not self._reward_on_gpu + ), + ) + ) + torch.cuda.synchronize() + t_sample_end = time.perf_counter() + if self._reward_on_gpu: + t_reward_clear = time.perf_counter() + clear_reward_models(self._reward_cfg) + torch.cuda.synchronize() + logger.info( + "[rewards] clear_after_sample=%.1fs", + time.perf_counter() - t_reward_clear, + ) + + # 2. Prepare samples (advantages). + t_adv_start = time.perf_counter() + samples = self._prepare_samples( + samples, epoch, iteration + ) + t_adv_end = time.perf_counter() + + # 3. PPO training. + torch.cuda.synchronize() + t_ppo_start = time.perf_counter() + ppo_metrics = self._ppo_train( + samples, epoch, iteration + ) + torch.cuda.synchronize() + t_ppo_end = time.perf_counter() + all_metrics.update(ppo_metrics) + + if ( + self._eval_every_steps > 0 + and iteration % self._eval_every_steps == 0 + ): + t_eval_start = time.perf_counter() + eval_ctx = ( + reward_models_on_device( + self._reward_cfg, device + ) + if self._reward_on_gpu + else contextlib.nullcontext() + ) + with eval_ctx: + eval_metrics = eval_once( + model=self.student, + scheduler=self._scheduler, + test_dataloader=self._test_dataloader, + text_encoder=self.student.text_encoder, + tokenizer=self.student.tokenizer, + sample_neg_prompt_embeds=( + self._sample_neg_embeds + ), + eval_reward_fn=self._reward_fn, + global_step=iteration, + ema_callback=None, + eval_num_steps=self._eval_num_steps, + eval_guidance_scale=( + self._eval_guidance_scale + ), + height=self._height, + width=self._width, + num_frames=self._num_frames, + device=device, + world_size=self._world_size, + rank=self._rank, + is_main_process=self._is_main, + tracker=None, + max_batches=self._eval_num_batches, + seed=self._seed + 1_000_000, + ) + if self._reward_on_gpu: + clear_reward_models(self._reward_cfg) + all_metrics.update(eval_metrics) + torch.cuda.synchronize() + all_metrics["time/eval_sec"] = ( + time.perf_counter() - t_eval_start + ) + + logger.info( + "[GenRL step %d] TIMING: " + "sample=%.1fs advantages=%.1fs " + "ppo_train=%.1fs total=%.1fs", + iteration, + t_sample_end - t_sample_start, + t_adv_end - t_adv_start, + t_ppo_end - t_ppo_start, + t_ppo_end - t_sample_start, + ) + all_metrics["time/sample_sec"] = ( + t_sample_end - t_sample_start + ) + all_metrics["time/advantages_sec"] = ( + t_adv_end - t_adv_start + ) + all_metrics["time/ppo_train_sec"] = ( + t_ppo_end - t_ppo_start + ) + + # Build outputs dict with sampled videos for + # logging callbacks. + outputs: dict[str, Any] = {} + if self._is_main and batch_videos: + vids = batch_videos[0] + outputs["sample_videos"] = ( + (vids * 255) + .clamp(0, 255) + .to(torch.uint8) + .cpu() + ) + outputs["sample_prompts"] = ( + batch_prompts[0] if batch_prompts else [] + ) + + # Return dummy loss (everything is internal). + dummy_loss = torch.zeros( + (), device=device, requires_grad=False + ) + return ( + {"total_loss": dummy_loss}, + outputs, + all_metrics, + ) + + def backward( + self, + loss_map: dict[str, torch.Tensor], + outputs: dict[str, Any], + *, + grad_accum_rounds: int = 1, + ) -> None: + pass # Handled internally. + + def get_optimizers( + self, + iteration: int, + ) -> list[torch.optim.Optimizer]: + return [] # Handled internally. + + def get_lr_schedulers( + self, + iteration: int, + ) -> list[Any]: + return [] + + @property + def _optimizer_dict( + self, + ) -> dict[str, torch.optim.Optimizer]: + return {"student": self._optimizer} + + @property + def _lr_scheduler_dict(self) -> dict[str, Any]: + return {"student": self._lr_scheduler} + + def get_grad_clip_targets( + self, + iteration: int, + ) -> dict[str, torch.nn.Module]: + return {} # We clip internally. + + def on_train_start(self) -> None: + """Seed RNG and call student on_train_start.""" + from fastvideo.utils import set_random_seed + + set_random_seed(self._seed) + self.cuda_generator = torch.Generator( + device=self.student.device + ).manual_seed(self._seed + self._rank) + self.student.on_train_start() + + # ------------------------------------------------------------------ + # Prepare samples + # ------------------------------------------------------------------ + + def _prepare_samples( + self, + samples: list[dict[str, Any]], + epoch: int, + global_step: int, + ) -> dict[str, torch.Tensor]: + """Collate, compute advantages, and filter.""" + device = self.student.device + + # Collate list of per-batch dicts into one dict. + collated: dict[str, Any] = {} + for k in samples[0]: + first = samples[0][k] + if isinstance(first, dict): + collated[k] = { + sk: torch.cat( + [s[k][sk] for s in samples], + dim=0, + ) + for sk in first + } + else: + collated[k] = torch.cat( + [s[k] for s in samples], dim=0 + ) + samples_t = collated + + # Apply KL penalty. + samples_t["rewards"]["ori_avg"] = samples_t[ + "rewards" + ]["avg"] + kl_penalty = ( + self._kl_reward * samples_t["kl"] + ) + samples_t["rewards"]["avg"] = ( + samples_t["rewards"]["avg"].unsqueeze(-1) + - kl_penalty + ) + + # Broadcast raw rewards to timestep dimension. + num_timesteps = samples_t["kl"].shape[1] + for rn in self._reward_cfg: + raw_key = f"{rn}_raw" + samples_t["rewards"][f"ori_{raw_key}"] = ( + samples_t["rewards"][raw_key] + ) + samples_t["rewards"][raw_key] = ( + samples_t["rewards"][raw_key] + .unsqueeze(-1) + .expand(-1, num_timesteps) + ) + + # Gather rewards / KL across processes. + gathered_rewards = { + k: _gather_tensor(v, self._world_size) + .cpu() + .numpy() + for k, v in samples_t["rewards"].items() + } + gathered_kl = ( + _gather_tensor(samples_t["kl"], self._world_size) + .cpu() + .numpy() + ) + + # Log reward stats. + raw_keys = [ + k + for k in gathered_rewards + if k.endswith("_raw") + and not k.startswith("ori_") + ] + reward_logs = { + f"reward_{k}": float( + gathered_rewards[k].mean() + ) + for k in raw_keys + } + kl_mean = float(gathered_kl.mean()) + reward_logs["kl"] = kl_mean + reward_logs["kl_abs"] = float( + np.abs(gathered_kl).mean() + ) + if self._is_main and self.tracker is not None: + self.tracker.log(reward_logs, global_step) + + # Decode prompts for advantage computation. + prompts = None + if self._per_prompt_stat_tracking: + prompt_ids = ( + _gather_tensor( + samples_t["prompt_ids"], + self._world_size, + ) + .cpu() + .numpy() + ) + prompts = ( + self.student.tokenizer.batch_decode( + prompt_ids, + skip_special_tokens=True, + ) + ) + + # Compute advantages. + advantages, adv_log = compute_advantages( + reward_fn_cfg=self._reward_cfg, + weight_advantages=self._weight_advantages, + per_prompt_stat_tracking=( + self._per_prompt_stat_tracking + ), + kl_reward=self._kl_reward, + samples=samples_t, + gathered_rewards=gathered_rewards, + gathered_kl=gathered_kl, + prompts=prompts, + stat_tracker=self._stat_tracker, + reward_stat_trackers=( + self._reward_stat_trackers + ), + kl_stat_tracker=self._kl_stat_tracker, + ) + if adv_log and self._is_main and self.tracker: + self.tracker.log(adv_log, global_step) + + # Shard advantages back to this rank. + advantages = torch.as_tensor(advantages) + advantages = advantages.reshape( + self._world_size, + -1, + advantages.shape[-1], + )[self._rank].to(device) + samples_t["advantages"] = advantages + + # Cleanup unused fields. + del samples_t["rewards"] + del samples_t["prompt_ids"] + + # Filter zero-advantage samples. + mask = ( + samples_t["advantages"].abs().sum(dim=1) != 0 + ) + num_batches = self._num_batches_per_epoch + true_count = mask.sum() + if true_count == 0: + samples_t["advantages"] = ( + samples_t["advantages"] + ADVANTAGE_EPSILON + ) + mask = ( + samples_t["advantages"].abs().sum(dim=1) + != 0 + ) + true_count = mask.sum() + if true_count % num_batches != 0: + false_idx = torch.where(~mask)[0] + need = num_batches - ( + true_count % num_batches + ) + if len(false_idx) >= need: + g = torch.Generator(device=device) + g.manual_seed(self._seed + epoch) + perm = torch.randperm( + len(false_idx), + device=device, + generator=g, + )[:need] + mask[false_idx[perm]] = True + + global_count = ( + _gather_tensor( + mask.sum().view(1), self._world_size + ) + .sum() + .item() + ) + actual_batch_size = ( + global_count + / (num_batches * self._world_size) + ) + if self._is_main and self.tracker: + self.tracker.log( + {"actual_batch_size": actual_batch_size}, + global_step, + ) + + samples_t = { + k: v[mask] for k, v in samples_t.items() + } + return samples_t + + # ------------------------------------------------------------------ + # PPO training + # ------------------------------------------------------------------ + + def _ppo_train( + self, + samples: dict[str, torch.Tensor], + epoch: int, + global_step: int, + ) -> dict[str, LogScalar]: + """Run inner PPO training loop.""" + device = self.student.device + total_batch_size, num_timesteps = samples[ + "timesteps" + ].shape + num_ts = self._num_train_timesteps + all_info: dict[str, list[float]] = defaultdict( + list + ) + + for inner_epoch in range(self._num_inner_epochs): + # Shuffle samples. + g = torch.Generator(device=device) + g.manual_seed( + self._seed + + epoch * SEED_EPOCH_STRIDE + + inner_epoch + ) + perm = torch.randperm( + total_batch_size, + device=device, + generator=g, + ) + samples = { + k: v[perm] for k, v in samples.items() + } + # Shuffle timestep dimension per-sample. + perms = torch.stack( + [ + torch.arange( + num_timesteps, device=device + ) + for _ in range(total_batch_size) + ] + ) + row_idx = torch.arange( + total_batch_size, device=device + )[:, None] + for key in [ + "timesteps", + "latents", + "next_latents", + "log_probs", + ]: + samples[key] = samples[key][ + row_idx, perms + ] + + # Batch into micro-batches. + micro = ( + total_batch_size + // self._num_batches_per_epoch + ) + batched = { + k: v.reshape( + -1, micro, *v.shape[1:] + ) + for k, v in samples.items() + } + batched_list = [ + dict(zip(batched, x, strict=False)) + for x in zip( + *batched.values(), strict=False + ) + ] + + self.student.transformer.train() + info: dict[str, list] = defaultdict(list) + _ppo_batch_times: list[float] = [] + + if ( + self._accumulate_ppo_microbatches + and len(batched_list) > 1 + ): + probe_args = None + for j in self._train_timesteps: + if self._optimizer_step_per_timestep: + self._optimizer.zero_grad() + + for sample in batched_list: + torch.cuda.synchronize() + _ppo_batch_t0 = time.perf_counter() + embeds = sample["prompt_embeds"] + neg_embeds = ( + self._train_neg_embeds[ + : len(embeds) + ] + if self._use_cfg + else None + ) + loss, loss_metrics = ( + self._compute_ppo_loss_and_metrics( + sample, j, embeds, neg_embeds + ) + ) + loss_scale = len(batched_list) + if not self._optimizer_step_per_timestep: + loss_scale *= num_ts + + timestep_j = sample["timesteps"][:, j] + with set_forward_context( + current_timestep=timestep_j, + attn_metadata=None, + ): + (loss / loss_scale).backward() + + for key, value in loss_metrics.items(): + info[key].append(value) + if probe_args is None: + probe_args = ( + sample, + j, + embeds, + neg_embeds, + ) + torch.cuda.synchronize() + _ppo_batch_times.append( + time.perf_counter() - _ppo_batch_t0 + ) + + if self._optimizer_step_per_timestep: + grad_norm = clip_grad_norm_if_needed( + self.student.transformer, + self._max_grad_norm, + ) + self._optimizer.step() + self._lr_scheduler.step() + info["grad_norm"].append(grad_norm) + info["learning_rate"].append( + float( + self._optimizer.param_groups[0]["lr"] + ) + ) + if probe_args is not None: + self._log_post_update_probe( + info, *probe_args + ) + self._optimizer.zero_grad() + + if not self._optimizer_step_per_timestep: + grad_norm = clip_grad_norm_if_needed( + self.student.transformer, + self._max_grad_norm, + ) + self._optimizer.step() + self._lr_scheduler.step() + info["grad_norm"].append(grad_norm) + info["learning_rate"].append( + float(self._optimizer.param_groups[0]["lr"]) + ) + if probe_args is not None: + self._log_post_update_probe( + info, *probe_args + ) + self._optimizer.zero_grad() + + if _ppo_batch_times: + logger.info( + "[GenRL PPO] inner_epoch=%d " + "accumulated_micro_batches=%d " + "micro_batch_times=%s " + "total=%.1fs", + inner_epoch, + len(batched_list), + [f"{t:.1f}s" for t in _ppo_batch_times], + sum(_ppo_batch_times), + ) + for k, v in info.items(): + all_info[k].extend(v) + continue + + for sample in batched_list: + torch.cuda.synchronize() + _ppo_batch_t0 = time.perf_counter() + # Get embeddings. + embeds = sample["prompt_embeds"] + neg_embeds = ( + self._train_neg_embeds[ + : len(embeds) + ] + if self._use_cfg + else None + ) + + self._optimizer.zero_grad() + + for j in self._train_timesteps: + if self._optimizer_step_per_timestep: + self._optimizer.zero_grad() + + loss, loss_metrics = ( + self._compute_ppo_loss_and_metrics( + sample, j, embeds, neg_embeds + ) + ) + + # Backward with gradient accumulation. + timestep_j = sample["timesteps"][ + :, j + ] + with set_forward_context( + current_timestep=timestep_j, + attn_metadata=None, + ): + if self._optimizer_step_per_timestep: + loss.backward() + else: + (loss / num_ts).backward() + + for key, value in loss_metrics.items(): + info[key].append(value) + + if self._optimizer_step_per_timestep: + grad_norm = clip_grad_norm_if_needed( + self.student.transformer, + self._max_grad_norm, + ) + self._optimizer.step() + self._lr_scheduler.step() + info["grad_norm"].append(grad_norm) + info["learning_rate"].append( + float( + self._optimizer.param_groups[0]["lr"] + ) + ) + self._log_post_update_probe( + info, + sample, + j, + embeds, + neg_embeds, + ) + self._optimizer.zero_grad() + + if not self._optimizer_step_per_timestep: + # Clip + step after accumulating all train timesteps. + grad_norm = clip_grad_norm_if_needed( + self.student.transformer, + self._max_grad_norm, + ) + self._optimizer.step() + self._lr_scheduler.step() + info["grad_norm"].append(grad_norm) + info["learning_rate"].append( + float(self._optimizer.param_groups[0]["lr"]) + ) + self._optimizer.zero_grad() + torch.cuda.synchronize() + _ppo_batch_times.append( + time.perf_counter() - _ppo_batch_t0 + ) + + if _ppo_batch_times: + logger.info( + "[GenRL PPO] inner_epoch=%d " + "micro_batch_times=%s " + "total=%.1fs", + inner_epoch, + [f"{t:.1f}s" for t in _ppo_batch_times], + sum(_ppo_batch_times), + ) + # Aggregate info for this inner epoch. + for k, v in info.items(): + all_info[k].extend(v) + + # Aggregate across inner epochs. + metrics: dict[str, LogScalar] = {} + for k, vals in all_info.items(): + metrics[k] = float(np.mean(vals)) + metrics["inner_epochs"] = float( + self._num_inner_epochs + ) + return metrics + + def _log_post_update_probe( + self, + info: dict[str, list], + sample: dict[str, torch.Tensor], + j: int, + embeds: torch.Tensor, + neg_embeds: torch.Tensor | None, + ) -> None: + """Recompute one log-prob after optimizer.step for diagnostics.""" + if not self._log_post_update_kl: + return + + with torch.no_grad(): + ( + _, + post_log_prob, + _, + _, + _, + _, + _, + ) = compute_log_prob( + self.student, + self._scheduler, + sample, + j, + embeds, + neg_embeds, + self._guidance_scale, + self._use_cfg, + self._noise_level, + self._sde_type, + self._diffusion_clip, + self._diffusion_clip_value, + ) + delta = post_log_prob - sample["log_probs"][:, j] + info["post_update_approx_kl"].append( + 0.5 * torch.mean(delta**2).detach().item() + ) + info["post_update_logprob_delta_abs"].append( + torch.mean(torch.abs(delta)).detach().item() + ) + + def _compute_ppo_loss_and_metrics( + self, + sample: dict[str, torch.Tensor], + j: int, + embeds: torch.Tensor, + neg_embeds: torch.Tensor | None, + ) -> tuple[torch.Tensor, dict[str, float]]: + """Compute one PPO/GRPO loss term and detached diagnostics.""" + prev_mean_ref = None + dt_sqrt_ref = None + if self._beta > 0: + ref_model, ref_ctx = self._get_reference_logprob_context() + if ref_model is not None: + with torch.no_grad(), ref_ctx: + ( + _, + _, + prev_mean_ref, + _, + dt_sqrt_ref, + _, + _, + ) = compute_log_prob( + ref_model, + self._scheduler, + sample, + j, + embeds, + neg_embeds, + self._guidance_scale, + self._use_cfg, + self._noise_level, + self._sde_type, + self._diffusion_clip, + self._diffusion_clip_value, + ) + + ( + _prev_sample, + log_prob, + prev_sample_mean, + std_dev_t, + dt_sqrt, + sigma, + sigma_max, + ) = compute_log_prob( + self.student, + self._scheduler, + sample, + j, + embeds, + neg_embeds, + self._guidance_scale, + self._use_cfg, + self._noise_level, + self._sde_type, + self._diffusion_clip, + self._diffusion_clip_value, + ) + + advantages = torch.clamp( + sample["advantages"][:, j], + -self._adv_clip_max, + self._adv_clip_max, + ) + logprob_delta = log_prob - sample["log_probs"][:, j] + ratio = torch.exp(logprob_delta) + unclipped = -advantages * ratio + clipped = -advantages * torch.clamp( + ratio, + 1.0 - self._clip_range, + 1.0 + self._clip_range, + ) + policy_loss = torch.mean( + torch.maximum(unclipped, clipped) + ) + + rw_scale, rw_scale_kl = self._compute_reweight_scales( + sigma=sigma, + sigma_max=sigma_max, + dt_sqrt=dt_sqrt, + ) + + metrics = { + "approx_kl": ( + 0.5 * torch.mean(logprob_delta**2).detach().item() + ), + "logprob_delta_abs": ( + torch.mean(torch.abs(logprob_delta)).detach().item() + ), + "advantage_abs": ( + torch.mean(torch.abs(advantages)).detach().item() + ), + "rw_scale": float( + rw_scale.detach().item() + if isinstance(rw_scale, torch.Tensor) + else rw_scale + ), + "clip_frac": ( + torch.mean( + (torch.abs(ratio - 1.0) > self._clip_range).float() + ) + .detach() + .item() + ), + "clip_frac_gt_one": ( + torch.mean( + (ratio - 1.0 > self._clip_range).float() + ) + .detach() + .item() + ), + "clip_frac_lt_one": ( + torch.mean( + (1.0 - ratio > self._clip_range).float() + ) + .detach() + .item() + ), + "policy_loss": policy_loss.detach().item(), + } + + if self._beta > 0 and prev_mean_ref is not None: + if self._sde_type == "flow_sde": + kl_denom = (std_dev_t * dt_sqrt_ref) ** 2 + elif self._sde_type == "flow_cps": + kl_denom = 0.5 + else: + msg = f"Unknown sde_type: {self._sde_type}" + raise ValueError(msg) + kl_loss = ((prev_sample_mean - prev_mean_ref) ** 2).mean( + dim=(1, 2, 3), + keepdim=True, + ) / (2 * kl_denom) + kl_loss = torch.mean(kl_loss) + loss = ( + rw_scale * policy_loss + + self._beta * kl_loss * rw_scale_kl + ) + metrics["kl_loss"] = kl_loss.detach().item() + else: + loss = rw_scale * policy_loss + + metrics["loss"] = loss.detach().item() + return loss, metrics + + def _compute_reweight_scales( + self, + *, + sigma: torch.Tensor, + sigma_max: float, + dt_sqrt: torch.Tensor, + ) -> tuple[torch.Tensor | float, torch.Tensor | float]: + """Return policy/KL loss scales for the selected timestep.""" + if self._loss_reweighting == "flash_tgr": + return 1.0, 1.0 + + if ( + self._loss_reweighting != "longcat" + or self._sde_type != "flow_sde" + ): + return 1.0, 1.0 + + safe_sigma = torch.where( + sigma == 1, + torch.tensor( + sigma_max, + device=sigma.device, + dtype=sigma.dtype, + ), + sigma, + ) + rw_scale = torch.sqrt( + sigma / (1 - safe_sigma) + ) / dt_sqrt + rw_scale = torch.mean(rw_scale) + if self._loss_reweighting_clip is not None: + rw_scale = torch.clamp( + rw_scale, + max=self._loss_reweighting_clip, + ) + return rw_scale, rw_scale**2 + + # ------------------------------------------------------------------ + # Reference model + # ------------------------------------------------------------------ + + def _get_ref_model(self): + """Get reference model for KL computation.""" + if self._reference is not None: + return self._reference + # LoRA case: caller should use disable_adapter. + return None + + def _get_lora_ref_transformer(self): + """Get LoRA transformer that can disable adapters for sampling KL.""" + transformer = getattr(self.student, "transformer", None) + if transformer is None: + return None + if hasattr(transformer, "disable_adapter"): + return transformer + return None + + def _has_reference_policy(self) -> bool: + return ( + self._reference is not None + or self._get_lora_ref_transformer() is not None + ) + + def _get_reference_logprob_context(self): + """Return model/context for reference log-prob computation.""" + if self._reference is not None: + return self._reference, contextlib.nullcontext() + + lora_ref = self._get_lora_ref_transformer() + if lora_ref is not None: + return self.student, lora_ref.disable_adapter() + + return None, contextlib.nullcontext() diff --git a/fastvideo/train/methods/rl/reward/HPSv3 b/fastvideo/train/methods/rl/reward/HPSv3 new file mode 160000 index 0000000000..a2eb2ef2c7 --- /dev/null +++ b/fastvideo/train/methods/rl/reward/HPSv3 @@ -0,0 +1 @@ +Subproject commit a2eb2ef2c7b5d91a566347a5825cf6d872122149 diff --git a/fastvideo/train/methods/rl/reward/VideoAlign b/fastvideo/train/methods/rl/reward/VideoAlign new file mode 160000 index 0000000000..aba26b658f --- /dev/null +++ b/fastvideo/train/methods/rl/reward/VideoAlign @@ -0,0 +1 @@ +Subproject commit aba26b658fec7d9fd30c295187b548ea673c8769 diff --git a/fastvideo/train/methods/rl/reward/__init__.py b/fastvideo/train/methods/rl/reward/__init__.py new file mode 100644 index 0000000000..9c37af3a5c --- /dev/null +++ b/fastvideo/train/methods/rl/reward/__init__.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reward functions for RL-based video generation training.""" + +from fastvideo.train.methods.rl.reward.hpsv3 import ( + hpsv3_general_score, + hpsv3_percentile_score, +) +from fastvideo.train.methods.rl.reward.ocr import ( + video_ocr_score, +) +from fastvideo.train.methods.rl.reward.videoalign import ( + videoalign_mq_score, + videoalign_ta_score, +) + +__all__ = [ + "hpsv3_general_score", + "hpsv3_percentile_score", + "video_ocr_score", + "videoalign_mq_score", + "videoalign_ta_score", +] diff --git a/fastvideo/train/methods/rl/reward/hpsv3.py b/fastvideo/train/methods/rl/reward/hpsv3.py new file mode 100644 index 0000000000..a32402b329 --- /dev/null +++ b/fastvideo/train/methods/rl/reward/hpsv3.py @@ -0,0 +1,313 @@ +# SPDX-License-Identifier: Apache-2.0 +"""HPSv3 reward functions for visual quality assessment.""" + +from __future__ import annotations + +import os +import sys +import tempfile +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from fastvideo.logger import init_logger +from fastvideo.train.methods.rl.reward.utils import ( + prepare_images, +) + +logger = init_logger(__name__) + +# Prefer local HPSv3 submodule over site-packages. +_HPSV3_ROOT = Path(__file__).resolve().parent / "HPSv3" +if _HPSV3_ROOT.exists(): + _hpsv3_path = str(_HPSV3_ROOT) + if _hpsv3_path not in sys.path: + sys.path.insert(0, _hpsv3_path) + +# Global cache of HPSv3 inferencers keyed by device. +_HPSV3_INFERENCERS: dict[str, Any] = {} +_HPSV3_LOAD_PATCHED = False + + +def _patch_transformers_video_input_alias() -> None: + """Keep HPSv3 compatible with newer transformers releases. + + HPSv3 imports ``VideoInput`` from ``transformers.image_utils`` for type + annotations. Some transformers versions used by FastVideo no longer + export that alias, even though the runtime image utilities HPSv3 needs are + still present. + """ + from transformers import image_utils + + if not hasattr(image_utils, "VideoInput"): + image_utils.VideoInput = image_utils.ImageInput + + +def _remap_hpsv3_state_dict(state_dict: dict[str, Any]) -> dict[str, Any]: + """Adapt HPSv3 checkpoints saved with older Qwen2-VL key names.""" + remapped = {} + for key, value in state_dict.items(): + if key.startswith("visual."): + key = f"model.{key}" + elif key.startswith("model.layers."): + key = f"model.language_model.{key[len('model.'):]}" + elif key.startswith("model.embed_tokens."): + key = f"model.language_model.{key[len('model.'):]}" + elif key.startswith("model.norm."): + key = f"model.language_model.{key[len('model.'):]}" + + key = key.replace( + "base_model.model.visual.", + "base_model.model.model.visual.", + 1, + ) + key = key.replace( + "base_model.model.model.layers.", + "base_model.model.model.language_model.layers.", + 1, + ) + key = key.replace( + "base_model.model.model.embed_tokens.", + "base_model.model.model.language_model.embed_tokens.", + 1, + ) + key = key.replace( + "base_model.model.model.norm.", + "base_model.model.model.language_model.norm.", + 1, + ) + remapped[key] = value + return remapped + + +def _walk_model_graph(model: Any): + """Yield common wrapper/base model objects without importing PEFT.""" + stack = [model] + seen = set() + while stack: + current = stack.pop() + if current is None or id(current) in seen: + continue + seen.add(id(current)) + yield current + for attr in ("base_model", "model"): + child = getattr(current, attr, None) + if child is not None: + stack.append(child) + + +def _patch_load_state_dict(cls: Any) -> None: + """Patch a model class to accept old Qwen2-VL checkpoint keys.""" + if getattr(cls, "_fastvideo_qwen2vl_key_remap", False): + return + + original_load_state_dict = cls.load_state_dict + + def load_state_dict_with_key_remap( + self, + state_dict, + strict=True, + assign=False, + ): + state_dict = _remap_hpsv3_state_dict(state_dict) + return original_load_state_dict( + self, + state_dict, + strict=strict, + assign=assign, + ) + + cls.load_state_dict = load_state_dict_with_key_remap + cls._fastvideo_qwen2vl_key_remap = True + + +def _patch_hpsv3_state_dict_loader() -> None: + """Patch HPSv3 reward model loading for transformers key drift.""" + global _HPSV3_LOAD_PATCHED + if _HPSV3_LOAD_PATCHED: + return + + from hpsv3.model.qwen2vl_trainer import Qwen2VLRewardModelBT + + _patch_load_state_dict(Qwen2VLRewardModelBT) + try: + from peft import PeftModel + except ImportError: + PeftModel = None + if PeftModel is not None: + _patch_load_state_dict(PeftModel) + _HPSV3_LOAD_PATCHED = True + + +def _patch_hpsv3_runtime_model(model: Any) -> None: + """Add aliases expected by HPSv3's older Qwen2-VL forward.""" + for candidate in _walk_model_graph(model): + language_model = getattr(candidate, "language_model", None) + if ( + language_model is not None + and not hasattr(candidate, "embed_tokens") + and hasattr(language_model, "embed_tokens") + ): + candidate.embed_tokens = language_model.embed_tokens + + +def _normalize_device(device) -> str: + if isinstance(device, torch.device): + return str(device) + return str(torch.device(device)) + + +def _move_hpsv3_inferencer(inferencer: Any, device) -> None: + """Move an HPSv3 inferencer across devices. + + HPSv3RewardInferencer does not expose ``.to()``, but it stores its torch + module on ``.model`` and reads ``.device`` when preparing batches. + """ + device_str = _normalize_device(device) + model = getattr(inferencer, "model", None) + if model is not None and hasattr(model, "to"): + model.to(device) + inferencer.device = device_str + + +def set_hpsv3_device(device) -> None: + """Move cached HPSv3 inferencer to given device.""" + key = _normalize_device(device) + if key in _HPSV3_INFERENCERS: + return + # Move from any existing device. + for old_key, inf in list(_HPSV3_INFERENCERS.items()): + if old_key != key: + _move_hpsv3_inferencer(inf, device) + _HPSV3_INFERENCERS[key] = inf + del _HPSV3_INFERENCERS[old_key] + return + + +def _get_hpsv3_inferencer(device): + """Get or create HPSv3 inferencer for device.""" + key = _normalize_device(device) + if key not in _HPSV3_INFERENCERS: + try: + _patch_transformers_video_input_alias() + from hpsv3 import HPSv3RewardInferencer + _patch_hpsv3_state_dict_loader() + except ImportError as exc: + msg = ( + "Failed to import HPSv3. Ensure the HPSv3 submodule is " + "checked out under fastvideo/train/methods/rl/reward/HPSv3 " + "and that its transformers dependencies are compatible." + ) + raise ImportError(msg) from exc + inf = HPSv3RewardInferencer(device=device) + _patch_hpsv3_runtime_model(inf.model) + _HPSV3_INFERENCERS[key] = inf + return _HPSV3_INFERENCERS[key] + + +def _save_frame_to_temp(frame: np.ndarray) -> str: + """Save a frame to a temporary PNG file.""" + from PIL import Image + + fd, path = tempfile.mkstemp(suffix=".png") + os.close(fd) + Image.fromarray(frame).save(path) + return path + + +def _extract_reward_scalar(result) -> float: + """Extract a float from HPSv3 result.""" + if isinstance(result, torch.Tensor): + return float(result.item()) + if isinstance(result, (float, int)): + return float(result) + if isinstance(result, (list, np.ndarray)): + return float(np.mean(result)) + return float(result) + + +def hpsv3_general_score(device): + """Return a reward fn that scores frames with + 'A high-quality image' as prompt. + + Returns mean score across all frames.""" + + def _score(images, prompts, metadata, only_strict=False): + inf = _get_hpsv3_inferencer(device) + images_np = prepare_images(images) + batch_scores = [] + + for b in range(len(images_np)): + frames = images_np[b] + if frames.ndim == 3: + frames = frames[np.newaxis] + frame_scores = [] + for frame in frames: + path = _save_frame_to_temp(frame) + try: + rewards = inf.reward( + ["A high-quality image"], [path] + ) + frame_scores.append( + _extract_reward_scalar(rewards[0][0]) + ) + finally: + os.remove(path) + batch_scores.append(np.mean(frame_scores)) + + reward = torch.tensor( + batch_scores, device=device + ).float() + return {"avg": reward}, {} + + return _score + + +def hpsv3_percentile_score(device): + """Return a reward fn that scores frames with per-prompt + text. Returns mean of top 30% frame scores.""" + + def _score(images, prompts, metadata, only_strict=False): + inf = _get_hpsv3_inferencer(device) + images_np = prepare_images(images) + batch_scores = [] + + for b in range(len(images_np)): + frames = images_np[b] + if frames.ndim == 3: + frames = frames[np.newaxis] + prompt = ( + prompts[b] + if prompts and b < len(prompts) + else "A high-quality image" + ) + frame_scores = [] + for frame in frames: + path = _save_frame_to_temp(frame) + try: + rewards = inf.reward( + [prompt], [path] + ) + frame_scores.append( + _extract_reward_scalar(rewards[0][0]) + ) + finally: + os.remove(path) + # Top 30% percentile. + if frame_scores: + k = max(1, int(len(frame_scores) * 0.3)) + top_k = sorted( + frame_scores, reverse=True + )[:k] + batch_scores.append(np.mean(top_k)) + else: + batch_scores.append(0.0) + + reward = torch.tensor( + batch_scores, device=device + ).float() + return {"avg": reward}, {} + + return _score diff --git a/fastvideo/train/methods/rl/reward/ocr.py b/fastvideo/train/methods/rl/reward/ocr.py new file mode 100644 index 0000000000..ca0ce4d142 --- /dev/null +++ b/fastvideo/train/methods/rl/reward/ocr.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +"""OCR-based reward for video-text alignment.""" + +from __future__ import annotations + +import re + +import numpy as np +import torch + +from fastvideo.logger import init_logger +from fastvideo.train.methods.rl.reward.utils import ( + prepare_images, +) + +logger = init_logger(__name__) + + +def _levenshtein_distance(s1: str, s2: str) -> int: + """Compute Levenshtein edit distance between strings.""" + if len(s1) < len(s2): + return _levenshtein_distance(s2, s1) + if len(s2) == 0: + return len(s1) + prev_row = list(range(len(s2) + 1)) + for i, c1 in enumerate(s1): + curr_row = [i + 1] + for j, c2 in enumerate(s2): + insertions = prev_row[j + 1] + 1 + deletions = curr_row[j] + 1 + substitutions = prev_row[j] + (c1 != c2) + curr_row.append( + min(insertions, deletions, substitutions) + ) + prev_row = curr_row + return prev_row[-1] + + +def _extract_text_from_prompt(prompt: str) -> str: + """Extract expected text from prompt (within quotes).""" + match = re.search(r'["\'](.+?)["\']', prompt) + if match: + return match.group(1) + return prompt + + +def video_ocr_score(): + """Return an OCR-based reward function (CPU).""" + try: + from paddleocr import PaddleOCR + except ImportError as exc: + msg = ( + "paddleocr not installed. " + "Install via: pip install paddleocr" + ) + raise ImportError(msg) from exc + + ocr = PaddleOCR( + use_angle_cls=True, lang="en", use_gpu=False + ) + + def _score(images, prompts, metadata, only_strict=False): + images_np = prepare_images(images) + batch_scores = [] + + for b in range(len(images_np)): + frames = images_np[b] + if frames.ndim == 3: + frames = frames[np.newaxis] + + expected = _extract_text_from_prompt( + prompts[b] if prompts else "" + ).lower() + + # Sample every 4th frame. + sample_indices = list( + range(0, len(frames), 4) + ) + if not sample_indices: + sample_indices = [0] + + best_score = 0.0 + for idx in sample_indices: + frame = frames[idx] + result = ocr.ocr(frame, cls=True) + detected = "" + if result and result[0]: + texts = [ + line[1][0] + for line in result[0] + if line[1] + ] + detected = " ".join(texts).lower() + + if not expected: + score = 1.0 if detected else 0.0 + elif not detected: + score = 0.0 + else: + dist = _levenshtein_distance( + detected, expected + ) + max_len = max( + len(detected), len(expected) + ) + score = 1.0 - (dist / max_len) + best_score = max(best_score, score) + + batch_scores.append(best_score) + + reward = torch.tensor(batch_scores).float() + return {"avg": reward}, {} + + return _score diff --git a/fastvideo/train/methods/rl/reward/utils.py b/fastvideo/train/methods/rl/reward/utils.py new file mode 100644 index 0000000000..661da33028 --- /dev/null +++ b/fastvideo/train/methods/rl/reward/utils.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Utility functions for reward computation.""" + +from __future__ import annotations + +import numpy as np +import torch + + +def prepare_images( + images: torch.Tensor | np.ndarray, +) -> np.ndarray: + """Convert tensor images to uint8 numpy (NHWC or NFHWC). + + Accepts: + - (N, C, H, W) or (N, H, W, C) tensors/arrays + - (N, F, C, H, W) or (N, F, H, W, C) video tensors + Returns: + uint8 numpy array in HWC/FHWC layout. + """ + if isinstance(images, torch.Tensor): + images = images.detach().cpu().numpy() + images = np.asarray(images) + + if images.ndim == 4: + # Image batch: (N, C, H, W) or (N, H, W, C) + if images.shape[-1] in (1, 3): + pass + elif images.shape[1] in (1, 3): + images = images.transpose(0, 2, 3, 1) + elif images.ndim == 5: + # Video batch: (N, F, H, W, C), (N, F, C, H, W), + # or (N, C, F, H, W). Check channel-last first because + # one-frame videos have shape[1] == 1. + if images.shape[-1] in (1, 3): + pass + elif images.shape[2] in (1, 3): + # (N, F, C, H, W) -> (N, F, H, W, C) + images = images.transpose(0, 1, 3, 4, 2) + elif images.shape[1] in (1, 3): + # (N, C, F, H, W) -> (N, F, H, W, C) + images = images.transpose(0, 2, 3, 4, 1) + + if images.dtype == np.float32 or images.dtype == np.float64: + images = np.clip(images * 255, 0, 255).astype( + np.uint8 + ) + return images diff --git a/fastvideo/train/methods/rl/reward/videoalign.py b/fastvideo/train/methods/rl/reward/videoalign.py new file mode 100644 index 0000000000..9e41a91b14 --- /dev/null +++ b/fastvideo/train/methods/rl/reward/videoalign.py @@ -0,0 +1,471 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VideoAlign reward functions for motion quality and +text-video alignment.""" + +from __future__ import annotations + +import os +import sys +import tempfile +from importlib import import_module, util +from typing import Any + +import numpy as np +import torch + +from fastvideo.logger import init_logger +from fastvideo.train.methods.rl.reward.utils import ( + prepare_images, +) + +logger = init_logger(__name__) + +# Add VideoAlign submodule to path for importing. +_VIDEOALIGN_ROOT = os.path.join( + os.path.dirname(__file__), "VideoAlign" +) +if os.path.isdir(_VIDEOALIGN_ROOT): + if _VIDEOALIGN_ROOT not in sys.path: + sys.path.insert(0, _VIDEOALIGN_ROOT) + +# Global cache of VideoAlign inferencers. +_VIDEOALIGN_INFERENCERS: dict[str, Any] = {} +_VIDEOALIGN_PATCHED = False + + +def _normalize_device_str(device) -> str: + if isinstance(device, torch.device): + return str(device) + return str(torch.device(device)) + + +def _move_videoalign_inferencer(inferencer: Any, device) -> None: + """Move a VideoAlign inferencer across devices.""" + device_str = _normalize_device_str(device) + model = getattr(inferencer, "model", None) + if model is not None and hasattr(model, "to"): + model.to(device) + inferencer.device = device_str + + +def _remap_qwen2vl_state_dict_keys( + state_dict: dict[str, Any], +) -> dict[str, Any]: + """Adapt checkpoints saved with older Qwen2-VL key names.""" + remapped = {} + for key, value in state_dict.items(): + if key.startswith("visual."): + key = f"model.{key}" + elif key.startswith("model.layers."): + key = f"model.language_model.{key[len('model.'):]}" + elif key.startswith("model.embed_tokens."): + key = f"model.language_model.{key[len('model.'):]}" + elif key.startswith("model.norm."): + key = f"model.language_model.{key[len('model.'):]}" + + key = key.replace( + "base_model.model.visual.", + "base_model.model.model.visual.", + 1, + ) + key = key.replace( + "base_model.model.model.layers.", + "base_model.model.model.language_model.layers.", + 1, + ) + key = key.replace( + "base_model.model.model.embed_tokens.", + "base_model.model.model.language_model.embed_tokens.", + 1, + ) + key = key.replace( + "base_model.model.model.norm.", + "base_model.model.model.language_model.norm.", + 1, + ) + remapped[key] = value + return remapped + + +def _walk_model_graph(model: Any): + """Yield common wrapper/base model objects without importing PEFT.""" + stack = [model] + seen = set() + while stack: + current = stack.pop() + if current is None or id(current) in seen: + continue + seen.add(id(current)) + yield current + for attr in ("base_model", "model"): + child = getattr(current, attr, None) + if child is not None: + stack.append(child) + + +def _patch_load_state_dict(cls: Any) -> None: + """Patch a model class to accept old VideoAlign checkpoint keys.""" + if getattr(cls, "_fastvideo_qwen2vl_key_remap", False): + return + + original_load_state_dict = cls.load_state_dict + + def load_state_dict_with_key_remap( + self, + state_dict, + strict=True, + assign=False, + ): + state_dict = _remap_qwen2vl_state_dict_keys(state_dict) + if not assign: + try: + assign = any( + getattr(param, "is_meta", False) + for param in self.parameters() + ) + except Exception: + assign = False + return original_load_state_dict( + self, + state_dict, + strict=strict, + assign=assign, + ) + + cls.load_state_dict = load_state_dict_with_key_remap + cls._fastvideo_qwen2vl_key_remap = True + + +def _select_videoalign_frame_indices( + vision_mod: Any, + ele: dict[str, Any], + total_frames: int, + video_fps: float, +) -> list[int]: + sample_type = ele.get("sample_type", "uniform") + if sample_type == "uniform": + nframes = vision_mod.smart_nframes( + ele, + total_frames=total_frames, + video_fps=video_fps, + ) + return torch.linspace( + 0, + total_frames - 1, + nframes, + ).round().long().tolist() + if sample_type == "multi_pts": + frames_each_pts = 6 + num_pts = 4 + fps = 8 + nframes = max( + frames_each_pts, + int(total_frames * fps // video_fps), + ) + frame_idx = torch.linspace( + 0, + total_frames - 1, + nframes, + ).round().long().tolist() + start_pt = int(frames_each_pts // 2) + end_pt = int(nframes - frames_each_pts // 2 - 1) + pts = torch.linspace( + start_pt, + end_pt, + num_pts, + ).round().long().tolist() + idx = [] + for pt in pts: + idx.extend( + frame_idx[ + pt - frames_each_pts // 2: + pt + frames_each_pts // 2 + ] + ) + return idx + raise ValueError(f"Unsupported VideoAlign sample_type: {sample_type}") + + +def _read_video_opencv( + vision_mod: Any, + ele: dict[str, Any], +) -> torch.Tensor: + """Read local MP4s without relying on torchvision.io.read_video.""" + import cv2 + + video_path = ele["video"] + if video_path.startswith("file://"): + video_path = video_path[7:] + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise ValueError(f"Could not open video: {video_path}") + + video_fps = float(cap.get(cv2.CAP_PROP_FPS) or 30.0) + total_file_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0) + start_frame = max( + 0, + int(round(float(ele.get("video_start", 0.0) or 0.0) * video_fps)), + ) + end_sec = ele.get("video_end") + if end_sec is None: + end_frame = total_file_frames if total_file_frames > 0 else None + else: + end_frame = int(round(float(end_sec) * video_fps)) + if total_file_frames > 0: + end_frame = min(end_frame, total_file_frames) + + cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame) + frames = [] + current_frame = start_frame + while end_frame is None or current_frame < end_frame: + ok, frame = cap.read() + if not ok: + break + frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + current_frame += 1 + cap.release() + + if not frames: + raise ValueError(f"No frames were read from video: {video_path}.") + + idx = _select_videoalign_frame_indices( + vision_mod, + ele, + total_frames=len(frames), + video_fps=video_fps, + ) + video = np.stack([frames[i] for i in idx], axis=0) + return torch.from_numpy(video).permute(0, 3, 1, 2) + + +def _torchvision_read_video_available() -> bool: + try: + torchvision_io = import_module("torchvision.io") + except Exception: + return False + return hasattr(torchvision_io, "read_video") + + +def _patch_videoalign_video_reader() -> None: + """Register an OpenCV reader for torchvision builds without read_video.""" + vision_mod = import_module("vision_process") + if "opencv" not in vision_mod.VIDEO_READER_BACKENDS: + + def read_video_opencv(ele): + return _read_video_opencv(vision_mod, ele) + + vision_mod.VIDEO_READER_BACKENDS["opencv"] = read_video_opencv + + if _torchvision_read_video_available(): + return + + vision_mod.FORCE_QWENVL_VIDEO_READER = "opencv" + if hasattr(vision_mod.get_video_reader_backend, "cache_clear"): + vision_mod.get_video_reader_backend.cache_clear() + + +def _patch_videoalign_modules() -> Any: + """Patch VideoAlign for the FastVideo dependency set.""" + global _VIDEOALIGN_PATCHED + + inference_mod = import_module("inference") + if _VIDEOALIGN_PATCHED: + return inference_mod + + train_reward_mod = import_module("train_reward") + trainer_mod = import_module("trainer") + _patch_videoalign_video_reader() + + if util.find_spec("flash_attn") is None: + for mod in (train_reward_mod, inference_mod): + original_create = mod.create_model_and_processor + + def create_model_and_processor_sdpa( + *args, + _original_create=original_create, + **kwargs, + ): + training_args = kwargs.get("training_args") + if training_args is not None: + training_args.disable_flash_attn2 = True + return _original_create(*args, **kwargs) + + mod.create_model_and_processor = create_model_and_processor_sdpa + + _patch_load_state_dict(trainer_mod.Qwen2VLRewardModelBT) + try: + peft_mod = import_module("peft") + except ImportError: + peft_mod = None + if peft_mod is not None: + _patch_load_state_dict(peft_mod.PeftModel) + + _VIDEOALIGN_PATCHED = True + return inference_mod + + +def _patch_videoalign_runtime_model(model: Any) -> None: + """Add aliases expected by VideoAlign's older Qwen2-VL forward.""" + for candidate in _walk_model_graph(model): + language_model = getattr(candidate, "language_model", None) + if ( + language_model is not None + and not hasattr(candidate, "embed_tokens") + and hasattr(language_model, "embed_tokens") + ): + candidate.embed_tokens = language_model.embed_tokens + + +def set_videoalign_device(device) -> None: + """Move cached VideoAlign inferencers to device.""" + key = _normalize_device_str(device) + for old_key, inf in list( + _VIDEOALIGN_INFERENCERS.items() + ): + if old_key != key and old_key.split(":")[0] != key: + new_key = inf._key_prefix + ":" + key + _move_videoalign_inferencer(inf, device) + _VIDEOALIGN_INFERENCERS[new_key] = inf + del _VIDEOALIGN_INFERENCERS[old_key] + + +def _get_inferencer( + device, + checkpoint_path: str | None = None, +): + """Get or create VideoAlign inferencer.""" + if checkpoint_path is None: + checkpoint_path = os.environ.get( + "VIDEOALIGN_CHECKPOINT_PATH", + os.path.join( + os.path.dirname(__file__), + "VideoAlign", + "checkpoints", + ), + ) + checkpoint_path = os.path.abspath(checkpoint_path) + + key = _normalize_device_str(device) + cache_key = f"{checkpoint_path}:{key}" + if cache_key not in _VIDEOALIGN_INFERENCERS: + try: + inference_mod = _patch_videoalign_modules() + VideoVLMRewardInference = inference_mod.VideoVLMRewardInference + except ImportError as exc: + msg = ( + "VideoAlign not found. Ensure the " + "VideoAlign submodule is checked out " + "under fastvideo/train/methods/rl/" + "reward/VideoAlign" + ) + raise ImportError(msg) from exc + + inf = VideoVLMRewardInference( + load_from_pretrained=checkpoint_path, + device=device, + ) + _patch_videoalign_runtime_model(inf.model) + inf._key_prefix = checkpoint_path or "default" + _VIDEOALIGN_INFERENCERS[cache_key] = inf + return _VIDEOALIGN_INFERENCERS[cache_key] + + +def _convert_to_grayscale( + frames: np.ndarray, +) -> np.ndarray: + """Convert FHWC frames to grayscale FHWC.""" + if frames.ndim == 4 and frames.shape[-1] == 3: + gray = np.mean(frames, axis=-1, keepdims=True) + return np.repeat(gray.astype(np.uint8), 3, axis=-1) + return frames + + +def _save_video_to_temp( + frames: np.ndarray, + fps: int = 8, +) -> str: + """Save frames to a temporary MP4 file.""" + import cv2 + + fd, path = tempfile.mkstemp(suffix=".mp4") + os.close(fd) + + h, w = frames.shape[1], frames.shape[2] + fourcc = cv2.VideoWriter.fourcc(*"mp4v") + writer = cv2.VideoWriter(path, fourcc, fps, (w, h)) + for frame in frames: + bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) + writer.write(bgr) + writer.release() + return path + + +def videoalign_mq_score( + device, + checkpoint_path: str | None = None, +): + """Return Motion Quality reward fn (grayscale).""" + + def _score(images, prompts, metadata, only_strict=False): + inf = _get_inferencer(device, checkpoint_path) + images_np = prepare_images(images) + batch_scores = [] + + for b in range(len(images_np)): + frames = images_np[b] + if frames.ndim == 3: + frames = frames[np.newaxis] + gray_frames = _convert_to_grayscale(frames) + path = _save_video_to_temp(gray_frames) + try: + results = inf.reward( + [path], [""], use_norm=True + ) + mq = float(results[0].get("MQ", 0)) + batch_scores.append(mq) + finally: + os.remove(path) + + reward = torch.tensor( + batch_scores, device=device + ).float() + return {"avg": reward}, {} + + return _score + + +def videoalign_ta_score( + device, + checkpoint_path: str | None = None, +): + """Return Text-Video Alignment reward fn (color).""" + + def _score(images, prompts, metadata, only_strict=False): + inf = _get_inferencer(device, checkpoint_path) + images_np = prepare_images(images) + batch_scores = [] + + for b in range(len(images_np)): + frames = images_np[b] + if frames.ndim == 3: + frames = frames[np.newaxis] + prompt = ( + prompts[b] if prompts and b < len(prompts) + else "" + ) + path = _save_video_to_temp(frames) + try: + results = inf.reward( + [path], [prompt], use_norm=True + ) + ta = float(results[0].get("TA", 0)) + batch_scores.append(ta) + finally: + os.remove(path) + + reward = torch.tensor( + batch_scores, device=device + ).float() + return {"avg": reward}, {} + + return _score diff --git a/fastvideo/train/methods/rl/utils/__init__.py b/fastvideo/train/methods/rl/utils/__init__.py new file mode 100644 index 0000000000..e4282003ae --- /dev/null +++ b/fastvideo/train/methods/rl/utils/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Utility modules for RL training.""" diff --git a/fastvideo/train/methods/rl/utils/advantages.py b/fastvideo/train/methods/rl/utils/advantages.py new file mode 100644 index 0000000000..b4d935af9f --- /dev/null +++ b/fastvideo/train/methods/rl/utils/advantages.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Advantage computation for RL training.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from fastvideo.logger import init_logger +from fastvideo.train.methods.rl.utils.stat_tracking import ( + EPSILON, + PerPromptStatTracker, +) + +logger = init_logger(__name__) + + +def _normalize_rewards( + rewards: np.ndarray, + epsilon: float = EPSILON, +) -> np.ndarray: + """Normalize rewards to zero mean and unit variance.""" + return (rewards - rewards.mean()) / ( + rewards.std() + epsilon + ) + + +def _compute_kl_advantages( + gathered_kl: np.ndarray, + kl_stat_tracker: PerPromptStatTracker | None, + prompts: list[str] | None, + use_per_prompt: bool, +) -> np.ndarray: + """Compute KL advantages (negative = penalty).""" + if use_per_prompt and kl_stat_tracker is not None: + return kl_stat_tracker.update( + prompts, -gathered_kl + ) + return _normalize_rewards(-gathered_kl) + + +def calculate_zero_std_ratio( + prompts, + gathered_rewards: dict[str, np.ndarray], + reward_key: str = "avg", +) -> float: + """Compute fraction of prompts with zero std.""" + prompts_arr = np.array(prompts) + rewards = gathered_rewards.get(reward_key) + if rewards is None: + return 0.0 + unique = np.unique(prompts_arr) + zero_count = 0 + for p in unique: + r = rewards[prompts_arr == p] + if np.std(r) < EPSILON: + zero_count += 1 + return zero_count / max(len(unique), 1) + + +def compute_advantages( + reward_fn_cfg: dict[str, float], + weight_advantages: bool, + per_prompt_stat_tracking: bool, + kl_reward: float, + samples: dict[str, Any], + gathered_rewards: dict[str, np.ndarray], + gathered_kl: np.ndarray, + prompts: list[str] | None, + stat_tracker: PerPromptStatTracker | None, + reward_stat_trackers: ( + dict[str, PerPromptStatTracker] | None + ), + kl_stat_tracker: PerPromptStatTracker | None, +) -> tuple[np.ndarray, dict[str, Any]]: + """Compute advantages from gathered rewards and KL. + + Supports two modes: + - Mode 1 (default): Weight rewards, then advantages. + - Mode 2 (weight_advantages=True): Per-reward + advantages, then weight. + + Returns: + (advantages, log_dict) + """ + log_dict: dict[str, Any] = {} + + if weight_advantages: + if per_prompt_stat_tracking: + if reward_stat_trackers is None: + msg = ( + "reward_stat_trackers required when " + "weight_advantages=True and " + "per_prompt_stat_tracking=True" + ) + raise ValueError(msg) + + weighted_list = [] + for reward_name in reward_fn_cfg: + raw_key = f"{reward_name}_raw" + adv = reward_stat_trackers[ + reward_name + ].update( + prompts, gathered_rewards[raw_key] + ) + weight = reward_fn_cfg[reward_name] + weighted_list.append(adv * weight) + + if kl_reward > 0: + if kl_stat_tracker is None: + msg = ( + "kl_stat_tracker required when " + "weight_advantages=True and " + "kl_reward > 0" + ) + raise ValueError(msg) + kl_adv = _compute_kl_advantages( + gathered_kl, + kl_stat_tracker, + prompts, + use_per_prompt=True, + ) + weighted_list.append( + kl_adv * kl_reward + ) + + advantages = sum(weighted_list) + + first_name = next(iter(reward_fn_cfg)) + group_size, trained_num = ( + reward_stat_trackers[ + first_name + ].get_stats() + ) + zero_std_ratios = {} + for rn in reward_fn_cfg: + raw_key = f"{rn}_raw" + zero_std_ratios[ + f"zero_std_ratio_{rn}" + ] = calculate_zero_std_ratio( + prompts, + gathered_rewards, + reward_key=f"ori_{raw_key}", + ) + log_dict = { + "group_size": group_size, + "trained_prompt_num": trained_num, + **zero_std_ratios, + } + for t in reward_stat_trackers.values(): + t.clear() + if kl_stat_tracker is not None: + kl_stat_tracker.clear() + else: + weighted_list = [] + for reward_name in reward_fn_cfg: + raw_key = f"{reward_name}_raw" + raw = gathered_rewards[raw_key] + adv = _normalize_rewards(raw) + weight = reward_fn_cfg[reward_name] + weighted_list.append(adv * weight) + + if kl_reward > 0: + kl_adv = _compute_kl_advantages( + gathered_kl, + None, + None, + use_per_prompt=False, + ) + weighted_list.append( + kl_adv * kl_reward + ) + + advantages = sum(weighted_list) + + elif per_prompt_stat_tracking: + if stat_tracker is None: + msg = ( + "stat_tracker required when " + "per_prompt_stat_tracking=True" + ) + raise ValueError(msg) + advantages = stat_tracker.update( + prompts, gathered_rewards["avg"] + ) + group_size, trained_num = ( + stat_tracker.get_stats() + ) + zero_std = calculate_zero_std_ratio( + prompts, gathered_rewards + ) + log_dict = { + "group_size": group_size, + "trained_prompt_num": trained_num, + "zero_std_ratio": zero_std, + } + stat_tracker.clear() + else: + advantages = _normalize_rewards( + gathered_rewards["avg"] + ) + + return advantages, log_dict diff --git a/fastvideo/train/methods/rl/utils/data.py b/fastvideo/train/methods/rl/utils/data.py new file mode 100644 index 0000000000..db54011d40 --- /dev/null +++ b/fastvideo/train/methods/rl/utils/data.py @@ -0,0 +1,254 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Text prompt datasets and samplers for RL training.""" + +from __future__ import annotations + +import json +import os + +import torch +from torch.utils.data import DataLoader, Dataset, Sampler + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +class TextPromptDataset(Dataset): + """Load plain text prompts from train.txt / test.txt.""" + + def __init__(self, dataset: str, split: str = "train"): + self.file_path = os.path.join( + dataset, f"{split}.txt" + ) + with open(self.file_path) as f: + self.prompts = [ + line.strip() for line in f.readlines() + ] + + def __len__(self) -> int: + return len(self.prompts) + + def __getitem__( + self, idx: int | tuple[int, int] + ) -> dict: + epoch_tag = None + if isinstance(idx, tuple): + epoch_tag, idx = idx + return { + "epoch": epoch_tag, + "prompt": self.prompts[idx], + "metadata": {}, + } + + @staticmethod + def collate_fn( + examples: list[dict], + ) -> tuple[int | None, list[str], list[dict]]: + epoch_tags = [ + example.get("epoch") for example in examples + ] + epoch_tag = ( + epoch_tags[0] + if all(tag == epoch_tags[0] for tag in epoch_tags) + else None + ) + prompts = [ + example["prompt"] for example in examples + ] + metadatas = [ + example["metadata"] for example in examples + ] + return epoch_tag, prompts, metadatas + + +class JsonPromptDataset(Dataset): + """Load prompts from JSONL files.""" + + def __init__(self, dataset: str, split: str = "train"): + self.file_path = os.path.join( + dataset, f"{split}.json" + ) + self._prompts: list[str] = [] + self._metadatas: list[dict] = [] + self._load_all_prompts() + + def _load_all_prompts(self): + with open(self.file_path, encoding="utf-8") as f: + for raw_line in f: + line = raw_line.strip() + if not line: + continue + try: + item = json.loads(line) + prompt = item.get("prompt", "") + if prompt: + self._prompts.append(prompt) + metadata = { + k: v + for k, v in item.items() + if k != "prompt" + } + self._metadatas.append(metadata) + except json.JSONDecodeError as e: + logger.warning( + "Skipping invalid JSON line: %s", + e, + ) + + def __len__(self) -> int: + return len(self._prompts) + + def __getitem__( + self, idx: int | tuple[int, int] + ) -> dict: + epoch_tag = None + if isinstance(idx, tuple): + epoch_tag, idx = idx + return { + "epoch": epoch_tag, + "prompt": self._prompts[idx], + "metadata": ( + self._metadatas[idx] + if self._metadatas + else {} + ), + } + + @staticmethod + def collate_fn( + examples: list[dict], + ) -> tuple[int | None, list[str], list[dict]]: + epoch_tags = [ + example.get("epoch") for example in examples + ] + epoch_tag = ( + epoch_tags[0] + if all(tag == epoch_tags[0] for tag in epoch_tags) + else None + ) + prompts = [ + example["prompt"] for example in examples + ] + metadatas = [ + example["metadata"] for example in examples + ] + return epoch_tag, prompts, metadatas + + +class DistributedKRepeatSampler(Sampler): + + def __init__( + self, + dataset: Dataset, + batch_size: int, + k: int, + num_replicas: int, + rank: int, + seed: int = 0, + ): + self.dataset = dataset + self.batch_size = batch_size + self.k = k + self.num_replicas = num_replicas + self.rank = rank + self.seed = seed + self.total_samples = num_replicas * batch_size + if self.batch_size % self.k != 0: + raise ValueError( + "batch_size must be divisible by k so each rank receives " + "whole prompt groups. Got " + f"batch_size={batch_size}, k={k}." + ) + assert self.total_samples % self.k == 0, ( + f"k cannot divide n*b: k={k} " + f"num_replicas={num_replicas} " + f"batch_size={batch_size}" + ) + self.m = self.total_samples // self.k + self.groups_per_rank = self.batch_size // self.k + self.epoch = 0 + + def __iter__(self): + while True: + g = torch.Generator() + g.manual_seed(self.seed + self.epoch) + indices = torch.randperm( + len(self.dataset), generator=g + )[: self.m].tolist() + start = self.rank * self.groups_per_rank + end = start + self.groups_per_rank + rank_groups = indices[start:end] + yield [ + (self.epoch, idx) + for idx in rank_groups + for _ in range(self.k) + ] + + def set_epoch(self, epoch: int): + self.epoch = epoch + + +def build_prompt_dataloaders( + prompt_dataset_path: str, + prompt_fn: str, + sample_batch_size: int, + eval_batch_size: int, + num_video_per_prompt: int, + num_processes: int, + process_index: int, + seed: int, +) -> tuple[DataLoader, DataLoader, DistributedKRepeatSampler]: + """Build train/eval prompt dataloaders. + + Returns: + (train_dataloader, test_dataloader, train_sampler) + """ + if prompt_fn == "general_ocr": + train_ds = TextPromptDataset( + prompt_dataset_path, "train" + ) + test_ds = TextPromptDataset( + prompt_dataset_path, "test" + ) + collate = TextPromptDataset.collate_fn + elif prompt_fn == "filtered_prompts": + train_ds = JsonPromptDataset( + prompt_dataset_path, "train" + ) + test_ds = JsonPromptDataset( + prompt_dataset_path, "test" + ) + collate = JsonPromptDataset.collate_fn + else: + msg = ( + f"Unsupported prompt_fn: {prompt_fn}. " + "Use 'general_ocr' or 'filtered_prompts'." + ) + raise NotImplementedError(msg) + + train_sampler = DistributedKRepeatSampler( + dataset=train_ds, + batch_size=sample_batch_size, + k=num_video_per_prompt, + num_replicas=num_processes, + rank=process_index, + seed=seed, + ) + + train_dl = DataLoader( + train_ds, + batch_sampler=train_sampler, + num_workers=1, + collate_fn=collate, + prefetch_factor=1, + persistent_workers=False, + ) + test_dl = DataLoader( + test_ds, + batch_size=eval_batch_size, + collate_fn=collate, + shuffle=False, + num_workers=8, + ) + return train_dl, test_dl, train_sampler diff --git a/fastvideo/train/methods/rl/utils/diffusion.py b/fastvideo/train/methods/rl/utils/diffusion.py new file mode 100644 index 0000000000..e43a8e4382 --- /dev/null +++ b/fastvideo/train/methods/rl/utils/diffusion.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Single diffusion step for PPO training phase.""" + +from __future__ import annotations + +import torch + +from fastvideo.train.methods.rl.utils.sde import ( + sde_step_with_logprob, +) + + +def compute_log_prob( + model, + scheduler, + sample: dict[str, torch.Tensor], + j: int, + embeds: torch.Tensor, + negative_embeds: torch.Tensor | None, + guidance_scale: float, + use_cfg: bool, + noise_level: float, + sde_type: str, + diffusion_clip: bool = False, + diffusion_clip_value: float = 0.45, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + float, +]: + """Run one diffusion step and return log-probability. + + Uses model.forward_transformer_raw() for the forward + pass. + + Args: + model: WanModel instance. + scheduler: Noise scheduler. + sample: Dict with latents, next_latents, timesteps. + j: Timestep index within the trajectory. + embeds: Conditional text embeddings. + negative_embeds: Unconditional embeddings (or None). + guidance_scale: CFG scale. + use_cfg: Whether to use classifier-free guidance. + noise_level: SDE noise level. + sde_type: 'flow_sde' or 'flow_cps'. + diffusion_clip: Clip SDE variance. + diffusion_clip_value: Clip threshold. + + Returns: + (prev_sample, log_prob, prev_sample_mean, + std_dev_t, dt_sqrt, sigma, sigma_max) + """ + dtype = embeds.dtype + latents_j = sample["latents"][:, j] + timestep_j = sample["timesteps"][:, j] + + if use_cfg and negative_embeds is not None: + noise_pred_text = model.forward_transformer_raw( + latents_j.to(dtype), + timestep_j, + embeds, + ) + noise_pred_uncond = model.forward_transformer_raw( + latents_j.to(dtype), + timestep_j, + negative_embeds, + ) + noise_pred = ( + noise_pred_uncond + + guidance_scale + * (noise_pred_text - noise_pred_uncond) + ) + else: + noise_pred = model.forward_transformer_raw( + latents_j.to(dtype), + timestep_j, + embeds, + ) + + ( + prev_sample, + log_prob, + prev_sample_mean, + std_dev_t, + dt_sqrt, + sigma, + sigma_max, + ) = sde_step_with_logprob( + scheduler, + noise_pred.float(), + timestep_j, + latents_j.float(), + noise_level=noise_level, + prev_sample=sample["next_latents"][:, j].float(), + sde_type=sde_type, + diffusion_clip=diffusion_clip, + diffusion_clip_value=diffusion_clip_value, + return_sqrt_dt_and_std_dev_t=True, + ) + return ( + prev_sample, + log_prob, + prev_sample_mean, + std_dev_t, + dt_sqrt, + sigma, + sigma_max, + ) diff --git a/fastvideo/train/methods/rl/utils/embeddings.py b/fastvideo/train/methods/rl/utils/embeddings.py new file mode 100644 index 0000000000..719b3c35cf --- /dev/null +++ b/fastvideo/train/methods/rl/utils/embeddings.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Text embedding utilities for RL training.""" + +from __future__ import annotations + +import torch + + +def compute_text_embeddings( + prompts: list[str], + text_encoder, + tokenizer, + max_sequence_length: int = 512, + device: torch.device | str = "cuda", +) -> torch.Tensor: + """Encode text prompts into embeddings using T5. + + Args: + prompts: List of text prompts. + text_encoder: T5 text encoder model. + tokenizer: T5 tokenizer. + max_sequence_length: Max token length. + device: Target device. + + Returns: + Tensor of shape (B, L, D). + """ + text_inputs = tokenizer( + prompts, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids.to(device) + attention_mask = text_inputs.attention_mask.to(device) + with torch.no_grad(): + prompt_embeds = text_encoder( + text_input_ids, attention_mask=attention_mask + ).last_hidden_state + + # make padding token 0 + prompt_embeds[attention_mask == 0] = 0 + return prompt_embeds diff --git a/fastvideo/train/methods/rl/utils/evaluation.py b/fastvideo/train/methods/rl/utils/evaluation.py new file mode 100644 index 0000000000..82f5090bd9 --- /dev/null +++ b/fastvideo/train/methods/rl/utils/evaluation.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Evaluation loop for RL training.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import torch + +from fastvideo.logger import init_logger +from fastvideo.train.methods.rl.utils.embeddings import ( + compute_text_embeddings, +) +from fastvideo.train.methods.rl.utils.pipeline import ( + wan_denoising_with_logprob, +) + +logger = init_logger(__name__) + + +def eval_once( + model, + scheduler, + test_dataloader, + text_encoder, + tokenizer, + sample_neg_prompt_embeds: torch.Tensor, + eval_reward_fn: Callable, + global_step: int, + ema_callback, + *, + eval_num_steps: int, + eval_guidance_scale: float, + height: int, + width: int, + num_frames: int, + device: torch.device, + world_size: int, + rank: int, + is_main_process: bool, + tracker: Any | None = None, + max_batches: int | None = None, + seed: int = 0, +) -> dict[str, float]: + """Run evaluation on test set. + + Args: + ema_callback: An ``EMACallback`` instance (or + ``None``). Used to temporarily swap EMA + weights into the transformer for evaluation. + + Returns: + Dict of aggregated eval metrics. + """ + model.transformer.eval() + all_rewards: dict[str, list[float]] = {} + + # Use EMA context manager if available. + if ema_callback is not None: + ctx = ema_callback.ema_context( + model.transformer + ) + else: + from contextlib import nullcontext + + ctx = nullcontext() + + with ctx: + for batch_idx, ( + _epoch_tag, + prompts, + metadata, + ) in enumerate(test_dataloader): + if max_batches is not None and batch_idx >= max_batches: + break + prompt_embeds = compute_text_embeddings( + prompts, + text_encoder, + tokenizer, + max_sequence_length=512, + device=device, + ) + neg_prompt_embeds = sample_neg_prompt_embeds[ + : len(prompts) + ] + + with torch.no_grad(): + generator = torch.Generator(device=device) + generator.manual_seed(seed + batch_idx) + ( + videos, + _latents, + _log_probs, + _kls, + _timesteps, + ) = wan_denoising_with_logprob( + model, + scheduler, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=( + neg_prompt_embeds + ), + num_inference_steps=eval_num_steps, + guidance_scale=eval_guidance_scale, + height=height, + width=width, + num_frames=num_frames, + generator=generator, + deterministic=True, + sde_type="flow_sde", + ) + + rewards, _ = eval_reward_fn( + videos, prompts, metadata + ) + for key, val in rewards.items(): + if key not in all_rewards: + all_rewards[key] = [] + if isinstance(val, torch.Tensor): + all_rewards[key].extend( + val.detach().cpu().tolist() + ) + else: + all_rewards[key].append(float(val)) + + # Aggregate metrics. + metrics = {} + for key, vals in all_rewards.items(): + avg = sum(vals) / max(len(vals), 1) + metrics[f"eval_{key}"] = avg + + if is_main_process and tracker is not None: + tracker.log(metrics, global_step) + + return metrics diff --git a/fastvideo/train/methods/rl/utils/pipeline.py b/fastvideo/train/methods/rl/utils/pipeline.py new file mode 100644 index 0000000000..a4a79260e2 --- /dev/null +++ b/fastvideo/train/methods/rl/utils/pipeline.py @@ -0,0 +1,364 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Multi-step denoising with log-probability tracking. + +Replaces diffusers' WanPipeline.__call__ for RL training. +Uses WanModel's forward_transformer_raw() instead of +the diffusers pipeline. +""" + +from __future__ import annotations + +import contextlib +import random +import time +from typing import Any + +import torch + +from fastvideo.logger import init_logger + +_pipeline_logger = init_logger(__name__) + +from fastvideo.train.methods.rl.utils.sde import ( + sde_step_with_logprob, +) + + +def wan_denoising_with_logprob( + model, + scheduler, + prompt_embeds: torch.Tensor, + negative_prompt_embeds: torch.Tensor | None = None, + num_inference_steps: int = 50, + guidance_scale: float = 5.0, + height: int = 480, + width: int = 832, + num_frames: int = 81, + generator: torch.Generator | None = None, + noise_level: float = 0.7, + sde_type: str = "flow_sde", + deterministic: bool = False, + diffusion_clip: bool = False, + diffusion_clip_value: float = 0.45, + sde_window_size: int = 0, + sde_window_range: tuple[int, int] | None = None, + kl_reward: float = 0.0, + ref_transformer: torch.nn.Module | None = None, + lora_model: Any | None = None, +) -> tuple[ + torch.Tensor, + list[torch.Tensor], + list[torch.Tensor], + list[torch.Tensor], + list[torch.Tensor], +]: + """Run full denoising loop, collecting latent + trajectories and log-probabilities at each step. + + Args: + model: WanModel (or GenRLWanModel) with + forward_transformer_raw and vae. + scheduler: Noise scheduler (UniPC/Euler). + prompt_embeds: (B, L, D) text embeddings. + negative_prompt_embeds: (B, L, D) or None. + num_inference_steps: Number of denoising steps. + guidance_scale: CFG scale. + height, width, num_frames: Video dimensions. + generator: RNG for reproducibility. + noise_level: SDE noise level. + sde_type: 'flow_sde' or 'flow_cps'. + deterministic: If True, no SDE noise. + diffusion_clip: Clip SDE variance. + diffusion_clip_value: Clip threshold. + sde_window_size: Window size for SDE training. + sde_window_range: (start, end) range for window. + kl_reward: KL penalty weight (>0 enables KL). + ref_transformer: Reference model for KL. + lora_model: LoRA model with disable_adapter(). + + Returns: + (videos, all_latents, all_log_probs, + all_kl, all_timesteps) + """ + device = model.device + batch_size = prompt_embeds.shape[0] + dtype = prompt_embeds.dtype + + do_cfg = ( + guidance_scale > 1.0 + and negative_prompt_embeds is not None + ) + + # Prepare initial noise. + vae_config = model.vae.config + vae_scale_temporal = getattr( + vae_config, "temporal_compression_ratio", 4 + ) + vae_scale_spatial = getattr( + vae_config, "spatial_compression_ratio", 8 + ) + num_channels = getattr(vae_config, "z_dim", 16) + + latent_frames = (num_frames - 1) // vae_scale_temporal + 1 + latent_h = height // vae_scale_spatial + latent_w = width // vae_scale_spatial + + latent_shape = ( + 1, num_channels, latent_frames, latent_h, latent_w, + ) + if isinstance(generator, list): + latents = torch.cat([ + torch.randn( + *latent_shape, + generator=generator[i], + device=device, + dtype=torch.float32, + ) + for i in range(batch_size) + ]) + else: + latents = torch.randn( + batch_size, + num_channels, + latent_frames, + latent_h, + latent_w, + generator=generator, + device=device, + dtype=torch.float32, + ) + + # Setup scheduler. + scheduler.set_timesteps( + num_inference_steps, device=device + ) + timesteps = scheduler.timesteps + + # Window setup. + use_window = ( + sde_window_size > 0 + and sde_window_range is not None + ) + if use_window: + if ( + sde_window_range[1] - sde_window_range[0] + < sde_window_size + ): + msg = ( + f"sde_window_range span " + f"({sde_window_range[1] - sde_window_range[0]}) " + f"must be >= sde_window_size " + f"({sde_window_size})" + ) + raise ValueError(msg) + if generator is not None: + gen = ( + generator[0] + if isinstance(generator, list) + else generator + ) + max_start = ( + sde_window_range[1] - sde_window_size + ) + start = torch.randint( + sde_window_range[0], + max_start + 1, + (1,), + generator=gen, + device=device, + ).item() + else: + start = random.randint( + sde_window_range[0], + sde_window_range[1] - sde_window_size, + ) + end = start + sde_window_size + sde_window = (start, end) + all_latents: list[torch.Tensor] = [] + else: + sde_window = None + all_latents: list[torch.Tensor] = [latents] + + all_log_probs: list[torch.Tensor] = [] + all_kl: list[torch.Tensor] = [] + all_timesteps: list[torch.Tensor] = [] + + _denoise_fwd_time = 0.0 + _denoise_sde_time = 0.0 + _denoise_kl_time = 0.0 + + for i, t in enumerate(timesteps): + latents_ori = latents.clone() + timestep = t.expand(batch_size) + + # Conditional prediction. + torch.cuda.synchronize() + _fwd_t0 = time.perf_counter() + noise_pred = model.forward_transformer_raw( + latents.to(dtype), + timestep, + prompt_embeds, + ) + noise_pred = noise_pred.to(dtype) + + # CFG. + if do_cfg: + noise_uncond = model.forward_transformer_raw( + latents.to(dtype), + timestep, + negative_prompt_embeds, + ) + noise_pred = noise_uncond + guidance_scale * ( + noise_pred - noise_uncond + ) + torch.cuda.synchronize() + _denoise_fwd_time += time.perf_counter() - _fwd_t0 + + # Determine noise level for this step. + if use_window: + if i < sde_window[0]: + cur_noise_level = 0.0 + elif i == sde_window[0]: + cur_noise_level = noise_level + all_latents.append(latents) + elif sde_window[0] < i < sde_window[1]: + cur_noise_level = noise_level + else: + cur_noise_level = 0.0 + else: + cur_noise_level = noise_level + + # SDE step. + _sde_t0 = time.perf_counter() + ( + latents, + log_prob, + prev_latents_mean, + std_dev_t, + sigma, + sigma_max, + ) = sde_step_with_logprob( + scheduler, + noise_pred.float(), + t.unsqueeze(0), + latents.float(), + noise_level=cur_noise_level, + sde_type=sde_type, + deterministic=deterministic, + diffusion_clip=diffusion_clip, + diffusion_clip_value=diffusion_clip_value, + ) + _denoise_sde_time += time.perf_counter() - _sde_t0 + prev_latents = latents.clone() + + # Record. + in_window = ( + use_window + and sde_window[0] <= i < sde_window[1] + ) + should_record = (not use_window) or in_window + + if should_record: + all_latents.append(latents) + all_log_probs.append(log_prob) + all_timesteps.append(t) + + # KL computation. + _kl_t0 = time.perf_counter() + if should_record and kl_reward > 0 and not deterministic: + ref_model = ref_transformer + ref_ctx: Any = contextlib.nullcontext() + if ref_model is None and lora_model is not None: + ref_model = lora_model + ref_ctx = lora_model.disable_adapter() + + if ref_model is not None: + with ref_ctx: + ref_noise = ref_model( + hidden_states=latents_ori.to(dtype), + timestep=timestep, + encoder_hidden_states=prompt_embeds, + return_dict=False, + ) + ref_noise = ref_noise.to(dtype) + if do_cfg: + ref_uncond = ref_model( + hidden_states=latents_ori.to( + dtype + ), + timestep=timestep, + encoder_hidden_states=( + negative_prompt_embeds + ), + return_dict=False, + ) + ref_noise = ( + ref_uncond + + guidance_scale + * (ref_noise - ref_uncond) + ) + + ( + _, + _ref_log_prob, + ref_prev_mean, + ref_std, + _ref_sigma, + _ref_sigma_max, + ) = sde_step_with_logprob( + scheduler, + ref_noise.float(), + t.unsqueeze(0), + latents_ori.float(), + noise_level=noise_level, + sde_type=sde_type, + prev_sample=prev_latents.float(), + deterministic=deterministic, + diffusion_clip=diffusion_clip, + diffusion_clip_value=diffusion_clip_value, + ) + kl = (prev_latents_mean - ref_prev_mean) ** 2 / ( + 2 * std_dev_t**2 + ) + kl = kl.mean( + dim=tuple(range(1, kl.ndim)) + ) + all_kl.append(kl) + else: + all_kl.append( + torch.zeros( + batch_size, device=device + ) + ) + elif should_record: + all_kl.append( + torch.zeros(batch_size, device=device) + ) + torch.cuda.synchronize() + _denoise_kl_time += time.perf_counter() - _kl_t0 + + # Decode to video. + torch.cuda.synchronize() + _vae_t0 = time.perf_counter() + videos = model.decode_latents(latents) + + torch.cuda.synchronize() + _vae_done = time.perf_counter() + _pipeline_logger.info( + "[denoising] %d steps: " + "transformer_fwd=%.1fs sde_step=%.1fs " + "kl=%.1fs vae_decode=%.1fs", + len(timesteps), + _denoise_fwd_time, + _denoise_sde_time, + _denoise_kl_time, + _vae_done - _vae_t0, + ) + + return ( + videos, + all_latents, + all_log_probs, + all_kl, + all_timesteps, + ) diff --git a/fastvideo/train/methods/rl/utils/rewards.py b/fastvideo/train/methods/rl/utils/rewards.py new file mode 100644 index 0000000000..d7017c2055 --- /dev/null +++ b/fastvideo/train/methods/rl/utils/rewards.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reward function loading and composition for RL training.""" + +from __future__ import annotations + +import gc +import importlib +import inspect +import time +from collections.abc import Callable +from contextlib import contextmanager + +import torch + +from fastvideo.logger import init_logger +from fastvideo.train.methods.rl.reward import ( + hpsv3_general_score, + hpsv3_percentile_score, + video_ocr_score, + videoalign_mq_score, + videoalign_ta_score, +) +from fastvideo.train.methods.rl.reward.hpsv3 import ( + _HPSV3_INFERENCERS, + set_hpsv3_device, +) +from fastvideo.train.methods.rl.reward.videoalign import ( + _VIDEOALIGN_INFERENCERS, + set_videoalign_device, +) + +logger = init_logger(__name__) + +_BUILTIN_REWARDS: dict[str, Callable] = { + "video_ocr": video_ocr_score, + "hpsv3_general": hpsv3_general_score, + "hpsv3_percentile": hpsv3_percentile_score, + "videoalign_mq": videoalign_mq_score, + "videoalign_ta": videoalign_ta_score, +} + +_GPU_REWARD_NAMES = { + "hpsv3_general", + "hpsv3_percentile", + "videoalign_mq", + "videoalign_ta", +} + + +def load_reward_fn( + name: str, + device, + module_path: str | None = None, +): + """Load a reward function by name.""" + if module_path: + mod = importlib.import_module(module_path) + fn = getattr(mod, f"{name}_score", None) + if fn is None: + msg = ( + f"Reward {name}_score not found " + f"in {module_path}" + ) + raise ValueError(msg) + return fn(device) if callable(fn) else fn + + if name in _BUILTIN_REWARDS: + fn = _BUILTIN_REWARDS[name] + sig = inspect.signature(fn) + accepts_device = any( + p.name in {"device", "dev"} + for p in sig.parameters.values() + ) + return fn(device) if accepts_device else fn() + + def _zero_fn(images, prompts, metadata, only_strict=False): + batch = ( + len(prompts) if prompts is not None else 1 + ) + zeros = torch.zeros(batch, device=device) + return {"avg": zeros}, {} + + return _zero_fn + + +def multi_score( + device, + reward_cfg: dict[str, float], + module_path: str | None = None, + return_raw_scores: bool = False, +): + """Compose multiple reward heads. + + Args: + device: Device for reward computation. + reward_cfg: Dict mapping reward name to weight. + module_path: Optional custom module path. + return_raw_scores: If True, include raw scores. + + Returns: + A callable (images, prompts, metadata, only_strict) + -> (scores_dict, metadata_dict). + """ + reward_fns = {} + weights = {} + for name, weight in reward_cfg.items(): + reward_fns[name] = load_reward_fn( + name, device, module_path + ) + weights[name] = weight + + def _fn(images, prompts, metadata, only_strict=True): + scores = {} + for name, fn in reward_fns.items(): + out, _meta = fn(images, prompts, metadata) + if isinstance(out, dict): + val = out.get( + "avg", out.get("reward", out) + ) + else: + val = out + if return_raw_scores: + scores[f"{name}_raw"] = val + scores[name] = val * weights[name] + stacked = torch.stack( + [scores[name] for name in reward_cfg], dim=0 + ) + scores["avg"] = stacked.mean(0) + return scores, {} + + return _fn + + +def _has_reward(reward_cfg, names) -> bool: + if not reward_cfg: + return False + return any(name in reward_cfg for name in names) + + +def _device_type(device) -> str: + if isinstance(device, torch.device): + return device.type + return torch.device(device).type + + +def move_reward_models(reward_cfg, device) -> None: + """Move GPU-backed reward models to device.""" + if _has_reward( + reward_cfg, + {"hpsv3_general", "hpsv3_percentile"}, + ): + set_hpsv3_device(device) + if _has_reward( + reward_cfg, + {"videoalign_mq", "videoalign_ta"}, + ): + set_videoalign_device(device) + + +def clear_reward_models(reward_cfg) -> None: + """Drop cached GPU-backed reward models before PPO training.""" + cleared = False + if _has_reward( + reward_cfg, + {"hpsv3_general", "hpsv3_percentile"}, + ): + _HPSV3_INFERENCERS.clear() + cleared = True + if _has_reward( + reward_cfg, + {"videoalign_mq", "videoalign_ta"}, + ): + _VIDEOALIGN_INFERENCERS.clear() + cleared = True + if cleared: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + + +@contextmanager +def reward_models_on_device(reward_cfg, device): + """Temporarily move reward models to device.""" + if _has_reward(reward_cfg, _GPU_REWARD_NAMES): + use_cuda = _device_type(device) == "cuda" + _t0 = time.perf_counter() + move_reward_models(reward_cfg, device) + if use_cuda: + torch.cuda.synchronize() + _t1 = time.perf_counter() + logger.info( + "[rewards] move_to_device=%.1fs", _t1 - _t0 + ) + try: + yield + finally: + _t2 = time.perf_counter() + move_reward_models(reward_cfg, "cpu") + if use_cuda: + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + _t3 = time.perf_counter() + logger.info( + "[rewards] move_to_cpu+gc=%.1fs", + _t3 - _t2, + ) + else: + yield diff --git a/fastvideo/train/methods/rl/utils/sampling.py b/fastvideo/train/methods/rl/utils/sampling.py new file mode 100644 index 0000000000..09a6359501 --- /dev/null +++ b/fastvideo/train/methods/rl/utils/sampling.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Sampling epoch for RL training — generate videos and +compute rewards.""" + +from __future__ import annotations + +import hashlib +import time +from collections.abc import Callable +from typing import Any + +import torch + +from fastvideo.logger import init_logger +from fastvideo.train.methods.rl.utils.embeddings import ( + compute_text_embeddings, +) +from fastvideo.train.methods.rl.utils.pipeline import ( + wan_denoising_with_logprob, +) + +logger = init_logger(__name__) + +SEED_EPOCH_STRIDE = 10_000 + + +def create_generator( + prompts: list[str], + base_seed: int, + device: torch.device, +) -> list[torch.Generator]: + """Create deterministic generators seeded by prompt.""" + generators = [] + for prompt in prompts: + prompt_seed = int.from_bytes( + hashlib.blake2b( + prompt.encode("utf-8"), + digest_size=8, + ).digest(), + "big", + ) + g = torch.Generator(device=device) + g.manual_seed(base_seed + prompt_seed % (2**31)) + generators.append(g) + return generators + + +def sample_epoch( + model, + scheduler, + train_sampler, + train_iter, + reward_fn: Callable, + sample_neg_prompt_embeds: torch.Tensor, + text_encoder, + tokenizer, + executor, + epoch: int, + global_step: int, + *, + # Config values passed explicitly. + sample_batch_size: int, + num_batches_per_epoch: int, + num_inference_steps: int, + guidance_scale: float, + height: int, + width: int, + num_frames: int, + noise_level: float, + sde_type: str, + diffusion_clip: bool, + diffusion_clip_value: float, + sde_window_size: int, + sde_window_range: tuple[int, int] | None, + kl_reward: float, + same_latent: bool, + seed: int, + device: torch.device, + is_main_process: bool, + ref_transformer: torch.nn.Module | None = None, + lora_model: Any | None = None, + tracker: Any | None = None, + async_reward_scoring: bool = True, +) -> tuple[ + list[dict[str, Any]], + list[torch.Tensor], + list[list[str]], +]: + """Run one sampling epoch: generate videos and compute rewards. + + Returns: + Tuple of (samples, all_videos, all_prompts): + - samples: list of sample dicts with prompt_ids, + prompt_embeds, latents, log_probs, kl, + timesteps, rewards. + - all_videos: list of decoded video tensors + per batch, each (B, 3, T, H, W) in [0, 1]. + - all_prompts: list of prompt string lists + per batch. + """ + samples = [] + all_videos: list[torch.Tensor] = [] + all_prompts: list[list[str]] = [] + + for i in range(num_batches_per_epoch): + current_epoch_tag = ( + epoch * num_batches_per_epoch + i + ) + train_sampler.set_epoch(current_epoch_tag) + + # Drain until epoch tag matches. + while True: + epoch_tag, prompts, prompt_metadata = next( + train_iter + ) + if epoch_tag == current_epoch_tag: + break + + torch.cuda.synchronize() + _t_batch_start = time.perf_counter() + + _t_embed = time.perf_counter() + prompt_embeds = compute_text_embeddings( + prompts, + text_encoder, + tokenizer, + max_sequence_length=512, + device=device, + ) + prompt_ids = tokenizer( + prompts, + padding="max_length", + max_length=512, + truncation=True, + return_tensors="pt", + ).input_ids.to(device) + + # Generator setup. + if same_latent: + gen = create_generator( + prompts, + base_seed=seed + epoch * SEED_EPOCH_STRIDE + i, + device=device, + ) + else: + gen = torch.Generator(device=device) + gen.manual_seed( + seed + epoch * SEED_EPOCH_STRIDE + i + ) + + torch.cuda.synchronize() + _t_embed_done = time.perf_counter() + + with torch.no_grad(): + ( + videos, + latents_list, + log_probs_list, + kls_list, + timesteps_list, + ) = wan_denoising_with_logprob( + model, + scheduler, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=( + sample_neg_prompt_embeds + ), + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + height=height, + width=width, + num_frames=num_frames, + generator=gen, + noise_level=noise_level, + sde_type=sde_type, + diffusion_clip=diffusion_clip, + diffusion_clip_value=diffusion_clip_value, + sde_window_size=sde_window_size, + sde_window_range=sde_window_range, + kl_reward=kl_reward, + ref_transformer=ref_transformer, + lora_model=lora_model, + ) + + torch.cuda.synchronize() + _t_denoise_done = time.perf_counter() + + latents = torch.stack(latents_list, dim=1) + log_probs = torch.stack(log_probs_list, dim=1) + kls = torch.stack(kls_list, dim=1) + kl = kls.detach() + + timesteps = ( + torch.stack(timesteps_list) + .unsqueeze(0) + .repeat(sample_batch_size, 1) + ) + + videos_cpu = videos.detach().cpu() + + # Collect decoded videos and prompts for logging. + all_videos.append(videos_cpu) + all_prompts.append(list(prompts)) + + if async_reward_scoring: + rewards = executor.submit( + reward_fn, + videos_cpu, + prompts, + prompt_metadata, + True, + ) + time.sleep(0) + else: + rewards = (videos_cpu, list(prompts), prompt_metadata) + + del videos + + logger.info( + "[sample_epoch] batch %d/%d: " + "text_embed=%.1fs denoise=%.1fs " + "batch_total=%.1fs", + i + 1, + num_batches_per_epoch, + _t_embed_done - _t_embed, + _t_denoise_done - _t_embed_done, + _t_denoise_done - _t_batch_start, + ) + + samples.append( + { + "prompt_ids": prompt_ids, + "prompt_embeds": prompt_embeds, + "negative_prompt_embeds": ( + sample_neg_prompt_embeds + ), + "timesteps": timesteps, + "latents": latents[:, :-1], + "next_latents": latents[:, 1:], + "log_probs": log_probs, + "kl": kl, + "rewards": rewards, + } + ) + + # Wait for all rewards. + torch.cuda.synchronize() + _t_reward_wait = time.perf_counter() + for sample in samples: + if async_reward_scoring: + rewards, _ = sample["rewards"].result() + else: + videos_cpu, prompts, prompt_metadata = sample["rewards"] + torch.cuda.empty_cache() + rewards, _ = reward_fn( + videos_cpu, prompts, prompt_metadata, True + ) + sample["rewards"] = { + key: torch.as_tensor(value, device=device).float() + for key, value in rewards.items() + } + _t_reward_done = time.perf_counter() + logger.info( + "[sample_epoch] reward_wait=%.1fs", + _t_reward_done - _t_reward_wait, + ) + + return samples, all_videos, all_prompts diff --git a/fastvideo/train/methods/rl/utils/sde.py b/fastvideo/train/methods/rl/utils/sde.py new file mode 100644 index 0000000000..7bbd86fb01 --- /dev/null +++ b/fastvideo/train/methods/rl/utils/sde.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: Apache-2.0 +"""SDE step with log-probability computation for RL training.""" + +from __future__ import annotations + +import math + +import torch +from diffusers.utils.torch_utils import randn_tensor + + +def sde_step_with_logprob( + scheduler, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor, + sample: torch.FloatTensor, + noise_level: float = 0.7, + prev_sample: torch.FloatTensor | None = None, + generator: torch.Generator | None = None, + sde_type: str | None = "flow_sde", + deterministic: bool = False, + return_sqrt_dt_and_std_dev_t: bool = False, + diffusion_clip: bool = False, + diffusion_clip_value: float = 0.45, +): + """Predict the sample from the previous timestep by reversing + the SDE, returning log-probability of the transition. + + Args: + scheduler: Noise scheduler with sigmas and timestep index. + model_output: Predicted noise/velocity. + timestep: Current timestep(s). + sample: Current latents. + noise_level: Noise level for SDE/CPS computation. + prev_sample: Optional precomputed previous sample. + generator: Optional RNG for sampling prev_sample. + sde_type: 'flow_sde' or 'flow_cps'. + deterministic: If True, no noise added. + return_sqrt_dt_and_std_dev_t: If True, return extra terms. + diffusion_clip: If True, clip std_dev_t. + diffusion_clip_value: Clipping threshold. + + Returns: + If return_sqrt_dt_and_std_dev_t: + (prev_sample, log_prob, prev_sample_mean, + std_dev_t, sqrt_neg_dt, sigma, sigma_max) + Else: + (prev_sample, log_prob, prev_sample_mean, + std_dev_t * sqrt_neg_dt, sigma, sigma_max) + """ + model_output = model_output.float() + sample = sample.float() + if prev_sample is not None: + prev_sample = prev_sample.float() + + step_index = [ + scheduler.index_for_timestep(t) for t in timestep + ] + prev_step_index = [step + 1 for step in step_index] + + scheduler.sigmas = scheduler.sigmas.to(sample.device) + sigma = scheduler.sigmas[step_index].view(-1, 1, 1, 1, 1) + sigma_prev = ( + scheduler.sigmas[prev_step_index].view(-1, 1, 1, 1, 1) + ) + sigma_max = scheduler.sigmas[1].item() + dt = sigma_prev - sigma + + if sde_type == "flow_sde": + std_dev_t = ( + torch.sqrt( + sigma + / ( + 1 + - torch.where( + sigma == 1, + torch.tensor( + sigma_max, + device=sigma.device, + dtype=sigma.dtype, + ), + sigma, + ) + ) + ) + * noise_level + ) + + if diffusion_clip: + max_std_dev_t = ( + diffusion_clip_value / torch.sqrt(-1 * dt) + ) + std_dev_t = torch.minimum(std_dev_t, max_std_dev_t) + + prev_sample_mean = ( + sample + * (1 + std_dev_t**2 / (2 * sigma) * dt) + + model_output + * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) + * dt + ) + + if prev_sample is None: + variance_noise = randn_tensor( + model_output.shape, + generator=generator, + device=model_output.device, + dtype=model_output.dtype, + ) + prev_sample = ( + prev_sample_mean + + std_dev_t * torch.sqrt(-1 * dt) * variance_noise + ) + + if deterministic: + prev_sample = sample + dt * model_output + + std_scale = std_dev_t * torch.sqrt(-1 * dt) + if torch.all(std_scale == 0): + log_prob = torch.zeros_like(prev_sample) + else: + std_scale = torch.clamp( + std_scale, + min=torch.finfo(std_scale.dtype).tiny, + ) + log_prob = ( + -((prev_sample.detach() - prev_sample_mean) ** 2) + / (2 * (std_scale**2)) + - torch.log(std_scale) + - torch.log( + torch.sqrt(2 * torch.as_tensor(math.pi)) + ) + ) + + elif sde_type == "flow_cps": + std_dev_t = sigma_prev * math.sin( + noise_level * math.pi / 2 + ) + pred_original_sample = sample - sigma * model_output + noise_estimate = ( + sample + model_output * (1 - sigma) + ) + prev_sample_mean = pred_original_sample * ( + 1 - sigma_prev + ) + noise_estimate * torch.sqrt( + sigma_prev**2 - std_dev_t**2 + ) + + if prev_sample is None: + variance_noise = randn_tensor( + model_output.shape, + generator=generator, + device=model_output.device, + dtype=model_output.dtype, + ) + prev_sample = ( + prev_sample_mean + std_dev_t * variance_noise + ) + + if deterministic: + prev_sample = ( + pred_original_sample * (1 - sigma_prev) + + noise_estimate * sigma_prev + ) + + log_prob = -( + (prev_sample.detach() - prev_sample_mean) ** 2 + ) + + else: + msg = ( + f"Unknown sde_type: {sde_type}. " + "Must be 'flow_sde' or 'flow_cps'." + ) + raise ValueError(msg) + + log_prob = log_prob.mean( + dim=tuple(range(1, log_prob.ndim)) + ) + + if return_sqrt_dt_and_std_dev_t: + return ( + prev_sample, + log_prob, + prev_sample_mean, + std_dev_t, + torch.sqrt(-1 * dt), + sigma, + sigma_max, + ) + return ( + prev_sample, + log_prob, + prev_sample_mean, + std_dev_t * torch.sqrt(-1 * dt), + sigma, + sigma_max, + ) diff --git a/fastvideo/train/methods/rl/utils/stat_tracking.py b/fastvideo/train/methods/rl/utils/stat_tracking.py new file mode 100644 index 0000000000..ea0ba4609e --- /dev/null +++ b/fastvideo/train/methods/rl/utils/stat_tracking.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Per-prompt statistics tracking for advantage computation.""" + +from __future__ import annotations + +import numpy as np +import torch + +EPSILON = 1e-4 + + +class PerPromptStatTracker: + """Track per-prompt reward history and compute advantages.""" + + def __init__( + self, + use_global_std: bool = False, + max_group_std: bool = False, + ): + self.use_global_std = use_global_std + self.max_group_std = max_group_std + self.stats: dict[str, np.ndarray] = {} + self.history_prompts: set[int] = set() + + def update( + self, + prompts, + rewards, + mode: str = "grpo", + ) -> np.ndarray: + """Update stats and compute advantages. + + Args: + prompts: Iterable of prompt strings. + rewards: Array-like rewards aligned with prompts. + mode: Advantage mode: grpo|rwr|sft|dpo. + + Returns: + Advantages array aligned with prompts. + """ + prompts = np.array(prompts) + rewards = np.array(rewards, dtype=np.float64) + unique = np.unique(prompts) + advantages = np.empty_like(rewards) * 0.0 + + for prompt in unique: + prompt_rewards = rewards[prompts == prompt] + if prompt not in self.stats: + self.stats[prompt] = [] + self.stats[prompt].extend(prompt_rewards) + self.history_prompts.add(hash(prompt)) + self.stats[prompt] = np.stack( + self.stats[prompt] + ) + + max_std = None + if self.max_group_std and len(unique) > 0: + prompt_stds = [] + for prompt in unique: + prompt_std = ( + np.std( + self.stats[prompt], + axis=0, + keepdims=True, + ) + + EPSILON + ) + prompt_stds.append(prompt_std) + max_std_value = max( + np.max(std) for std in prompt_stds + ) + max_std = np.full_like( + prompt_stds[0], max_std_value + ) + + for prompt in unique: + prompt_rewards = rewards[prompts == prompt] + mean = np.mean( + self.stats[prompt], axis=0, keepdims=True + ) + if self.use_global_std: + std = ( + np.std(rewards, axis=0, keepdims=True) + + EPSILON + ) + elif self.max_group_std: + std = max_std + else: + std = ( + np.std( + self.stats[prompt], + axis=0, + keepdims=True, + ) + + EPSILON + ) + if mode == "grpo": + advantages[prompts == prompt] = ( + prompt_rewards - mean + ) / std + elif mode == "rwr": + advantages[prompts == prompt] = ( + prompt_rewards + ) + elif mode == "sft": + advantages[prompts == prompt] = ( + ( + torch.tensor(prompt_rewards) + == torch.max( + torch.tensor(prompt_rewards) + ) + ) + .float() + .numpy() + ) + elif mode == "dpo": + pa = torch.tensor(prompt_rewards) + max_idx = torch.argmax(pa) + min_idx = torch.argmin(pa) + if max_idx == min_idx: + min_idx = 0 + max_idx = 1 + result = torch.zeros_like(pa).float() + result[max_idx] = 1.0 + result[min_idx] = -1.0 + advantages[prompts == prompt] = ( + result.numpy() + ) + return advantages + + def get_stats(self) -> tuple[float, int]: + """Return (avg_group_size, num_unique_prompts).""" + avg = ( + sum(len(v) for v in self.stats.values()) + / len(self.stats) + if self.stats + else 0 + ) + return avg, len(self.history_prompts) + + def clear(self): + """Clear stored statistics.""" + self.stats = {} diff --git a/fastvideo/train/models/wan/__init__.py b/fastvideo/train/models/wan/__init__.py index 9a8113ac14..afe592d649 100644 --- a/fastvideo/train/models/wan/__init__.py +++ b/fastvideo/train/models/wan/__init__.py @@ -5,3 +5,5 @@ WanModel as WanModel, ) from fastvideo.train.models.wan.wan_causal import ( WanCausalModel as WanCausalModel, ) +from fastvideo.train.models.wan.wan_genrl import ( + GenRLWanModel as GenRLWanModel, ) diff --git a/fastvideo/train/models/wan/wan.py b/fastvideo/train/models/wan/wan.py index 57d4c4e401..d490b87661 100644 --- a/fastvideo/train/models/wan/wan.py +++ b/fastvideo/train/models/wan/wan.py @@ -580,6 +580,100 @@ def _build_distill_input_kwargs( def _get_transformer(self, timestep: torch.Tensor) -> torch.nn.Module: return self.transformer + # ------------------------------------------------------------------ + # RL pipeline primitives + # ------------------------------------------------------------------ + + def forward_transformer_raw( + self, + latents: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + ) -> torch.Tensor: + """Direct transformer forward for RL pipelines. + + Bypasses batch preparation / attention-metadata. + Uses dense attention (attn_metadata=None). + + Args: + latents: (B, C, T, H, W) in diffusion space. + timestep: (B,) or scalar timestep. + encoder_hidden_states: (B, L, D) text embeddings. + + Returns: + Model output tensor (B, C, T, H, W). + """ + dtype = self._get_training_dtype() + device_type = self.device.type + with ( + torch.autocast(device_type, dtype=dtype), + set_forward_context( + current_timestep=timestep, + attn_metadata=None, + ), + ): + output = self.transformer( + hidden_states=latents.to(dtype), + timestep=timestep, + encoder_hidden_states=( + encoder_hidden_states.to(dtype) + ), + return_dict=False, + ) + return output + + def decode_latents( + self, + latents: torch.Tensor, + ) -> torch.Tensor: + """Decode latents to pixel-space video. + + Denormalizes from flow diffusion space and passes + through the VAE decoder. + + Args: + latents: (B, C, T, H, W) denoised latents. + + Returns: + Video tensor (B, 3, T_out, H_out, W_out) + in [0, 1] range. + """ + vae = self.vae + vae_config = vae.config + z_dim = getattr(vae_config, "z_dim", 16) + + latents_mean = ( + torch.tensor(vae_config.latents_mean) + .view(1, z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std_inv = ( + 1.0 + / torch.tensor(vae_config.latents_std).view( + 1, z_dim, 1, 1, 1 + ) + ).to(latents.device, latents.dtype) + latents = latents / latents_std_inv + latents_mean + + # Decode one sample at a time. + vae_dtype = next(vae.parameters()).dtype + videos = [] + with torch.no_grad(): + for idx in range(latents.shape[0]): + decoded = vae.decode( + latents[idx : idx + 1].to(vae_dtype) + ) + if isinstance(decoded, tuple): + decoded = decoded[0] + videos.append(decoded.float()) + video = torch.cat(videos, dim=0) + + # Normalize to [0, 1]. + video = (video / 2 + 0.5).clamp(0, 1) + return video + + # ------------------------------------------------------------------ + def _get_uncond_text_dict( self, batch: TrainingBatch, diff --git a/fastvideo/train/models/wan/wan_genrl.py b/fastvideo/train/models/wan/wan_genrl.py new file mode 100644 index 0000000000..dbde67b494 --- /dev/null +++ b/fastvideo/train/models/wan/wan_genrl.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wan model extended for GenRL (RL training with text prompts). + +Overrides ``init_preprocessors`` to load a T5 text encoder +and tokenizer instead of the standard parquet video +dataloader. Provides a trivial dummy dataloader so the +trainer's outer loop has something to iterate over — the +real prompt dataloaders are created by the GenRLMethod. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from types import MethodType +from typing import Any, TYPE_CHECKING + +import torch + +from fastvideo.distributed import ( + get_sp_group, + get_world_group, +) +from fastvideo.logger import init_logger +from fastvideo.train.models.wan.wan import WanModel +from fastvideo.train.utils.moduleloader import ( + load_module_from_path, +) + +if TYPE_CHECKING: + from fastvideo.train.utils.training_config import ( + TrainingConfig, + ) + +logger = init_logger(__name__) + + +def _is_lora_target( + module_name: str, + target_modules: list[str], +) -> bool: + return any( + module_name == target + or module_name.endswith(f".{target}") + or target in module_name + for target in target_modules + ) + + +def _apply_fastvideo_lora( + transformer: Any, + *, + lora_rank: int, + lora_alpha: int, + target_modules: list[str], + init_weights: str, +) -> int: + from fastvideo.layers.lora.linear import ( + get_lora_layer, + replace_submodule, + ) + + transformer.requires_grad_(False) + converted_count = 0 + for name, layer in list(transformer.named_modules()): + if not _is_lora_target(name, target_modules): + continue + lora_layer = get_lora_layer( + layer, + lora_rank=lora_rank, + lora_alpha=lora_alpha, + training_mode=True, + ) + if lora_layer is None: + continue + _init_lora_weights(lora_layer, init_weights, lora_rank) + replace_submodule(transformer, name, lora_layer) + converted_count += 1 + return converted_count + + +def _init_lora_weights( + lora_layer: Any, + init_weights: str, + lora_rank: int, +) -> None: + """Match PEFT's useful LoRA initialization modes.""" + init = init_weights.lower() + if init == "default": + return + + lora_A = getattr(lora_layer, "lora_A", None) + lora_B = getattr(lora_layer, "lora_B", None) + if lora_A is None or lora_B is None: + return + + if init == "gaussian": + torch.nn.init.normal_(lora_A, std=1 / max(1, lora_rank)) + torch.nn.init.zeros_(lora_B) + return + + raise ValueError( + "Unsupported GenRLWanModel LoRA init_weights=" + f"{init_weights!r}. Use 'gaussian' or 'default'." + ) + + +@contextmanager +def _disable_lora_adapters(transformer: Any): + """Temporarily run a LoRA-wrapped transformer as its frozen base model.""" + lora_layers = [ + module for module in transformer.modules() + if hasattr(module, "disable_lora") + ] + previous = [bool(module.disable_lora) for module in lora_layers] + try: + for module in lora_layers: + module.disable_lora = True + yield + finally: + for module, was_disabled in zip(lora_layers, previous, strict=True): + module.disable_lora = was_disabled + + +def _attach_disable_adapter(transformer: Any) -> None: + """Expose a PEFT-compatible disable_adapter context manager.""" + + def disable_adapter(self): + return _disable_lora_adapters(self) + + transformer.disable_adapter = MethodType( # type: ignore[attr-defined] + disable_adapter, + transformer, + ) + + +class _InfiniteDummyLoader: + """Trivial iterable that yields empty dicts forever.""" + + def __iter__(self): + while True: + yield {} + + +class GenRLWanModel(WanModel): + """Wan model with text encoder for RL training. + + Compared to the base :class:`WanModel`, this variant: + + * Loads the T5 text encoder and tokenizer from the + pretrained model path. + * Sets a dummy dataloader so the trainer can iterate + without blocking. + * Does **not** build the standard parquet video + dataloader. + """ + + def __init__( + self, + *, + init_from: str, + training_config: TrainingConfig, + trainable: bool = True, + use_lora: bool = False, + lora_r: int = 32, + lora_alpha: int = 64, + lora_target_modules: list[str] | None = None, + lora_path: str | None = None, + lora_init_weights: str = "gaussian", + disable_custom_init_weights: bool = False, + flow_shift: float = 3.0, + enable_gradient_checkpointing_type: str + | None = None, + transformer_override_safetensor: str + | None = None, + ) -> None: + super().__init__( + init_from=init_from, + training_config=training_config, + trainable=trainable, + disable_custom_init_weights=( + disable_custom_init_weights + ), + flow_shift=flow_shift, + enable_gradient_checkpointing_type=( + enable_gradient_checkpointing_type + ), + transformer_override_safetensor=( + transformer_override_safetensor + ), + ) + if use_lora: + if lora_target_modules is None: + raise ValueError( + "GenRLWanModel use_lora=True requires " + "lora_target_modules." + ) + if lora_path: + raise ValueError( + "GenRLWanModel lora_path is not supported for " + "FastVideo LoRA training yet." + ) + converted_count = _apply_fastvideo_lora( + self.transformer, + lora_rank=int(lora_r), + lora_alpha=int(lora_alpha), + target_modules=lora_target_modules, + init_weights=lora_init_weights, + ) + if converted_count == 0: + raise ValueError( + "GenRLWanModel use_lora=True did not match any " + f"FastVideo linear layers: {lora_target_modules}" + ) + logger.info( + "Converted %d GenRL Wan transformer layers to LoRA", + converted_count, + ) + _attach_disable_adapter(self.transformer) + self.text_encoder: Any = None + self.tokenizer: Any = None + + def disable_adapter(self): + """PEFT-compatible context manager for reference KL with LoRA.""" + return _disable_lora_adapters(self.transformer) + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + def init_preprocessors( + self, + training_config: TrainingConfig, + ) -> None: # type: ignore[override] + """Load VAE, text encoder, and tokenizer.""" + # Load VAE. + self.vae = load_module_from_path( + model_path=str(training_config.model_path), + module_type="vae", + training_config=training_config, + ) + + self.world_group = get_world_group() + self.sp_group = get_sp_group() + self._init_timestep_mechanics() + + # Load text encoder and tokenizer. + model_path = str(training_config.model_path) + self._load_text_encoder( + model_path, training_config + ) + + # Dummy dataloader for the trainer's outer loop. + self.dataloader = _InfiniteDummyLoader() + self.start_step = 0 + + def _load_text_encoder( + self, + model_path: str, + training_config: TrainingConfig, + ) -> None: + from transformers import AutoTokenizer + + logger.info( + "Loading tokenizer from %s", model_path + ) + self.tokenizer = AutoTokenizer.from_pretrained( + model_path, subfolder="tokenizer" + ) + + logger.info( + "Loading text encoder from %s", model_path + ) + self.text_encoder = load_module_from_path( + model_path=model_path, + module_type="text_encoder", + training_config=training_config, + ) + self.text_encoder.requires_grad_(False) + + def on_train_start(self) -> None: + """Skip negative conditioning (handled by method).""" diff --git a/fastvideo/train/trainer.py b/fastvideo/train/trainer.py index d814a6cfe6..1c8128383e 100644 --- a/fastvideo/train/trainer.py +++ b/fastvideo/train/trainer.py @@ -127,6 +127,8 @@ def run( metric_sums: dict[str, float] = {} for accum_iter in range(grad_accum): batch = next(data_stream) + # TODO: have method.single_train_step return a single dict of outputs directly + # TODO: ask single_train_step to do backward and remove explicity method.backward call loss_map, outputs, step_metrics = method.single_train_step( batch, step, @@ -156,6 +158,7 @@ def run( self.callbacks.on_before_optimizer_step( method, iteration=step, + outputs=outputs, ) method.optimizers_schedulers_step(step) method.optimizers_zero_grad(step)