Skip to content

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

Merged
Rockdu merged 37 commits into
mainfrom
kangrui/hps-reward-pool-placement
Sep 6, 2026
Merged

feat(reward): per-reward colocate, mixture RM, verified SD3.5 HPS recipe#234
Rockdu merged 37 commits into
mainfrom
kangrui/hps-reward-pool-placement

Conversation

@Rockdu

@Rockdu Rockdu commented Sep 6, 2026

Copy link
Copy Markdown
Collaborator

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 force-pushed the kangrui/hps-reward-pool-placement branch from 3bcccae to 972553a Compare September 6, 2026 01:22
@Rockdu Rockdu added the run-ci-e2e Run e2e metric-regression tests on this PR label Sep 6, 2026
@Rockdu
Rockdu force-pushed the kangrui/hps-reward-pool-placement branch from a1ec27c to c0492b8 Compare September 6, 2026 01:57
@Rockdu
Rockdu force-pushed the kangrui/hps-reward-pool-placement branch 2 times, most recently from 0a8f6b9 to 74dd434 Compare September 6, 2026 05:48
@Rockdu
Rockdu force-pushed the kangrui/hps-reward-pool-placement branch 2 times, most recently from 8412b00 to 9d56096 Compare September 6, 2026 06:00
@Rockdu
Rockdu force-pushed the kangrui/hps-reward-pool-placement branch 5 times, most recently from 888fa37 to 8f906c7 Compare September 6, 2026 06:37
@Rockdu
Rockdu force-pushed the kangrui/hps-reward-pool-placement branch from 4c27a14 to ff0d81b Compare September 6, 2026 06:58
@Rockdu
Rockdu force-pushed the kangrui/hps-reward-pool-placement branch from afd40f3 to d6df53c Compare September 6, 2026 07:12
@Rockdu
Rockdu force-pushed the kangrui/hps-reward-pool-placement branch from d6df53c to 6b52349 Compare September 6, 2026 07:13
@Rockdu

Rockdu commented Sep 6, 2026

Copy link
Copy Markdown
Collaborator Author

@JingwenGu0829 Thanks for your contribution to this feature! Since this newer PR has been merged, I will close the original one.

@Rockdu
Rockdu merged commit 368240b into main Sep 6, 2026
34 of 35 checks passed
Rockdu added a commit that referenced this pull request Sep 6, 2026
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.

2 participants