[fsdp,trainer,vllm_omni,algo] feat: support FlowGRPO training for QwenImage

[zhtmike]

What does this PR do?

Follow-up Work for #4639

• A trainable script for the FlowGRPO algorithm for Qwen-Image is provided.
• Support for vLLM-Omni has been added to the rollout engine.
• Diffusers has been integrated as the training engine for the diffusion model.
• The scripts have been tested with training using cfg/non-cfg, w/o kl loss, lora/full-model, fsdp/fsdp2, collocated mode, and async reward with a standalone resource pool.
• Unit tests for diffusion rollout/reward and diffusers training engines have been added.

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here: ...
• Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
  • {modules} include fsdp, megatron, veomni, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data, cfg, reward
  • If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
  • {type} is in feat, fix, refactor, chore, test
  • If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
  • Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

Install Method

We provide a quick way to test with our examples codes. The formal support will be released together with vllm-omni 0.17.

pip install vllm==0.16
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni && git reset --hard 92e1a544 && pip install . && cd ../

Experiments were conducted on 4 NVIDIA A100 GPUs using the OCR reward (Levenshtein distance) with Qwen3-VL-8B-Instruct model.

Test Dataset: OCR Dataset
Training Sample Size: 19653
Test Sample Size: 1018

Training GPU hours required to reach and maintain a validation reward score of approximately 0.9:

Repo	Model	Algorithm	Hybrid Engine	# Cards	Reward Fn	# GPUs for Actor	# GPUs for Rollout	# GPUs for Async Reward	Batch Size	rollout.n	Learning Rate	# Val Samples	Throughput (Samples / Seconds)	# GPU Hour
verl-omni (ours)	Qwen-Image	Flow-GRPO-Fast	True	4+1	qwenvl-ocr-vllm	4	4	1	32	16	3e-4	1k (full set)	0.04	49
Flow-GRPO	Qwen-Image	Flow-GRPO-Fast	True	4+1	qwenvl-ocr-vllm	4	4	1	32	16	3e-4	1k (full set)	0.03	65

• Result reported based on the script examples/flowgrpo_trainer/run_flowgrpo_async_reward.sh

Some results from wandb logging:

Loss:

Validation Score

Performance:

Visualization Result:

Prompt: A medieval knight holds a shield emblazoned with the crest "Fortis et Fidelis", standing resolutely in a sunlit courtyard, surrounded by ancient stone walls and banners fluttering in the breeze.

Before RL	After RL

For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc.

API and Usage Example

Demonstrate how the API changes if any, and provide usage example(s) if possible.

# Add code snippet or script demonstrating how to use this

Design & Code Changes

Please refer to #4639

Demonstrate the high-level design if this PR is complex, and list the specific changes.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update the documentation.
• Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
• Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)
• If your PR is related to the recipe submodule, please also update the reference to the submodule commit via git submodule update --remote or cd recipe && git pull origin main.

[lrq619]

Hi thanks for the great PR, I'm also trying out RL postraining for flowGRPO
Is it possible for you to post the reward curve during training steps? I want to validate whether my training is correct, thank you!

[zhtmike]

Hi thanks for the great PR, I'm also trying out RL postraining for flowGRPO Is it possible for you to post the reward curve during training steps? I want to validate whether my training is correct, thank you!

The reward curve can be found in the chart val-aux/flow_grpo/ocr/score/mean@1 mentioned in the PR description. It was run based on the script examples/flowgrpo_trainer/run_flowgrpo.sh with the training with CFG :)

[lrq619]

Hi thanks for the great PR, I'm also trying out RL postraining for flowGRPO Is it possible for you to post the reward curve during training steps? I want to validate whether my training is correct, thank you!

The reward curve can be found in the chart val-aux/flow_grpo/ocr/score/mean@1 mentioned in the PR description. It was run based on the script examples/flowgrpo_trainer/run_flowgrpo.sh with the training with CFG :)

Thanks! Is it possible to get the reward on training set e.g. (critic/rewards/mean)? My reward on training set cannot converge. Is it normal that in the initial steps, the reward starts from 0.2? Because I see your validation set reward is around 0.9 even in step 0

[zhtmike]

Thanks! Is it possible to get the reward on training set e.g. (critic/rewards/mean)? My reward on training set cannot converge. Is it normal that in the initial steps, the reward starts from 0.2? Because I see your validation set reward is around 0.9 even in step 0

This is the reward on training set.

If you are training Qwen-Image with CFG, the score should be sufficiently high because the model already performs quite well on OCR tasks. If you are training without CFG, a relatively low score at the beginning should be resonable.

[lrq619]

Thanks! Is it possible to get the reward on training set e.g. (critic/rewards/mean)? My reward on training set cannot converge. Is it normal that in the initial steps, the reward starts from 0.2? Because I see your validation set reward is around 0.9 even in step 0

This is the reward on training set.

If you are training Qwen-Image with CFG, the score should be sufficiently high because the model already performs quite well on OCR tasks. If you are training without CFG, a relatively low score at the beginning should be resonable.

Thanks for the information, that's really helpful. I have fixed the model and it starts from ~0.8 now.
I have a question regarding the performance. I see that your performance is around 500s/step under 4 A100 GPUs. But I tried my performance is around 2000s/step even under 4 H800 GPUs. May I know did you use the full denoising steps during rollout phase (50 steps) as in the configuration?

[zhtmike]

Thanks for the information, that's really helpful. I have fixed the model and it starts from ~0.8 now. I have a question regarding the performance. I see that your performance is around 500s/step under 4 A100 GPUs. But I tried my performance is around 2000s/step even under 4 H800 GPUs. May I know did you use the full denoising steps during rollout phase (50 steps) as in the configuration?

We follow the paper setting, use 10 steps at the rollout phase at training stage, 50 steps at the evaluation stage.

[BesmingY]

Hello, may I ask which repository should be cloned to directly reproduce the training curves you reported? Additionally, could you please specify the exact versions of vllm and vllm_omnithat were used? Thank you for your help.

[zhtmike]

Hello, may I ask which repository should be cloned to directly reproduce the training curves you reported? Additionally, could you please specify the exact versions of vllm and vllm_omnithat were used? Thank you for your help.

Hi, you can refer to the PR branch for validating. And the version of the vllm and vllm-omni is listed in the Install Method under Test in PR description.

[BesmingY]

Hello, may I ask which repository should be cloned to directly reproduce the training curves you reported? Additionally, could you please specify the exact versions of vllm and vllm_omnithat were used? Thank you for your help.

Hi, you can refer to the PR branch for validating. And the version of the vllm and vllm-omni is listed in the Install Method under Test in PR description.

Hello, may I ask which repository should be cloned to directly reproduce the training curves you reported? Additionally, could you please specify the exact versions of vllm and vllm_omnithat were used? Thank you for your help.

Hi, you can refer to the PR branch for validating. And the version of the vllm and vllm-omni is listed in the Install Method under Test in PR description.

Hi, thank you for your reply! I noticed that your PR repository is still being updated. Could you please confirm whether I should use the latest vllm-omni-pr​ or the vllm-omni-20260211​ branch to reproduce the results?

[zhtmike]

Hi, thank you for your reply! I noticed that your PR repository is still being updated. Could you please confirm whether I should use the latest vllm-omni-pr​ or the vllm-omni-20260211​ branch to reproduce the results?

This PR will be no longer updated except small bugs/typo fixes (which will not affect the reproduction result). You can use the latest version of vllm-omni-pr​ for your test.