[reward, trainer] feat: add managed multi-reward deployments - #524
Draft
Sky-Trigger wants to merge 1 commit into
Draft
[reward, trainer] feat: add managed multi-reward deployments#524Sky-Trigger wants to merge 1 commit into
Sky-Trigger wants to merge 1 commit into
Conversation
Collaborator
Author
|
@chenyingshu PTAL. I've been working on this recently, but it still needs some time to be further refined. |
Sky-Trigger
force-pushed
the
MultiRewardDeployments
branch
from
September 4, 2026 07:35
0784651 to
5b5b8a5
Compare
Sky-Trigger
force-pushed
the
MultiRewardDeployments
branch
2 times, most recently
from
September 7, 2026 04:49
b0407c9 to
63e0760
Compare
4 tasks
Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Trigger <Meng.Bo.Wang@outlook.com>
Sky-Trigger
force-pushed
the
MultiRewardDeployments
branch
from
September 7, 2026 09:07
63e0760 to
8bf963d
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.
What does this PR do?
Related to #432.
This PR adds named, managed reward-model deployments to the reward loop. It
supports two explicit backends under one
MultiRewardModelManager:engine: one upstreamverl.RewardModelManagerper deployment, retainingvLLM/vLLM-Omni router, replica, tensor-parallel, wake-up, and sleep behavior.
native: a worker-local Transformers or custom scorer for models that cannotyet run through an engine. It has explicit wake-up, score, sleep, and
accelerator-cache cleanup behavior.
The trainer still chooses either the existing
global_poolorreward_pool.MultiRewardModelManagerthen partitions that parent pool into disjointengine subpools and one native parent subpool. Each native deployment claims
an explicit, disjoint list of bundle indices inside that native subpool; it no
longer relies on the global
reward.num_workerssetting.Reward terms use
reward.reward_functions.<term>.deployment=<name>. Existingpath/name reward functions remain supported, and terms are combined by the
existing weighted-sum behavior. A failure now propagates instead of silently
contributing a zero score.
For native PickScore, the PR preserves the existing maximum-16-request batch
forward. The batching queue is scoped to one
PickScoreNativeScorerinstance,so native executor sleep can first drain scoring work, then stop the consumer
and release the CLIP model before actor update.
Architecture
Ownership and control plane
MultiVisualRewardManageris not a third model backend. It stays inside eachreward-loop worker and evaluates only the terms assigned to that worker group.
For an engine term it uses an
EngineRewardExecutor; for a native term it usesthe group's
NativeRewardExecutor.Parent pool split and native placement
placement.devicesis a list of native-subpool-relative placement-groupbundle indices, not host-global CUDA/NPU IDs and not process-local
cuda:Nindices. That makes the configuration stable when Ray rewrites visible-device
variables for an actor.
The native subpool reserves
max(placement.devices) + 1bundles. Each listedindex creates one full model replica, not a tensor-parallel shard. Sparse
layouts are allowed but leave unused bundles idle. Validation rejects an empty
list, negative/non-integer values, duplicate values, cross-deployment overlap,
and engine-only settings such as
rollout,replicas,n_gpus_per_node,nnodes, TP, DP, or PP.Reward-path data flow
Each group uses the existing
compute_score_batch()path. Its localMultiVisualRewardManagerapplies term weights first; the manager then sumsall group outputs. The resulting semantics remain:
Intermediate group-local
reward/combinedvalues are discarded during themerge. Individual
reward/<term>values are retained and there is one finalreward/combined. Native deployments use the normal batch reward phase ratherthan the streaming-agent path, because the upstream streaming interface accepts
only one worker list and cannot safely fan a request to several native groups.
API and Usage Example
Use a native PickScore deployment by naming the model deployment and binding a
reward term to it:
+reward.deployments.pickscore.backend=native \ +reward.deployments.pickscore.adapter=pickscore \ +reward.deployments.pickscore.model_path=yuvalkirstain/PickScore_v1 \ +reward.deployments.pickscore.placement.devices='[0,1,2,3]' \ +reward.reward_functions.pickscore.deployment=pickscoreplacement.devicesis mandatory for a named native deployment. The fourentries above create four full PickScore replicas on the first four native
subpool bundles; they are not tensor-parallel ranks.
weight=1.0andreward.num_workers=4are unnecessary in this example: a term's defaultweight is
1.0, and native worker count islen(placement.devices).The native scorer wakes for the reward batch and sleeps after it. It reuses the
existing
RewardLoopWorker.compute_score_batch()path; this PR does not add anative router, replica scheduler, or native tensor parallel runtime.
A multi-native deployment configuration can be expressed directly as a clear
dictionary:
The current Qwen-Image-Edit CUDA and V1-NPU launchers are migrated to this
named-native form. They expose
PICKSCORE_MODEL_PATHandNATIVE_REWARD_DEVICESinstead of the legacy PickScore reward-function path,reward.accelerator_workers.enabled, andreward.num_workerssettings.For an engine-backed reward, set
backend=engineand supply the deploymentresource and rollout configuration. Existing path/name reward functions can
then consume that deployment's
reward_router_addressandmodel_name.Design & Code Changes
reward.deploymentsconfiguration and early validation for deploymentreferences, backend-specific terms, and requested resource sizes.
MultiRewardModelManager,EngineRewardDeployment, andNativeRewardDeployment; engine deployments wrap the existing upstreamRewardModelManagerinstead of duplicating engine lifecycle logic.placement.devicesindices owned by its deployment. Retainreward.accelerator_workers.enabledand the deprecatedcustom_reward_function.use_acceleratoralias for existing custom-rewardlaunchers.
EngineRewardExecutoradapters for existing router-based rewardfunctions and the current PickScore embedding endpoint path.
NativeRewardExecutorlifecycle management. It waits for active scoringbefore closing a scorer and clearing accelerator cache, but does not impose a
per-scorer serialization lock on a scorer that implements its own batching.
consumer so wake-up/sleep can load and release the model deterministically.
Qwen-Image-Edit CUDA and V1-NPU PickScore recipes to named native
deployments.
worker-group isolation, explicit bundle scheduling, cross-group score merge,
executor lifecycle, and native PickScore batching.
Scope
This PR does not:
runtime;
vLLM/vLLM-Omni reward-model semantic;
to native deployments.