ViperEkura
95c43368ae
- RolloutGenerator now delegates prefill/decode to InferenceScheduler.run_batch (sync API, no background thread), sharing one KV-cache code path with the inference server and eliminating O(n^2) recompute in rollout - Add sample(return_logprobs=) and Executor.execute_decode(return_logprobs=) to expose behaviour-policy log-probs through the engine; Task gains output_logprobs - RolloutResult now subclasses RawRollout (adds rewards only), removing duplicated fields - RolloutRunner.__call__ returns (result, is_fresh) instead of relying on object identity, removing the fragile refresh-detection contract - Remove O(n^2) generate_responses helper and dead code (_tokenize_prompts, unused old_model arg) - train_context.py wires InferenceScheduler directly instead of hand-rolling SamplingPipeline - Tests: +11 covering return_logprobs, run_batch, and KV-cache-backed rollout semantics; 404 pass
266 lines
9.2 KiB
Python
266 lines
9.2 KiB
Python
"""Online rollout runner for RL training.
|
|
|
|
Provides:
|
|
- :class:`RawRollout` — generation output container (no reward yet)
|
|
- :class:`RolloutResult` — a :class:`RawRollout` with rewards attached
|
|
- :class:`BaseRewardModel` — pluggable reward interface
|
|
- :class:`RolloutGenerator` — KV-cache-backed generation of grouped
|
|
responses + decoding (no reward); delegates the generation loop to
|
|
:class:`~astrai.inference.core.scheduler.InferenceScheduler.run_batch`
|
|
so rollout and the production inference server share one code path
|
|
- :class:`RolloutRunner` — orchestrates generation + scoring with a
|
|
step-driven cache; its ``__call__`` returns ``(RolloutResult, is_fresh)``
|
|
so callers do not need to rely on object identity to detect refreshes.
|
|
"""
|
|
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass, field
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch import Tensor
|
|
|
|
from astrai.inference.core.scheduler import InferenceScheduler
|
|
|
|
|
|
@dataclass(kw_only=True)
|
|
class RawRollout:
|
|
"""Generation output before reward scoring.
|
|
|
|
Produced by :class:`RolloutGenerator`; consumed by :class:`RolloutRunner`
|
|
to assemble a :class:`RolloutResult` once rewards are attached.
|
|
|
|
Fields are designed to cover all common RL algorithms:
|
|
GRPO, PPO, Online DPO, Rejection Sampling, etc.
|
|
"""
|
|
|
|
prompts: Tensor
|
|
"""Tokenized prompts, shape ``[B, P_len]``."""
|
|
|
|
responses: Tensor
|
|
"""Generated response token IDs, shape ``[B, G, R_max]``."""
|
|
|
|
response_mask: Tensor
|
|
"""Boolean mask for real (non-pad) response tokens, shape ``[B, G, R_max]``."""
|
|
|
|
logprobs_old: Tensor
|
|
"""Per-token log-probs under the behaviour policy, shape ``[B, G, R_max]``."""
|
|
|
|
prompt_texts: List[str] = field(default_factory=list)
|
|
"""Decoded prompt strings (for reward models that need text)."""
|
|
|
|
response_texts: List[List[str]] = field(default_factory=list)
|
|
"""Decoded response strings, shape ``[B, G]`` (for reward models)."""
|
|
|
|
|
|
@dataclass(kw_only=True)
|
|
class RolloutResult(RawRollout):
|
|
"""A :class:`RawRollout` with reward scoring attached.
|
|
|
|
Produced by :class:`RolloutRunner` once the :class:`BaseRewardModel`
|
|
has scored the decoded responses.
|
|
"""
|
|
|
|
rewards: Tensor
|
|
"""Reward per response, shape ``[B, G]``."""
|
|
|
|
|
|
class BaseRewardModel(ABC):
|
|
"""Pluggable reward model interface.
|
|
|
|
Subclasses should implement ``score()`` to return a ``[B, G]`` float
|
|
tensor of rewards. Implementations can be:
|
|
* A loaded reward model (e.g. ArmoRM, Skywork-Reward)
|
|
* An external API call
|
|
* A rule-based function (format, length, keyword matching)
|
|
"""
|
|
|
|
@abstractmethod
|
|
def score(self, prompts: List[str], responses: List[List[str]]) -> Tensor:
|
|
"""Score each generated response.
|
|
|
|
Args:
|
|
prompts: Raw prompt strings, length ``B``.
|
|
responses: Generated response strings, shape ``[B, G]``.
|
|
|
|
Returns:
|
|
Float tensor of shape ``[B, G]``.
|
|
"""
|
|
...
|
|
|
|
|
|
_PAD = 0
|
|
|
|
|
|
class RolloutGenerator:
|
|
"""Pure generation + decoding for a group of responses per prompt.
|
|
|
|
Delegates the prefill/decode loop to
|
|
:meth:`~astrai.inference.core.scheduler.InferenceScheduler.run_batch`,
|
|
which uses a real KV cache (no O(n²) recompute). Has no dependency
|
|
on any reward model; can be reused in isolation for offline
|
|
generation, qualitative sampling, or eval pipelines.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
scheduler: InferenceScheduler,
|
|
tokenizer,
|
|
max_tokens: int = 1024,
|
|
group_size: int = 8,
|
|
temperature: float = 1.0,
|
|
top_k: int = 0,
|
|
top_p: float = 1.0,
|
|
frequency_penalty: float = 0.0,
|
|
rep_window: int = 64,
|
|
):
|
|
self.scheduler = scheduler
|
|
self.tokenizer = tokenizer
|
|
self.max_tokens = max_tokens
|
|
self.group_size = group_size
|
|
self.temperature = temperature
|
|
self.top_k = top_k
|
|
self.top_p = top_p
|
|
self.frequency_penalty = frequency_penalty
|
|
self.rep_window = rep_window
|
|
|
|
@torch.no_grad()
|
|
def generate(self, batch: Dict[str, Tensor]) -> RawRollout:
|
|
"""Expand prompts by ``group_size`` and generate one response each."""
|
|
prompt_ids = batch["input_ids"] if "input_ids" in batch else batch["prompts"]
|
|
prompt_mask = (
|
|
batch["attention_mask"] if "attention_mask" in batch else (prompt_ids != 0)
|
|
)
|
|
B, _ = prompt_ids.shape
|
|
G = self.group_size
|
|
|
|
prompt_texts: List[str] = []
|
|
flat_prompt_ids: List[List[int]] = []
|
|
for i in range(B):
|
|
ids = prompt_ids[i, prompt_mask[i]].tolist()
|
|
text = self.tokenizer.decode(ids, skip_special_tokens=True)
|
|
for _ in range(G):
|
|
flat_prompt_ids.append(list(ids))
|
|
prompt_texts.append(text)
|
|
|
|
results = self.scheduler.run_batch(
|
|
flat_prompt_ids,
|
|
max_tokens=self.max_tokens,
|
|
temperature=self.temperature,
|
|
top_k=self.top_k,
|
|
top_p=self.top_p,
|
|
frequency_penalty=self.frequency_penalty,
|
|
rep_window=self.rep_window,
|
|
return_logprobs=True,
|
|
)
|
|
|
|
# Each element is (token_ids, logprobs); pad to max length.
|
|
max_len = 0
|
|
for token_ids, _lp in results:
|
|
max_len = max(max_len, len(token_ids))
|
|
max_len = max(max_len, 1)
|
|
|
|
device = prompt_ids.device
|
|
responses = torch.full((B, G, max_len), _PAD, dtype=torch.long, device=device)
|
|
response_mask = torch.zeros((B, G, max_len), dtype=torch.bool, device=device)
|
|
logprobs_old = torch.zeros((B, G, max_len), dtype=torch.float, device=device)
|
|
|
|
flat_idx = 0
|
|
response_texts: List[List[str]] = [[] for _ in range(B)]
|
|
for i in range(B):
|
|
for g in range(G):
|
|
token_ids, lps = results[flat_idx]
|
|
flat_idx += 1
|
|
n = len(token_ids)
|
|
if n:
|
|
responses[i, g, :n] = torch.tensor(
|
|
token_ids, dtype=torch.long, device=device
|
|
)
|
|
response_mask[i, g, :n] = True
|
|
logprobs_old[i, g, :n] = torch.tensor(
|
|
lps, dtype=torch.float, device=device
|
|
)
|
|
response_texts[i].append(
|
|
self.tokenizer.decode(token_ids, skip_special_tokens=True)
|
|
)
|
|
|
|
return RawRollout(
|
|
prompts=prompt_ids,
|
|
responses=responses,
|
|
response_mask=response_mask,
|
|
logprobs_old=logprobs_old,
|
|
prompt_texts=prompt_texts,
|
|
response_texts=response_texts,
|
|
)
|
|
|
|
|
|
class RolloutRunner:
|
|
"""Produces :class:`RolloutResult` from a prompt batch.
|
|
|
|
Composes a :class:`RolloutGenerator` (generation + decoding) with a
|
|
:class:`BaseRewardModel` (scoring). Maintains an internal cache so
|
|
the same batch prompt can be replayed for multiple gradient steps.
|
|
A new rollout is triggered every ``rollout_interval`` calls to
|
|
:meth:`step` (or after :meth:`clear_cache`).
|
|
|
|
The ``__call__`` contract returns a ``(RolloutResult, is_fresh)``
|
|
tuple — callers must use the boolean to detect a refreshed rollout
|
|
rather than relying on object identity.
|
|
|
|
Usage::
|
|
|
|
generator = RolloutGenerator(policy, tokenizer, pipeline, ...)
|
|
runner = RolloutRunner(generator, reward_model, rollout_interval=512)
|
|
result, is_fresh = runner(prompt_batch)
|
|
if is_fresh:
|
|
... # e.g. sync behaviour policy
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
generator: RolloutGenerator,
|
|
reward_model: BaseRewardModel,
|
|
rollout_interval: int = 512,
|
|
):
|
|
self.generator = generator
|
|
self.reward_model = reward_model
|
|
self.rollout_interval = rollout_interval
|
|
|
|
self._cache: Optional[RolloutResult] = None
|
|
self._steps_since_rollout: int = 0
|
|
|
|
def step(self):
|
|
"""Advance the internal counter (call once per optimizer step)."""
|
|
self._steps_since_rollout += 1
|
|
|
|
def clear_cache(self):
|
|
"""Force next call to re-run rollout."""
|
|
self._cache = None
|
|
|
|
def _score(self, raw: RawRollout) -> RolloutResult:
|
|
rewards = self.reward_model.score(raw.prompt_texts, raw.response_texts)
|
|
device = raw.prompts.device
|
|
return RolloutResult(
|
|
prompts=raw.prompts,
|
|
responses=raw.responses,
|
|
response_mask=raw.response_mask,
|
|
rewards=rewards.to(device=device),
|
|
logprobs_old=raw.logprobs_old,
|
|
prompt_texts=raw.prompt_texts,
|
|
response_texts=raw.response_texts,
|
|
)
|
|
|
|
def __call__(self, batch: Dict[str, Tensor]) -> Tuple[RolloutResult, bool]:
|
|
"""Return ``(cached or fresh) RolloutResult`` plus an ``is_fresh`` flag.
|
|
|
|
Triggers a new rollout when ``_steps_since_rollout >= rollout_interval``
|
|
or when the cache is empty.
|
|
"""
|
|
if self._cache is None or self._steps_since_rollout >= self.rollout_interval:
|
|
raw = self.generator.generate(batch)
|
|
self._cache = self._score(raw)
|
|
self._steps_since_rollout = 0
|
|
return self._cache, True
|
|
return self._cache, False
|