Question about GRPO training setup for Qwen-Image-Edit-2509 or 2511

[Weistrass]

Hi, thanks for sharing this great repository.

I noticed that your repo supports GRPO training for Qwen-Image-Edit-2509, and that it allows choosing between LoRA and full-parameter training modes. I have a few questions regarding training setup and customization:

Multi-reference image training

(1) I would like to train the model to generate a target image conditioned on 4–5 reference images simultaneously. In this case, do you think a configuration like 8 × 140 is sufficient for either LoRA or full training?

(2) Are there any practical differences in feasibility or stability between LoRA and full training for this multi-reference setting?

Configuration considerations

For the above setup, which configuration aspects should I pay special attention to? For example:

(1) Image resolution / sequence length

(2) GRPO-specific hyperparameters (e.g., rollout length, reward normalization)

(3) Any model- or data-related constraints specific to Qwen-Image-Edit-2509 or 2511

Adding custom reward functions:

(1) If I want to add custom reward functions (e.g., for reference consistency or visual alignment). Which part of the codebase should I modify or extend?

(2) Is there a recommended interface or example for registering new reward functions in the GRPO pipeline?

Thanks a lot for your time and for open-sourcing this work. Any guidance would be greatly appreciated.

[Jayce-Ping]

Hi! Thanks for checking out this repo!

Multi-reference image training

(1) I would like to train the model to generate a target image conditioned on 4–5 reference images simultaneously. In this case, do you think a configuration like 8 × 140 is sufficient for either LoRA or full training?

That will be sufficient, I think. I haven't tried these many condition images yet. But for LoRA tuning, it should be enough. As for full, maybe it requires lower-resolution for condition images.

(2) Are there any practical differences in feasibility or stability between LoRA and full training for this multi-reference setting?

I suggest to train a LoRA with some easy reward and see if the reward goes well as expected. Then, turn to full-fine tuning. There is no much existing known feasibility/stability differences.

Configuration considerations

(1) Image resolution / sequence length

I suggest to set condition_image_size to a small number such as 256 if there are too many reference images. Since the Qwen-Image-Edit-2509 concatenate the images on the sequence dimension, it may cause OOM if the resolution and image numbers are both large.

(2) GRPO-specific hyperparameters (e.g., rollout length, reward normalization)

