NVIDIA-NeMo / NVIDIA-NeMo/RL

RFC: image/video diffusion (flow-matching) RL — flow-GRPO training

Open
#3,233 0 comments 0 reactions 1 assignee Claimed by @terrykong View on GitHub
accuracy community-request enhancement Feature waiting-on-maintainers
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Is your feature request related to a problem? Please describe.**

NeMo RL currently covers LLMs and VLMs (understanding), but has no support for RL post-training of image or video generation models. Flow-matching base models such as Qwen-Image benefit substantially from online RL against task rewards (OCR accuracy, human preference) — see [Flow-GRPO](https://arxiv.org/abs/2505.05470). Today users must leave the NeMo RL ecosystem to do this.

**Describe the solution you'd like**
Add a diffusion-GRPO training path: convert the flow-matching ODE into an equivalent SDE so per-step transition log-probs are computable in closed form, then apply the standard GRPO recipe (group-relative advantages, clipped policy gradients, optional Gaussian KL vs. the LoRA-disabled reference).

Naming note: to avoid confusion with diffusion language model work (e.g. #3167), everything here is visual generation (continuous flow-matching); the two efforts are independent and do not overlap in code. The initial work is scoped to text-to-image; video flow-matching models are a natural follow-up — they reuse the same SDE/GRPO machinery through the pipeline-adapter protocol.

I have a working implementation (Qwen-Image via Diffusers, LoRA + Ray DP workers, pluggable image reward environment, OCR task converging) and will upstream it as a series of focused, self-contained PRs. Each PR introduces its own scope, notes how it relates to the previously merged ones, and links back to this issue.

**Describe alternatives you've considered**
A single monolithic PR (~6.6k lines) — rejected in favor of reviewable, independently-testable increments.
Building on an inference engine for rollouts — not applicable: the training framework itself samples the SDE trajectory (log-probs must match the training graph), so no vLLM/refit machinery is involved.

**Additional context**
Reference implementation parity: hyperparameters and SDE semantics mirror the Flow-GRPO authors' repo and verl-omni's flowgrpo trainer; the exemplar config reproduces their Qwen-Image OCR LoRA recipe.
Scope/limitations of the initial series: Qwen-Image (Diffusers) only, LoRA + single-GPU-per-worker data parallelism (multi-node via Ray), no FSDP/Megatron sharding. The pipeline adapter is a Protocol, so other flow-matching pipelines (image or video) can be added incrementally.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.