feat(reward): per-reward colocate, mixture RM, verified SD3.5 HPS recipe - #237
Merged
Conversation
Co-authored-by: Jingwen Gu <75733630+JingwenGu0829@users.noreply.github.com>
Rockdu
force-pushed
the
kangrui/reapply-234
branch
from
September 6, 2026 18:45
d322f85 to
ae5a6e3
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
1fc7cf0are new here).What
--colocate-rewardbecomes--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.RolloutManagerowns the slot ledger and seats every colocated pool before the first rollout; the module-global placement state is gone.--custom-rm-args(opaque string handed to the--custom-rm-pathfunction) and the shipped examplerm_hub/weighted_mixture_rm.py:--custom-rm-args "hps=0.7,pickscore=0.3" --reward-key weightedsums 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 ownrollout/reward/<name>_mean(andeval/<dataset>/<name>) panel, so component curves stay visible while the sum is optimized.generated_outputitself. 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_rmscores its rewards concurrently.--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_mean0.284 → 0.349 (mean of the last 100 rollouts), peak 10-rollout moving average 0.362, max single step 0.384. Training reward onhpdv2prompts (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-commitpasses on every touched file (--all-filesnot run)pytest -xover the whole suite not run)python3 train.py --helpnot re-run after the flag changesdocs/user-guide/rewards.md