feat: add online rollout framework for RL strategies
- RolloutRunner: generate + score responses with cached re-rollout trigger - BaseStrategy.__call__ switches online/offline via runner injection - GRPO/DPO implement prepare_from_rollout; aliases online_grpo/online_dpo - TrainConfig + train.py add rollout params and CLI flags - Tests cover generate_responses, RolloutRunner cache, shared __call__
This commit is contained in:
@@ -138,6 +138,32 @@ class TrainConfig(BaseConfig):
|
||||
metadata={"help": "NEFTune noise alpha (0=disabled, typical: 5.0)."},
|
||||
)
|
||||
|
||||
# online rollout
|
||||
rollout_interval: int = field(
|
||||
default=512,
|
||||
metadata={"help": "Number of optimizer steps between online rollouts."},
|
||||
)
|
||||
rollout_temperature: float = field(
|
||||
default=0.7, metadata={"help": "Sampling temperature for online rollout."}
|
||||
)
|
||||
rollout_top_k: int = field(
|
||||
default=0, metadata={"help": "Top-k filtering for online rollout (0=disable)."}
|
||||
)
|
||||
rollout_top_p: float = field(
|
||||
default=0.9,
|
||||
metadata={"help": "Top-p (nucleus) filtering for online rollout."},
|
||||
)
|
||||
rollout_max_tokens: int = field(
|
||||
default=1024,
|
||||
metadata={"help": "Maximum generated tokens per response in rollout."},
|
||||
)
|
||||
reward_model_fn: Optional[Callable] = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": "Factory for reward model (required for online RL strategies)."
|
||||
},
|
||||
)
|
||||
|
||||
executor_kwargs: Dict[str, Any] = field(
|
||||
default_factory=dict,
|
||||
metadata={"help": "Extra kwargs passed to ExecutorFactory.create()."},
|
||||
|
||||
@@ -0,0 +1,272 @@
|
||||
"""Online rollout runner for RL training.
|
||||
|
||||
Provides:
|
||||
- :class:`RolloutResult` — universal data container for online sampling
|
||||
- :class:`BaseRewardModel` — pluggable reward interface
|
||||
- :class:`RolloutRunner` — generates + scores batches for any RL strategy
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from astrai.inference.sample import SamplingPipeline
|
||||
|
||||
|
||||
@dataclass
|
||||
class RolloutResult:
|
||||
"""Universal container produced by :class:`RolloutRunner`.
|
||||
|
||||
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]``."""
|
||||
|
||||
rewards: Tensor
|
||||
"""Reward per response, shape ``[B, G]``."""
|
||||
|
||||
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)."""
|
||||
|
||||
|
||||
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]``.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
def generate_responses(
|
||||
model: nn.Module,
|
||||
input_ids: Tensor,
|
||||
attention_mask: Tensor,
|
||||
max_new_tokens: int,
|
||||
sampling_pipeline: SamplingPipeline,
|
||||
stop_ids: List[int],
|
||||
) -> Dict[str, Tensor]:
|
||||
"""Autoregressive generation with log-prob tracking.
|
||||
|
||||
Args:
|
||||
model: Policy model (``forward`` returns ``{"logits": ...}``).
|
||||
input_ids: ``[B, P_len]`` prompt token IDs.
|
||||
attention_mask: ``[B, P_len]`` boolean mask.
|
||||
max_new_tokens: Maximum tokens to generate.
|
||||
sampling_pipeline: Composed sampling strategies.
|
||||
stop_ids: Token IDs that stop generation (eos, etc.).
|
||||
|
||||
Returns:
|
||||
``dict`` with keys:
|
||||
- ``generated_ids``: ``[B, max_new_tokens]`` (padded to same length)
|
||||
- ``generated_mask``: ``[B, max_new_tokens]``
|
||||
- ``logprobs``: ``[B, max_new_tokens]`` per-token log-probs
|
||||
"""
|
||||
_PAD = 0
|
||||
B, P_len = input_ids.shape
|
||||
device = input_ids.device
|
||||
stop_ids_set = set(stop_ids)
|
||||
done = torch.zeros(B, dtype=torch.bool, device=device)
|
||||
all_ids = input_ids.clone()
|
||||
all_mask = attention_mask.clone()
|
||||
logprob_list: List[Tensor] = []
|
||||
|
||||
for _ in range(max_new_tokens):
|
||||
outputs = model(input_ids=all_ids, input_mask=all_mask)
|
||||
logits = outputs["logits"][:, -1, :].float()
|
||||
log_probs = F.log_softmax(logits, dim=-1)
|
||||
|
||||
logits = sampling_pipeline.apply(logits, input_ids=all_ids, input_mask=all_mask)
|
||||
probs = torch.softmax(logits, dim=-1)
|
||||
next_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)
|
||||
|
||||
next_tokens[done] = _PAD
|
||||
chosen_logprobs = torch.gather(log_probs, -1, next_tokens.unsqueeze(-1))
|
||||
logprob_list.append(chosen_logprobs)
|
||||
|
||||
all_ids = torch.cat([all_ids, next_tokens.unsqueeze(1)], dim=-1)
|
||||
all_mask = torch.cat([all_mask, (~done).unsqueeze(1)], dim=-1)
|
||||
|
||||
done = done | torch.tensor(
|
||||
[t.item() in stop_ids_set for t in next_tokens],
|
||||
device=device,
|
||||
)
|
||||
if done.all():
|
||||
break
|
||||
|
||||
logprobs = torch.cat(logprob_list, dim=-1)
|
||||
if logprobs.size(1) < max_new_tokens:
|
||||
pad_len = max_new_tokens - logprobs.size(1)
|
||||
logprobs = F.pad(logprobs, (0, pad_len), value=0.0)
|
||||
|
||||
generated_ids = all_ids[:, P_len:]
|
||||
if generated_ids.size(1) < max_new_tokens:
|
||||
pad_len = max_new_tokens - generated_ids.size(1)
|
||||
generated_ids = F.pad(generated_ids, (0, pad_len), value=_PAD)
|
||||
|
||||
generated_mask = generated_ids != _PAD
|
||||
|
||||
return {
|
||||
"generated_ids": generated_ids,
|
||||
"generated_mask": generated_mask,
|
||||
"logprobs": logprobs,
|
||||
}
|
||||
|
||||
|
||||
class RolloutRunner:
|
||||
"""Produces :class:`RolloutResult` from a prompt batch.
|
||||
|
||||
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`.
|
||||
|
||||
Usage::
|
||||
|
||||
runner = RolloutRunner(policy, old_policy, tokenizer,
|
||||
reward_model, sampling_pipeline, config)
|
||||
result = runner(prompt_batch)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
policy_model: nn.Module,
|
||||
old_model: Optional[nn.Module],
|
||||
tokenizer,
|
||||
reward_model: BaseRewardModel,
|
||||
sampling_pipeline: SamplingPipeline,
|
||||
max_tokens: int = 1024,
|
||||
group_size: int = 8,
|
||||
rollout_interval: int = 512,
|
||||
):
|
||||
self.policy_model = policy_model
|
||||
self.old_model = old_model
|
||||
self.tokenizer = tokenizer
|
||||
self.reward_model = reward_model
|
||||
self.sampling_pipeline = sampling_pipeline
|
||||
self.max_tokens = max_tokens
|
||||
self.group_size = group_size
|
||||
self.rollout_interval = rollout_interval
|
||||
self.stop_ids = getattr(tokenizer, "stop_ids", []) or []
|
||||
|
||||
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 _tokenize_prompts(self, raw_texts: List[str]) -> Dict[str, Tensor]:
|
||||
ids_list = self.tokenizer.encode(raw_texts, out_ids=True)
|
||||
B = len(ids_list)
|
||||
P_max = max(len(ids) for ids in ids_list) if ids_list else 0
|
||||
input_ids = torch.zeros(B, P_max, dtype=torch.long)
|
||||
for i, ids in enumerate(ids_list):
|
||||
input_ids[i, : len(ids)] = torch.tensor(ids[:P_max], dtype=torch.long)
|
||||
attention_mask = input_ids != 0
|
||||
return {"input_ids": input_ids, "attention_mask": attention_mask}
|
||||
|
||||
def _decode(self, token_ids: Tensor, mask: Tensor) -> List[List[str]]:
|
||||
B, G, _ = token_ids.shape
|
||||
texts = []
|
||||
for i in range(B):
|
||||
group_texts = []
|
||||
for g in range(G):
|
||||
ids = token_ids[i, g, mask[i, g]].tolist()
|
||||
group_texts.append(self.tokenizer.decode(ids, skip_special_tokens=True))
|
||||
texts.append(group_texts)
|
||||
return texts
|
||||
|
||||
@torch.no_grad()
|
||||
def _run(self, batch: Dict[str, Tensor]) -> RolloutResult:
|
||||
"""Execute the actual generation + reward scoring."""
|
||||
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, P_len = prompt_ids.shape
|
||||
G = self.group_size
|
||||
device = prompt_ids.device
|
||||
|
||||
prompt_texts: List[str] = []
|
||||
for i in range(B):
|
||||
ids = prompt_ids[i, prompt_mask[i]].tolist()
|
||||
prompt_texts.append(self.tokenizer.decode(ids, skip_special_tokens=True))
|
||||
|
||||
expanded_ids = prompt_ids.unsqueeze(1).expand(-1, G, -1).reshape(B * G, P_len)
|
||||
expanded_mask = prompt_mask.unsqueeze(1).expand(-1, G, -1).reshape(B * G, P_len)
|
||||
|
||||
gen_out = generate_responses(
|
||||
model=self.policy_model,
|
||||
input_ids=expanded_ids,
|
||||
attention_mask=expanded_mask,
|
||||
max_new_tokens=self.max_tokens,
|
||||
sampling_pipeline=self.sampling_pipeline,
|
||||
stop_ids=self.stop_ids,
|
||||
)
|
||||
|
||||
gen_ids = gen_out["generated_ids"].reshape(B, G, -1)
|
||||
gen_mask = gen_out["generated_mask"].reshape(B, G, -1)
|
||||
gen_logprobs = gen_out["logprobs"].reshape(B, G, -1)
|
||||
|
||||
response_texts = self._decode(gen_ids, gen_mask)
|
||||
reward_tensor = self.reward_model.score(prompt_texts, response_texts)
|
||||
rewards = reward_tensor.to(device=device)
|
||||
|
||||
return RolloutResult(
|
||||
prompts=prompt_ids,
|
||||
responses=gen_ids,
|
||||
response_mask=gen_mask,
|
||||
rewards=rewards,
|
||||
logprobs_old=gen_logprobs,
|
||||
prompt_texts=prompt_texts,
|
||||
response_texts=response_texts,
|
||||
)
|
||||
|
||||
def __call__(self, batch: Dict[str, Tensor]) -> RolloutResult:
|
||||
"""Return cached or fresh :class:`RolloutResult`.
|
||||
|
||||
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:
|
||||
self._cache = self._run(batch)
|
||||
self._steps_since_rollout = 0
|
||||
return self._cache
|
||||
+101
-3
@@ -9,6 +9,7 @@ import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from astrai.factory import BaseFactory
|
||||
from astrai.trainer.rollout import RolloutResult
|
||||
|
||||
|
||||
def create_ref_model(
|
||||
@@ -87,7 +88,15 @@ def make_doc_boundary_mask(position_ids: Tensor) -> Tensor:
|
||||
|
||||
|
||||
class BaseStrategy(ABC):
|
||||
"""Abstract base class for training strategies."""
|
||||
"""Abstract base class for training strategies.
|
||||
|
||||
When a :class:`~astrai.trainer.rollout.RolloutRunner` is injected via
|
||||
:meth:`set_rollout_runner`, the strategy transparently switches to
|
||||
online mode: each ``__call__`` produces a :class:`RolloutResult`,
|
||||
converts it to a training batch via :meth:`prepare_from_rollout`, and
|
||||
then computes the loss. Without a runner the strategy runs in
|
||||
offline mode and consumes the batch directly.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -99,6 +108,8 @@ class BaseStrategy(ABC):
|
||||
self.device = device
|
||||
self.executor = kwargs.pop("executor", None)
|
||||
self.extra_kwargs = kwargs
|
||||
self._rollout_runner = None
|
||||
self._prev_rollout_result = None
|
||||
|
||||
@abstractmethod
|
||||
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
|
||||
@@ -112,9 +123,51 @@ class BaseStrategy(ABC):
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def supports_online(self) -> bool:
|
||||
"""Whether this strategy can operate with a rollout runner.
|
||||
|
||||
Base implementation returns ``False``; strategies that implement
|
||||
:meth:`prepare_from_rollout` should override to return ``True``.
|
||||
"""
|
||||
return False
|
||||
|
||||
def set_rollout_runner(self, runner):
|
||||
"""Inject a :class:`RolloutRunner` to enable online rollout mode."""
|
||||
self._rollout_runner = runner
|
||||
|
||||
def prepare_from_rollout(self, result: RolloutResult) -> Dict[str, Tensor]:
|
||||
"""Map a :class:`RolloutResult` to the batch layout expected by
|
||||
:meth:`compute_loss`.
|
||||
|
||||
Strategies that return ``True`` from :meth:`supports_online` must
|
||||
override this. Default raises :class:`NotImplementedError`.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
f"{type(self).__name__} does not support online rollout"
|
||||
)
|
||||
|
||||
def _on_rollout_refresh(self):
|
||||
"""Hook fired when a fresh rollout result is produced.
|
||||
|
||||
Override to refresh stale state (e.g. syncing the behaviour
|
||||
policy). Default is a no-op.
|
||||
"""
|
||||
pass
|
||||
|
||||
def __call__(self, batch: Dict[str, Tensor]) -> Tensor:
|
||||
"""Allow calling strategy directly as a callable."""
|
||||
return self.compute_loss(batch)
|
||||
"""Run offline or online forward depending on runner injection."""
|
||||
if self._rollout_runner is None:
|
||||
return self.compute_loss(batch)
|
||||
|
||||
result = self._rollout_runner(batch)
|
||||
if result is not self._prev_rollout_result:
|
||||
self._on_rollout_refresh()
|
||||
self._prev_rollout_result = result
|
||||
if self.executor and self.executor.sync_gradients:
|
||||
self._rollout_runner.step()
|
||||
|
||||
train_batch = self.prepare_from_rollout(result)
|
||||
return self.compute_loss(train_batch)
|
||||
|
||||
|
||||
class StrategyFactory(BaseFactory["BaseStrategy"]):
|
||||
@@ -260,6 +313,29 @@ class DPOStrategy(BaseStrategy):
|
||||
|
||||
return dpo_loss
|
||||
|
||||
def supports_online(self) -> bool:
|
||||
return True
|
||||
|
||||
def prepare_from_rollout(self, result: RolloutResult) -> Dict[str, Tensor]:
|
||||
"""Pick best/worst response per prompt by reward as chosen/rejected."""
|
||||
rewards = result.rewards
|
||||
responses = result.responses
|
||||
masks = result.response_mask
|
||||
best = rewards.argmax(dim=-1)
|
||||
worst = rewards.argmin(dim=-1)
|
||||
B = responses.shape[0]
|
||||
idx = torch.arange(B, device=responses.device)
|
||||
chosen = responses[idx, best]
|
||||
chosen_mask = masks[idx, best].float()
|
||||
rejected = responses[idx, worst]
|
||||
rejected_mask = masks[idx, worst].float()
|
||||
return {
|
||||
"chosen": chosen,
|
||||
"chosen_mask": chosen_mask,
|
||||
"rejected": rejected,
|
||||
"rejected_mask": rejected_mask,
|
||||
}
|
||||
|
||||
|
||||
@StrategyFactory.register("grpo")
|
||||
class GRPOStrategy(BaseStrategy):
|
||||
@@ -371,3 +447,25 @@ class GRPOStrategy(BaseStrategy):
|
||||
total_loss = policy_loss + kl_penalty
|
||||
|
||||
return total_loss
|
||||
|
||||
def supports_online(self) -> bool:
|
||||
return True
|
||||
|
||||
def prepare_from_rollout(self, result: RolloutResult) -> Dict[str, Tensor]:
|
||||
return {
|
||||
"prompts": result.prompts,
|
||||
"responses": result.responses,
|
||||
"masks": result.response_mask,
|
||||
"rewards": result.rewards,
|
||||
}
|
||||
|
||||
def _on_rollout_refresh(self):
|
||||
"""Sync the behaviour policy whenever a fresh rollout arrives."""
|
||||
self.sync_old_model()
|
||||
|
||||
|
||||
# Factory aliases: online variants use the same strategy class; the
|
||||
# ``RolloutRunner`` is injected by ``TrainContextBuilder`` to enable
|
||||
# online mode, so no separate subclass is needed.
|
||||
StrategyFactory._entries["online_grpo"] = GRPOStrategy
|
||||
StrategyFactory._entries["online_dpo"] = DPOStrategy
|
||||
|
||||
@@ -8,11 +8,19 @@ from torch.utils.data import DataLoader, random_split
|
||||
|
||||
from astrai.config.train_config import TrainConfig
|
||||
from astrai.dataset import RDSampler
|
||||
from astrai.inference.sample import (
|
||||
SamplingPipeline,
|
||||
TemperatureStrategy,
|
||||
TopKStrategy,
|
||||
TopPStrategy,
|
||||
)
|
||||
from astrai.model.components.lora import inject_lora
|
||||
from astrai.parallel.executor import BaseExecutor, ExecutorFactory
|
||||
from astrai.parallel.setup import get_current_device, get_rank, get_world_size
|
||||
from astrai.protocols import OptimizerProtocol, SchedulerProtocol
|
||||
from astrai.serialization import Checkpoint, load_json
|
||||
from astrai.tokenize import AutoTokenizer
|
||||
from astrai.trainer.rollout import RolloutRunner
|
||||
from astrai.trainer.strategy import BaseStrategy, StrategyFactory, create_ref_model
|
||||
|
||||
|
||||
@@ -27,7 +35,6 @@ class TrainContext:
|
||||
config: TrainConfig = field(default=None)
|
||||
model_config: dict = field(default_factory=dict)
|
||||
executor: BaseExecutor = field(default=None)
|
||||
|
||||
epoch: int = field(default=0)
|
||||
consumed_samples: int = field(default=0)
|
||||
loss: float = field(default=0.0)
|
||||
@@ -194,13 +201,22 @@ class TrainContextBuilder:
|
||||
|
||||
strategy_kwargs = dict(cfg.extra_kwargs)
|
||||
|
||||
if cfg.strategy in ("dpo", "grpo"):
|
||||
needs_ref = cfg.strategy in (
|
||||
"dpo",
|
||||
"grpo",
|
||||
"online_grpo",
|
||||
"online_dpo",
|
||||
)
|
||||
needs_old = cfg.strategy in ("grpo", "online_grpo")
|
||||
|
||||
if needs_ref:
|
||||
ref_model = create_ref_model(
|
||||
cfg.model_fn, executor.unwrap_model(context.model)
|
||||
).to(device=device)
|
||||
strategy_kwargs["ref_model"] = ref_model
|
||||
|
||||
if cfg.strategy == "grpo":
|
||||
old_model = None
|
||||
if needs_old:
|
||||
old_model = create_ref_model(
|
||||
cfg.model_fn, executor.unwrap_model(context.model)
|
||||
).to(device=device)
|
||||
@@ -214,4 +230,37 @@ class TrainContextBuilder:
|
||||
**strategy_kwargs,
|
||||
)
|
||||
|
||||
# Enable online rollout when the train_type is an ``online_*`` variant.
|
||||
is_online = cfg.strategy.startswith("online_")
|
||||
if is_online:
|
||||
if not context.strategy.supports_online():
|
||||
raise ValueError(
|
||||
f"Strategy '{cfg.strategy}' does not support online rollout"
|
||||
)
|
||||
if cfg.reward_model_fn is None:
|
||||
raise ValueError("reward_model_fn is required for online RL strategies")
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(self._param_path)
|
||||
reward_model = cfg.reward_model_fn()
|
||||
|
||||
pipeline = SamplingPipeline(
|
||||
[
|
||||
TemperatureStrategy(cfg.rollout_temperature),
|
||||
TopKStrategy(cfg.rollout_top_k),
|
||||
TopPStrategy(cfg.rollout_top_p),
|
||||
]
|
||||
)
|
||||
|
||||
runner = RolloutRunner(
|
||||
policy_model=context.model,
|
||||
old_model=old_model,
|
||||
tokenizer=tokenizer,
|
||||
reward_model=reward_model,
|
||||
sampling_pipeline=pipeline,
|
||||
max_tokens=cfg.rollout_max_tokens,
|
||||
group_size=strategy_kwargs.get("group_size", 8),
|
||||
rollout_interval=cfg.rollout_interval,
|
||||
)
|
||||
context.strategy.set_rollout_runner(runner)
|
||||
|
||||
return context
|
||||
|
||||
+59
-3
@@ -1,7 +1,7 @@
|
||||
import argparse
|
||||
import os
|
||||
from functools import partial
|
||||
from typing import Any, Dict
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
|
||||
import torch
|
||||
import torch.optim as optim
|
||||
@@ -12,6 +12,7 @@ from astrai.dataset import DatasetFactory, dpo_collate_fn, grpo_collate_fn
|
||||
from astrai.model import AutoRegressiveLM
|
||||
from astrai.model.components.decoder_block import DecoderBlock
|
||||
from astrai.trainer import SchedulerFactory, Trainer
|
||||
from astrai.trainer.rollout import BaseRewardModel
|
||||
|
||||
|
||||
class MuonMix(optim.Optimizer):
|
||||
@@ -101,7 +102,7 @@ def parse_args() -> argparse.Namespace:
|
||||
"--train_type",
|
||||
type=str,
|
||||
required=True,
|
||||
choices=["seq", "sft", "dpo", "grpo"],
|
||||
choices=["seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo"],
|
||||
help="Train type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -217,6 +218,39 @@ def parse_args() -> argparse.Namespace:
|
||||
default=0.0,
|
||||
help="cross_entropy function label smoothing parameter",
|
||||
)
|
||||
|
||||
# online rollout
|
||||
parser.add_argument(
|
||||
"--rollout_interval",
|
||||
type=int,
|
||||
default=512,
|
||||
help="Number of optimizer steps between online rollouts.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rollout_temperature",
|
||||
type=float,
|
||||
default=0.7,
|
||||
help="Sampling temperature for online rollout.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rollout_top_k",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Top-k filtering for online rollout (0=disable).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rollout_top_p",
|
||||
type=float,
|
||||
default=0.9,
|
||||
help="Top-p (nucleus) filtering for online rollout.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rollout_max_tokens",
|
||||
type=int,
|
||||
default=1024,
|
||||
help="Maximum generated tokens per response in rollout.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--gradient_checkpointing",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
@@ -428,7 +462,14 @@ def train(
|
||||
decay_steps: int,
|
||||
**kwargs,
|
||||
):
|
||||
assert train_type in ["seq", "sft", "dpo", "grpo"]
|
||||
assert train_type in [
|
||||
"seq",
|
||||
"sft",
|
||||
"dpo",
|
||||
"grpo",
|
||||
"online_grpo",
|
||||
"online_dpo",
|
||||
]
|
||||
assert os.path.exists(param_path)
|
||||
if nprocs > 1 and parallel_mode == "none":
|
||||
raise ValueError("--nprocs > 1 requires --parallel_mode to be 'ddp' or 'fsdp'")
|
||||
@@ -449,6 +490,13 @@ def train(
|
||||
"group_size": kwargs.pop("group_size"),
|
||||
}
|
||||
|
||||
rollout_interval = kwargs.pop("rollout_interval", 512)
|
||||
rollout_temperature = kwargs.pop("rollout_temperature", 0.7)
|
||||
rollout_top_k = kwargs.pop("rollout_top_k", 0)
|
||||
rollout_top_p = kwargs.pop("rollout_top_p", 0.9)
|
||||
rollout_max_tokens = kwargs.pop("rollout_max_tokens", 1024)
|
||||
reward_model_fn: Optional[Callable[[], BaseRewardModel]] = None
|
||||
|
||||
executor_kwargs = {}
|
||||
if parallel_mode == "ddp":
|
||||
executor_kwargs.update(
|
||||
@@ -512,6 +560,8 @@ def train(
|
||||
collate_fn = dpo_collate_fn
|
||||
elif train_type == "grpo":
|
||||
collate_fn = grpo_collate_fn
|
||||
elif train_type in ("online_grpo", "online_dpo"):
|
||||
collate_fn = None
|
||||
|
||||
train_config = TrainConfig(
|
||||
model_fn=model_fn,
|
||||
@@ -546,6 +596,12 @@ def train(
|
||||
extra_kwargs=strategy_kwargs,
|
||||
neftune_alpha=neftune_alpha,
|
||||
collate_fn=collate_fn,
|
||||
rollout_interval=rollout_interval,
|
||||
rollout_temperature=rollout_temperature,
|
||||
rollout_top_k=rollout_top_k,
|
||||
rollout_top_p=rollout_top_p,
|
||||
rollout_max_tokens=rollout_max_tokens,
|
||||
reward_model_fn=reward_model_fn,
|
||||
)
|
||||
|
||||
trainer = Trainer(train_config)
|
||||
|
||||
@@ -0,0 +1,337 @@
|
||||
"""Unit tests for online rollout integration in :class:`BaseStrategy`.
|
||||
|
||||
Covers the shared rollout-trigger logic in ``BaseStrategy.__call__``
|
||||
(runner injection, cache-driven refresh hook, ``step()`` callback) and
|
||||
the per-strategy ``prepare_from_rollout`` mappings for both
|
||||
:class:`GRPOStrategy` and :class:`DPOStrategy`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from astrai.config.model_config import AutoRegressiveLMConfig
|
||||
from astrai.model.transformer import AutoRegressiveLM
|
||||
from astrai.trainer.rollout import RolloutResult
|
||||
from astrai.trainer.strategy import (
|
||||
DPOStrategy,
|
||||
GRPOStrategy,
|
||||
StrategyFactory,
|
||||
)
|
||||
|
||||
|
||||
class _FakeExecutor:
|
||||
"""Executor stub tracking ``sync_gradients`` and providing unwrap_model."""
|
||||
|
||||
def __init__(self, sync_gradients=True):
|
||||
self._sync_gradients = sync_gradients
|
||||
|
||||
@property
|
||||
def sync_gradients(self):
|
||||
return self._sync_gradients
|
||||
|
||||
def unwrap_model(self, model):
|
||||
return model.state_dict()
|
||||
|
||||
|
||||
def _make_config(vocab_size=200, max_len=64):
|
||||
return AutoRegressiveLMConfig(
|
||||
vocab_size=vocab_size,
|
||||
dim=16,
|
||||
n_heads=2,
|
||||
n_kv_heads=1,
|
||||
dim_ffn=32,
|
||||
max_len=max_len,
|
||||
n_layers=2,
|
||||
norm_eps=1e-5,
|
||||
)
|
||||
|
||||
|
||||
def _make_model(device):
|
||||
cfg = _make_config()
|
||||
return AutoRegressiveLM(cfg).to(device=device), cfg
|
||||
|
||||
|
||||
def _make_frozen(model, device):
|
||||
cfg = _make_config()
|
||||
copy = AutoRegressiveLM(cfg).to(device=device)
|
||||
copy.load_state_dict(model.state_dict())
|
||||
copy.requires_grad_(False)
|
||||
copy.eval()
|
||||
return copy
|
||||
|
||||
|
||||
def _make_rollout_result(B=2, G=4, P=6, R=8, device="cpu"):
|
||||
return RolloutResult(
|
||||
prompts=torch.randint(3, 200, (B, P), device=device),
|
||||
responses=torch.randint(3, 200, (B, G, R), device=device),
|
||||
response_mask=torch.ones(B, G, R, dtype=torch.bool, device=device),
|
||||
rewards=torch.randn(B, G, device=device),
|
||||
logprobs_old=torch.zeros(B, G, R, device=device),
|
||||
)
|
||||
|
||||
|
||||
class _RecordingRunner:
|
||||
"""Fake RolloutRunner that returns a fixed result and tracks calls."""
|
||||
|
||||
def __init__(self, result):
|
||||
self.result = result
|
||||
self.calls = 0
|
||||
self.step_calls = 0
|
||||
|
||||
def __call__(self, batch):
|
||||
self.calls += 1
|
||||
return self.result
|
||||
|
||||
def step(self):
|
||||
self.step_calls += 1
|
||||
|
||||
def swap_result(self, result):
|
||||
self.result = result
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def device():
|
||||
return "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
|
||||
def _make_grpo(device, executor=None):
|
||||
model, _ = _make_model(device)
|
||||
old_model = _make_frozen(model, device)
|
||||
ref_model = _make_frozen(model, device)
|
||||
return GRPOStrategy(
|
||||
model=model,
|
||||
device=device,
|
||||
old_model=old_model,
|
||||
ref_model=ref_model,
|
||||
clip_eps=0.2,
|
||||
kl_coef=0.01,
|
||||
group_size=4,
|
||||
model_fn=lambda c=_make_config(): AutoRegressiveLM(c).to(device=device),
|
||||
executor=executor or _FakeExecutor(),
|
||||
)
|
||||
|
||||
|
||||
def _make_dpo(device, executor=None):
|
||||
model, _ = _make_model(device)
|
||||
ref_model = _make_frozen(model, device)
|
||||
return DPOStrategy(
|
||||
model=model,
|
||||
device=device,
|
||||
ref_model=ref_model,
|
||||
beta=0.1,
|
||||
reduction="sum",
|
||||
model_fn=lambda c=_make_config(): AutoRegressiveLM(c).to(device=device),
|
||||
executor=executor or _FakeExecutor(),
|
||||
)
|
||||
|
||||
|
||||
def test_factory_registers_online_aliases():
|
||||
assert StrategyFactory.is_registered("online_grpo")
|
||||
assert StrategyFactory.is_registered("online_dpo")
|
||||
assert StrategyFactory._entries["online_grpo"] is GRPOStrategy
|
||||
assert StrategyFactory._entries["online_dpo"] is DPOStrategy
|
||||
|
||||
|
||||
def test_grpo_supports_online(device):
|
||||
assert _make_grpo(device).supports_online() is True
|
||||
|
||||
|
||||
def test_dpo_supports_online(device):
|
||||
assert _make_dpo(device).supports_online() is True
|
||||
|
||||
|
||||
def test_base_strategy_prepare_from_rollout_raises_by_default(device):
|
||||
from astrai.trainer.strategy import BaseStrategy
|
||||
|
||||
class _Offline(BaseStrategy):
|
||||
def compute_loss(self, batch):
|
||||
return torch.tensor(0.0)
|
||||
|
||||
strat = _Offline(model=torch.nn.Linear(1, 1), device="cpu")
|
||||
with pytest.raises(NotImplementedError):
|
||||
strat.prepare_from_rollout(_make_rollout_result(device="cpu"))
|
||||
|
||||
|
||||
def test_base_strategy_supports_online_default_false():
|
||||
from astrai.trainer.strategy import BaseStrategy
|
||||
|
||||
class _Offline(BaseStrategy):
|
||||
def compute_loss(self, batch):
|
||||
return torch.tensor(0.0)
|
||||
|
||||
strat = _Offline(model=torch.nn.Linear(1, 1), device="cpu")
|
||||
assert strat.supports_online() is False
|
||||
|
||||
|
||||
def test_grpo_prepare_from_rollout_mapping(device):
|
||||
strat = _make_grpo(device)
|
||||
r = _make_rollout_result(device=device)
|
||||
batch = strat.prepare_from_rollout(r)
|
||||
assert batch["prompts"] is r.prompts
|
||||
assert batch["responses"] is r.responses
|
||||
assert batch["masks"] is r.response_mask
|
||||
assert batch["rewards"] is r.rewards
|
||||
|
||||
|
||||
def test_dpo_prepare_from_rollout_picks_best_worst(device):
|
||||
strat = _make_dpo(device)
|
||||
r = _make_rollout_result(B=3, G=4, R=5, device=device)
|
||||
batch = strat.prepare_from_rollout(r)
|
||||
assert batch["chosen"].shape == (3, 5)
|
||||
assert batch["rejected"].shape == (3, 5)
|
||||
assert batch["chosen_mask"].shape == (3, 5)
|
||||
assert batch["rejected_mask"].shape == (3, 5)
|
||||
idx = torch.arange(3, device=device)
|
||||
expected_best = r.responses[idx, r.rewards.argmax(dim=-1)]
|
||||
expected_worst = r.responses[idx, r.rewards.argmin(dim=-1)]
|
||||
assert torch.equal(batch["chosen"], expected_best)
|
||||
assert torch.equal(batch["rejected"], expected_worst)
|
||||
|
||||
|
||||
def test_call_without_runner_falls_back_to_compute_loss_grpo(device):
|
||||
strat = _make_grpo(device)
|
||||
batch = {
|
||||
"prompts": torch.randint(3, 200, (2, 4), device=device),
|
||||
"responses": torch.randint(3, 200, (2, 4, 6), device=device),
|
||||
"masks": torch.ones(2, 4, 6, device=device),
|
||||
"rewards": torch.randn(2, 4, device=device),
|
||||
}
|
||||
loss = strat(batch)
|
||||
assert torch.isfinite(loss).item()
|
||||
|
||||
|
||||
def test_call_with_runner_returns_finite_loss_grpo(device):
|
||||
strat = _make_grpo(device)
|
||||
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
|
||||
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
assert torch.isfinite(loss).item()
|
||||
|
||||
|
||||
def test_call_with_runner_returns_finite_loss_dpo(device):
|
||||
strat = _make_dpo(device)
|
||||
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
|
||||
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
assert torch.isfinite(loss).item()
|
||||
|
||||
|
||||
def test_call_invokes_runner_each_time(device):
|
||||
strat = _make_grpo(device)
|
||||
runner = _RecordingRunner(_make_rollout_result(device=device))
|
||||
strat.set_rollout_runner(runner)
|
||||
strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
assert runner.calls == 2
|
||||
|
||||
|
||||
def test_grpo_syncs_old_model_on_first_rollout(device):
|
||||
strat = _make_grpo(device)
|
||||
runner = _RecordingRunner(_make_rollout_result(device=device))
|
||||
strat.set_rollout_runner(runner)
|
||||
with torch.no_grad():
|
||||
for p in strat.model.parameters():
|
||||
p.add_(0.1)
|
||||
old_before = {k: v.clone() for k, v in strat.old_model.state_dict().items()}
|
||||
strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
old_after = strat.old_model.state_dict()
|
||||
synced = any(
|
||||
not torch.allclose(old_before[k], old_after[k])
|
||||
for k in old_before
|
||||
if k in old_after
|
||||
)
|
||||
assert synced
|
||||
|
||||
|
||||
def test_grpo_no_resync_when_same_cached_result(device):
|
||||
strat = _make_grpo(device)
|
||||
runner = _RecordingRunner(_make_rollout_result(device=device))
|
||||
strat.set_rollout_runner(runner)
|
||||
strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
assert runner.calls == 2
|
||||
assert runner.step_calls == 1
|
||||
|
||||
|
||||
def test_grpo_resync_when_new_rollout_result(device):
|
||||
strat = _make_grpo(device)
|
||||
runner = _RecordingRunner(_make_rollout_result(device=device))
|
||||
strat.set_rollout_runner(runner)
|
||||
strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
runner.swap_result(_make_rollout_result(device=device))
|
||||
strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
assert runner.calls == 2
|
||||
assert runner.step_calls == 2
|
||||
|
||||
|
||||
def test_dpo_no_sync_hook_when_new_rollout_result(device):
|
||||
"""DPO has no old_model, so ``_on_rollout_refresh`` must be a no-op.
|
||||
|
||||
We verify by ensuring no AttributeError is raised (DPO has no
|
||||
old_model) and that step is still called.
|
||||
"""
|
||||
strat = _make_dpo(device)
|
||||
runner = _RecordingRunner(_make_rollout_result(device=device))
|
||||
strat.set_rollout_runner(runner)
|
||||
strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
runner.swap_result(_make_rollout_result(device=device))
|
||||
strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
assert runner.step_calls == 2
|
||||
|
||||
|
||||
def test_step_not_called_when_sync_gradients_false(device):
|
||||
executor = _FakeExecutor(sync_gradients=False)
|
||||
strat = _make_grpo(device, executor=executor)
|
||||
runner = _RecordingRunner(_make_rollout_result(device=device))
|
||||
strat.set_rollout_runner(runner)
|
||||
strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
assert runner.step_calls == 0
|
||||
|
||||
|
||||
def test_step_called_when_sync_gradients_true(device):
|
||||
executor = _FakeExecutor(sync_gradients=True)
|
||||
strat = _make_grpo(device, executor=executor)
|
||||
runner = _RecordingRunner(_make_rollout_result(device=device))
|
||||
strat.set_rollout_runner(runner)
|
||||
strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
assert runner.step_calls == 1
|
||||
|
||||
|
||||
def test_loss_is_differentiable_grpo(device):
|
||||
strat = _make_grpo(device)
|
||||
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
|
||||
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
loss.backward()
|
||||
has_grad = any(
|
||||
p.grad is not None and p.grad.abs().sum() > 0 for p in strat.model.parameters()
|
||||
)
|
||||
assert has_grad
|
||||
|
||||
|
||||
def test_loss_is_differentiable_dpo(device):
|
||||
strat = _make_dpo(device)
|
||||
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
|
||||
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
loss.backward()
|
||||
has_grad = any(
|
||||
p.grad is not None and p.grad.abs().sum() > 0 for p in strat.model.parameters()
|
||||
)
|
||||
assert has_grad
|
||||
|
||||
|
||||
def test_ref_and_old_model_not_updated_by_backward_grpo(device):
|
||||
strat = _make_grpo(device)
|
||||
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
|
||||
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
loss.backward()
|
||||
for p in strat.ref_model.parameters():
|
||||
assert p.grad is None
|
||||
for p in strat.old_model.parameters():
|
||||
assert p.grad is None
|
||||
|
||||
|
||||
def test_ref_model_not_updated_by_backward_dpo(device):
|
||||
strat = _make_dpo(device)
|
||||
strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
|
||||
loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
|
||||
loss.backward()
|
||||
for p in strat.ref_model.parameters():
|
||||
assert p.grad is None
|
||||
@@ -0,0 +1,264 @@
|
||||
"""Unit tests for the online rollout module.
|
||||
|
||||
Covers :class:`RolloutResult`, :class:`BaseRewardModel`,
|
||||
:func:`generate_responses`, and :class:`RolloutRunner` including
|
||||
its internal cache and rollout-interval trigger logic.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from astrai.config.model_config import AutoRegressiveLMConfig
|
||||
from astrai.inference.sample import (
|
||||
SamplingPipeline,
|
||||
TemperatureStrategy,
|
||||
TopKStrategy,
|
||||
TopPStrategy,
|
||||
)
|
||||
from astrai.model.transformer import AutoRegressiveLM
|
||||
from astrai.trainer.rollout import (
|
||||
BaseRewardModel,
|
||||
RolloutResult,
|
||||
RolloutRunner,
|
||||
generate_responses,
|
||||
)
|
||||
|
||||
|
||||
class FakeTokenizer:
|
||||
"""Minimal char-level tokenizer stub for rollout tests.
|
||||
|
||||
Vocab: 0 = pad, 1..255 = byte values. ``stop_ids = [2]`` (a fake
|
||||
EOS) so tests can verify early-stopping behaviour.
|
||||
"""
|
||||
|
||||
stop_ids = [2]
|
||||
|
||||
def encode(self, texts, out_ids=True, **_):
|
||||
if isinstance(texts, str):
|
||||
texts = [texts]
|
||||
return [[b for b in t.encode("utf-8")] for t in texts]
|
||||
|
||||
def decode(self, ids, skip_special_tokens=True):
|
||||
out = bytes(b for b in ids if b > 2 or not skip_special_tokens).decode(
|
||||
"utf-8", errors="ignore"
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
class ConstantRewardModel(BaseRewardModel):
|
||||
"""Returns a constant reward for every response."""
|
||||
|
||||
def __init__(self, value: float = 1.0):
|
||||
self.value = value
|
||||
|
||||
def score(self, prompts, responses):
|
||||
B = len(prompts)
|
||||
G = len(responses[0]) if B else 0
|
||||
return torch.full((B, G), float(self.value))
|
||||
|
||||
|
||||
class _FakeOldModel:
|
||||
"""Placeholder old-model; RolloutRunner stores but never calls it."""
|
||||
|
||||
|
||||
def _make_config(vocab_size=200, max_len=128):
|
||||
return AutoRegressiveLMConfig(
|
||||
vocab_size=vocab_size,
|
||||
dim=16,
|
||||
n_heads=2,
|
||||
n_kv_heads=1,
|
||||
dim_ffn=32,
|
||||
max_len=max_len,
|
||||
n_layers=2,
|
||||
norm_eps=1e-5,
|
||||
)
|
||||
|
||||
|
||||
def _make_model(device):
|
||||
cfg = _make_config()
|
||||
m = AutoRegressiveLM(cfg).to(device=device)
|
||||
m.eval()
|
||||
return m, cfg
|
||||
|
||||
|
||||
def _make_pipeline():
|
||||
return SamplingPipeline(
|
||||
[TemperatureStrategy(1.0), TopKStrategy(0), TopPStrategy(1.0)]
|
||||
)
|
||||
|
||||
|
||||
def _make_prompt_batch(batch_size=2, prompt_len=6, device="cpu"):
|
||||
ids = torch.randint(3, 200, (batch_size, prompt_len), device=device)
|
||||
mask = torch.ones(batch_size, prompt_len, dtype=torch.bool, device=device)
|
||||
return {"input_ids": ids, "attention_mask": mask}
|
||||
|
||||
|
||||
def test_rollout_result_fields():
|
||||
r = RolloutResult(
|
||||
prompts=torch.zeros(2, 4, dtype=torch.long),
|
||||
responses=torch.zeros(2, 3, 5, dtype=torch.long),
|
||||
response_mask=torch.ones(2, 3, 5, dtype=torch.bool),
|
||||
rewards=torch.zeros(2, 3),
|
||||
logprobs_old=torch.zeros(2, 3, 5),
|
||||
)
|
||||
assert r.prompts.shape == (2, 4)
|
||||
assert r.responses.shape == (2, 3, 5)
|
||||
assert r.prompt_texts == []
|
||||
assert r.response_texts == []
|
||||
|
||||
|
||||
def test_base_reward_model_is_abstract():
|
||||
with pytest.raises(TypeError):
|
||||
BaseRewardModel()
|
||||
|
||||
|
||||
def test_constant_reward_model_shape():
|
||||
rm = ConstantRewardModel(0.5)
|
||||
out = rm.score(["a", "b"], [["x", "y", "z"], ["p", "q", "r"]])
|
||||
assert out.shape == (2, 3)
|
||||
assert torch.all(out == 0.5)
|
||||
|
||||
|
||||
def test_generate_responses_shapes():
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
model, _ = _make_model(device)
|
||||
pipeline = _make_pipeline()
|
||||
ids = torch.randint(3, 200, (2, 4), device=device)
|
||||
mask = torch.ones(2, 4, dtype=torch.bool, device=device)
|
||||
|
||||
out = generate_responses(
|
||||
model=model,
|
||||
input_ids=ids,
|
||||
attention_mask=mask,
|
||||
max_new_tokens=8,
|
||||
sampling_pipeline=pipeline,
|
||||
stop_ids=[],
|
||||
)
|
||||
assert out["generated_ids"].shape == (2, 8)
|
||||
assert out["generated_mask"].shape == (2, 8)
|
||||
assert out["logprobs"].shape == (2, 8)
|
||||
|
||||
|
||||
def test_generate_responses_stops_on_stop_id():
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
model, _ = _make_model(device)
|
||||
pipeline = _make_pipeline()
|
||||
ids = torch.randint(3, 200, (1, 3), device=device)
|
||||
mask = torch.ones(1, 3, dtype=torch.bool, device=device)
|
||||
|
||||
out = generate_responses(
|
||||
model=model,
|
||||
input_ids=ids,
|
||||
attention_mask=mask,
|
||||
max_new_tokens=16,
|
||||
sampling_pipeline=pipeline,
|
||||
stop_ids=[7],
|
||||
)
|
||||
gen = out["generated_ids"][0]
|
||||
mask = out["generated_mask"][0]
|
||||
# If a 7 appeared, all tokens after it must be pad (mask False).
|
||||
nonzero_stop = (gen == 7).nonzero()
|
||||
if nonzero_stop.numel():
|
||||
first = nonzero_stop[0].item()
|
||||
assert mask[first + 1 :].sum() == 0
|
||||
|
||||
|
||||
def test_generate_responses_logprobs_match_tokens():
|
||||
"""logprobs[i] must be the logprob of generated_ids[i]."""
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
model, _ = _make_model(device)
|
||||
pipeline = _make_pipeline()
|
||||
ids = torch.randint(3, 200, (1, 2), device=device)
|
||||
mask = torch.ones(1, 2, dtype=torch.bool, device=device)
|
||||
|
||||
out = generate_responses(
|
||||
model=model,
|
||||
input_ids=ids,
|
||||
attention_mask=mask,
|
||||
max_new_tokens=4,
|
||||
sampling_pipeline=pipeline,
|
||||
stop_ids=[],
|
||||
)
|
||||
gen = out["generated_ids"][0]
|
||||
lp = out["logprobs"][0]
|
||||
for i in range(4):
|
||||
if gen[i] == 0 and not out["generated_mask"][0, i]:
|
||||
continue
|
||||
assert lp[i] <= 0.0
|
||||
|
||||
|
||||
def _make_runner(device, **kw):
|
||||
model, _ = _make_model(device)
|
||||
rm = ConstantRewardModel(1.0)
|
||||
return RolloutRunner(
|
||||
policy_model=model,
|
||||
old_model=_FakeOldModel(),
|
||||
tokenizer=FakeTokenizer(),
|
||||
reward_model=rm,
|
||||
sampling_pipeline=_make_pipeline(),
|
||||
max_tokens=kw.get("max_tokens", 8),
|
||||
group_size=kw.get("group_size", 2),
|
||||
rollout_interval=kw.get("rollout_interval", 2),
|
||||
), model
|
||||
|
||||
|
||||
def test_rollout_runner_shapes():
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
runner, _ = _make_runner(device, group_size=3, max_tokens=5)
|
||||
batch = _make_prompt_batch(batch_size=2, prompt_len=4, device=device)
|
||||
r = runner(batch)
|
||||
assert r.prompts.shape == (2, 4)
|
||||
assert r.responses.shape == (2, 3, 5)
|
||||
assert r.response_mask.shape == (2, 3, 5)
|
||||
assert r.rewards.shape == (2, 3)
|
||||
assert r.logprobs_old.shape == (2, 3, 5)
|
||||
assert len(r.prompt_texts) == 2
|
||||
assert len(r.response_texts) == 2
|
||||
assert len(r.response_texts[0]) == 3
|
||||
|
||||
|
||||
def test_rollout_runner_cache_returns_same_object():
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
runner, _ = _make_runner(device, rollout_interval=10)
|
||||
batch = _make_prompt_batch(device=device)
|
||||
r1 = runner(batch)
|
||||
r2 = runner(batch)
|
||||
assert r1 is r2
|
||||
|
||||
|
||||
def test_rollout_runner_step_triggers_new_rollout():
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
runner, _ = _make_runner(device, rollout_interval=2)
|
||||
batch = _make_prompt_batch(device=device)
|
||||
r1 = runner(batch)
|
||||
runner.step()
|
||||
# interval=2 means trigger when _steps_since_rollout >= 2; 1 step not enough
|
||||
r2 = runner(batch)
|
||||
assert r1 is r2
|
||||
runner.step()
|
||||
# Now _steps_since_rollout == 2 -> re-rollout
|
||||
r3 = runner(batch)
|
||||
assert r3 is not r1
|
||||
|
||||
|
||||
def test_rollout_runner_clear_cache_forces_rerun():
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
runner, _ = _make_runner(device, rollout_interval=100)
|
||||
batch = _make_prompt_batch(device=device)
|
||||
r1 = runner(batch)
|
||||
runner.clear_cache()
|
||||
r2 = runner(batch)
|
||||
assert r2 is not r1
|
||||
|
||||
|
||||
def test_rollout_runner_step_resets_counter():
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
runner, _ = _make_runner(device, rollout_interval=1)
|
||||
batch = _make_prompt_batch(device=device)
|
||||
r1 = runner(batch)
|
||||
runner.step()
|
||||
r2 = runner(batch)
|
||||
assert r2 is not r1
|
||||
# Counter reset after rollout; second call w/o step should be cached.
|
||||
r3 = runner(batch)
|
||||
assert r3 is r2
|
||||
Reference in New Issue
Block a user