[feat] GenRL: add runtime and memory stability helpers (#1402)

Co-authored-by: Davids048 <jundasu@ucsd.edu>

Author: Abecid

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[*]:-}

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 "

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.
 
@@ -65,20 +67,27 @@ def eval_once(
         ctx = nullcontext()
 
     with ctx:
+        generator = torch.Generator(device=device)
         for batch_idx, (
             _epoch_tag,
             prompts,
             metadata,
         ) in enumerate(test_dataloader):
+            if max_batches is not None and batch_idx >= max_batches:
+                break
             prompt_embeds = compute_text_embeddings(
                 prompts,
                 text_encoder,
                 tokenizer,
                 max_sequence_length=512,
                 device=device,
             )
+            neg_prompt_embeds = sample_neg_prompt_embeds[
+                : len(prompts)
+            ]
 
             with torch.no_grad():
+                generator.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",
                 )

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)
+                        )
 
                 (
                     _,

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,15 +239,22 @@ def sample_epoch(
                 "next_latents": latents[:, 1:],
                 "log_probs": log_probs,
                 "kl": kl,
-                "rewards": rewards_future,
+                "rewards": rewards,
             }
         )
 
     # Wait for all rewards.
     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()

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=math.sqrt(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(