Skip to content

feat(reward): per-reward colocate, mixture RM, verified SD3.5 HPS recipe - #237

Merged
Rockdu merged 1 commit into
mainfrom
kangrui/reapply-234
Sep 6, 2026
Merged

feat(reward): per-reward colocate, mixture RM, verified SD3.5 HPS recipe#237
Rockdu merged 1 commit into
mainfrom
kangrui/reapply-234

Conversation

@Rockdu

@Rockdu Rockdu commented Sep 6, 2026

Copy link
Copy Markdown
Collaborator

Re-applies #234 (reverted in #236) so the squash commit carries the co-author credit for the 24 commits that came from #207 (Jingwen Gu). Same tree as 368240b; CI already passed on it.


Stacked on #207 (its commits are included; only the commits after 1fc7cf0 are new here).

What

  • Per-reward colocate placement. --colocate-reward becomes --pickscore-reward-colocate / --hps-reward-colocate; each colocated pool takes one slot per placement-group bundle, sharing the GPU with the train actor and rollout engine. RolloutManager owns the slot ledger and seats every colocated pool before the first rollout; the module-global placement state is gone.
  • Mixture reward. New --custom-rm-args (opaque string handed to the --custom-rm-path function) and the shipped example rm_hub/weighted_mixture_rm.py: --custom-rm-args "hps=0.7,pickscore=0.3" --reward-key weighted sums built-in rewards, each keeping its own placement flags. The example returns a dict per sample and every entry of a dict reward now gets its own rollout/reward/<name>_mean (and eval/<dataset>/<name>) panel, so component curves stay visible while the sum is optimized.
  • Rewards receive generated_output itself. The rm functions hand the float tensors to the pools and every reward actor quantises to uint8 on its own terms (HPS/OCR round, PickScore truncates, unchanged); no frame decoding on the manager. weighted_mixture_rm scores its rewards concurrently.
  • SD3.5 Flow-GRPO + HPSv2.1 verified. The recipe now ships the settings of the verified run (--diffusion-kl-beta 0.01, group-wise advantage std) and is marked 📈 V.

Validation

600 rollouts (2 optimizer steps each), 2×H200, run_diffusion_grpo_sd3_hps_sglang.py: rollout/reward/raw_mean 0.284 → 0.349 (mean of the last 100 rollouts), peak 10-rollout moving average 0.362, max single step 0.384. Training reward on hpdv2 prompts (T=10 SDE); community SD3.5-M + HPSv2.1 numbers (DanceGRPO 0.356, MixGRPO 0.365) are eval rewards on held-out prompts, so they are only a rough reference.

CPU tests (all registered to stage-a-cpu): test_reward_pool.py (deal order + slot ledger), test_reward_pool_placement.py (one slot per colocated worker; unseated colocate pool rejected), test_weighted_mixture_rm.py (weights applied once per reward; unknown name rejected), test_hps_reward.py (HPS rounds, PickScore truncates and averages frames, per-pool queue depth).

Checklist

  • pre-commit passes on every touched file (--all-files not run)
  • Added/updated tests for new behaviour
  • New and adjacent fast tests are green (pytest -x over the whole suite not run)
  • python3 train.py --help not re-run after the flag changes
  • New public flags are in the CLI reference docs
  • The mixture example has a walkthrough in docs/user-guide/rewards.md

@Rockdu Rockdu added the run-ci-e2e Run e2e metric-regression tests on this PR label Sep 6, 2026
Co-authored-by: Jingwen Gu <75733630+JingwenGu0829@users.noreply.github.com>
@Rockdu
Rockdu force-pushed the kangrui/reapply-234 branch from d322f85 to ae5a6e3 Compare September 6, 2026 18:45
@Rockdu
Rockdu merged commit 5e90ffe into main Sep 6, 2026
6 of 8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci-e2e Run e2e metric-regression tests on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant