improve ema

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

fastvideo/train/methods/rl/ema.py

@@ -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,
-        }

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():

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)