UniRL/examples/diffusion/qwen_image/qwen_image_trainside.yaml at main · Tencent-Hunyuan/UniRL · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
# @package _global_
# Qwen-Image + trainside (FSDP) colocate, v2 trainer, GRPO.
#
# Trainside = rollout runs in-process on the FSDP-wrapped training pipeline
# (no separate rollout actor, no weight sync). Strongest colocate form.

num_devices: 8
batch_size: 48                # prompts_per_rollout
adv_use_global_std: true     # advantage: divide by ONE batch-wide std (v1 parity), not per-group
num_rollouts: 10000

logging:
  report_to_wandb: false        # flip to true to enable wandb (rank-0/driver only)
  project_name: unirl
  run_name: null
  entity: null
  tags: null

bundle:
  _target_: unirl.models.qwen_image.bundle.QwenImageBundle.from_config
  config:
    _target_: unirl.models.qwen_image.config.QwenImagePipelineConfig
    pretrained_model_ckpt_path: ${oc.env:QWEN_IMAGE_PATH,Qwen/Qwen-Image}
    model_precision: bf16
    max_sequence_length: 512
    # Static-shift fallback; the real checkpoint enables dynamic shifting via
    # scheduler/scheduler_config.json and this value is ignored at runtime.
    shift: 3.0

# Pipeline relies on the optional-stages constructor (mirrors SD3): the
# trainer auto-injects `bundle=self.bundle`; text_embed/diffusion/vae_decode
# are built from the bundle here using the precision policy below.
pipeline:
  _target_: unirl.models.qwen_image.pipeline.QwenImagePipeline
  shift: 3.0
  autocast_precision: bf16
  trajectory_precision: bf16
  logprob_precision: fp32
  max_sequence_length: 512
  strategy:
    _target_: unirl.sde.kernels.FlowSDEStrategy

backend:
  _target_: unirl.train.backend.fsdp.FSDPBackend
  block_class_names: ["QwenImageTransformerBlock"]
  trainable_attr: transformer
  fsdp_cfg:
    _target_: unirl.train.configs.FSDPConfig
    param_dtype: bf16
    cpu_offload: false
    mixed_precision: true
    fsdp_mode: full
    reshard_after_forward: true
    # Qwen-Image transformer + Qwen2.5-VL text encoder together exceed 96GB
    # per H20 under full-graph replay without AC; v1 uses bf16_full_ac.
    activation_checkpointing: true
    use_torch_compile: false
  optimizer_cfg:
    _target_: unirl.train.backend.base.OptimizerConfig
    learning_rate: 3.0e-4
    adam_beta1: 0.9
    adam_beta2: 0.999
    adam_epsilon: 1.0e-8
    weight_decay: 0.0
  scheduler_cfg:
    _target_: unirl.train.backend.base.LrSchedulerConfig
    type: constant
    warmup_steps: 0
    total_steps: 10000
  lora_cfg:
    _target_: unirl.train.configs.LoraConfig
    rank: 64
    alpha: 128
    dropout: 0.0
    bias: none
    task_type: FEATURE_EXTRACTION
    # JointAttnProcessor naming — image stream (to_q/k/v, to_out.0) plus
    # text stream (add_*_proj, to_add_out). Verified against the diffusers
    # QwenImageTransformer2DModel attention layer.
    target_modules:
      - attn.add_k_proj
      - attn.add_q_proj
      - attn.add_v_proj
      - attn.to_add_out
      - attn.to_k
      - attn.to_out.0
      - attn.to_q
      - attn.to_v

rollout:
  _target_: unirl.rollout.engine.trainside.engine.TrainsideRolloutEngine
  stage_attrs: [diffusion]
  # Qwen-Image VAE + text encoder are heavy; chunk forward to 1 sample at a
  # time. Mirrors v1 `rollout.plan.forward_batch_size: 1`.
  forward_batch_size: 1

reward:
  _target_: unirl.reward.service.RewardService
  backend:
    _target_: unirl.reward.local.pickscore.PickScoreRewardScorer
    base_device: cuda
    config:
      _target_: unirl.reward.local.pickscore.PickScoreSpec
      batch_size: 8
      device: auto
      processor_id: laion/CLIP-ViT-H-14-laion2B-s32B-b79K
      model_id: yuvalkirstain/PickScore_v1

algorithm:
  _target_: unirl.algorithms.flowgrpo.FlowGRPO
  stage_attr: diffusion
  clip_range: 1.0e-4
  clip_schedule: constant
  conditions_cls:
    _target_: hydra.utils.get_class
    path: unirl.models.qwen_image.conditions.QwenImageConditions
  params: ${sampling}

stack:
  _target_: unirl.train.stack.TrainStack
  micro_batch_size: 1
  max_grad_norm: 1.0
  num_updates_per_batch: 2      # PPO mini-batches per rollout (π_old frozen once); v1 parity

data_source:
  _target_: unirl.data.data_source.MultimodalRLDataSource
  args:
    run:
      data_path: datasets/pickscore/train.txt
      eval_data_path: datasets/pickscore/test.txt
      seed: 42
    algorithm:
      prompts_per_rollout: ${batch_size}

sampling:
  _target_: unirl.types.sampling.DiffusionSamplingParams
  num_inference_steps: 12       # qwen_384 preset default
  guidance_scale: 1.0
  height: 384
  width: 384
  eta: 0.7
  samples_per_prompt: 16
  seed: 42
  init_same_noise: false
  scheduler:
    _target_: unirl.sde.index_schedule.AllSDEScheduler
    num_timesteps: ${..num_inference_steps}
    num_sde_steps: 3
    # SDE-step window: [0,0.5] confines SDE noise to the early high-σ steps
    # (the validated reproduce setting); [0,1] spreads it across all steps.
    timestep_fraction: [0, 0.5]