z_image_trainside.yaml

# @package _global_
# Z-Image (base) + trainside (FSDP) colocate, v2 trainer, FlowGRPO.
#
# Trainside = rollout runs in-process on the FSDP-wrapped training pipeline
# (no separate rollout actor, no weight sync). Strongest colocate form.
#
# This recipe targets the BASE checkpoint Tongyi-MAI/Z-Image (not the distilled
# Turbo): multi-step sampling and static shift=6.0. The S3-DiT single-stream
# transformer + Qwen3 text encoder are loaded by the bundle; the VAE is the
# flux-style 16-channel AutoencoderKL.
#
# CFG is intentionally OFF here (guidance_scale=0 -> single forward, no negative
# prompt). Empirically CFG slows the reward curve: it sharpens the per-prompt
# distribution (less in-group diversity -> weaker GRPO advantage) AND doubles
# the per-step forward cost. Off = more exploration + ~2x faster rollouts.
# (To re-enable CFG: set sampling.guidance_scale>0; the pipeline then auto-adds
# empty-string negatives. To run Turbo: set Z_IMAGE_PATH=.../Z-Image-Turbo and
# override shift=3.0, sampling.num_inference_steps=8.)
#
# Scale is set by the launcher (num_devices = cluster GPU count); this same
# recipe runs on 2 (H20 smoke), 8 (1-node) and 32 (4x8) GPUs. batch_size=64
# divides 2/8/32 evenly. For a quick smoke, shrink via CLI overrides, e.g.:
#   batch_size=2 sampling.samples_per_prompt=4 num_rollouts=2

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

logging:
  # Env-driven so multi-node runs enable W&B with REPORT_TO_WANDB=true (the
  # launcher exports it). rank-0/driver only; no-op when false.
  report_to_wandb: ${oc.decode:${oc.env:REPORT_TO_WANDB,false}}
  project_name: ${oc.env:WANDB_PROJECT,unirl-zimage}
  run_name: ${oc.env:WANDB_RUN_NAME,z_image_base_trainside}
  entity: ${oc.env:WANDB_ENTITY,null}
  tags: null

bundle:
  _target_: unirl.models.z_image.bundle.ZImageBundle.from_config
  config:
    _target_: unirl.models.z_image.config.ZImagePipelineConfig
    pretrained_model_ckpt_path: ${oc.env:Z_IMAGE_PATH,Tongyi-MAI/Z-Image}
    model_precision: bf16
    max_sequence_length: 512
    # Static FlowMatch shift; base Z-Image's scheduler_config.json declares
    # use_dynamic_shifting=false + shift=6.0 (read by build_schedule_policy).
    shift: 6.0

# Pipeline relies on the optional-stages constructor (mirrors SD3/Qwen-Image):
# 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.z_image.pipeline.ZImagePipeline
  shift: 6.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: ["ZImageTransformerBlock"]
  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
    # 6B S3-DiT + Qwen3 text encoder under full-graph replay: keep AC on.
    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
    # Z-Image single-stream attention (diffusers Attention inside
    # ZImageTransformerBlock): query/key/value projections + output proj.
    # Matches every block (main layers + noise/context refiners).
    target_modules:
      - to_k
      - to_out.0
      - to_q
      - to_v

rollout:
  _target_: unirl.rollout.engine.trainside.engine.TrainsideRolloutEngine
  stage_attrs: [diffusion]
  # 6B DiT + Qwen3 encoder are heavy; chunk the generate() forward so the
  # rollout/decode peak stays bounded on an H20.
  forward_batch_size: 8

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
  # Recompute pi_old at the exact train micro-geometry so the on-policy ratio is
  # exactly 1 (bf16 forwards are batch-shape sensitive). Correct & safe to keep
  # even now that micro_batch_size happens to equal forward_batch_size, and
  # required again the moment they diverge (see SD3 recipe).
  old_logp_source: replay
  conditions_cls:
    _target_: hydra.utils.get_class
    path: unirl.models.z_image.conditions.ZImageConditions
  params: ${sampling}

stack:
  _target_: unirl.train.stack.TrainStack
  # Train-side micro-batch (pure speed/memory knob — it only sets grad-accum
  # granularity, NOT the effective batch, so it never changes the optimization
  # math). Observed: mbs 1->4 barely moved HBM (~42 -> ~40 GB) because resident
  # frozen weights (transformer + Qwen3 encoder + VAE) dominate and activations
  # don't persist under AC. So push it hard. On the 4x8 run the per-GPU
  # mini-batch is 16 (batch_size*spp/num_devices/num_updates_per_batch =
  # 64*16/32/2), so mbs=16 means zero accumulation (max SM feed); 8 = 2 accum
  # steps. Start at 8, bump to 16 via `stack.micro_batch_size=16` if HBM allows.
  micro_batch_size: 8
  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       # base Z-Image: multi-step (RL-reduced from ~30-50;
                                # 12 ≈ Qwen-Image recipe — trims rollout cost, which
                                # CFG already doubles per step). Set 8 only for Turbo.
  guidance_scale: 0.0           # CFG OFF (single forward, no negative). Faster
                                # reward curve: more in-group diversity for GRPO +
                                # ~2x cheaper rollouts. Set >0 to re-enable CFG.
  height: 512
  width: 512
  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: confine SDE noise to the early high-σ steps.
    timestep_fraction: [0, 0.5]