From 0a01c12b54384956039365ad9c692165c023decd Mon Sep 17 00:00:00 2001 From: Peiyuan Zhang Date: Tue, 10 Mar 2026 17:44:50 +0000 Subject: [PATCH 1/7] first edit --- .../genrl_wan2.1_t2v_1.3B_longcat.yaml | 125 ++ fastvideo/train/methods/__init__.py | 4 + fastvideo/train/methods/rl/__init__.py | 2 + fastvideo/train/methods/rl/advantages.py | 204 ++++ fastvideo/train/methods/rl/data.py | 263 +++++ fastvideo/train/methods/rl/diffusion.py | 113 ++ fastvideo/train/methods/rl/ema.py | 121 ++ fastvideo/train/methods/rl/embeddings.py | 46 + fastvideo/train/methods/rl/evaluation.py | 122 ++ fastvideo/train/methods/rl/genrl.py | 1033 +++++++++++++++++ fastvideo/train/methods/rl/pipeline.py | 318 +++++ fastvideo/train/methods/rl/reward/__init__.py | 22 + fastvideo/train/methods/rl/reward/hpsv3.py | 162 +++ fastvideo/train/methods/rl/reward/ocr.py | 114 ++ fastvideo/train/methods/rl/reward/utils.py | 42 + .../train/methods/rl/reward/videoalign.py | 170 +++ fastvideo/train/methods/rl/rewards.py | 172 +++ fastvideo/train/methods/rl/sampling.py | 204 ++++ fastvideo/train/methods/rl/sde.py | 192 +++ fastvideo/train/methods/rl/stat_tracking.py | 143 +++ fastvideo/train/models/wan/__init__.py | 2 + fastvideo/train/models/wan/wan.py | 94 ++ fastvideo/train/models/wan/wan_genrl.py | 140 +++ 23 files changed, 3808 insertions(+) create mode 100644 examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml create mode 100644 fastvideo/train/methods/rl/__init__.py create mode 100644 fastvideo/train/methods/rl/advantages.py create mode 100644 fastvideo/train/methods/rl/data.py create mode 100644 fastvideo/train/methods/rl/diffusion.py create mode 100644 fastvideo/train/methods/rl/ema.py create mode 100644 fastvideo/train/methods/rl/embeddings.py create mode 100644 fastvideo/train/methods/rl/evaluation.py create mode 100644 fastvideo/train/methods/rl/genrl.py create mode 100644 fastvideo/train/methods/rl/pipeline.py create mode 100644 fastvideo/train/methods/rl/reward/__init__.py create mode 100644 fastvideo/train/methods/rl/reward/hpsv3.py create mode 100644 fastvideo/train/methods/rl/reward/ocr.py create mode 100644 fastvideo/train/methods/rl/reward/utils.py create mode 100644 fastvideo/train/methods/rl/reward/videoalign.py create mode 100644 fastvideo/train/methods/rl/rewards.py create mode 100644 fastvideo/train/methods/rl/sampling.py create mode 100644 fastvideo/train/methods/rl/sde.py create mode 100644 fastvideo/train/methods/rl/stat_tracking.py create mode 100644 fastvideo/train/models/wan/wan_genrl.py 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..00dfde7f72 --- /dev/null +++ b/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml @@ -0,0 +1,125 @@ +# GenRL / Video GRPO: Wan 2.1 T2V 1.3B with LongCat reweighting. +# +# Ported from GenRL/config/longcat.yaml. +# +# - Student: trainable (LoRA or full finetune) +# - Reference: frozen copy for KL penalty (optional, only for full finetune with beta > 0) +# +# Usage: +# torchrun --nnodes=1 --nproc_per_node=8 \ +# fastvideo/train/entrypoint/train.py \ +# --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 + 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 + + # ---- Data ---- + prompt_dataset_path: datasets/filtered_prompts + prompt_fn: filtered_prompts + + # ---- 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 + beta: 3.0e-4 + use_cfg: true + loss_reweighting: longcat + weight_advantages: true + max_grad_norm: 1.0 + seed: 42 + + # ---- Advantage computation ---- + per_prompt_stat_tracking: true + global_std: false + max_group_std: true + + # ---- EMA ---- + use_ema: true + ema_decay: 0.9 + ema_update_interval: 8 + +training: + distributed: + num_gpus: 8 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 8 + + 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: + 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 == 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 + run_name: wan_2_1_t2v_1_3b_longcat + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + max_grad_norm: 0.0 # Disabled; GenRLMethod clips internally. + +pipeline: + flow_shift: 3.0 diff --git a/fastvideo/train/methods/__init__.py b/fastvideo/train/methods/__init__.py index 61fd6ef2e7..8385298bfe 100644 --- a/fastvideo/train/methods/__init__.py +++ b/fastvideo/train/methods/__init__.py @@ -8,6 +8,7 @@ "FineTuneMethod", "SelfForcingMethod", "DiffusionForcingSFTMethod", + "GenRLMethod", ] @@ -24,4 +25,7 @@ def __getattr__(name: str) -> object: if name == "DiffusionForcingSFTMethod": from fastvideo.train.methods.fine_tuning.dfsft import DiffusionForcingSFTMethod return DiffusionForcingSFTMethod + if name == "GenRLMethod": + from fastvideo.train.methods.rl.genrl import GenRLMethod + return GenRLMethod 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/advantages.py b/fastvideo/train/methods/rl/advantages.py new file mode 100644 index 0000000000..f8fb09c562 --- /dev/null +++ b/fastvideo/train/methods/rl/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.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/data.py b/fastvideo/train/methods/rl/data.py new file mode 100644 index 0000000000..8ca0607d68 --- /dev/null +++ b/fastvideo/train/methods/rl/data.py @@ -0,0 +1,263 @@ +# 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): + """Repeat each prompt k times per global batch and + shard across ranks.""" ++ + 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 + 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.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() + repeated = [ + idx + for idx in indices + for _ in range(self.k) + ] + shuffled_idx = torch.randperm( + len(repeated), generator=g + ).tolist() + shuffled = [ + repeated[i] for i in shuffled_idx + ] + per_card = [] + for i in range(self.num_replicas): + start = i * self.batch_size + end = start + self.batch_size + per_card.append( + [ + (self.epoch, idx) + for idx in shuffled[start:end] + ] + ) + yield per_card[self.rank] + + 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/diffusion.py b/fastvideo/train/methods/rl/diffusion.py new file mode 100644 index 0000000000..3fc47639cf --- /dev/null +++ b/fastvideo/train/methods/rl/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.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/ema.py b/fastvideo/train/methods/rl/ema.py new file mode 100644 index 0000000000..20f4bcdf64 --- /dev/null +++ b/fastvideo/train/methods/rl/ema.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Exponential Moving Average wrapper for RL training.""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch + + +class EMAModuleWrapper: + """Maintains EMA copies of model parameters.""" + + def __init__( + self, + parameters: Iterable[torch.nn.Parameter], + decay: float = 0.9999, + update_step_interval: int = 1, + device: torch.device | None = None, + ): + parameters = list(parameters) + self.ema_parameters = [ + p.clone().detach().to(device) for p in parameters + ] + self.temp_stored_parameters = None + self.decay = decay + self.update_step_interval = update_step_interval + self.device = device + + def get_current_decay(self, optimization_step) -> float: + return min( + (1 + optimization_step) + / (10 + optimization_step), + self.decay, + ) + + @torch.no_grad() + def step( + self, + parameters: Iterable[torch.nn.Parameter], + optimization_step, + ): + parameters = list(parameters) + one_minus_decay = ( + 1 - self.get_current_decay(optimization_step) + ) + + if ( + optimization_step + 1 + ) % self.update_step_interval == 0: + for ema_p, p in zip( + self.ema_parameters, + parameters, + strict=True, + ): + if p.requires_grad: + if ema_p.device == p.device: + ema_p.add_( + one_minus_decay * (p - ema_p) + ) + else: + p_copy = p.detach().to(ema_p.device) + p_copy.sub_(ema_p) + p_copy.mul_(one_minus_decay) + ema_p.add_(p_copy) + del p_copy + + def to( + self, + device: torch.device = None, + dtype: torch.dtype = None, + ) -> None: + self.device = device + self.ema_parameters = [ + ( + p.to(device=device, dtype=dtype) + if p.is_floating_point() + else p.to(device=device) + ) + for p in self.ema_parameters + ] + + def copy_ema_to( + self, + parameters: Iterable[torch.nn.Parameter], + store_temp: bool = True, + ) -> None: + if store_temp: + self.temp_stored_parameters = [ + p.detach().cpu() for p in parameters + ] + parameters = list(parameters) + for ema_p, p in zip( + self.ema_parameters, parameters, strict=True + ): + p.data.copy_(ema_p.to(p.device).data) + + def copy_temp_to( + self, + parameters: Iterable[torch.nn.Parameter], + ) -> None: + for temp_p, p in zip( + self.temp_stored_parameters, + parameters, + strict=True, + ): + p.data.copy_(temp_p.data) + self.temp_stored_parameters = None + + def load_state_dict(self, state_dict: dict) -> None: + self.decay = state_dict.get("decay", self.decay) + self.ema_parameters = state_dict.get( + "ema_parameters" + ) + self.to(self.device) + + def state_dict(self) -> dict: + return { + "decay": self.decay, + "ema_parameters": self.ema_parameters, + } diff --git a/fastvideo/train/methods/rl/embeddings.py b/fastvideo/train/methods/rl/embeddings.py new file mode 100644 index 0000000000..a978fdc36a --- /dev/null +++ b/fastvideo/train/methods/rl/embeddings.py @@ -0,0 +1,46 @@ +# 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) + + with torch.no_grad(): + prompt_embeds = text_encoder( + text_input_ids + )[0] + + prompt_embeds = prompt_embeds.to( + dtype=text_encoder.dtype, device=device + ) + return prompt_embeds diff --git a/fastvideo/train/methods/rl/evaluation.py b/fastvideo/train/methods/rl/evaluation.py new file mode 100644 index 0000000000..37a657b6f4 --- /dev/null +++ b/fastvideo/train/methods/rl/evaluation.py @@ -0,0 +1,122 @@ +# 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.embeddings import ( + compute_text_embeddings, +) +from fastvideo.train.methods.rl.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, + transformer_params, + *, + 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, +) -> dict[str, float]: + """Run evaluation on test set. + + Returns: + Dict of aggregated eval metrics. + """ + model.transformer.eval() + + # Apply EMA weights if available. + if ema is not None: + ema.copy_ema_to(transformer_params, store_temp=True) + + all_rewards: dict[str, list[float]] = {} + + try: + for batch_idx, ( + _epoch_tag, + prompts, + metadata, + ) in enumerate(test_dataloader): + prompt_embeds = compute_text_embeddings( + prompts, + text_encoder, + tokenizer, + max_sequence_length=512, + device=device, + ) + + with torch.no_grad(): + ( + videos, + _latents, + _log_probs, + _kls, + _timesteps, + ) = wan_denoising_with_logprob( + model, + scheduler, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=( + sample_neg_prompt_embeds + ), + num_inference_steps=eval_num_steps, + guidance_scale=eval_guidance_scale, + height=height, + width=width, + num_frames=num_frames, + 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)) + + finally: + # Restore original weights. + if ema is not None: + ema.copy_temp_to(transformer_params) + + # 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/genrl.py b/fastvideo/train/methods/rl/genrl.py new file mode 100644 index 0000000000..600ff11f59 --- /dev/null +++ b/fastvideo/train/methods/rl/genrl.py @@ -0,0 +1,1033 @@ +# 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 copy +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.advantages import ( + compute_advantages, +) +from fastvideo.train.methods.rl.data import ( + build_prompt_dataloaders, +) +from fastvideo.train.methods.rl.diffusion import ( + compute_log_prob, +) +from fastvideo.train.methods.rl.embeddings import ( + compute_text_embeddings, +) +from fastvideo.train.methods.rl.ema import ( + EMAModuleWrapper, +) +from fastvideo.train.methods.rl.rewards import ( + multi_score, + reward_models_on_device, +) +from fastvideo.train.methods.rl.sampling import ( + sample_epoch, +) +from fastvideo.train.methods.rl.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) + + # 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_fn = multi_score( + torch.device("cpu"), + self._reward_cfg, + mc.get("reward_module"), + return_raw_scores=True, + ) + + # Prompt dataloaders. + wg = get_world_group() + self._world_size = wg.world_size + self._rank = wg.rank + self._is_main = wg.rank == 0 + + 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() + + # EMA. + self._init_ema() + + # 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._weight_advantages = bool( + mc.get("weight_advantages", False) + ) + self._max_grad_norm = float( + mc.get("max_grad_norm", 1.0) + ) + self._train_batch_size = int( + mc.get("train_batch_size", 8) + ) + + # 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 + + # EMA config. + self._use_ema = bool(mc.get("use_ema", True)) + self._ema_decay = float( + mc.get("ema_decay", 0.9) + ) + self._ema_update_interval = int( + mc.get("ema_update_interval", 8) + ) + + # Reward config. + self._reward_cfg = dict(mc.get("reward_fn", {})) + + # ------------------------------------------------------------------ + # 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 + ] + 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 _init_ema(self) -> None: + self._ema = None + if self._use_ema: + self._ema = EMAModuleWrapper( + self._transformer_params, + decay=self._ema_decay, + update_step_interval=( + self._ema_update_interval + ), + device=self.student.device, + ) + + 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. + self.student.transformer.eval() + with reward_models_on_device( + self._reward_cfg, device + ): + samples = 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 + ), + tracker=self.tracker, + ) + + # 2. Prepare samples (advantages). + samples = self._prepare_samples( + samples, epoch, iteration + ) + + # 3. PPO training. + ppo_metrics = self._ppo_train( + samples, epoch, iteration + ) + all_metrics.update(ppo_metrics) + + # Return dummy loss (everything is internal). + dummy_loss = torch.zeros( + (), device=device, requires_grad=False + ) + return ( + {"total_loss": dummy_loss}, + {}, + 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 + + 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) + + for sample in batched_list: + # 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: + # Reference model output (for KL). + prev_mean_ref = None + dt_sqrt_ref = None + if self._beta > 0: + ref_model = self._get_ref_model() + if ref_model is not None: + with torch.no_grad(): + ( + _, + _, + 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, + ) + + # Policy forward. + ( + _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, + ) + + # PPO loss. + advantages = torch.clamp( + sample["advantages"][:, j], + -self._adv_clip_max, + self._adv_clip_max, + ) + ratio = torch.exp( + log_prob + - sample["log_probs"][:, j] + ) + 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) + ) + + # Loss reweighting. + rw_scale = 1.0 + rw_scale_kl = 1.0 + if ( + self._loss_reweighting + == "longcat" + and self._sde_type == "flow_sde" + ): + rw_scale = ( + torch.sqrt( + sigma + / ( + 1 + - torch.where( + sigma == 1, + torch.tensor( + sigma_max, + device=( + sigma.device + ), + dtype=( + sigma.dtype + ), + ), + sigma, + ) + ) + ) + / dt_sqrt + ) + rw_scale = torch.mean(rw_scale) + rw_scale_kl = rw_scale**2 + + # KL loss. + 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 = ( + "Unknown sde_type: " + f"{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 + ) + else: + loss = rw_scale * policy_loss + + # Backward with gradient accumulation. + timestep_j = sample["timesteps"][ + :, j + ] + with set_forward_context( + current_timestep=timestep_j, + attn_metadata=None, + ): + (loss / num_ts).backward() + + # Track. + info["approx_kl"].append( + 0.5 + * torch.mean( + ( + log_prob + - sample["log_probs"][ + :, j + ] + ) + ** 2 + ) + .detach() + .item() + ) + info["clip_frac"].append( + torch.mean( + ( + torch.abs(ratio - 1.0) + > self._clip_range + ).float() + ) + .detach() + .item() + ) + info["policy_loss"].append( + policy_loss.detach().item() + ) + if self._beta > 0 and prev_mean_ref is not None: + info["kl_loss"].append( + kl_loss.detach().item() + ) + info["loss"].append( + loss.detach().item() + ) + + # Clip + step after all timesteps. + clip_grad_norm_if_needed( + self.student.transformer, + self._max_grad_norm, + ) + self._optimizer.step() + self._optimizer.zero_grad() + + # EMA. + if self._ema is not None: + self._ema.step( + self._transformer_params, + global_step, + ) + + # 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 + + # ------------------------------------------------------------------ + # 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 diff --git a/fastvideo/train/methods/rl/pipeline.py b/fastvideo/train/methods/rl/pipeline.py new file mode 100644 index 0000000000..dd35fc2d68 --- /dev/null +++ b/fastvideo/train/methods/rl/pipeline.py @@ -0,0 +1,318 @@ +# 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 +from typing import Any + +import torch + +from fastvideo.train.methods.rl.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 + + 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] = [] + + for i, t in enumerate(timesteps): + latents_ori = latents.clone() + timestep = t.expand(batch_size) + + # Conditional prediction. + 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 + ) + + # 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. + ( + 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, + ) + 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. + 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: + with ref_ctx: + 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) + ) + + # Decode to video. + videos = model.decode_latents(latents) + + return ( + videos, + all_latents, + all_log_probs, + all_kl, + all_timesteps, + ) 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..0551e6f1c4 --- /dev/null +++ b/fastvideo/train/methods/rl/reward/hpsv3.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +"""HPSv3 reward functions for visual quality assessment.""" + +from __future__ import annotations + +import os +import tempfile +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__) + +# Global cache of HPSv3 inferencers keyed by device. +_HPSV3_INFERENCERS: dict[str, Any] = {} + + +def _normalize_device(device) -> str: + if isinstance(device, torch.device): + return str(device) + return str(torch.device(device)) + + +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: + inf.to(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: + from hpsv3 import HPSv3Inferencer + except ImportError as exc: + msg = ( + "hpsv3 package not installed. " + "Install via: pip install hpsv3" + ) + raise ImportError(msg) from exc + inf = HPSv3Inferencer(device=device) + _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: + score = inf.score( + path, "A high-quality image" + ) + frame_scores.append( + _extract_reward_scalar(score) + ) + 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: + score = inf.score(path, prompt) + frame_scores.append( + _extract_reward_scalar(score) + ) + 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..91f9cef02c --- /dev/null +++ b/fastvideo/train/methods/rl/reward/utils.py @@ -0,0 +1,42 @@ +# 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): + images = images.transpose(0, 2, 3, 1) + elif images.ndim == 5: + # Video batch: (N, F, C, H, W) or (N, C, F, H, W) + if 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..9d93e6a369 --- /dev/null +++ b/fastvideo/train/methods/rl/reward/videoalign.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +"""VideoAlign reward functions for motion quality and +text-video alignment.""" + +from __future__ import annotations + +import os +import tempfile +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__) + +# Global cache of VideoAlign inferencers. +_VIDEOALIGN_INFERENCERS: dict[str, Any] = {} + + +def _normalize_device_str(device) -> str: + if isinstance(device, torch.device): + return str(device) + return str(torch.device(device)) + + +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 + inf.to(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.""" + key = _normalize_device_str(device) + cache_key = f"{checkpoint_path or 'default'}:{key}" + if cache_key not in _VIDEOALIGN_INFERENCERS: + try: + from videoalign import VideoAlignInferencer + except ImportError as exc: + msg = ( + "videoalign not installed. " + "Install from VideoAlign repo." + ) + raise ImportError(msg) from exc + + kwargs = {"device": device} + if checkpoint_path: + kwargs["checkpoint_path"] = checkpoint_path + inf = VideoAlignInferencer(**kwargs) + 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: + result = inf.score_video(path) + mq = float( + result.get("mq", result.get("avg", 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: + result = inf.score_video( + path, prompt=prompt + ) + ta = float( + result.get("ta", result.get("avg", 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/rewards.py b/fastvideo/train/methods/rl/rewards.py new file mode 100644 index 0000000000..b5f178c33f --- /dev/null +++ b/fastvideo/train/methods/rl/rewards.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Reward function loading and composition for RL training.""" + +from __future__ import annotations + +import importlib +import inspect +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 ( + set_hpsv3_device, +) +from fastvideo.train.methods.rl.reward.videoalign import ( + 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) + + +@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" + move_reward_models(reward_cfg, device) + try: + yield + finally: + move_reward_models(reward_cfg, "cpu") + if use_cuda: + import gc + + gc.collect() + torch.cuda.empty_cache() + else: + yield diff --git a/fastvideo/train/methods/rl/sampling.py b/fastvideo/train/methods/rl/sampling.py new file mode 100644 index 0000000000..6d8aea0bbe --- /dev/null +++ b/fastvideo/train/methods/rl/sampling.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Sampling epoch for RL training — generate videos and +compute rewards.""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from typing import Any + +import torch + +from fastvideo.logger import init_logger +from fastvideo.train.methods.rl.embeddings import ( + compute_text_embeddings, +) +from fastvideo.train.methods.rl.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: + g = torch.Generator(device=device) + g.manual_seed(base_seed + hash(prompt) % (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, +) -> list[dict[str, Any]]: + """Run one sampling epoch: generate videos, compute + rewards asynchronously. + + Returns: + List of sample dicts with prompt_ids, + prompt_embeds, latents, log_probs, kl, + timesteps, rewards. + """ + samples = [] + + 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 + + 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=epoch * SEED_EPOCH_STRIDE + i, + device=device, + ) + else: + gen = torch.Generator(device=device) + gen.manual_seed( + seed + epoch * SEED_EPOCH_STRIDE + i + ) + + 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, + ) + + 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) + ) + + # Async reward computation. + rewards_future = executor.submit( + reward_fn, + videos, + prompts, + prompt_metadata, + True, + ) + time.sleep(0) + + 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_future, + } + ) + + # Wait for all rewards. + for sample in samples: + rewards, _ = sample["rewards"].result() + sample["rewards"] = { + key: torch.as_tensor(value, device=device).float() + for key, value in rewards.items() + } + + return samples diff --git a/fastvideo/train/methods/rl/sde.py b/fastvideo/train/methods/rl/sde.py new file mode 100644 index 0000000000..1f676f5981 --- /dev/null +++ b/fastvideo/train/methods/rl/sde.py @@ -0,0 +1,192 @@ +# 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 + + log_prob = ( + -( + (prev_sample.detach() - prev_sample_mean) ** 2 + ) + / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2)) + - torch.log(std_dev_t * torch.sqrt(-1 * dt)) + - 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/stat_tracking.py b/fastvideo/train/methods/rl/stat_tracking.py new file mode 100644 index 0000000000..ea0ba4609e --- /dev/null +++ b/fastvideo/train/methods/rl/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..5304f46de5 --- /dev/null +++ b/fastvideo/train/models/wan/wan_genrl.py @@ -0,0 +1,140 @@ +# 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 typing import Any, TYPE_CHECKING + +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__) + + +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, + 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 + ), + ) + self.text_encoder: Any = None + self.tokenizer: Any = None + + # ------------------------------------------------------------------ + # 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) + + # Dummy dataloader for the trainer's outer loop. + self.dataloader = _InfiniteDummyLoader() + self.start_step = 0 + + def _load_text_encoder(self, model_path: str) -> None: + from transformers import ( + AutoTokenizer, + T5EncoderModel, + ) + + logger.info( + "Loading tokenizer from %s", model_path + ) + self.tokenizer = AutoTokenizer.from_pretrained( + model_path, subfolder="tokenizer" + ) + + logger.info( + "Loading T5 text encoder from %s", model_path + ) + dtype = self._get_training_dtype() + self.text_encoder = T5EncoderModel.from_pretrained( + model_path, + subfolder="text_encoder", + torch_dtype=dtype, + ) + self.text_encoder.to(self.device) + self.text_encoder.requires_grad_(False) + self.text_encoder.eval() + + def on_train_start(self) -> None: + """Skip negative conditioning (handled by method).""" From c678b9cfbad19a92c8be9a096230e21fa761e277 Mon Sep 17 00:00:00 2001 From: Peiyuan Zhang Date: Tue, 10 Mar 2026 19:04:50 +0000 Subject: [PATCH 2/7] improve ema --- AGENTS.md | 27 ++++ GenRL | 1 + .../genrl_wan2.1_t2v_1.3B_longcat.yaml | 10 +- fastvideo/train/methods/rl/ema.py | 121 ------------------ fastvideo/train/methods/rl/evaluation.py | 28 ++-- fastvideo/train/methods/rl/genrl.py | 34 ----- 6 files changed, 48 insertions(+), 173 deletions(-) create mode 160000 GenRL delete mode 100644 fastvideo/train/methods/rl/ema.py 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 index 00dfde7f72..8c5d8f62ab 100644 --- a/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml +++ b/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml @@ -29,7 +29,7 @@ method: reward_module: null # ---- Data ---- - prompt_dataset_path: datasets/filtered_prompts + prompt_dataset_path: data/filtered_prompts prompt_fn: filtered_prompts # ---- Sampling ---- @@ -70,11 +70,6 @@ method: global_std: false max_group_std: true - # ---- EMA ---- - use_ema: true - ema_decay: 0.9 - ema_update_interval: 8 - training: distributed: num_gpus: 8 @@ -120,6 +115,9 @@ training: 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/fastvideo/train/methods/rl/ema.py b/fastvideo/train/methods/rl/ema.py deleted file mode 100644 index 20f4bcdf64..0000000000 --- a/fastvideo/train/methods/rl/ema.py +++ /dev/null @@ -1,121 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Exponential Moving Average wrapper for RL training.""" - -from __future__ import annotations - -from collections.abc import Iterable - -import torch - - -class EMAModuleWrapper: - """Maintains EMA copies of model parameters.""" - - def __init__( - self, - parameters: Iterable[torch.nn.Parameter], - decay: float = 0.9999, - update_step_interval: int = 1, - device: torch.device | None = None, - ): - parameters = list(parameters) - self.ema_parameters = [ - p.clone().detach().to(device) for p in parameters - ] - self.temp_stored_parameters = None - self.decay = decay - self.update_step_interval = update_step_interval - self.device = device - - def get_current_decay(self, optimization_step) -> float: - return min( - (1 + optimization_step) - / (10 + optimization_step), - self.decay, - ) - - @torch.no_grad() - def step( - self, - parameters: Iterable[torch.nn.Parameter], - optimization_step, - ): - parameters = list(parameters) - one_minus_decay = ( - 1 - self.get_current_decay(optimization_step) - ) - - if ( - optimization_step + 1 - ) % self.update_step_interval == 0: - for ema_p, p in zip( - self.ema_parameters, - parameters, - strict=True, - ): - if p.requires_grad: - if ema_p.device == p.device: - ema_p.add_( - one_minus_decay * (p - ema_p) - ) - else: - p_copy = p.detach().to(ema_p.device) - p_copy.sub_(ema_p) - p_copy.mul_(one_minus_decay) - ema_p.add_(p_copy) - del p_copy - - def to( - self, - device: torch.device = None, - dtype: torch.dtype = None, - ) -> None: - self.device = device - self.ema_parameters = [ - ( - p.to(device=device, dtype=dtype) - if p.is_floating_point() - else p.to(device=device) - ) - for p in self.ema_parameters - ] - - def copy_ema_to( - self, - parameters: Iterable[torch.nn.Parameter], - store_temp: bool = True, - ) -> None: - if store_temp: - self.temp_stored_parameters = [ - p.detach().cpu() for p in parameters - ] - parameters = list(parameters) - for ema_p, p in zip( - self.ema_parameters, parameters, strict=True - ): - p.data.copy_(ema_p.to(p.device).data) - - def copy_temp_to( - self, - parameters: Iterable[torch.nn.Parameter], - ) -> None: - for temp_p, p in zip( - self.temp_stored_parameters, - parameters, - strict=True, - ): - p.data.copy_(temp_p.data) - self.temp_stored_parameters = None - - def load_state_dict(self, state_dict: dict) -> None: - self.decay = state_dict.get("decay", self.decay) - self.ema_parameters = state_dict.get( - "ema_parameters" - ) - self.to(self.device) - - def state_dict(self) -> dict: - return { - "decay": self.decay, - "ema_parameters": self.ema_parameters, - } diff --git a/fastvideo/train/methods/rl/evaluation.py b/fastvideo/train/methods/rl/evaluation.py index 37a657b6f4..d019c591b9 100644 --- a/fastvideo/train/methods/rl/evaluation.py +++ b/fastvideo/train/methods/rl/evaluation.py @@ -28,8 +28,7 @@ def eval_once( sample_neg_prompt_embeds: torch.Tensor, eval_reward_fn: Callable, global_step: int, - ema, - transformer_params, + ema_callback, *, eval_num_steps: int, eval_guidance_scale: float, @@ -44,18 +43,28 @@ def eval_once( ) -> 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]] = {} - # Apply EMA weights if available. - if ema is not None: - ema.copy_ema_to(transformer_params, store_temp=True) + # Use EMA context manager if available. + if ema_callback is not None: + ctx = ema_callback.ema_context( + model.transformer + ) + else: + from contextlib import nullcontext - all_rewards: dict[str, list[float]] = {} + ctx = nullcontext() - try: + with ctx: for batch_idx, ( _epoch_tag, prompts, @@ -105,11 +114,6 @@ def eval_once( else: all_rewards[key].append(float(val)) - finally: - # Restore original weights. - if ema is not None: - ema.copy_temp_to(transformer_params) - # Aggregate metrics. metrics = {} for key, vals in all_rewards.items(): diff --git a/fastvideo/train/methods/rl/genrl.py b/fastvideo/train/methods/rl/genrl.py index 600ff11f59..7931c0def3 100644 --- a/fastvideo/train/methods/rl/genrl.py +++ b/fastvideo/train/methods/rl/genrl.py @@ -41,9 +41,6 @@ from fastvideo.train.methods.rl.embeddings import ( compute_text_embeddings, ) -from fastvideo.train.methods.rl.ema import ( - EMAModuleWrapper, -) from fastvideo.train.methods.rl.rewards import ( multi_score, reward_models_on_device, @@ -156,9 +153,6 @@ def __init__( # Optimizer and scheduler. self._init_optimizer() - # EMA. - self._init_ema() - # Async reward executor. self._executor = futures.ThreadPoolExecutor( max_workers=8 @@ -255,15 +249,6 @@ def _parse_config(self, mc: dict[str, Any]) -> None: if self._num_video_per_prompt == 1: self._per_prompt_stat_tracking = False - # EMA config. - self._use_ema = bool(mc.get("use_ema", True)) - self._ema_decay = float( - mc.get("ema_decay", 0.9) - ) - self._ema_update_interval = int( - mc.get("ema_update_interval", 8) - ) - # Reward config. self._reward_cfg = dict(mc.get("reward_fn", {})) @@ -363,18 +348,6 @@ def _init_optimizer(self) -> None: scheduler_name=str(tc.optimizer.lr_scheduler), ) - def _init_ema(self) -> None: - self._ema = None - if self._use_ema: - self._ema = EMAModuleWrapper( - self._transformer_params, - decay=self._ema_decay, - update_step_interval=( - self._ema_update_interval - ), - device=self.student.device, - ) - def _compute_train_timesteps(self) -> None: if self._sde_window_size > 0: num_ts = self._sde_window_size @@ -1001,13 +974,6 @@ def _ppo_train( self._optimizer.step() self._optimizer.zero_grad() - # EMA. - if self._ema is not None: - self._ema.step( - self._transformer_params, - global_step, - ) - # Aggregate info for this inner epoch. for k, v in info.items(): all_info[k].extend(v) From b7dee6d5f009a804750527fe5792126ce56c2e6b Mon Sep 17 00:00:00 2001 From: Peiyuan Zhang Date: Tue, 10 Mar 2026 21:51:08 +0000 Subject: [PATCH 3/7] gen RL runing --- .../genrl_wan2.1_t2v_1.3B_longcat.yaml | 7 +- .../configs/genrl_wan2.1_t2v_1.3B_ocr.yaml | 112 ++++++++++++++++++ fastvideo/train/methods/__init__.py | 19 --- fastvideo/train/methods/rl/data.py | 4 +- fastvideo/train/methods/rl/pipeline.py | 32 +++-- fastvideo/train/methods/rl/reward/HPSv3 | 1 + fastvideo/train/methods/rl/reward/VideoAlign | 1 + fastvideo/train/methods/rl/reward/hpsv3.py | 18 ++- .../train/methods/rl/reward/videoalign.py | 35 ++++-- 9 files changed, 183 insertions(+), 46 deletions(-) create mode 100644 examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml create mode 160000 fastvideo/train/methods/rl/reward/HPSv3 create mode 160000 fastvideo/train/methods/rl/reward/VideoAlign 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 index 8c5d8f62ab..4752d3e01b 100644 --- a/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml +++ b/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml @@ -72,11 +72,12 @@ method: training: distributed: - num_gpus: 8 + + num_gpus: 4 sp_size: 1 tp_size: 1 - hsdp_replicate_dim: 1 - hsdp_shard_dim: 8 + hsdp_replicate_dim: 4 + hsdp_shard_dim: 1 data: # Not used by GenRL (prompt dataloaders are in method config) 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..25b06d54c9 --- /dev/null +++ b/examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml @@ -0,0 +1,112 @@ +# GenRL / Video GRPO: Wan 2.1 T2V 1.3B — OCR reward, full finetune, 4 GPUs. +# +# Usage: +# torchrun --nnodes=1 --nproc_per_node=4 \ +# fastvideo/train/entrypoint/train.py \ +# --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: data/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 + beta: 3.0e-4 + 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/fastvideo/train/methods/__init__.py b/fastvideo/train/methods/__init__.py index 8385298bfe..0a5a8f99ac 100644 --- a/fastvideo/train/methods/__init__.py +++ b/fastvideo/train/methods/__init__.py @@ -10,22 +10,3 @@ "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 - if name == "GenRLMethod": - from fastvideo.train.methods.rl.genrl import GenRLMethod - return GenRLMethod - raise AttributeError(name) diff --git a/fastvideo/train/methods/rl/data.py b/fastvideo/train/methods/rl/data.py index 8ca0607d68..b0e0a787bd 100644 --- a/fastvideo/train/methods/rl/data.py +++ b/fastvideo/train/methods/rl/data.py @@ -137,9 +137,7 @@ def collate_fn( class DistributedKRepeatSampler(Sampler): - """Repeat each prompt k times per global batch and - shard across ranks.""" -+ + def __init__( self, dataset: Dataset, diff --git a/fastvideo/train/methods/rl/pipeline.py b/fastvideo/train/methods/rl/pipeline.py index dd35fc2d68..9c1e8eb54a 100644 --- a/fastvideo/train/methods/rl/pipeline.py +++ b/fastvideo/train/methods/rl/pipeline.py @@ -98,16 +98,30 @@ def wan_denoising_with_logprob( latent_h = height // vae_scale_spatial latent_w = width // vae_scale_spatial - latents = torch.randn( - batch_size, - num_channels, - latent_frames, - latent_h, - latent_w, - generator=generator, - device=device, - dtype=torch.float32, + 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( 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/hpsv3.py b/fastvideo/train/methods/rl/reward/hpsv3.py index 0551e6f1c4..a1578dd8c5 100644 --- a/fastvideo/train/methods/rl/reward/hpsv3.py +++ b/fastvideo/train/methods/rl/reward/hpsv3.py @@ -4,7 +4,9 @@ from __future__ import annotations import os +import sys import tempfile +from pathlib import Path from typing import Any import numpy as np @@ -17,6 +19,13 @@ 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] = {} @@ -46,14 +55,15 @@ def _get_hpsv3_inferencer(device): key = _normalize_device(device) if key not in _HPSV3_INFERENCERS: try: - from hpsv3 import HPSv3Inferencer + from hpsv3 import HPSv3RewardInferencer except ImportError as exc: msg = ( - "hpsv3 package not installed. " - "Install via: pip install hpsv3" + "hpsv3 package not found. Ensure the " + "HPSv3 submodule is checked out under " + "fastvideo/train/methods/rl/reward/HPSv3" ) raise ImportError(msg) from exc - inf = HPSv3Inferencer(device=device) + inf = HPSv3RewardInferencer(device=device) _HPSV3_INFERENCERS[key] = inf return _HPSV3_INFERENCERS[key] diff --git a/fastvideo/train/methods/rl/reward/videoalign.py b/fastvideo/train/methods/rl/reward/videoalign.py index 9d93e6a369..90e53957ee 100644 --- a/fastvideo/train/methods/rl/reward/videoalign.py +++ b/fastvideo/train/methods/rl/reward/videoalign.py @@ -5,6 +5,7 @@ from __future__ import annotations import os +import sys import tempfile from typing import Any @@ -18,6 +19,14 @@ 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] = {} @@ -46,22 +55,32 @@ def _get_inferencer( checkpoint_path: str | None = None, ): """Get or create VideoAlign inferencer.""" + if checkpoint_path is None: + checkpoint_path = os.path.join( + os.path.dirname(__file__), + "..", "..", "..", "..", "..", + "data", "VideoReward", + ) + checkpoint_path = os.path.abspath(checkpoint_path) + key = _normalize_device_str(device) - cache_key = f"{checkpoint_path or 'default'}:{key}" + cache_key = f"{checkpoint_path}:{key}" if cache_key not in _VIDEOALIGN_INFERENCERS: try: - from videoalign import VideoAlignInferencer + from inference import VideoVLMRewardInference except ImportError as exc: msg = ( - "videoalign not installed. " - "Install from VideoAlign repo." + "VideoAlign not found. Ensure the " + "VideoAlign submodule is checked out " + "under fastvideo/train/methods/rl/" + "reward/VideoAlign" ) raise ImportError(msg) from exc - kwargs = {"device": device} - if checkpoint_path: - kwargs["checkpoint_path"] = checkpoint_path - inf = VideoAlignInferencer(**kwargs) + inf = VideoVLMRewardInference( + load_from_pretrained=checkpoint_path, + device=device, + ) inf._key_prefix = checkpoint_path or "default" _VIDEOALIGN_INFERENCERS[cache_key] = inf return _VIDEOALIGN_INFERENCERS[cache_key] From 0fb70ad9e67792a3e70781cfe87da1f9aca86639 Mon Sep 17 00:00:00 2001 From: Peiyuan Zhang Date: Tue, 10 Mar 2026 23:49:08 +0000 Subject: [PATCH 4/7] time profiling and log sampled videos --- .../genrl_wan2.1_t2v_1.3B_longcat.yaml | 5 + fastvideo/train/callbacks/__init__.py | 3 + fastvideo/train/callbacks/callback.py | 2 + fastvideo/train/callbacks/grad_clip.py | 3 +- fastvideo/train/callbacks/log_rl_samples.py | 127 ++++++++++++ fastvideo/train/methods/rl/genrl.py | 187 +++++++++++++----- fastvideo/train/methods/rl/pipeline.py | 33 ++++ fastvideo/train/methods/rl/reward/hpsv3.py | 12 +- .../train/methods/rl/reward/videoalign.py | 14 +- fastvideo/train/methods/rl/rewards.py | 14 ++ fastvideo/train/methods/rl/sampling.py | 53 ++++- fastvideo/train/models/wan/wan_genrl.py | 4 +- fastvideo/train/trainer.py | 3 + 13 files changed, 390 insertions(+), 70 deletions(-) create mode 100644 fastvideo/train/callbacks/log_rl_samples.py 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 index 4752d3e01b..73783b0230 100644 --- a/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml +++ b/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml @@ -27,6 +27,7 @@ method: videoalign_mq: 1.0 videoalign_ta: 1.0 reward_module: null + reward_on_gpu: true # ---- Data ---- prompt_dataset_path: data/filtered_prompts @@ -119,6 +120,10 @@ callbacks: ema: decay: 0.9 start_iter: 0 + log_rl_samples: + every_steps: 1 + max_videos: 4 + fps: 16 pipeline: flow_shift: 3.0 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/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/rl/genrl.py b/fastvideo/train/methods/rl/genrl.py index 7931c0def3..73eba7c430 100644 --- a/fastvideo/train/methods/rl/genrl.py +++ b/fastvideo/train/methods/rl/genrl.py @@ -13,7 +13,9 @@ from __future__ import annotations +import contextlib import copy +import time from collections import defaultdict from concurrent import futures from typing import Any @@ -42,6 +44,7 @@ compute_text_embeddings, ) from fastvideo.train.methods.rl.rewards import ( + move_reward_models, multi_score, reward_models_on_device, ) @@ -108,12 +111,24 @@ def __init__( ) # 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( - torch.device("cpu"), + 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() @@ -383,74 +398,132 @@ def single_train_step( } # 1. Sample epoch. + torch.cuda.synchronize() + t_sample_start = time.perf_counter() self.student.transformer.eval() - with reward_models_on_device( - self._reward_cfg, device - ): - samples = 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 - ), - tracker=self.tracker, + 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 + ), + tracker=self.tracker, + ) ) + torch.cuda.synchronize() + t_sample_end = time.perf_counter() # 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) + 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, ) @@ -756,8 +829,11 @@ def _ppo_train( self.student.transformer.train() info: dict[str, list] = defaultdict(list) + _ppo_batch_times: list[float] = [] for sample in batched_list: + torch.cuda.synchronize() + _ppo_batch_t0 = time.perf_counter() # Get embeddings. embeds = sample["prompt_embeds"] neg_embeds = ( @@ -973,7 +1049,20 @@ def _ppo_train( ) self._optimizer.step() 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) diff --git a/fastvideo/train/methods/rl/pipeline.py b/fastvideo/train/methods/rl/pipeline.py index 9c1e8eb54a..59ee42ca70 100644 --- a/fastvideo/train/methods/rl/pipeline.py +++ b/fastvideo/train/methods/rl/pipeline.py @@ -10,10 +10,15 @@ 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.sde import ( sde_step_with_logprob, ) @@ -178,11 +183,17 @@ def wan_denoising_with_logprob( 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, @@ -200,6 +211,8 @@ def wan_denoising_with_logprob( 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: @@ -216,6 +229,7 @@ def wan_denoising_with_logprob( cur_noise_level = noise_level # SDE step. + _sde_t0 = time.perf_counter() ( latents, log_prob, @@ -234,6 +248,7 @@ def wan_denoising_with_logprob( diffusion_clip=diffusion_clip, diffusion_clip_value=diffusion_clip_value, ) + _denoise_sde_time += time.perf_counter() - _sde_t0 prev_latents = latents.clone() # Record. @@ -249,6 +264,7 @@ def wan_denoising_with_logprob( 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() @@ -319,10 +335,27 @@ def wan_denoising_with_logprob( 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, diff --git a/fastvideo/train/methods/rl/reward/hpsv3.py b/fastvideo/train/methods/rl/reward/hpsv3.py index a1578dd8c5..b017653353 100644 --- a/fastvideo/train/methods/rl/reward/hpsv3.py +++ b/fastvideo/train/methods/rl/reward/hpsv3.py @@ -108,11 +108,11 @@ def _score(images, prompts, metadata, only_strict=False): for frame in frames: path = _save_frame_to_temp(frame) try: - score = inf.score( - path, "A high-quality image" + rewards = inf.reward( + ["A high-quality image"], [path] ) frame_scores.append( - _extract_reward_scalar(score) + _extract_reward_scalar(rewards[0][0]) ) finally: os.remove(path) @@ -148,9 +148,11 @@ def _score(images, prompts, metadata, only_strict=False): for frame in frames: path = _save_frame_to_temp(frame) try: - score = inf.score(path, prompt) + rewards = inf.reward( + [prompt], [path] + ) frame_scores.append( - _extract_reward_scalar(score) + _extract_reward_scalar(rewards[0][0]) ) finally: os.remove(path) diff --git a/fastvideo/train/methods/rl/reward/videoalign.py b/fastvideo/train/methods/rl/reward/videoalign.py index 90e53957ee..abfc2d31c5 100644 --- a/fastvideo/train/methods/rl/reward/videoalign.py +++ b/fastvideo/train/methods/rl/reward/videoalign.py @@ -134,10 +134,10 @@ def _score(images, prompts, metadata, only_strict=False): gray_frames = _convert_to_grayscale(frames) path = _save_video_to_temp(gray_frames) try: - result = inf.score_video(path) - mq = float( - result.get("mq", result.get("avg", 0)) + results = inf.reward( + [path], [""], use_norm=True ) + mq = float(results[0].get("MQ", 0)) batch_scores.append(mq) finally: os.remove(path) @@ -171,12 +171,10 @@ def _score(images, prompts, metadata, only_strict=False): ) path = _save_video_to_temp(frames) try: - result = inf.score_video( - path, prompt=prompt - ) - ta = float( - result.get("ta", result.get("avg", 0)) + results = inf.reward( + [path], [prompt], use_norm=True ) + ta = float(results[0].get("TA", 0)) batch_scores.append(ta) finally: os.remove(path) diff --git a/fastvideo/train/methods/rl/rewards.py b/fastvideo/train/methods/rl/rewards.py index b5f178c33f..a1e24b012a 100644 --- a/fastvideo/train/methods/rl/rewards.py +++ b/fastvideo/train/methods/rl/rewards.py @@ -5,6 +5,7 @@ import importlib import inspect +import time from collections.abc import Callable from contextlib import contextmanager @@ -158,15 +159,28 @@ 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: import gc gc.collect() torch.cuda.empty_cache() + _t3 = time.perf_counter() + logger.info( + "[rewards] move_to_cpu+gc=%.1fs", + _t3 - _t2, + ) else: yield diff --git a/fastvideo/train/methods/rl/sampling.py b/fastvideo/train/methods/rl/sampling.py index 6d8aea0bbe..cc663d0ef1 100644 --- a/fastvideo/train/methods/rl/sampling.py +++ b/fastvideo/train/methods/rl/sampling.py @@ -72,16 +72,27 @@ def sample_epoch( ref_transformer: torch.nn.Module | None = None, lora_model: Any | None = None, tracker: Any | None = None, -) -> list[dict[str, Any]]: +) -> tuple[ + list[dict[str, Any]], + list[torch.Tensor], + list[list[str]], +]: """Run one sampling epoch: generate videos, compute rewards asynchronously. Returns: - List of sample dicts with prompt_ids, - prompt_embeds, latents, log_probs, kl, - timesteps, rewards. + 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 = ( @@ -97,6 +108,10 @@ def sample_epoch( 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, @@ -125,6 +140,9 @@ def sample_epoch( seed + epoch * SEED_EPOCH_STRIDE + i ) + torch.cuda.synchronize() + _t_embed_done = time.perf_counter() + with torch.no_grad(): ( videos, @@ -156,6 +174,9 @@ def sample_epoch( 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) @@ -167,6 +188,10 @@ def sample_epoch( .repeat(sample_batch_size, 1) ) + # Collect decoded videos and prompts for logging. + all_videos.append(videos) + all_prompts.append(list(prompts)) + # Async reward computation. rewards_future = executor.submit( reward_fn, @@ -177,6 +202,17 @@ def sample_epoch( ) time.sleep(0) + 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, @@ -194,11 +230,18 @@ def sample_epoch( ) # Wait for all rewards. + torch.cuda.synchronize() + _t_reward_wait = time.perf_counter() for sample in samples: rewards, _ = sample["rewards"].result() 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 + return samples, all_videos, all_prompts diff --git a/fastvideo/train/models/wan/wan_genrl.py b/fastvideo/train/models/wan/wan_genrl.py index 5304f46de5..18a6010858 100644 --- a/fastvideo/train/models/wan/wan_genrl.py +++ b/fastvideo/train/models/wan/wan_genrl.py @@ -113,7 +113,7 @@ def init_preprocessors( def _load_text_encoder(self, model_path: str) -> None: from transformers import ( AutoTokenizer, - T5EncoderModel, + UMT5EncoderModel, ) logger.info( @@ -127,7 +127,7 @@ def _load_text_encoder(self, model_path: str) -> None: "Loading T5 text encoder from %s", model_path ) dtype = self._get_training_dtype() - self.text_encoder = T5EncoderModel.from_pretrained( + self.text_encoder = UMT5EncoderModel.from_pretrained( model_path, subfolder="text_encoder", torch_dtype=dtype, 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) From d9890ab3aeea2871944da389585c64f564b011c4 Mon Sep 17 00:00:00 2001 From: Peiyuan Zhang Date: Wed, 11 Mar 2026 00:57:06 +0000 Subject: [PATCH 5/7] sampled video looks correct --- fastvideo/train/methods/rl/embeddings.py | 13 +++++------ fastvideo/train/models/wan/wan_genrl.py | 28 ++++++++++++------------ 2 files changed, 20 insertions(+), 21 deletions(-) diff --git a/fastvideo/train/methods/rl/embeddings.py b/fastvideo/train/methods/rl/embeddings.py index a978fdc36a..719b3c35cf 100644 --- a/fastvideo/train/methods/rl/embeddings.py +++ b/fastvideo/train/methods/rl/embeddings.py @@ -34,13 +34,12 @@ def compute_text_embeddings( 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 - )[0] - - prompt_embeds = prompt_embeds.to( - dtype=text_encoder.dtype, device=device - ) + 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/models/wan/wan_genrl.py b/fastvideo/train/models/wan/wan_genrl.py index 18a6010858..f01d97695c 100644 --- a/fastvideo/train/models/wan/wan_genrl.py +++ b/fastvideo/train/models/wan/wan_genrl.py @@ -104,17 +104,20 @@ def init_preprocessors( # Load text encoder and tokenizer. model_path = str(training_config.model_path) - self._load_text_encoder(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) -> None: - from transformers import ( - AutoTokenizer, - UMT5EncoderModel, - ) + def _load_text_encoder( + self, + model_path: str, + training_config: TrainingConfig, + ) -> None: + from transformers import AutoTokenizer logger.info( "Loading tokenizer from %s", model_path @@ -124,17 +127,14 @@ def _load_text_encoder(self, model_path: str) -> None: ) logger.info( - "Loading T5 text encoder from %s", model_path + "Loading text encoder from %s", model_path ) - dtype = self._get_training_dtype() - self.text_encoder = UMT5EncoderModel.from_pretrained( - model_path, - subfolder="text_encoder", - torch_dtype=dtype, + self.text_encoder = load_module_from_path( + model_path=model_path, + module_type="text_encoder", + training_config=training_config, ) - self.text_encoder.to(self.device) self.text_encoder.requires_grad_(False) - self.text_encoder.eval() def on_train_start(self) -> None: """Skip negative conditioning (handled by method).""" From 78cbbebc7b3f71be76b24cf4ced7d0394fa325d1 Mon Sep 17 00:00:00 2001 From: Peiyuan Zhang Date: Wed, 11 Mar 2026 02:03:08 +0000 Subject: [PATCH 6/7] mv to utils --- .../configs/genrl_wan2.1_t2v_1.3B_longcat.yaml | 3 +-- fastvideo/train/methods/rl/genrl.py | 14 +++++++------- fastvideo/train/methods/rl/utils/__init__.py | 2 ++ .../train/methods/rl/{ => utils}/advantages.py | 2 +- fastvideo/train/methods/rl/{ => utils}/data.py | 0 .../train/methods/rl/{ => utils}/diffusion.py | 2 +- .../train/methods/rl/{ => utils}/embeddings.py | 0 .../train/methods/rl/{ => utils}/evaluation.py | 4 ++-- fastvideo/train/methods/rl/{ => utils}/pipeline.py | 2 +- fastvideo/train/methods/rl/{ => utils}/rewards.py | 0 fastvideo/train/methods/rl/{ => utils}/sampling.py | 4 ++-- fastvideo/train/methods/rl/{ => utils}/sde.py | 0 .../train/methods/rl/{ => utils}/stat_tracking.py | 0 13 files changed, 17 insertions(+), 16 deletions(-) create mode 100644 fastvideo/train/methods/rl/utils/__init__.py rename fastvideo/train/methods/rl/{ => utils}/advantages.py (99%) rename fastvideo/train/methods/rl/{ => utils}/data.py (100%) rename fastvideo/train/methods/rl/{ => utils}/diffusion.py (98%) rename fastvideo/train/methods/rl/{ => utils}/embeddings.py (100%) rename fastvideo/train/methods/rl/{ => utils}/evaluation.py (96%) rename fastvideo/train/methods/rl/{ => utils}/pipeline.py (99%) rename fastvideo/train/methods/rl/{ => utils}/rewards.py (100%) rename fastvideo/train/methods/rl/{ => utils}/sampling.py (98%) rename fastvideo/train/methods/rl/{ => utils}/sde.py (100%) rename fastvideo/train/methods/rl/{ => utils}/stat_tracking.py (100%) 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 index 73783b0230..3965f8a722 100644 --- a/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml +++ b/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml @@ -115,8 +115,7 @@ training: enable_gradient_checkpointing_type: full callbacks: - grad_clip: - max_grad_norm: 0.0 # Disabled; GenRLMethod clips internally. + # Gradnorm call back Disabled; GenRLMethod clips internally. ema: decay: 0.9 start_iter: 0 diff --git a/fastvideo/train/methods/rl/genrl.py b/fastvideo/train/methods/rl/genrl.py index 73eba7c430..24cc80243b 100644 --- a/fastvideo/train/methods/rl/genrl.py +++ b/fastvideo/train/methods/rl/genrl.py @@ -31,27 +31,27 @@ LogScalar, TrainingMethod, ) -from fastvideo.train.methods.rl.advantages import ( +from fastvideo.train.methods.rl.utils.advantages import ( compute_advantages, ) -from fastvideo.train.methods.rl.data import ( +from fastvideo.train.methods.rl.utils.data import ( build_prompt_dataloaders, ) -from fastvideo.train.methods.rl.diffusion import ( +from fastvideo.train.methods.rl.utils.diffusion import ( compute_log_prob, ) -from fastvideo.train.methods.rl.embeddings import ( +from fastvideo.train.methods.rl.utils.embeddings import ( compute_text_embeddings, ) -from fastvideo.train.methods.rl.rewards import ( +from fastvideo.train.methods.rl.utils.rewards import ( move_reward_models, multi_score, reward_models_on_device, ) -from fastvideo.train.methods.rl.sampling import ( +from fastvideo.train.methods.rl.utils.sampling import ( sample_epoch, ) -from fastvideo.train.methods.rl.stat_tracking import ( +from fastvideo.train.methods.rl.utils.stat_tracking import ( PerPromptStatTracker, ) from fastvideo.train.models.base import ModelBase 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/advantages.py b/fastvideo/train/methods/rl/utils/advantages.py similarity index 99% rename from fastvideo/train/methods/rl/advantages.py rename to fastvideo/train/methods/rl/utils/advantages.py index f8fb09c562..b4d935af9f 100644 --- a/fastvideo/train/methods/rl/advantages.py +++ b/fastvideo/train/methods/rl/utils/advantages.py @@ -8,7 +8,7 @@ import numpy as np from fastvideo.logger import init_logger -from fastvideo.train.methods.rl.stat_tracking import ( +from fastvideo.train.methods.rl.utils.stat_tracking import ( EPSILON, PerPromptStatTracker, ) diff --git a/fastvideo/train/methods/rl/data.py b/fastvideo/train/methods/rl/utils/data.py similarity index 100% rename from fastvideo/train/methods/rl/data.py rename to fastvideo/train/methods/rl/utils/data.py diff --git a/fastvideo/train/methods/rl/diffusion.py b/fastvideo/train/methods/rl/utils/diffusion.py similarity index 98% rename from fastvideo/train/methods/rl/diffusion.py rename to fastvideo/train/methods/rl/utils/diffusion.py index 3fc47639cf..e43a8e4382 100644 --- a/fastvideo/train/methods/rl/diffusion.py +++ b/fastvideo/train/methods/rl/utils/diffusion.py @@ -5,7 +5,7 @@ import torch -from fastvideo.train.methods.rl.sde import ( +from fastvideo.train.methods.rl.utils.sde import ( sde_step_with_logprob, ) diff --git a/fastvideo/train/methods/rl/embeddings.py b/fastvideo/train/methods/rl/utils/embeddings.py similarity index 100% rename from fastvideo/train/methods/rl/embeddings.py rename to fastvideo/train/methods/rl/utils/embeddings.py diff --git a/fastvideo/train/methods/rl/evaluation.py b/fastvideo/train/methods/rl/utils/evaluation.py similarity index 96% rename from fastvideo/train/methods/rl/evaluation.py rename to fastvideo/train/methods/rl/utils/evaluation.py index d019c591b9..1c10bf64b6 100644 --- a/fastvideo/train/methods/rl/evaluation.py +++ b/fastvideo/train/methods/rl/utils/evaluation.py @@ -9,10 +9,10 @@ import torch from fastvideo.logger import init_logger -from fastvideo.train.methods.rl.embeddings import ( +from fastvideo.train.methods.rl.utils.embeddings import ( compute_text_embeddings, ) -from fastvideo.train.methods.rl.pipeline import ( +from fastvideo.train.methods.rl.utils.pipeline import ( wan_denoising_with_logprob, ) diff --git a/fastvideo/train/methods/rl/pipeline.py b/fastvideo/train/methods/rl/utils/pipeline.py similarity index 99% rename from fastvideo/train/methods/rl/pipeline.py rename to fastvideo/train/methods/rl/utils/pipeline.py index 59ee42ca70..197ec7989b 100644 --- a/fastvideo/train/methods/rl/pipeline.py +++ b/fastvideo/train/methods/rl/utils/pipeline.py @@ -19,7 +19,7 @@ _pipeline_logger = init_logger(__name__) -from fastvideo.train.methods.rl.sde import ( +from fastvideo.train.methods.rl.utils.sde import ( sde_step_with_logprob, ) diff --git a/fastvideo/train/methods/rl/rewards.py b/fastvideo/train/methods/rl/utils/rewards.py similarity index 100% rename from fastvideo/train/methods/rl/rewards.py rename to fastvideo/train/methods/rl/utils/rewards.py diff --git a/fastvideo/train/methods/rl/sampling.py b/fastvideo/train/methods/rl/utils/sampling.py similarity index 98% rename from fastvideo/train/methods/rl/sampling.py rename to fastvideo/train/methods/rl/utils/sampling.py index cc663d0ef1..8f68730b23 100644 --- a/fastvideo/train/methods/rl/sampling.py +++ b/fastvideo/train/methods/rl/utils/sampling.py @@ -11,10 +11,10 @@ import torch from fastvideo.logger import init_logger -from fastvideo.train.methods.rl.embeddings import ( +from fastvideo.train.methods.rl.utils.embeddings import ( compute_text_embeddings, ) -from fastvideo.train.methods.rl.pipeline import ( +from fastvideo.train.methods.rl.utils.pipeline import ( wan_denoising_with_logprob, ) diff --git a/fastvideo/train/methods/rl/sde.py b/fastvideo/train/methods/rl/utils/sde.py similarity index 100% rename from fastvideo/train/methods/rl/sde.py rename to fastvideo/train/methods/rl/utils/sde.py diff --git a/fastvideo/train/methods/rl/stat_tracking.py b/fastvideo/train/methods/rl/utils/stat_tracking.py similarity index 100% rename from fastvideo/train/methods/rl/stat_tracking.py rename to fastvideo/train/methods/rl/utils/stat_tracking.py From bb5f60e0a1a75cd65cbf82f61dae21749af6544c Mon Sep 17 00:00:00 2001 From: Adam Lee Date: Sat, 23 May 2026 18:12:06 -0700 Subject: [PATCH 7/7] [bugfix]: stabilize GenRL reward and PPO training --- .../genrl_wan2.1_t2v_1.3B_longcat.yaml | 62 +- .../configs/genrl_wan2.1_t2v_1.3B_ocr.yaml | 7 +- examples/train/run.sh | 2 +- examples/train/run_slurm.sh | 2 +- fastvideo/layers/lora/linear.py | 20 +- fastvideo/train/callbacks/ema.py | 7 +- fastvideo/train/methods/rl/genrl.py | 759 +++++++++++++----- fastvideo/train/methods/rl/reward/hpsv3.py | 147 +++- fastvideo/train/methods/rl/reward/utils.py | 12 +- .../train/methods/rl/reward/videoalign.py | 296 ++++++- fastvideo/train/methods/rl/utils/data.py | 33 +- .../train/methods/rl/utils/evaluation.py | 12 +- fastvideo/train/methods/rl/utils/pipeline.py | 15 +- fastvideo/train/methods/rl/utils/rewards.py | 28 +- fastvideo/train/methods/rl/utils/sampling.py | 53 +- fastvideo/train/methods/rl/utils/sde.py | 22 +- fastvideo/train/models/wan/wan_genrl.py | 141 ++++ 17 files changed, 1343 insertions(+), 275 deletions(-) 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 index 3965f8a722..c102752298 100644 --- a/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml +++ b/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml @@ -2,12 +2,14 @@ # # Ported from GenRL/config/longcat.yaml. # -# - Student: trainable (LoRA or full finetune) -# - Reference: frozen copy for KL penalty (optional, only for full finetune with beta > 0) +# - 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=8 \ -# fastvideo/train/entrypoint/train.py \ +# 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: @@ -15,6 +17,19 @@ models: _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: @@ -30,13 +45,20 @@ method: reward_on_gpu: true # ---- Data ---- - prompt_dataset_path: data/filtered_prompts + prompt_dataset_path: GenRL/datasets/filtered_prompts prompt_fn: filtered_prompts # ---- Sampling ---- sample_batch_size: 4 eval_batch_size: 2 - num_batches_per_epoch: 1 + # 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 @@ -57,12 +79,21 @@ method: # ---- PPO training ---- train_batch_size: 4 num_inner_epochs: 1 - clip_range: 1.0e-3 + clip_range: 1.0e-4 adv_clip_max: 5.0 - beta: 3.0e-4 + # 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 - loss_reweighting: longcat + # 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 @@ -77,8 +108,9 @@ training: num_gpus: 4 sp_size: 1 tp_size: 1 - hsdp_replicate_dim: 4 - hsdp_shard_dim: 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) @@ -91,7 +123,9 @@ training: num_frames: 81 optimizer: - learning_rate: 1.0e-4 + # 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 @@ -109,7 +143,8 @@ training: tracker: project_name: VideoRL - run_name: wan_2_1_t2v_1_3b_longcat + # Leave blank so W&B auto-generates a unique display name per run. + run_name: "" model: enable_gradient_checkpointing_type: full @@ -119,6 +154,7 @@ callbacks: ema: decay: 0.9 start_iter: 0 + update_interval: 8 log_rl_samples: every_steps: 1 max_videos: 4 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 index 25b06d54c9..1b1c67db91 100644 --- a/examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml +++ b/examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml @@ -2,7 +2,7 @@ # # Usage: # torchrun --nnodes=1 --nproc_per_node=4 \ -# fastvideo/train/entrypoint/train.py \ +# -m fastvideo.train.entrypoint.train \ # --config examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml models: @@ -21,7 +21,7 @@ method: reward_module: null # ---- Data ---- - prompt_dataset_path: data/ocr + prompt_dataset_path: GenRL/datasets/ocr prompt_fn: general_ocr # ---- Sampling ---- @@ -50,7 +50,8 @@ method: num_inner_epochs: 1 clip_range: 1.0e-3 adv_clip_max: 5.0 - beta: 3.0e-4 + # No frozen reference model is configured in this launch. + beta: 0.0 use_cfg: true loss_reweighting: longcat weight_advantages: false 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/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/methods/rl/genrl.py b/fastvideo/train/methods/rl/genrl.py index 24cc80243b..824257d754 100644 --- a/fastvideo/train/methods/rl/genrl.py +++ b/fastvideo/train/methods/rl/genrl.py @@ -43,7 +43,11 @@ 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, @@ -101,6 +105,7 @@ def __init__( # Parse RL config. self._parse_config(mc) + self._validate_config() # Init student preprocessors (VAE, text encoder). self.student.init_preprocessors(tc) @@ -135,6 +140,14 @@ def __init__( 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( @@ -241,15 +254,44 @@ def _parse_config(self, mc: dict[str, Any]) -> None: 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)) @@ -267,6 +309,51 @@ def _parse_config(self, mc: dict[str, Any]) -> None: # 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 # ------------------------------------------------------------------ @@ -348,6 +435,18 @@ def _init_optimizer(self) -> None: 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, @@ -459,11 +558,23 @@ def single_train_step( 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() @@ -482,6 +593,54 @@ def single_train_step( 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 " @@ -746,6 +905,23 @@ def _prepare_samples( )[: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() } @@ -831,6 +1007,107 @@ def _ppo_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() @@ -847,154 +1124,14 @@ def _ppo_train( self._optimizer.zero_grad() for j in self._train_timesteps: - # Reference model output (for KL). - prev_mean_ref = None - dt_sqrt_ref = None - if self._beta > 0: - ref_model = self._get_ref_model() - if ref_model is not None: - with torch.no_grad(): - ( - _, - _, - 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, - ) - - # Policy forward. - ( - _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, - ) + if self._optimizer_step_per_timestep: + self._optimizer.zero_grad() - # PPO loss. - advantages = torch.clamp( - sample["advantages"][:, j], - -self._adv_clip_max, - self._adv_clip_max, - ) - ratio = torch.exp( - log_prob - - sample["log_probs"][:, j] - ) - 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) - ) - - # Loss reweighting. - rw_scale = 1.0 - rw_scale_kl = 1.0 - if ( - self._loss_reweighting - == "longcat" - and self._sde_type == "flow_sde" - ): - rw_scale = ( - torch.sqrt( - sigma - / ( - 1 - - torch.where( - sigma == 1, - torch.tensor( - sigma_max, - device=( - sigma.device - ), - dtype=( - sigma.dtype - ), - ), - sigma, - ) - ) - ) - / dt_sqrt - ) - rw_scale = torch.mean(rw_scale) - rw_scale_kl = rw_scale**2 - - # KL loss. - 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 = ( - "Unknown sde_type: " - f"{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 + loss, loss_metrics = ( + self._compute_ppo_loss_and_metrics( + sample, j, embeds, neg_embeds ) - else: - loss = rw_scale * policy_loss + ) # Backward with gradient accumulation. timestep_j = sample["timesteps"][ @@ -1004,51 +1141,49 @@ def _ppo_train( current_timestep=timestep_j, attn_metadata=None, ): - (loss / num_ts).backward() - - # Track. - info["approx_kl"].append( - 0.5 - * torch.mean( - ( - log_prob - - sample["log_probs"][ - :, j - ] + 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"] ) - ** 2 ) - .detach() - .item() - ) - info["clip_frac"].append( - torch.mean( - ( - torch.abs(ratio - 1.0) - > self._clip_range - ).float() + self._log_post_update_probe( + info, + sample, + j, + embeds, + neg_embeds, ) - .detach() - .item() - ) - info["policy_loss"].append( - policy_loss.detach().item() + 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, ) - if self._beta > 0 and prev_mean_ref is not None: - info["kl_loss"].append( - kl_loss.detach().item() - ) - info["loss"].append( - loss.detach().item() + self._optimizer.step() + self._lr_scheduler.step() + info["grad_norm"].append(grad_norm) + info["learning_rate"].append( + float(self._optimizer.param_groups[0]["lr"]) ) - - # Clip + step after all timesteps. - clip_grad_norm_if_needed( - self.student.transformer, - self._max_grad_norm, - ) - self._optimizer.step() - self._optimizer.zero_grad() + self._optimizer.zero_grad() torch.cuda.synchronize() _ppo_batch_times.append( time.perf_counter() - _ppo_batch_t0 @@ -1076,6 +1211,232 @@ def _ppo_train( ) 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 # ------------------------------------------------------------------ @@ -1086,3 +1447,29 @@ def _get_ref_model(self): 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.py b/fastvideo/train/methods/rl/reward/hpsv3.py index b017653353..a32402b329 100644 --- a/fastvideo/train/methods/rl/reward/hpsv3.py +++ b/fastvideo/train/methods/rl/reward/hpsv3.py @@ -28,6 +28,129 @@ # 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: @@ -36,6 +159,19 @@ def _normalize_device(device) -> str: 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) @@ -44,7 +180,7 @@ def set_hpsv3_device(device) -> None: # Move from any existing device. for old_key, inf in list(_HPSV3_INFERENCERS.items()): if old_key != key: - inf.to(device) + _move_hpsv3_inferencer(inf, device) _HPSV3_INFERENCERS[key] = inf del _HPSV3_INFERENCERS[old_key] return @@ -55,15 +191,18 @@ def _get_hpsv3_inferencer(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 = ( - "hpsv3 package not found. Ensure the " - "HPSv3 submodule is checked out under " - "fastvideo/train/methods/rl/reward/HPSv3" + "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] diff --git a/fastvideo/train/methods/rl/reward/utils.py b/fastvideo/train/methods/rl/reward/utils.py index 91f9cef02c..661da33028 100644 --- a/fastvideo/train/methods/rl/reward/utils.py +++ b/fastvideo/train/methods/rl/reward/utils.py @@ -24,11 +24,17 @@ def prepare_images( if images.ndim == 4: # Image batch: (N, C, H, W) or (N, H, W, C) - if images.shape[1] in (1, 3): + 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, C, H, W) or (N, C, F, H, W) - if images.shape[2] in (1, 3): + # 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): diff --git a/fastvideo/train/methods/rl/reward/videoalign.py b/fastvideo/train/methods/rl/reward/videoalign.py index abfc2d31c5..9e41a91b14 100644 --- a/fastvideo/train/methods/rl/reward/videoalign.py +++ b/fastvideo/train/methods/rl/reward/videoalign.py @@ -7,6 +7,7 @@ import os import sys import tempfile +from importlib import import_module, util from typing import Any import numpy as np @@ -29,6 +30,7 @@ # Global cache of VideoAlign inferencers. _VIDEOALIGN_INFERENCERS: dict[str, Any] = {} +_VIDEOALIGN_PATCHED = False def _normalize_device_str(device) -> str: @@ -37,6 +39,283 @@ def _normalize_device_str(device) -> str: 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) @@ -45,7 +324,7 @@ def set_videoalign_device(device) -> None: ): if old_key != key and old_key.split(":")[0] != key: new_key = inf._key_prefix + ":" + key - inf.to(device) + _move_videoalign_inferencer(inf, device) _VIDEOALIGN_INFERENCERS[new_key] = inf del _VIDEOALIGN_INFERENCERS[old_key] @@ -56,10 +335,13 @@ def _get_inferencer( ): """Get or create VideoAlign inferencer.""" if checkpoint_path is None: - checkpoint_path = os.path.join( - os.path.dirname(__file__), - "..", "..", "..", "..", "..", - "data", "VideoReward", + checkpoint_path = os.environ.get( + "VIDEOALIGN_CHECKPOINT_PATH", + os.path.join( + os.path.dirname(__file__), + "VideoAlign", + "checkpoints", + ), ) checkpoint_path = os.path.abspath(checkpoint_path) @@ -67,7 +349,8 @@ def _get_inferencer( cache_key = f"{checkpoint_path}:{key}" if cache_key not in _VIDEOALIGN_INFERENCERS: try: - from inference import VideoVLMRewardInference + inference_mod = _patch_videoalign_modules() + VideoVLMRewardInference = inference_mod.VideoVLMRewardInference except ImportError as exc: msg = ( "VideoAlign not found. Ensure the " @@ -81,6 +364,7 @@ def _get_inferencer( 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] diff --git a/fastvideo/train/methods/rl/utils/data.py b/fastvideo/train/methods/rl/utils/data.py index b0e0a787bd..db54011d40 100644 --- a/fastvideo/train/methods/rl/utils/data.py +++ b/fastvideo/train/methods/rl/utils/data.py @@ -154,12 +154,19 @@ def __init__( 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): @@ -169,28 +176,14 @@ def __iter__(self): indices = torch.randperm( len(self.dataset), generator=g )[: self.m].tolist() - repeated = [ - idx - for idx in indices + 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) ] - shuffled_idx = torch.randperm( - len(repeated), generator=g - ).tolist() - shuffled = [ - repeated[i] for i in shuffled_idx - ] - per_card = [] - for i in range(self.num_replicas): - start = i * self.batch_size - end = start + self.batch_size - per_card.append( - [ - (self.epoch, idx) - for idx in shuffled[start:end] - ] - ) - yield per_card[self.rank] def set_epoch(self, epoch: int): self.epoch = epoch diff --git a/fastvideo/train/methods/rl/utils/evaluation.py b/fastvideo/train/methods/rl/utils/evaluation.py index 1c10bf64b6..82f5090bd9 100644 --- a/fastvideo/train/methods/rl/utils/evaluation.py +++ b/fastvideo/train/methods/rl/utils/evaluation.py @@ -40,6 +40,8 @@ def eval_once( 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. @@ -70,6 +72,8 @@ def eval_once( 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, @@ -77,8 +81,13 @@ def eval_once( 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, @@ -90,13 +99,14 @@ def eval_once( scheduler, prompt_embeds=prompt_embeds, negative_prompt_embeds=( - sample_neg_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", ) diff --git a/fastvideo/train/methods/rl/utils/pipeline.py b/fastvideo/train/methods/rl/utils/pipeline.py index 197ec7989b..a4a79260e2 100644 --- a/fastvideo/train/methods/rl/utils/pipeline.py +++ b/fastvideo/train/methods/rl/utils/pipeline.py @@ -280,9 +280,8 @@ def wan_denoising_with_logprob( encoder_hidden_states=prompt_embeds, return_dict=False, ) - ref_noise = ref_noise.to(dtype) - if do_cfg: - with ref_ctx: + ref_noise = ref_noise.to(dtype) + if do_cfg: ref_uncond = ref_model( hidden_states=latents_ori.to( dtype @@ -293,11 +292,11 @@ def wan_denoising_with_logprob( ), return_dict=False, ) - ref_noise = ( - ref_uncond - + guidance_scale - * (ref_noise - ref_uncond) - ) + ref_noise = ( + ref_uncond + + guidance_scale + * (ref_noise - ref_uncond) + ) ( _, diff --git a/fastvideo/train/methods/rl/utils/rewards.py b/fastvideo/train/methods/rl/utils/rewards.py index a1e24b012a..d7017c2055 100644 --- a/fastvideo/train/methods/rl/utils/rewards.py +++ b/fastvideo/train/methods/rl/utils/rewards.py @@ -3,6 +3,7 @@ from __future__ import annotations +import gc import importlib import inspect import time @@ -20,9 +21,11 @@ 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, ) @@ -154,6 +157,28 @@ def move_reward_models(reward_cfg, device) -> None: 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.""" @@ -173,10 +198,9 @@ def reward_models_on_device(reward_cfg, device): _t2 = time.perf_counter() move_reward_models(reward_cfg, "cpu") if use_cuda: - import gc - gc.collect() torch.cuda.empty_cache() + torch.cuda.ipc_collect() _t3 = time.perf_counter() logger.info( "[rewards] move_to_cpu+gc=%.1fs", diff --git a/fastvideo/train/methods/rl/utils/sampling.py b/fastvideo/train/methods/rl/utils/sampling.py index 8f68730b23..09a6359501 100644 --- a/fastvideo/train/methods/rl/utils/sampling.py +++ b/fastvideo/train/methods/rl/utils/sampling.py @@ -4,6 +4,7 @@ from __future__ import annotations +import hashlib import time from collections.abc import Callable from typing import Any @@ -31,8 +32,15 @@ def create_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 + hash(prompt) % (2**31)) + g.manual_seed(base_seed + prompt_seed % (2**31)) generators.append(g) return generators @@ -72,13 +80,13 @@ def sample_epoch( 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, compute - rewards asynchronously. + """Run one sampling epoch: generate videos and compute rewards. Returns: Tuple of (samples, all_videos, all_prompts): @@ -131,7 +139,7 @@ def sample_epoch( if same_latent: gen = create_generator( prompts, - base_seed=epoch * SEED_EPOCH_STRIDE + i, + base_seed=seed + epoch * SEED_EPOCH_STRIDE + i, device=device, ) else: @@ -188,19 +196,25 @@ def sample_epoch( .repeat(sample_batch_size, 1) ) + videos_cpu = videos.detach().cpu() + # Collect decoded videos and prompts for logging. - all_videos.append(videos) + all_videos.append(videos_cpu) all_prompts.append(list(prompts)) - # Async reward computation. - rewards_future = executor.submit( - reward_fn, - videos, - prompts, - prompt_metadata, - True, - ) - time.sleep(0) + 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: " @@ -225,7 +239,7 @@ def sample_epoch( "next_latents": latents[:, 1:], "log_probs": log_probs, "kl": kl, - "rewards": rewards_future, + "rewards": rewards, } ) @@ -233,7 +247,14 @@ def sample_epoch( torch.cuda.synchronize() _t_reward_wait = time.perf_counter() for sample in samples: - rewards, _ = sample["rewards"].result() + 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() diff --git a/fastvideo/train/methods/rl/utils/sde.py b/fastvideo/train/methods/rl/utils/sde.py index 1f676f5981..7bbd86fb01 100644 --- a/fastvideo/train/methods/rl/utils/sde.py +++ b/fastvideo/train/methods/rl/utils/sde.py @@ -115,16 +115,22 @@ def sde_step_with_logprob( if deterministic: prev_sample = sample + dt * model_output - log_prob = ( - -( - (prev_sample.detach() - prev_sample_mean) ** 2 + 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, ) - / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2)) - - torch.log(std_dev_t * torch.sqrt(-1 * dt)) - - torch.log( - torch.sqrt(2 * torch.as_tensor(math.pi)) + 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( diff --git a/fastvideo/train/models/wan/wan_genrl.py b/fastvideo/train/models/wan/wan_genrl.py index f01d97695c..dbde67b494 100644 --- a/fastvideo/train/models/wan/wan_genrl.py +++ b/fastvideo/train/models/wan/wan_genrl.py @@ -10,8 +10,12 @@ 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, @@ -30,6 +34,105 @@ 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.""" @@ -57,6 +160,12 @@ def __init__( 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 @@ -79,9 +188,41 @@ def __init__( 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 # ------------------------------------------------------------------