[bugfix]: stabilize GenRL reward and PPO training - #1391
Conversation
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Welcome to FastVideo! Thanks for your first pull request.
How our CI works:
PRs run a two-tier CI system:
- Pre-commit — formatting (yapf), linting (ruff), type checking (mypy). Runs immediately on every PR.
- Fastcheck — core GPU tests (encoders, VAEs, transformers, kernels, unit tests). Runs automatically via Buildkite on relevant file changes (~10-15 min).
- Full Suite — integration tests, training pipelines, SSIM regression. Runs only when a reviewer adds the
readylabel.
Before your PR is reviewed:
-
pre-commit run --all-filespasses locally - You've added or updated tests for your changes
- The PR description explains what and why
If pre-commit fails, a bot comment will explain how to fix it. Fastcheck and Full Suite results appear in the Checks section below.
Useful links:
mergify
Bot
added
type: bugfix
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces LoRA support for Wan-based video RL training, implements PPO microbatch accumulation, and adds an evaluation suite during training. It also includes extensive compatibility patches for reward models, improves deterministic sampling, and optimizes GPU memory management. Review feedback identified several issues in the training logic, such as incorrect diagnostic capturing in the microbatch path, missing gradient resets for certain optimizer settings, and a validation mismatch regarding batch size divisibility. It was also suggested to replace assertions in the LoRA forward pass with explicit runtime errors to improve robustness.
| if probe_args is None: | ||
| probe_args = ( | ||
| sample, | ||
| j, | ||
| embeds, | ||
| neg_embeds, | ||
| ) |
There was a problem hiding this comment.
In the accumulated microbatch path, probe_args is captured only once at the beginning of the first timestep (j=0). When _log_post_update_probe is called inside the j loop (line 1074), it will always recompute diagnostics for timestep 0, regardless of which timestep j was actually updated. This makes the post-update KL diagnostics incorrect for all timesteps except the first one. probe_args should be updated for each timestep j to use a sample corresponding to that timestep.
probe_sample = batched_list[0]
probe_args = (
probe_sample,
j,
probe_sample["prompt_embeds"],
self._train_neg_embeds[:len(probe_sample["prompt_embeds"])] if self._use_cfg else None,
)| and len(batched_list) > 1 | ||
| ): | ||
| probe_args = None | ||
| for j in self._train_timesteps: |
There was a problem hiding this comment.
When accumulate_ppo_microbatches is True and optimizer_step_per_timestep is False, the optimizer gradients are never zeroed at the start of the inner epoch. This can lead to accumulating gradients from previous steps or epochs. An explicit self._optimizer.zero_grad() should be added before starting the timestep/microbatch loops.
| for j in self._train_timesteps: | |
| if not self._optimizer_step_per_timestep: | |
| self._optimizer.zero_grad() | |
| for j in self._train_timesteps: |
| 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 |
There was a problem hiding this comment.
Using assert in the forward pass is generally discouraged for production code as it can be optimized away with the -O flag. It is better to use a proper conditional check and raise a ValueError or RuntimeError if the LoRA parameters are missing when they are expected.
| assert lora_A is not None and lora_B is not None | |
| if lora_A is None or lora_B is None: | |
| raise RuntimeError("LoRA weights (lora_A, lora_B) must be initialized before forward pass when LoRA is enabled.") |
| 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}." | ||
| ) |
There was a problem hiding this comment.
The divisibility check here uses total_samples (global batch size), but the DistributedKRepeatSampler (in fastvideo/train/methods/rl/utils/data.py) enforces a stricter requirement: the per-rank batch_size must be divisible by num_video_per_prompt to ensure whole prompt groups are kept on the same rank. To avoid a ValueError later in the sampler initialization, this check should be updated to match.
if self._sample_batch_size % self._num_video_per_prompt != 0:
raise ValueError(
"sample_batch_size (per rank) must be divisible by "
"num_video_per_prompt for DistributedKRepeatSampler. Got "
f"{self._sample_batch_size} and {self._num_video_per_prompt}."
)There was a problem hiding this comment.
this is tied to the first change. Keep repeated samples from the same prompt together in DistributedKRepeatSampler.
should patch to enforce the rule
| 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() |
There was a problem hiding this comment.
The _log_post_update_probe call is missing in the non-accumulated path when optimizer_step_per_timestep is False. This results in missing post-update KL diagnostics for this specific configuration, which is inconsistent with the other branches.
Davids048
left a comment
Davids048
left a comment
There was a problem hiding this comment.
Currently the PR combines runtime/OOM fixes, reward-model compatibility patches, sampler behavior, PPO/GRPO training-loop changes, LoRA support, and config/default changes. That makes it hard to attribute regressions in the future.
I would recommend splitting this into a small stack of PRs. For example:
- OOM fixes
- Reward model compatibility fixes
- prompt group placement
- training loop changes
- LoRA support for GenRL
(Don't need to follow this exactly tho.)
There is an agent skill that may help: link
The Gemini comments look valid to me and should be addressed before merge.
3 additional issues I think are important:
1. LoRA is inserted after the Wan transformer has already gone through the FSDP/HSDP loader path
In fastvideo/train/models/wan/wan_genrl.py, GenRLWanModel.__init__() calls the parent WanModel init first, which loads the transformer through load_module_from_path(). After that, when use_lora=True, _apply_fastvideo_lora() replaces modules and creates new LoRA parameters.
That means the LoRA parameters are introduced after the transformer has already gone through the distributed loader/sharding path. Those newly inserted parameters may not be managed with the same FSDP/HSDP device mesh or checkpointing assumptions as the rest of the model.
2. Synchronous GPU reward mode still retains decoded CPU videos for all rollout batches before scoring
With reward_on_gpu: true, async reward scoring is disabled, but sample_epoch() still stores decoded videos_cpu for each rollout batch and delays reward scoring until after all batches are sampled. With num_batches_per_epoch: 4, this can retain multiple large decoded video tensors on CPU at once.
3. Config naming
examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml is confusing because FastVideo already has an actual LongCat-Video model. The "longcat" should be more specific and point to the reward method(s) used.Please rename it to something more explicit.
It is fine to document inside the file that the recipe is derived from GenRL's LongCat config.
Test
We should add testing (CI/unit tests) on the roadmap for this development.
|
Thanks for the review. I split this mega PR into six focused PRs so each concern can be reviewed independently:
I also addressed the Gemini comments in the split PRs:
For your three additional issues:
I’m going to treat #1391 as superseded by this stack unless you prefer I keep it open for reference. |
|
This PR has merge conflicts with the base branch. Please rebase: git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease |
Extracted from hao-ai-lab#1391. Renames the LongCat-labelled recipe to reflect the reward setup. GenRL-Stack: 6/6
|
closing for now, feel free to repo |
4 participants
Purpose
Stabilize the GenRL / diffusion RL training path on top of
py/add_rl.This PR focuses on getting the RL pipeline to run more reliably for Wan T2V training with multi-reward scoring, especially HPSv3 + VideoAlign rewards, and reducing PPO/GRPO instability from noisy rollout grouping and reward-model memory pressure.
Fixes #
Changes
Keep repeated samples from the same prompt together in
DistributedKRepeatSampler.sample_batch_sizeis divisible bynum_video_per_prompt.Add safer PPO/GRPO training behavior.
accumulate_ppo_microbatches.flash_tgrloss reweighting mode to avoid overly large LongCat timestep reweighting during full-parameter training.1e-4.Improve reward stack robustness.
Add/adjust GenRL training configs.
Running Experiments
Example local/multi-GPU launch:
Experiment Environment Notes
I verified the GenRL run in a Modal environment roughly equivalent to:
nvidia/cuda:12.8.1-devel-ubuntu22.043.12ffmpeg,libgl1,libglib2.0-0,build-essential,ninja-build,cmake,git-lfsdatasets==3.6.0matplotlib==3.10.3peft==0.10.0qwen-vl-utils==0.0.11safetensors==0.5.3timm==1.0.15trl==0.8.6numpy==1.26.4scipy==1.15.2flash_attn-2.8.3+cu128torch2.10prebuilt wheelThe Modal launcher itself is not included in this PR; these notes are only to document the environment used for testing/debugging.
Performance
Confirmed the following reward metrics going up over 40 iterations: