[RFC] Verl recipe for image generation model

[Franklin-Zhang0]

Motivation

Reinforcement learning for autoregressive image generation with textual chain-of-though is attracting growing interest. A recent project, ReasonGen-R1, employs GRPO via Verl to enhance an image generation model’s ability to follow instructions. We plan to contribute this training pipeline as a Verl recipe to benefit the wider community.

Overall Structure

/recipe/image_generation
      config/
           image_generation_rl.yaml
           image_generation_sft.yaml
      main.py
      ray_trainer.py
      sft_trainer.py
      fsdp_worker.py
      dp_actor.py
      hf_rollout.py
      datasets/
            rl_datasets.py
            sft_datasets.py
      Janus/

Proposed Major Changes

New Function or Classes

ImageRewardModelWorker in fsdp_worker.py: We need to implement a reward model that takes in the generated images and prompts to assess the image quality and instruction following. The overall structure is similar to the current RewardModelWorker.
Janus-model: An example model class for RL training. Need to implement inside the recipe for interleaved image-text generation and official genrate and forward function is not released by deepseek.
AdaptiveEntropyCoefficient in dp_actor.py: Adaptive entorpy loss coefficient for stable training in text-image interleaved RL training. It updates using the target entropy and the entropy of output logits.

Function or Classes that needs modification:

FSDPSFTTrainer in FSDPSFTTrainer: support sft training for image generation model
HFSFTDataset in sft_dataset.py and RLDataset in rl_dataset.py: support data loading and formating for image generation
_build_model_optimizer in ActorRolloutRefWorker: modify to support janus loading
update_actor in DataParallelPPOActor: handle the adaptive entropy and seperate entropy computation for text and image
hf_rollout: return generated images and generated texts as seperate output.
fit in ray_trainer.py: support group_filtering in dapo

CC

@eric-haibin-lin

[vermouth1992]

Could you give an example of image generation model?

[Franklin-Zhang0]

An example image generation model we want to integrate is Janus-Pro, which is also the base model in ReasonGen-R1

[eric-haibin-lin]

Thanks! What code goes in Janus-Pro? Is it possible to reuse from huggingface/transformers instead of copying all its code?

[Franklin-Zhang0]

I just checked the latest huggingface/transformers and found that they updated the janus-pro model last week.
So, yes. We can reuse code from there. For pure image generation or pure text generation, reusing their code would be enough.

We might still need to customize some generation functions for generating interleaved text and image, which is not originally supported by janus-pro.

[eric-haibin-lin]

Thanks. For rollout is only hf rollout supported? What's the issue with vllm/sglang?

[Franklin-Zhang0]

By the time I implement it, janus is not supported by vllm/sglang. I think sglang support janus now. However, I think it still doesn't support interleaved generation with text and images. But maybe I can do some customization to integrate vllm/sglang.

[eric-haibin-lin]

I'm okay with proceeding with the current plan, and having hf rollout as the initial PR to land this recipe. sglang/vllm rollout can be added as separate enhance PRs.
cc @zhaochenyang20 for sglang feature visibility

[zhaochenyang20]

By the time I implement it, janus is not supported by vllm/sglang. I think sglang support janus now. However, I think it still doesn't support interleaved generation with text and images. But maybe I can do some customization to integrate vllm/sglang.

SGLang support Janus. What do you need us to support? I can ask @mickqian

[mickqian]

FYI, image-generation capabilities are not supported in sglang, nor are they in sglang's roadmap. Not sure if it conflicts with sglang's goal.

[pimang62]

Is there any updated feature now? I'm hopefully waiting for RL training in image generation task. 🙏