Skip to content

[reward, trainer] feat: add managed multi-reward deployments - #524

Draft
Sky-Trigger wants to merge 1 commit into
verl-project:mainfrom
Sky-Trigger:MultiRewardDeployments
Draft

[reward, trainer] feat: add managed multi-reward deployments#524
Sky-Trigger wants to merge 1 commit into
verl-project:mainfrom
Sky-Trigger:MultiRewardDeployments

Conversation

@Sky-Trigger

@Sky-Trigger Sky-Trigger commented Sep 3, 2026

Copy link
Copy Markdown
Collaborator

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 upstream verl.RewardModelManager per deployment, retaining
    vLLM/vLLM-Omni router, replica, tensor-parallel, wake-up, and sleep behavior.
  • native: a worker-local Transformers or custom scorer for models that cannot
    yet run through an engine. It has explicit wake-up, score, sleep, and
    accelerator-cache cleanup behavior.

The trainer still chooses either the existing global_pool or reward_pool.
MultiRewardModelManager then partitions that parent pool into disjoint
engine 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_workers setting.

Reward terms use reward.reward_functions.<term>.deployment=<name>. Existing
path/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 PickScoreNativeScorer instance,
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

Trainer
│
├─ selects the Role.RewardModel parent pool
│    ├─ reward.reward_model.enable_resource_pool=false -> global_pool
│    └─ reward.reward_model.enable_resource_pool=true  -> reward_pool
│
└─ OmniRewardLoopManager
   │
   ├─ MultiRewardModelManager
   │    ├─ validates deployments, term references, and native placement
   │    ├─ partitions the selected parent pool
   │    │
   │    ├─ EngineRewardDeployment[name]
   │    │    └─ upstream verl.RewardModelManager
   │    │         └─ router + vLLM/vLLM-Omni replicas + wake_up/sleep
   │    │
   │    └─ NativeRewardDeployment[name]
   │         └─ static NativeRewardExecutor spec; no model is loaded here
   │
   └─ deployment-aware reward worker groups
        ├─ shared group (optional): rule terms and engine terms
        └─ one native group for each native deployment
             └─ workers own only that deployment's NativeRewardExecutor

MultiVisualRewardManager is not a third model backend. It stays inside each
reward-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 uses
the group's NativeRewardExecutor.

Parent pool split and native placement

Trainer-selected parent pool: global_pool or reward_pool
│
├─ engine_resource_pool[ocr]
│    └─ EngineRewardDeployment[ocr]
│         └─ engine replicas; size = replicas × TP × DP × PP
│
├─ engine_resource_pool[clip]
│    └─ EngineRewardDeployment[clip]
│         └─ independent engine replicas; own TP/DP/PP topology
│
└─ native_resource_pool
     ├─ bundle 0 -> native_pickscore: complete PickScore model
     ├─ bundle 1 -> native_pickscore: complete PickScore model
     ├─ bundle 2 -> native_pickscore: complete PickScore model
     ├─ bundle 3 -> native_pickscore: complete PickScore model
     ├─ bundle 4 -> native_hpsv3:     complete HPSv3 model
     ├─ bundle 5 -> native_hpsv3:     complete HPSv3 model
     ├─ bundle 6 -> native_hpsv3:     complete HPSv3 model
     └─ bundle 7 -> native_hpsv3:     complete HPSv3 model

placement.devices is a list of native-subpool-relative placement-group
bundle indices, not host-global CUDA/NPU IDs and not process-local cuda:N
indices. That makes the configuration stable when Ray rewrites visible-device
variables for an actor.

The native subpool reserves max(placement.devices) + 1 bundles. Each listed
index 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

rollout batch
  │
  └─ OmniRewardLoopManager.compute_rm_score(batch)
       │
       ├─ wake engine deployments
       ├─ shared group: rule / engine terms -> weighted partial scores
       ├─ native_pickscore: batch chunks on bundles [0, 1, 2, 3]
       ├─ native_hpsv3: batch chunks on bundles [4, 5, 6, 7]
       │
       ├─ final_reward[i] = sum(group_partial_score[i])
       ├─ merge term metrics; emit exactly one reward/combined
       └─ sleep engine deployments

Each group uses the existing compute_score_batch() path. Its local
MultiVisualRewardManager applies term weights first; the manager then sums
all group outputs. The resulting semantics remain:

final_reward[i] = Σ_term (term.weight × term.score[i])

Intermediate group-local reward/combined values are discarded during the
merge. Individual reward/<term> values are retained and there is one final
reward/combined. Native deployments use the normal batch reward phase rather
than 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=pickscore

