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 all commits
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
30 changes: 26 additions & 4 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 Expand Up @@ -88,10 +92,28 @@ async def batched_async_rm(
sample.reward = reward
return None

if args.custom_rm_path is not None:
samples_with_overrides = [(index, sample) for index, sample in enumerate(samples) if sample.custom_rm_path]
if not samples_with_overrides and 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)
tasks = [async_rm(args, sample, **kwargs) for sample in samples]
rewards = await asyncio.gather(*tasks)
return rewards

samples_without_overrides = [(index, sample) for index, sample in enumerate(samples) if not sample.custom_rm_path]
override_rewards = asyncio.gather(*(async_rm(args, sample, **kwargs) for _, sample in samples_with_overrides))
if args.custom_rm_path is not None and samples_without_overrides:
rm_function = load_function(args.custom_rm_path)
default_rewards = rm_function(args, [sample for _, sample in samples_without_overrides], **kwargs)
else:
default_rewards = asyncio.gather(
*(async_rm(args, sample, **kwargs) for _, sample in samples_without_overrides)
)

override_rewards, default_rewards = await asyncio.gather(override_rewards, default_rewards)
rewards_by_index = dict(
zip(
[index for index, _ in samples_with_overrides] + [index for index, _ in samples_without_overrides],
[*override_rewards, *default_rewards],
strict=True,
)
)
return [rewards_by_index[index] for index in range(len(samples))]
Comment on lines +101 to +119

@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.

medium

The current implementation of batched_async_rm has a few potential issues:

  1. Nested asyncio.gather: It nests asyncio.gather calls (e.g., gathering override_rewards and default_rewards where both are already futures returned by asyncio.gather). This adds unnecessary overhead and complexity.
  2. Synchronous RM Compatibility: It assumes that the custom batch reward function (rm_function) is always asynchronous. If a user provides a synchronous custom batch reward function, calling it will return a list of rewards directly, and passing this list to asyncio.gather will raise a TypeError.

We can simplify the logic by flattening the asyncio.gather calls and defensively checking if the batch task is awaitable using inspect.isawaitable. This ensures compatibility with both synchronous and asynchronous custom reward functions while maintaining optimal concurrency.

    samples_without_overrides = [(index, sample) for index, sample in enumerate(samples) if not sample.custom_rm_path]
    override_tasks = [async_rm(args, sample, **kwargs) for _, sample in samples_with_overrides]
    if args.custom_rm_path is not None and samples_without_overrides:
        rm_function = load_function(args.custom_rm_path)
        batch_task = rm_function(args, [sample for _, sample in samples_without_overrides], **kwargs)
        import inspect
        if inspect.isawaitable(batch_task):
            results = await asyncio.gather(*override_tasks, batch_task)
            override_rewards = results[:len(override_tasks)]
            default_rewards = results[len(override_tasks)]
        else:
            results = await asyncio.gather(*override_tasks)
            override_rewards = results
            default_rewards = batch_task
    else:
        non_override_tasks = [async_rm(args, sample, **kwargs) for _, sample in samples_without_overrides]
        results = await asyncio.gather(*override_tasks, *non_override_tasks)
        override_rewards = results[:len(override_tasks)]
        default_rewards = results[len(override_tasks):]

    rewards_by_index = dict(
        zip(
            [index for index, _ in samples_with_overrides] + [index for index, _ in samples_without_overrides],
            [*override_rewards, *default_rewards],
            strict=True,
        )
    )
    return [rewards_by_index[index] for index in range(len(samples))]

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
42 changes: 42 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,22 @@

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):
if isinstance(sample, list):
return [999] * len(sample)
return 999


@pytest.fixture
def mock_args():
Expand Down Expand Up @@ -95,8 +109,36 @@ 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:
def test_per_sample_custom_rm_takes_priority(self, mock_args):
mock_args.custom_rm_path = GLOBAL_RM_PATH
samples = [
Sample(prompt="", response="", label="", custom_rm_path=PER_SAMPLE_RM_PATH),
Sample(prompt="", response="", label=""),
]
with function_registry.temporary(PER_SAMPLE_RM_PATH, _per_sample_rm), function_registry.temporary(
GLOBAL_RM_PATH, _global_rm
):
rewards = run(batched_async_rm(mock_args, samples))
assert rewards == [111, 999]

@pytest.mark.parametrize(
"rm_type,samples_data,expected",
[
Expand Down