feat(reward): support HPSv2 and refine reward scheduling

[JingwenGu0829]

Summary

• Add HPSv2 as a batched GPU reward model in rm_hub, with v2.0/v2.1 checkpoints and aligned prompt-image scoring.
• Add a two-GPU colocated SD3.5 Flow-GRPO recipe, CLI flags, tests, and reward documentation.
• Stage the HPS checkpoint on CPU with weights_only=True to avoid a second checkpoint copy in GPU memory during actor initialization.

Design notes / known issues

1. Documentation and verification status. The generic HPS reward API is documented under Rewards and the runnable recipe is listed on the SD3 page. The preferred final placement/wording may need maintainer guidance. Because only short end-to-end runs have been performed and no complete reference curve or deterministic gate is committed, the recipe is marked ○ NV rather than Verified/FG.

2. Only one GPU reward pool may use --colocate-reward. A colocated train actor uses 0.70 GPU and the rollout engine uses 0.25, leaving one 0.05 slot per placement-group bundle. HPS and PickScore singleton pools previously both selected the first bundle, so constructing both pools could leave the second actor pending forever. rm_hub now claims the first colocated GPU pool and fails fast when a different pool tries to start. Mixed GPU reward types must use dedicated reward GPUs until bundle capacity is allocated across pools.

if _colocated_reward_pool_name not in (None, name):
    raise RuntimeError(
        f"Only one GPU reward pool can use --colocate-reward; {_colocated_reward_pool_name} is already active, "
        f"so {name} cannot start. Use one GPU reward type or dedicated reward GPUs."
    )
_colocated_reward_pool_name = name

3. SD3.5 text-encoder fallback. With the official SD3.5 Diffusers snapshot, the SGLang customized text-encoder loader sees duplicate tensor names across model.safetensors and model.fp16.safetensors, rejects the ambiguous load, and falls back to the native Transformers CLIPTextModel. The fallback completes and the end-to-end HPS run proceeds; the Miles-d team considers this a Diffusers/checkpoint-layout issue with limited performance impact. This PR does not change that upstream loader behavior.

[Rockdu]

A quick discussion here: the mental model of "colocate reward" is actually "reward slots colocate with rollout&train actors". In most cases, the reward models cannot take up all the slots (e.g., in wan2.2 17gpu pickscore, the ratio of rollout engines:reward is 16:4). Maybe we should arrange these reward models in the slots available instead of giving different reward models separate pools

[JingwenGu0829]

#207 (comment) Design here. Upstreamed

[Rockdu]

This function has already been absorbed into processing_utils.py, maybe let's reuse _sample_to_rgb_hwc_uint8_frames

[JingwenGu0829]

I took another look at this. I think keeping the thin HPS adapter separate is cleaner here.

The shared conversion is already reused through cfhw_to_fhwc and image_or_video_to_uint8.

The remaining behavior is pretty RM-specific, including HPS rejects multi-frame outputs and uses round_normalized=True, while PickScore performs video-frame sampling and currently uses the default truncating conversion.

Reusing _sample_to_rgb_hwc_uint8_frames would therefore require another conversion flag plus a separate HPS frame-count check, while introducing an HPS-to-PickScore private-helper dependency.

It would remove very little logic, so I would prefer to keep the two small adapters separate. What do you think?

[Rockdu]

Oh, actually images are 1-frame videos in sgl-d's output so don't worry about that~

[JingwenGu0829]

Minimal design for colocated reward slot allocation

Mirror existing rollout placement semantics

Use exactly the same effective span as rollout actor creation:

span = min(num_gpus_per_engine, num_gpus_per_node)

When span > 0, Miles creates a floor-divided number of rollout actors. The
reward ordering mirrors that without validating divisibility:

covered = len(bundle_indices) // span * span

For a position inside covered:

position_in_actor = position % span
is_base = position_in_actor == 0

Positions after covered are real PG bundles that Miles did not assign to a
complete rollout actor span. They remain valid reward slots and are treated as
unassigned non-base bundles.

If span <= 0, treat every bundle as unassigned. Do not introduce a new error;
the existing rollout component retains responsibility for its own invalid
configuration behavior.

Reward-local dispersal order

Add no topology class. _dispersal_order operates directly on the existing
lists:

def _dispersal_order(
    bundle_indices: list[int],
    gpu_ids: list[int],
    num_gpus_per_node: int,
    num_gpus_per_engine: int,
) -> tuple[int, ...]:
    span = min(num_gpus_per_engine, num_gpus_per_node)
    covered = len(bundle_indices) // span * span if span > 0 else 0

    def key(item: tuple[int, int, int]) -> tuple[int, int, int, int]:
        position, _bundle_index, gpu_id = item
        if position >= covered:
            # Not occupied by a rollout actor: prefer the extra headroom.
            return (0, -1, gpu_id, position)

        position_in_actor = position % span
        if position_in_actor == 0:
            # The Ray rollout actor consumes 0.25 on this bundle; use it last.
            return (1, 0, gpu_id, position)
        return (0, position_in_actor, gpu_id, position)

    slots = zip(range(len(bundle_indices)), bundle_indices, gpu_ids)
    return tuple(bundle_index for _, bundle_index, _ in sorted(slots, key=key))

