feat(reward): per-reward colocate, mixture RM, verified SD3.5 HPS recipe - #234
Merged
Conversation
This reverts commit 41492cd.
Rockdu
force-pushed
the
kangrui/hps-reward-pool-placement
branch
from
September 6, 2026 01:22
3bcccae to
972553a
Compare
Rockdu
force-pushed
the
kangrui/hps-reward-pool-placement
branch
from
September 6, 2026 01:57
a1ec27c to
c0492b8
Compare
Rockdu
force-pushed
the
kangrui/hps-reward-pool-placement
branch
2 times, most recently
from
September 6, 2026 05:48
0a8f6b9 to
74dd434
Compare
Rockdu
force-pushed
the
kangrui/hps-reward-pool-placement
branch
2 times, most recently
from
September 6, 2026 06:00
8412b00 to
9d56096
Compare
…e custom-rm example
…ings of the verified run
Rockdu
force-pushed
the
kangrui/hps-reward-pool-placement
branch
5 times, most recently
from
September 6, 2026 06:37
888fa37 to
8f906c7
Compare
…mixture_rm return them
Rockdu
force-pushed
the
kangrui/hps-reward-pool-placement
branch
from
September 6, 2026 06:58
4c27a14 to
ff0d81b
Compare
…ated, and require PickScore paths for it
Rockdu
force-pushed
the
kangrui/hps-reward-pool-placement
branch
from
September 6, 2026 07:12
afd40f3 to
d6df53c
Compare
Rockdu
force-pushed
the
kangrui/hps-reward-pool-placement
branch
from
September 6, 2026 07:13
d6df53c to
6b52349
Compare
Collaborator
Author
|
@JingwenGu0829 Thanks for your contribution to this feature! Since this newer PR has been merged, I will close the original one. |
This was referenced Sep 6, 2026
Rockdu
added a commit
that referenced
this pull request
Sep 6, 2026
6 tasks
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.
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