placement.devices is mandatory for a named native deployment. The four
entries above create four full PickScore replicas on the first four native
subpool bundles; they are not tensor-parallel ranks. weight=1.0 and
reward.num_workers=4 are unnecessary in this example: a term's default
weight is 1.0, and native worker count is len(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 a
native router, replica scheduler, or native tensor parallel runtime.

A multi-native deployment configuration can be expressed directly as a clear
dictionary:

reward:
  deployments:
    native_pickscore:
      backend: native
      adapter: pickscore
      model_path: yuvalkirstain/PickScore_v1
      placement:
        devices: [0, 1, 2, 3]
    native_hpsv3:
      backend: native
      model_path: /models/HPSv3
      executor:
        scorer: my_package.hpsv3:HPSv3NativeScorer
      placement:
        devices: [4, 5, 6, 7]
  reward_functions:
    pickscore:
      deployment: native_pickscore
      weight: 0.7
    hpsv3:
      deployment: native_hpsv3
      weight: 0.3

The current Qwen-Image-Edit CUDA and V1-NPU launchers are migrated to this
named-native form. They expose PICKSCORE_MODEL_PATH and
NATIVE_REWARD_DEVICES instead of the legacy PickScore reward-function path,
reward.accelerator_workers.enabled, and reward.num_workers settings.

For an engine-backed reward, set backend=engine and supply the deployment
resource and rollout configuration. Existing path/name reward functions can
then consume that deployment's reward_router_address and model_name.

Design & Code Changes

Trainer
  -> global_pool or reward_pool
  -> MultiRewardModelManager
       -> EngineRewardDeployment[name] -> upstream RewardModelManager
       -> NativeRewardDeployment[name] -> executor metadata
  -> shared worker group: rule terms + EngineRewardExecutor[name]
  -> native worker group[name]: NativeRewardExecutor[name] only
  -> manager-level merge: one rm_scores tensor + one reward/combined metric
  • Add reward.deployments configuration and early validation for deployment
    references, backend-specific terms, and requested resource sizes.
  • Add MultiRewardModelManager, EngineRewardDeployment, and
    NativeRewardDeployment; engine deployments wrap the existing upstream
    RewardModelManager instead of duplicating engine lifecycle logic.
  • Bind each native reward-loop worker group to only the explicit
    placement.devices indices owned by its deployment. Retain
    reward.accelerator_workers.enabled and the deprecated
    custom_reward_function.use_accelerator alias for existing custom-reward
    launchers.
  • Add EngineRewardExecutor adapters for existing router-based reward
    functions and the current PickScore embedding endpoint path.
  • Add NativeRewardExecutor lifecycle management. It waits for active scoring
    before closing a scorer and clearing accelerator cache, but does not impose a
    per-scorer serialization lock on a scorer that implements its own batching.
  • Move native PickScore batching from module-global state to an instance-local
    consumer so wake-up/sleep can load and release the model deterministically.
  • Keep legacy non-deployment rewards on their original paths. Migrate the
    Qwen-Image-Edit CUDA and V1-NPU PickScore recipes to named native
    deployments.
  • Add focused tests for native placement validation, native subpool size,
    worker-group isolation, explicit bundle scheduling, cross-group score merge,
    executor lifecycle, and native PickScore batching.

Scope

This PR does not:

  • change rollout generation or actor-update algorithms;
  • implement native tensor parallelism or a generic native model-serving
    runtime;
  • provide a generic adapter-registration mechanism for every future
    vLLM/vLLM-Omni reward-model semantic;
  • claim numerical parity between vLLM CLIP pooling and Transformers PickScore;
  • migrate unrelated legacy reward utilities such as CLAP, HPSv3, or ImageBind
    to native deployments.

@Sky-Trigger

Copy link
Copy Markdown
Collaborator Author

@chenyingshu PTAL. I've been working on this recently, but it still needs some time to be further refined.

@Sky-Trigger Sky-Trigger added the ready-for-ci read for running CI label Sep 4, 2026
@Sky-Trigger
Sky-Trigger force-pushed the MultiRewardDeployments branch from 0784651 to 5b5b8a5 Compare September 4, 2026 07:35
@github-actions github-actions Bot removed the ready-for-ci read for running CI label Sep 4, 2026
@Sky-Trigger Sky-Trigger added the ready-for-ci read for running CI label Sep 4, 2026
@github-actions github-actions Bot removed the ready-for-ci read for running CI label Sep 7, 2026
@Sky-Trigger Sky-Trigger added the ready-for-ci read for running CI label Sep 7, 2026
@github-actions github-actions Bot removed the ready-for-ci read for running CI label Sep 7, 2026
@Sky-Trigger Sky-Trigger added the ready-for-ci read for running CI label Sep 7, 2026
@github-actions github-actions Bot removed the ready-for-ci read for running CI label Sep 7, 2026
@Sky-Trigger
Sky-Trigger force-pushed the MultiRewardDeployments branch 2 times, most recently from b0407c9 to 63e0760 Compare September 7, 2026 04:49
@Sky-Trigger Sky-Trigger added ready-for-ci read for running CI and removed ready-for-ci read for running CI labels Sep 7, 2026
@Sky-Trigger Sky-Trigger changed the title [draft, reward, trainer] feat: add managed multi-reward deployments [reward, trainer] feat: add managed multi-reward deployments Sep 7, 2026
@Sky-Trigger Sky-Trigger added the ready-for-ci read for running CI label Sep 7, 2026
@Sky-Trigger Sky-Trigger added ready-for-ci read for running CI and removed ready-for-ci read for running CI labels Sep 7, 2026
Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: Trigger <Meng.Bo.Wang@outlook.com>
@Sky-Trigger
Sky-Trigger force-pushed the MultiRewardDeployments branch from 63e0760 to 8bf963d Compare September 7, 2026 09:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant