perf(bagel): pack FlowGRPO replay steps

[zzhuoxin1508]

Summary

• add pipeline.batch_replay_steps as the Bagel replay-packing opt-in
• when enabled, pack every runtime-resolved SDE target (K = len(target)) as isolated Navit varlen sequences in one forward
• repeat the generation KV context for T2I/IT2I and vectorize the Flow SDE transition
• preserve per-sequence projection/MLP reduction boundaries so packed BF16 replay keeps serial-equivalent gradients

Scope

• Bagel FlowGRPO T2I/IT2I only
• no CFG support (cfg_text_scale=cfg_img_scale=1 is required)
• retain each recipe's existing rollout/replay anchor policy
• no unified-model, shared algorithm, benchmark, or profiling changes

Validation

• serial-vs-packed T2I and IT2I log-prob/mean/gradient parity
  • full-LoRA relative gradient error: ~0.4%
  • IT2I packed-vs-serial log-prob relative error: 6.5e-8
• Ruff and Python compile checks for changed first-party Bagel modules
• complete 300-rollout vLLM-Omni IT2I packed/main comparison
  • packed train time: 48.05 s vs main serial 53.24 s (10.8% faster)
  • packed end-to-end step time: 83.25 s vs 87.97 s (5.7% faster)
  • packed completed 300/300 without reward collapse; final reward 1.523
  • final-50 reward mean: packed 1.393 vs main serial 1.157
  • packed W&B: https://wandb.ai/zhouzhuoxin1508-tencent/bagel-editreward/runs/rt96ey1i
  • main serial W&B: https://wandb.ai/zhouzhuoxin1508-tencent/bagel-editreward/runs/mn48q5ud

[zzhuoxin1508]

bagel test result

[zzhuoxin1508]

BAGEL IT2I packed replay — 300-rollout result

• Setup: 8×H20, vLLM-Omni DP8, 10,732-sample real image-edit dataset, 14 diffusion steps, 3-step SDE window, 8 samples/prompt.
• Correctness: completed 300/300 rollouts without reward collapse; final reward 1.523.
• Gradient parity: full-LoRA packed-vs-serial relative error ~0.4%; log-prob relative error 6.5e-8.
• Performance (matched rollouts 5–300):
  • train: 48.05 s packed vs 53.24 s main serial (10.8% faster)
  • end-to-end: 83.25 s packed vs 87.97 s main serial (5.7% faster)
  • throughput: 43.24 vs 40.92 rollouts/hour
• Reward over final 50 rollouts: 1.393 packed vs 1.157 main serial.
• W&B: packed · main serial