diff --git a/GenRL b/GenRL new file mode 160000 index 000000000..82a867ed2 --- /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 00dfde7f7..8c5d8f62a 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 20f4bcdf6..000000000 --- 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 37a657b6f..d019c591b 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() - - # 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: + # Use EMA context manager if available. + if ema_callback is not None: + ctx = ema_callback.ema_context( + model.transformer + ) + else: + from contextlib import nullcontext + + ctx = nullcontext() + + with ctx: for batch_idx, ( _epoch_tag, prompts, @@ -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 600ff11f5..7931c0def 100644 --- a/fastvideo/train/methods/rl/genrl.py +++ b/fastvideo/train/methods/rl/genrl.py @@ -41,9 +41,6 @@ from fastvideo.train.methods.rl.diffusion import ( 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 @@ class GenRLMethod(TrainingMethod): # Optimizer and scheduler. self._init_optimizer() - # EMA. - self._init_ema() - # Async reward executor. self._executor = futures.ThreadPoolExecutor( max_workers=8 @@ -255,15 +249,6 @@ class GenRLMethod(TrainingMethod): 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 @@ class GenRLMethod(TrainingMethod): 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 @@ class GenRLMethod(TrainingMethod): 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)