The final position tie-break preserves the existing node-major order when
two slots otherwise have the same key. There is no isinstance, ragged-group
handling, node-boundary inference, or topology validation.

What the stable sort does

For normal homogeneous nodes, equal local GPU IDs appear once per node. Sorting
the node-major input by GPU ID transposes it:

input:  n0g0 n0g1 n0g2 n0g3 | n1g0 n1g1 n1g2 n1g3
output: n0g0 n1g0 n0g1 n1g1 n0g2 n1g2 n0g3 n1g3

position_in_actor adds the engine level, and the base phase moves rollout
actor bundles to the tail.

For two nodes × eight GPUs with four GPUs per rollout actor:

n0e0g1 n1e0g1 n0e1g1 n1e1g1
n0e0g2 n1e0g2 n0e1g2 n1e1g2
n0e0g3 n1e0g3 n0e1g3 n1e1g3
n0e0g0 n1e0g0 n0e1g0 n1e1g0

For single-GPU engines all occupied bundles are bases, so the order reduces to
the existing node round-robin behavior.

Graceful degradation

The ordering is total for every concrete PG view:

• 4-of-5 physical GPUs: all four PG bundles appear exactly once; the unused
  fifth physical GPU is irrelevant because it is not part of this PG.
• 2-of-3 physical GPUs: both PG bundles appear exactly once.
• Five bundles, span two: four bundles mirror two rollout actors; the fifth
  unassigned bundle is preferred and remains allocatable.
• Uneven node occupancy: missing GPU IDs are simply absent from the stable
  sort; every observed bundle is still retained.
• Disjoint local GPU-ID subsets across nodes: dispersion may be less even
  because the existing tuple no longer carries node identity. This affects
  balance only. It cannot duplicate a slot, lose capacity, or cause a Ray
  scheduling deadlock.

This tradeoff is deliberate: exact balance under an information-poor unusual
placement is not worth changing Miles' placement interfaces.

Shared runtime allocator

ColocatedRewardSlots owns the fixed order and process-lifetime claims:

class ColocatedRewardSlots:
    def __init__(self, order: tuple[int, ...]):
        self._order = order
        self._owners: dict[int, str] = {}
        self._pool_names: set[str] = set()

    def allocate(self, name: str, num_workers: int) -> list[int]:
        if name in self._pool_names:
            raise RuntimeError(f"{name} already owns reward slots")

        start = len(self._owners)
        remaining = len(self._order) - start
        if num_workers > remaining:
            raise RuntimeError(...actual capacity and ownership summary...)

        slots = list(self._order[start : start + num_workers])
        self._owners.update({slot: name for slot in slots})
        self._pool_names.add(name)
        return slots

The capacity check happens before mutation, so exhaustion cannot partially
claim a pool. _pool_names is separate from _owners so a zero-worker or
failed singleton construction cannot allocate again under the same name.

Contiguous slices continue from the phase where the previous pool stopped. On
homogeneous nodes, the order is cyclic at the node and engine levels, so every
slice reaches the integer lower bound (q or q + 1 workers per group). Across
pool boundaries, the remainder of one slice is naturally compensated by the
next.

For irregular placements, the proof is intentionally weaker: allocation
follows the deterministic best-effort order and consumes every real slot.

Integration

Only rm_hub needs the new ordering logic. RolloutManager supplies the two
existing configuration values used by rollout itself:

set_manager_placement_group(
    pg,
    num_gpus_per_node=args.num_gpus_per_node,
    num_gpus_per_engine=args.rollout_num_gpus_per_engine,
)

set_manager_placement_group reads all three existing tuple fields:

_pg, bundle_indices, gpu_ids = pg
_reward_slots = ColocatedRewardSlots(
    _dispersal_order(
        bundle_indices,
        gpu_ids,
        num_gpus_per_node,
        num_gpus_per_engine,
    )
)

The lazy AsyncRewardActorPool path remains unchanged after allocation: one
explicit PlacementGroupSchedulingStrategy per returned bundle.

Tests

Keep the tests reward-local and small:

1. exact SD3 single-GPU-engine order;
2. exact WAN two-node/four-GPU-engine order;
3. two pools receive disjoint contiguous slices with cross-pool compensation;
4. 4-of-5 and 2-of-3 partial-PG shapes retain every bundle without raising;
5. a five-bundle/span-two remainder is retained and selected before bases;
6. uneven GPU-ID lists retain every bundle exactly once;
7. exhaustion is atomic and reports current ownership;
8. duplicate pool names do not consume another slice.

No test asserts that a Miles configuration or Ray placement is invalid.

[Rockdu]

Oh, actually images are 1-frame videos in sgl-d's output so don't worry about that~