group_size can be set 16 or 24. This is enough for most of tasks (I don't see any task which needs larger yet). Don't worry too much about others. I will give you an example at the end of this response.

(3) Any model- or data-related constraints specific to Qwen-Image-Edit-2509 or 2511

No. They share the exact same architecture and just trained with different weights. If your task is not about like human face or realistism, using a random one is just fine.

Adding custom reward functions:

(1) If I want to add custom reward functions (e.g., for reference consistency or visual alignment). Which part of the codebase should I modify or extend?

(2) Is there a recommended interface or example for registering new reward functions in the GRPO pipeline?

Please check here for more infomation. Just implement your own logic of computing rewards in this file and set corresponding config (see below).

An example configuration file

Here I provide an example config file for you. Sorry that I don't have the GPU resources in the recent few days, I cannot verify it by myself. Please bring any feedback or issues. Thanks!

# Environment Configuration
launcher: "accelerate"  # Options: accelerate
config_file: config/accelerate_configs/fsdp2.yaml  # Use FSDP2 to shard the model as well. Switch to config/deepspeed/deepspeed_zero2.yaml if you have enough GPU memory.
num_processes: 8  # Number of processes to launch (overrides config file)
main_process_port: 29500
mixed_precision: "bf16"  # Options: no, fp16, bf16

run_name: null  # Run name (auto: {model_type}_{finetune_type}_{timestamp})
project: "Flow-Factory"  # Project name for logging
logging_backend: "wandb"  # Options: wandb, swanlab, none

# Data Configuration
data:
  dataset_dir: "dataset/sharegpt4o_image_mini"  # Path to dataset folder
  preprocessing_batch_size: 8  # Batch size for preprocessing
  dataloader_num_workers: 16  # Number of workers for DataLoader
  enable_preprocess: true  # Enable dataset preprocessing
  force_reprocess: true  # Force reprocessing of the dataset
  cache_dir: "~/.cache/flow_factory/datasets" # Cache directory for preprocessed datasets
  max_dataset_size: 1000  # Limit the maximum number of samples in the dataset

# Model Configuration
model:
  finetune_type: 'lora' # Options: full, lora
  lora_rank : 64
  lora_alpha : 128
  target_modules: "default" # Try default first, if OOM, then try ["to_k", "to_q", "to_v", "to_out.0"] for attention layers only.
  model_name_or_path: "Qwen/Qwen-Image-Edit-2509"  # Qwen/Qwen-Image-Edit-2509 or Qwen/Qwen-Image-Edit-2511
  model_type: "qwen-image-edit-plus"
  resume_path: null # Path to load previous checkpoint/lora adapter
  resume_training_state: false # Whether to resume training state, only effective when resume_path is a directory with full checkpoint

log:
  save_dir: "~/Flow-Factory"  # Directory to save model checkpoints and logs
  save_freq: 20  # Save frequency in epochs (0 to disable)
  save_model_only: true  # Save only the model weights (not optimizer, scheduler, etc.)

# Training Configuration
train:
  # Training settings
  trainer_type: 'grpo'
  enable_gradient_checkpointing: true  # Enable gradient checkpointing to save memory with extra compute. If OOM, set it `true`.

  # Image settings
  resolution: 512  # 512 is worth trying at first, if OOM, try 384 and then 256.
  condition_image_size: [512, 512]  # Keep as 512, or set the same as `resolution` above.
  
  # Batch and sampling
  per_device_batch_size: 1  # Qwen-Image-Edit-Plus accepts varying length of multi-image as condition, so the batch_size will always fallback to 1.
  group_size: 16  # Group size for GRPO sampling
  global_std: false  # Use global std for advantage normalization
  unique_sample_num_per_epoch: 48  # Unique samples per group
  gradient_step_per_epoch: 2  # Gradient steps per epoch
  
  # Clipping
  clip_range: 1.0e-4  # PPO/GRPO clipping range
  adv_clip_range: 5.0  # Advantage clipping range
  max_grad_norm: 1.0  # Max gradient norm for clipping
  
  # KL div
  kl_type: 'v-based' # Options: 'x-based', 'v-based'
  kl_beta: 0 # KL divergence beta. Set to 0 to disable it and save memory at first. If rewards grows as expected, set it to values like 0.04.
  ref_param_device: 'same_as_model' # Options: cpu, same_as_model

  # Denoising process
  num_inference_steps: 10  # Number of timesteps. For many tasks, 10 is good enough. Increase it to 20 if your task requires higher quality.
  guidance_scale: 4  # Guidance scale for sampling. 4 is following recommendation from Qwen-Image-Edit official model card.
  
  # Optimization
  seed: 42  # Random seed
  learning_rate: 3.0e-4  # Initial learning rate 3.0e-4 for LoRA and 1.0e-5 for full fine-tuning
  adam_weight_decay: 1.0e-4  # AdamW weight decay
  adam_betas: [0.9, 0.999]  # AdamW betas
  adam_epsilon: 1.0e-8  # AdamW epsilon

  # EMA
  ema_decay: 0.9  # EMA decay rate (0 to disable)
  ema_update_interval: 4  # EMA update interval (in epochs)

# Scheduler Configuration
scheduler:
  dynamics_type: "Flow-SDE"  # Options: Flow-SDE, Dance-SDE, CPS, ODE
  noise_level: 1.0  # Noise level for sampling
  num_train_steps: 1  # Number of noise steps
  train_steps: [1, 2, 3]  # Custom noise window, noise steps are randomly selected from this list during training
  seed: 42  # Scheduler seed (for noise step selection)

# Evaluation settings
eval:
  resolution: 512  # Evaluation resolution
  condition_image_size: [512, 512]  # Max condition image resolution, int or [height, width]
  auto_resize: true # Enable auto-resize to fit the condition images' aspect ratio during inference
  guidance_scale: 4  # Guidance scale for sampling
  num_inference_steps: 50  # Number of eval timesteps, 50 is recommended for better quality according to Qwen-Image-Edit official model card.
  per_device_batch_size: 1  # Eval batch size
  seed: 42  # Eval seed
  eval_freq: 20  # Eval frequency in epochs (0 to disable)

# Reward Model Configuration
rewards:
  - name: "visual_consistency"
    reward_model: "flow_factory.rewards.my_reward.VisualConsistencyRewardModel" # path to custom reward model
    batch_size: 16 # Batch size for reward model inference
    device: "cuda"
    dtype: bfloat16

# Optional Evaluation Reward Models
# eval_rewards:
#   - name: "text_alignment"
#     reward_model: "CLIP"
#     batch_size: 16
#     dtype: bfloat16
#     device: "cuda"

If you meet any issue, feel free to post here and I am happy to help🤗

[Weistrass]

感谢您的详细回复，目前我想快速验证一下这个框架对于qwen-edit-2509 grpo lora rank64的可行性，数据集我打算先用您提供的dataset/sharegpt4o_image_mini，然后奖励单纯就先用pickscore，不知道这种能否观察到您项目中提到的reward mean上升曲线。

[Jayce-Ping]

这个组合我已经验证过了，我把试验曲线贴到这里：

Training Reward

Eval Reward

还有一些Evaluation的例子，很明显可以看到图像是变得越来越符合PickScore的审美了。

Step 0

Step 80

Step 160

Step 240

配置也贴到这里：

# Environment Configuration
launcher: "accelerate"  # Options: accelerate
config_file: config/deepspeed/deepspeed_zero2.yaml  # Use FSDP2 to shard the model as well. Switch to config/deepspeed/deepspeed_zero2.yaml if you have enough GPU memory.
num_processes: 8  # Number of processes to launch (overrides config file)
main_process_port: 29500
mixed_precision: "bf16"  # Options: no, fp16, bf16

run_name: null  # Run name (auto: {model_type}_{finetune_type}_{timestamp})
project: "Flow-Factory"  # Project name for logging
logging_backend: "wandb"  # Options: wandb, swanlab, none

# Data Configuration
data:
  dataset_dir: "dataset/sharegpt4o_image_mini"  # Path to dataset folder
  preprocessing_batch_size: 8  # Batch size for preprocessing
  dataloader_num_workers: 16  # Number of workers for DataLoader
  enable_preprocess: true  # Enable dataset preprocessing
  force_reprocess: true  # Force reprocessing of the dataset
  cache_dir: "~/jcy/.cache/flow_factory/datasets" # Cache directory for preprocessed datasets
  max_dataset_size: 1000  # Limit the maximum number of samples in the dataset

# Model Configuration
model:
  finetune_type: 'lora' # Options: full, lora
  lora_rank : 64
  lora_alpha : 128
  target_modules: "default" # Options: all, default, or list of module names like ["to_k", "to_q", "to_v", "to_out.0"]
  model_name_or_path: "Qwen/Qwen-Image-Edit-2509"  # Qwen/Qwen-Image-Edit-2509 or Qwen/Qwen-Image-Edit-2511
  model_type: "qwen-image-edit-plus"
  resume_path: null # Path to load previous checkpoint/lora adapter
  resume_training_state: false # Whether to resume training state, only effective when resume_path is a directory with full checkpoint

log:
  save_dir: "~/jcy/Flow-Factory"  # Directory to save model checkpoints and logs
  save_freq: 20  # Save frequency in epochs (0 to disable)
  save_model_only: true  # Save only the model weights (not optimizer, scheduler, etc.)

# Training Configuration
train:
  # Training settings
  trainer_type: 'grpo'
  enable_gradient_checkpointing: false  # Enable gradient checkpointing to save memory with extra compute

  # Image settings
  resolution: 384  # Can be int or [height, width]
  auto_resize: true # Enable auto-resize to fit the condition images' aspect ratio during inference
  condition_image_size: [512, 512]  # Max condition image resolution, int or [height, width]
  
  # Batch and sampling
  per_device_batch_size: 1  # Qwen-Image-Edit-Plus accepts varying length of multi-image as condition, so the batch_size will always fallback to 1.
  group_size: 16  # Group size for GRPO sampling
  global_std: false  # Use global std for advantage normalization
  unique_sample_num_per_epoch: 48  # Unique samples per group
  gradient_step_per_epoch: 2  # Gradient steps per epoch
  
  # Clipping
  clip_range: 1.0e-4  # PPO/GRPO clipping range
  adv_clip_range: 5.0  # Advantage clipping range
  max_grad_norm: 1.0  # Max gradient norm for clipping
  
  # KL div
  kl_type: 'v-based' # Options: 'x-based', 'v-based'
  kl_beta: 0.04 # KL divergence beta
  ref_param_device: 'same_as_model' # Options: cpu, same_as_model

  # Denoising process
  num_inference_steps: 10  # Number of timesteps
  guidance_scale: 4  # Guidance scale for sampling
  
  # Optimization
  seed: 42  # Random seed
  learning_rate: 3.0e-4  # Initial learning rate
  adam_weight_decay: 1.0e-4  # AdamW weight decay
  adam_betas: [0.9, 0.999]  # AdamW betas
  adam_epsilon: 1.0e-8  # AdamW epsilon

  # EMA
  ema_decay: 0.9  # EMA decay rate (0 to disable)
  ema_update_interval: 4  # EMA update interval (in epochs)

# Scheduler Configuration
scheduler:
  dynamics_type: "Flow-SDE"  # Options: Flow-SDE, Dance-SDE, CPS, ODE
  noise_level: 1.0  # Noise level for sampling
  num_train_steps: 1  # Number of noise steps
  train_steps: [1, 2, 3]  # Custom noise window, noise steps are randomly selected from this list during training
  seed: 42  # Scheduler seed (for noise step selection)

# Evaluation settings
eval:
  resolution: 512  # Evaluation resolution
  condition_image_size: [512, 512]  # Max condition image resolution, int or [height, width]
  auto_resize: true # Enable auto-resize to fit the condition images' aspect ratio during inference
  guidance_scale: 4  # Guidance scale for sampling
  num_inference_steps: 40  # Number of eval timesteps
  per_device_batch_size: 1  # Eval batch size
  seed: 42  # Eval seed
  eval_freq: 20  # Eval frequency in epochs (0 to disable)

# Reward Model Configuration
rewards:
  - name: "pick_score"
    reward_model: "PickScore"
    batch_size: 16
    device: "cuda"
    dtype: bfloat16

当时我为了快，把训练分辨率设置在384，Eval的在512，和目前仓库中的example略有不同。因为分辨率较低，你可以看到step 0生成的图像有些扭曲，这好像是Qwen-Image-Edit在低分辨率自带的问题，不过RL在低分辨率能把这个问题修复过来。看到后面的训练图像都会好很多。这个问题在 X-GenGroup/PaCo-RL#2 也有讨论。你可以也试试把分辨率调低到512左右先跑，这样会快很多，可以很快验证能不能涨分

[Weistrass]

太感谢您的宝贵经验了，关于分辨率这个问题，qwen-edit-2509这个模型好像有过拟合现象的（拟合官方推理代码中默认的1024*1024分辨率），推理时，用1024x1024的分辨率得到的结果会比512x512好很多。不过我对你提到的这个讨论，在低分辨率，例如512训练完之后，能极大得缓解qwen-edit-2509这个模型的这种过拟合现象（即在非1024分辨率下的表现）比较感兴趣，我看你的建议是训练和采样时的分辨率应该一致，然后推理时的分辨率应小于等于训练时的分辨率？

[Jayce-Ping]

应该是：RL训练过程中 采样和优化时采用低分辨率（512），在推理时采用正常的高分辨率即可（如1024）。比如我上面给出的例子是在384下训练的+512上评估，可以看到后面在384上的表现会变好。评估时在512也涨分，我肯定在1024上也会变好。PaCo-RL这篇文章发现了这种机制——低分辨率RL的收益可以一并转移到高分辨率。这是一个非常简单有效的trick，可以显著加速训练。可能存在的问题就是太低分辨率导致图像细节的缺失会影响Reward Model打分，所以选一个差不多低的最好，但一定没必要跑1024训练.

[Weistrass]

好的，感谢，我会按照您的建议快速尝试一下，希望后面能继续和您交流

[Jayce-Ping]

我刚看到FLUX2-klein发布了，有4B和9B的两个版本，有统一的推理Pipeline，而且支持文生图，（多）图生图。很激动，我会尽快加上支持。后面这种多图任务就不用训练Qwen-Image-Edit-Plus和FLUX2-dev这两个大家伙了😃

[Weistrass]

我刚看到FLUX2-klein发布了，有4B和9B的两个版本，有统一的推理Pipeline，而且支持文生图，（多）图生图。很激动，我会尽快加上支持。后面这种多图任务就不用训练Qwen-Image-Edit-Plus和FLUX2-dev这两个大家伙了😃

那真是太好了，感谢您为这个领域做出的贡献，我也希望能有小模型能快速验证算法。我之前有特意调研过能支持多图参考训练的模型，ominigen2其实有专门为支持多张参考图片设计对应的位置编码，然后模型大小只有7B（可惜配套的lora微调或者强化学习不如qwen-edit-2509完善），其他的能支持多图的模型，例如mosaic，dreamomniv2都是作者额外改进位置编码做训练才实现的，比较麻烦，并且都没有开源训练代码，无法很好得进行修改。之所以选择qwen-edit-2509，是因为这个基模性能更好（我希望能在某个方面能做出比闭源模型，例如nano pro相当或者更好的效果），同时训练框架，例如对lora，或者flow-grpo的支持更好，能更快速的搭建起来，避免深陷工程的泥塘里，所以您的项目对于我来说非常有意义。

[Weistrass]

这是昨天我按照您给的配置训练的结果，训练和采样分辨率我都是512，但测试时是1024，总体趋势是上升的，所以应该没有太大问题。不过有几个问题需要和您请教，例如您之前的配置，训练384，测试512要训练多久呢，我看您给的图片是300 step eval/reward都还没有变得平缓，还有上升的空间。 我这边到160 step也同样没有收敛的意思，我看您写的是while循环，会一直训练下去，直到自己主动停止。我这边的算力资源最多能占用72小时，如果要加速，只能考虑多机多卡来训了。目前是单机8卡

[Jayce-Ping]

我这边是8*140G最多18个小时。所以最远我也不知道能到多少。实际上到200步多你可以看到图像质量已经有降低了，会出现过拟合PickScore偏好的风格，这个时候基本上就可以手动停止了。在配置文件可以为训练和评估分别采用两种奖励模型，比如用两个不同美学模型，刚开始可能会两个都涨分，直到评估开始掉分时，峰值的checkpoint应该就是最好的了。

[Jayce-Ping]

这个框架对模型封装的层次和diffusers几乎是一样的，我写代码时直接导入diffuers.pipeline来减少重复工作，不过这样会有明显滞后性。例如刚发布的FLUX2-klein，我需要等待最新版diffusers发布后才能兼容进来。另一个办法是直接copy 对应pipeline的代码拉进来改，我在考虑后面要不要都这么做。

[Weistrass]

我这边是8*140G最多18个小时。所以最远我也不知道能到多少。实际上到200步多你可以看到图像质量已经有降低了，会出现过拟合PickScore偏好的风格，这个时候基本上就可以手动停止了。在配置文件可以为训练和评估分别采用两种奖励模型，比如用两个不同美学模型，刚开始可能会两个都涨分，直到评估开始掉分时，峰值的checkpoint应该就是最好的了。

好的，感谢您的建议。

[Weistrass]

这个框架对模型封装的层次和diffusers几乎是一样的，我写代码时直接导入diffuers.pipeline来减少重复工作，不过这样会有明显滞后性。例如刚发布的FLUX2-klein，我需要等待最新版diffusers发布后才能兼容进来。另一个办法是直接copy 对应pipeline的代码拉进来改，我在考虑后面要不要都这么做。

我没有太多构建大型项目的经验，不过我认为可以先尝试一下工作量，仅先支持最热门的，并且diffusers还没支持的试试（尽量避免破坏原有的可拓展性）。我实际的科研经验来看，很多事可能很难一开始就想的特别清楚。

[Weistrass]

您好，我在训练模型支持更多参考图片的时候，发现如果训练和测试的分辨率都为512，group size为16，最多支持4张参考图片的训练，如果降低group size到8，可以支持最多到6张参考图片的训练。不过这个group size从16降低到8，对性能影响会有多大呢？希望能从您这得到一些经验。我是希望模型能够在更多参考图片的情况下仍然保持一个不错的性能，我不确定训练时最大参考图片为4的情况下，测试时泛化到7-8张参考图片的效果是否会好。所以面临一个参考图片数量和group size的权衡，分辨率目前不太能降低，担心在384的情况下训出来的模型在1024测试下表现不好，毕竟差距有点大。

[Jayce-Ping]

我不太明白为什么（分辨率+参考图像数量）和group_size会有冲突。分辨率+参考图像数量影响显存占用，而group_size不管多大都不影响峰值显存。但是增大group_size需要更多时间采样，训练效率更低，分辨率+参考图像数量也会影响训练效率。所以我猜你想表达的瓶颈应该是72h的训练时间。如果是训练时间的话，其实也不用很担心，可以先训72h，然后再开一个训练直接加载之前的检查点接着训练就行了。