Skip to content

feat(rm): allow per-sample custom_rm_path override #1637

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
EazyReal wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat(rm): allow per-sample custom_rm_path override #1637

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6ab3990
feat(rm): allow per-sample custom_rm_path override
EazyReal Jul 3, 2026
d964a78
fix(rm): honor sample overrides in batched rewards
EazyReal Jul 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions miles/rollout/generate_utils/sample_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,7 @@ def _merge_metadata():
status=b.status,
metadata=_merge_metadata(),
generate_function_path=_merge_equal_value("generate_function_path"),
custom_rm_path=_merge_equal_value("custom_rm_path"),
train_metadata=_merge_equal_value("train_metadata"),
session_id=_merge_equal_value("session_id"),
non_generation_time=_merge_equal_value("non_generation_time"),
Expand Down
4 changes: 4 additions & 0 deletions miles/rollout/rm_hub/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,10 @@ async def remote_rm(args, sample: Sample):


async def async_rm(args, sample: Sample, **kwargs):
if sample.custom_rm_path:
rm_function = load_function(sample.custom_rm_path)
return await rm_function(args, sample, **kwargs)
Comment on lines +32 to +34

@gemini-code-assist gemini-code-assist Bot Jul 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

While async_rm now correctly prioritizes sample.custom_rm_path over args.custom_rm_path, batched_async_rm (defined later in this file) does not.

Specifically, in batched_async_rm:

    if args.custom_rm_path is not None:
        # Ensure the custom reward function is implemented in batch mode
        rm_function = load_function(args.custom_rm_path)
        return await rm_function(args, samples, **kwargs)

If a global args.custom_rm_path is configured, batched_async_rm will unconditionally route all samples to that global custom RM, completely ignoring any per-sample sample.custom_rm_path overrides. This breaks the core objective of supporting mixed-task batches with different reward functions.

To fix this, batched_async_rm should group the samples by their effective custom RM path (either sample.custom_rm_path or args.custom_rm_path) and batch them accordingly, falling back to individual async_rm calls for samples that do not use any custom RM.

Here is a suggested implementation for batched_async_rm:

async def batched_async_rm(
    args,
    samples: list[Sample],
    inplace_set_reward_field: bool = False,
    **kwargs,
) -> list[int | float] | None:
    if inplace_set_reward_field:
        rewards = await batched_async_rm(args, samples, **kwargs)
        for sample, reward in zip(samples, rewards, strict=True):
            assert (
                sample.reward is None
            ), f"Overriding sample.reward from {sample.reward} to {reward}, is this intended?"
            sample.reward = reward
        return None

    # Group samples by their effective custom RM path to preserve batching where possible
    indexed_samples = list(enumerate(samples))
    groups = {}
    for idx, sample in indexed_samples:
        path = sample.custom_rm_path or args.custom_rm_path
        groups.setdefault(path, []).append((idx, sample))

    rewards = [None] * len(samples)
    tasks = []

    for path, group in groups.items():
        sub_indices = [item[0] for item in group]
        sub_samples = [item[1] for item in group]

        if path is not None:
            async def run_custom_batch(p=path, indices=sub_indices, s_samples=sub_samples):
                rm_function = load_function(p)
                batch_rewards = await rm_function(args, s_samples, **kwargs)
                for i, r in zip(indices, batch_rewards):
                    rewards[i] = r
            tasks.append(run_custom_batch())
        else:
            async def run_individual_fallback(indices=sub_indices, s_samples=sub_samples):
                individual_tasks = [async_rm(args, s, **kwargs) for s in s_samples]
                batch_rewards = await asyncio.gather(*individual_tasks)
                for i, r in zip(indices, batch_rewards):
                    rewards[i] = r
            tasks.append(run_individual_fallback())

    await asyncio.gather(*tasks)
    return rewards


if args.custom_rm_path is not None:
rm_function = load_function(args.custom_rm_path)
return await rm_function(args, sample, **kwargs)
Expand Down
1 change: 1 addition & 0 deletions miles/rollout/sglang_rollout.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -561,6 +561,7 @@ async def eval_rollout_single_dataset(
sample_index += 1
sample.metadata = dataset_cfg.inject_metadata(getattr(sample, "metadata", None))
sample.generate_function_path = getattr(dataset_cfg, "custom_generate_function_path", None)
sample.custom_rm_path = dataset_cfg.custom_rm_path
sampling_params = base_sampling_params
if getattr(args, "sglang_enable_deterministic_inference", False):
sampling_params = base_sampling_params.copy()
Expand Down
3 changes: 3 additions & 0 deletions miles/utils/eval_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,9 @@ class EvalDatasetConfig:
# per-dataset custom generate function (e.g., for tool calling)
custom_generate_function_path: str | None = None

# per-dataset custom reward function (overrides the global --custom-rm-path)
custom_rm_path: str | None = None

metadata_overrides: dict[str, Any] = field(default_factory=dict)

def __post_init__(self) -> None:
Expand Down
1 change: 1 addition & 0 deletions miles/utils/types.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ class Status(Enum):

metadata: dict = field(default_factory=dict)
generate_function_path: str | None = None
custom_rm_path: str | None = None
# metadata used during training, e.g., what loss to use for this sample.
train_metadata: dict | None = None

Expand Down
28 changes: 28 additions & 0 deletions tests/fast/rollout/rm_hub/test_rm_hub.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,20 @@

from miles.rollout.rm_hub import async_rm, batched_async_rm
from miles.utils.async_utils import run
from miles.utils.misc import function_registry
from miles.utils.types import Sample

PER_SAMPLE_RM_PATH = "test.per_sample.rm"
GLOBAL_RM_PATH = "test.global.rm"


async def _per_sample_rm(args, sample, **kwargs):
return 111


async def _global_rm(args, sample, **kwargs):
return 999


@pytest.fixture
def mock_args():
Expand Down Expand Up @@ -95,6 +107,22 @@ def test_invalid_rm_type_raises(self, mock_args, rm_type, match):
with pytest.raises(NotImplementedError, match=match):
run(async_rm(mock_args, sample))

def test_per_sample_custom_rm_takes_priority(self, mock_args):
mock_args.custom_rm_path = GLOBAL_RM_PATH
sample = Sample(prompt="", response="", label="", custom_rm_path=PER_SAMPLE_RM_PATH)
with function_registry.temporary(PER_SAMPLE_RM_PATH, _per_sample_rm), function_registry.temporary(
GLOBAL_RM_PATH, _global_rm
):
reward = run(async_rm(mock_args, sample))
assert reward == 111

def test_falls_back_to_global_custom_rm(self, mock_args):
mock_args.custom_rm_path = GLOBAL_RM_PATH
sample = Sample(prompt="", response="", label="")
with function_registry.temporary(GLOBAL_RM_PATH, _global_rm):
reward = run(async_rm(mock_args, sample))
assert reward == 999


class TestBatchedAsyncRm:
@pytest.mark.parametrize(
Expand Down