diff --git a/benchmarks/benchmark_stateless_executor.py b/benchmarks/benchmark_stateless_executor.py new file mode 100644 index 000000000..214852ed2 --- /dev/null +++ b/benchmarks/benchmark_stateless_executor.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import argparse +import csv +import json +import sys +from datetime import datetime, timezone +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import torch + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.executors.paged_kv_baseline import ( # noqa: E402 + PagedKVScoringBaseline, + PagedKVScoringConfig, +) +from rl_engine.executors.stateless_executor import ( # noqa: E402 + StatelessForwardConfig, + StatelessForwardExecutor, + StatelessForwardInputs, +) + +CSV_COLUMNS = [ + "timestamp", + "candidate", + "mode", + "stage", + "batch_size", + "seq_len", + "active_tokens", + "device", + "dtype", + "elapsed_ms", + "peak_allocated_mb", + "peak_reserved_mb", + "paged_kv_baseline_mb", + "kv_cache_output_mb", + "savings_vs_local_reservation_baseline_mb", + "status", + "notes", +] + +SCORING_MODES = {"reference", "reward", "both"} +DTYPE_CHOICES = { + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float32": torch.float32, +} + + +class TinyReferenceModel(torch.nn.Module): + """Small deterministic causal-LM-like module for executor smoke benchmarks.""" + + def __init__(self, vocab_size: int, hidden_dim: int): + super().__init__() + self.embedding = torch.nn.Embedding(vocab_size, hidden_dim) + self.proj = torch.nn.Linear(hidden_dim, vocab_size) + self.use_cache_calls: list[bool | None] = [] + + def forward(self, input_ids, attention_mask=None, use_cache=None): + del attention_mask + self.use_cache_calls.append(use_cache) + hidden = self.embedding(input_ids.long()) + return SimpleNamespace( + logits=self.proj(hidden), + past_key_values=_tiny_past_key_values(hidden, use_cache), + ) + + +class TinyRewardModel(torch.nn.Module): + """Small deterministic scalar-reward module for executor smoke benchmarks.""" + + def __init__(self, vocab_size: int, hidden_dim: int): + super().__init__() + self.embedding = torch.nn.Embedding(vocab_size, hidden_dim) + self.reward_head = torch.nn.Linear(hidden_dim, 1) + self.use_cache_calls: list[bool | None] = [] + + def forward(self, input_ids, attention_mask=None, use_cache=None): + self.use_cache_calls.append(use_cache) + hidden = self.embedding(input_ids.long()) + mask = attention_mask.to(device=hidden.device, dtype=hidden.dtype).unsqueeze(-1) + pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0) + return { + "logits": self.reward_head(pooled), + "past_key_values": _tiny_past_key_values(hidden, use_cache), + } + + +class TinyBothModel(torch.nn.Module): + """Tiny model with both LM logits and scalar rewards in one forward.""" + + def __init__(self, vocab_size: int, hidden_dim: int): + super().__init__() + self.embedding = torch.nn.Embedding(vocab_size, hidden_dim) + self.lm_head = torch.nn.Linear(hidden_dim, vocab_size) + self.reward_head = torch.nn.Linear(hidden_dim, 1) + self.use_cache_calls: list[bool | None] = [] + + def forward(self, input_ids, attention_mask=None, use_cache=None): + self.use_cache_calls.append(use_cache) + hidden = self.embedding(input_ids.long()) + mask = attention_mask.to(device=hidden.device, dtype=hidden.dtype).unsqueeze(-1) + pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0) + return { + "logits": self.lm_head(hidden), + "rewards": self.reward_head(pooled).squeeze(-1), + "past_key_values": _tiny_past_key_values(hidden, use_cache), + } + + +def _tiny_past_key_values(hidden: torch.Tensor, use_cache: bool | None): + if use_cache is not True: + return None + batch_size, seq_len, hidden_dim = hidden.shape + head_dim = max(1, min(hidden_dim, 8)) + key = torch.empty( + batch_size, + 1, + seq_len, + head_dim, + device=hidden.device, + dtype=hidden.dtype, + ) + value = torch.empty_like(key) + return ((key, value),) + + +def _timestamp() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _environment() -> str: + parts = [f"torch={torch.__version__}", f"cuda_available={torch.cuda.is_available()}"] + if torch.cuda.is_available(): + parts.append(f"cuda={torch.version.cuda}") + parts.append(f"gpu={torch.cuda.get_device_name(0)}") + return ";".join(parts) + + +def _append_row(path: Path, row: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + exists = path.exists() + with path.open("a", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=CSV_COLUMNS) + if not exists: + writer.writeheader() + writer.writerow({column: row.get(column, "") for column in CSV_COLUMNS}) + + +def _write_json(path: Path | None, payload: dict[str, Any]) -> None: + if path is None: + return + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8") + + +def _make_inputs(args: argparse.Namespace, device: torch.device) -> StatelessForwardInputs: + generator = torch.Generator(device=device) + generator.manual_seed(args.seed) + input_ids = torch.randint( + 0, + args.vocab_size, + (args.batch_size, args.seq_len), + device=device, + generator=generator, + ) + attention_mask = torch.ones_like(input_ids, dtype=torch.bool) + completion_mask = torch.zeros_like(input_ids, dtype=torch.bool) + completion_start = max(1, args.seq_len - args.completion_len) + completion_mask[:, completion_start:] = True + + if args.valid_density < 1.0: + keep = ( + torch.rand( + args.batch_size, + args.seq_len - completion_start, + device=device, + generator=generator, + ) + < args.valid_density + ) + keep[:, 0] = True + completion_mask[:, completion_start:] &= keep + prefix_positions = ( + torch.arange(args.seq_len, device=device).unsqueeze(0).lt(completion_start) + ) + attention_mask &= completion_mask | prefix_positions + + return StatelessForwardInputs( + input_ids=input_ids, + attention_mask=attention_mask, + completion_mask=completion_mask, + ) + + +def _make_model(args: argparse.Namespace, device: torch.device) -> torch.nn.Module: + mode = _scoring_mode(args) + if mode == "reference": + model: torch.nn.Module = TinyReferenceModel(args.vocab_size, args.hidden_dim) + elif mode == "reward": + model = TinyRewardModel(args.vocab_size, args.hidden_dim) + else: + model = TinyBothModel(args.vocab_size, args.hidden_dim) + return model.to(device=device).eval() + + +def _run(args: argparse.Namespace) -> dict[str, Any]: + if args.mode == "paged-kv-compare": + return _run_paged_kv_only(args) + + device = torch.device(args.device) + if device.type == "cuda" and not torch.cuda.is_available(): + return _blocked_row(args, "CUDA is not available") + + torch.manual_seed(args.seed) + inputs = _make_inputs(args, device) + model = _make_model(args, device) + scoring_mode = _scoring_mode(args) + executor = StatelessForwardExecutor( + model, + StatelessForwardConfig(mode=scoring_mode, return_token_scores=scoring_mode != "reward"), + ) + + result = executor.score(inputs) + peak_allocated = result.metrics.get("peak_allocated_mb", "") + peak_reserved = result.metrics.get("peak_reserved_mb", "") + + row = { + "timestamp": _timestamp(), + "candidate": "issue-47-candidate-1", + "mode": args.mode, + "stage": "evaluation", + "batch_size": args.batch_size, + "seq_len": args.seq_len, + "active_tokens": result.metrics["active_completion_tokens"], + "device": str(device), + "dtype": str(inputs.input_ids.dtype).replace("torch.", ""), + "elapsed_ms": f"{float(result.metrics['elapsed_ms']):.4f}", + "peak_allocated_mb": f"{float(peak_allocated):.4f}" if peak_allocated != "" else "", + "peak_reserved_mb": f"{float(peak_reserved):.4f}" if peak_reserved != "" else "", + "paged_kv_baseline_mb": "", + "kv_cache_output_mb": f"{float(result.metrics['kv_cache_output_mb']):.4f}", + "savings_vs_local_reservation_baseline_mb": "", + "status": "pass", + "notes": ( + f"use_cache_passed={result.metrics['use_cache_passed']};" + f"detached_outputs={result.metrics['detached_outputs']};" + f"zero_kv_cache={result.metrics['zero_kv_cache']};" + f"attention_backend={result.metrics['attention_backend']};" + f"hidden_dim={args.hidden_dim};vocab_size={args.vocab_size};{_environment()}" + ), + } + if args.compare_paged_kv: + paged_result = _score_paged_kv(args, inputs, model) + paged_metrics = paged_result.metrics + row["paged_kv_baseline_mb"] = f"{float(paged_metrics['paged_kv_cache_reserved_mb']):.4f}" + savings_mb = float(paged_metrics["total_kv_cache_mb"]) - float( + result.metrics["kv_cache_output_mb"] + ) + row["savings_vs_local_reservation_baseline_mb"] = f"{savings_mb:.4f}" + row["notes"] += ( + f";paged_kv_elapsed_ms={float(paged_metrics['elapsed_ms']):.4f};" + f"paged_kv_blocks={paged_metrics['paged_kv_blocks']};" + f"paged_kv_required_blocks={paged_metrics['paged_kv_required_blocks']};" + f"paged_kv_block_size={paged_metrics['paged_kv_block_size']};" + f"paged_kv_dtype={paged_metrics['kv_cache_dtype']};" + f"paged_kv_use_cache_passed={paged_metrics['use_cache_passed']};" + f"paged_kv_model_cache_mb={float(paged_metrics['model_kv_cache_output_mb']):.4f};" + f"paged_kv_total_cache_mb={float(paged_metrics['total_kv_cache_mb']):.4f}" + ) + correctness = _paged_kv_correctness_notes(result, paged_result) + if correctness: + row["notes"] += f";{correctness}" + return row + + +def _run_paged_kv_only(args: argparse.Namespace) -> dict[str, Any]: + device = torch.device(args.device) + if device.type == "cuda" and not torch.cuda.is_available(): + return _blocked_row(args, "CUDA is not available") + + torch.manual_seed(args.seed) + inputs = _make_inputs(args, device) + model = _make_model(args, device) + result = _score_paged_kv(args, inputs, model) + return _paged_kv_row(args, inputs, result) + + +def _score_paged_kv( + args: argparse.Namespace, + inputs: StatelessForwardInputs, + model: torch.nn.Module, +): + baseline = PagedKVScoringBaseline(model, _paged_kv_config(args)) + return baseline.score(inputs) + + +def _paged_kv_config(args: argparse.Namespace) -> PagedKVScoringConfig: + scoring_mode = _scoring_mode(args) + return PagedKVScoringConfig( + mode=scoring_mode, + num_layers=args.paged_kv_layers, + num_kv_heads=args.paged_kv_heads, + head_dim=args.paged_kv_head_dim, + block_size=args.paged_kv_block_size, + kv_cache_dtype=DTYPE_CHOICES[args.paged_kv_dtype], + kv_cache_blocks=args.paged_kv_blocks, + return_token_scores=scoring_mode != "reward", + ) + + +def _paged_kv_row( + args: argparse.Namespace, + inputs: StatelessForwardInputs, + result, +) -> dict[str, Any]: + metrics = result.metrics + peak_allocated = metrics.get("peak_allocated_mb", "") + peak_reserved = metrics.get("peak_reserved_mb", "") + return { + "timestamp": _timestamp(), + "candidate": "issue-47-candidate-1", + "mode": args.mode, + "stage": "evaluation", + "batch_size": args.batch_size, + "seq_len": args.seq_len, + "active_tokens": metrics["active_completion_tokens"], + "device": str(inputs.input_ids.device), + "dtype": str(inputs.input_ids.dtype).replace("torch.", ""), + "elapsed_ms": f"{float(metrics['elapsed_ms']):.4f}", + "peak_allocated_mb": f"{float(peak_allocated):.4f}" if peak_allocated != "" else "", + "peak_reserved_mb": f"{float(peak_reserved):.4f}" if peak_reserved != "" else "", + "paged_kv_baseline_mb": f"{float(metrics['paged_kv_cache_reserved_mb']):.4f}", + "kv_cache_output_mb": f"{float(metrics['model_kv_cache_output_mb']):.4f}", + "savings_vs_local_reservation_baseline_mb": "", + "status": "pass", + "notes": ( + f"baseline_kind={metrics['baseline_kind']};" + f"scoring_mode={metrics['mode']};" + f"use_cache_passed={metrics['use_cache_passed']};" + f"paged_kv_layers={metrics['paged_kv_layers']};" + f"paged_kv_heads={metrics['paged_kv_heads']};" + f"paged_kv_head_dim={metrics['paged_kv_head_dim']};" + f"paged_kv_block_size={metrics['paged_kv_block_size']};" + f"paged_kv_blocks={metrics['paged_kv_blocks']};" + f"paged_kv_required_blocks={metrics['paged_kv_required_blocks']};" + f"paged_kv_dtype={metrics['kv_cache_dtype']};" + f"model_kv_cache_output_present={metrics['model_kv_cache_output_present']};" + f"total_kv_cache_mb={float(metrics['total_kv_cache_mb']):.4f};" + f"hidden_dim={args.hidden_dim};vocab_size={args.vocab_size};{_environment()}" + ), + } + + +def _paged_kv_correctness_notes(stateless_result, paged_result) -> str: + notes: list[str] = [] + if stateless_result.reference_logps is not None and paged_result.reference_logps is not None: + reference_match = torch.allclose( + stateless_result.reference_logps, + paged_result.reference_logps, + ) + notes.append(f"paged_kv_reference_allclose={bool(reference_match)}") + if stateless_result.rewards is not None and paged_result.rewards is not None: + reward_match = torch.allclose( + stateless_result.rewards, + paged_result.rewards, + ) + notes.append(f"paged_kv_rewards_allclose={bool(reward_match)}") + return ";".join(notes) + + +def _scoring_mode(args: argparse.Namespace) -> str: + if args.mode in SCORING_MODES: + return args.mode + return "reference" + + +def _blocked_row(args: argparse.Namespace, reason: str) -> dict[str, Any]: + return { + "timestamp": _timestamp(), + "candidate": "issue-47-candidate-1", + "mode": args.mode, + "stage": "evaluation", + "batch_size": args.batch_size, + "seq_len": args.seq_len, + "device": args.device, + "dtype": "int64", + "status": "blocked", + "notes": f"{reason};{_environment()}", + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Stateless Reference/Reward executor benchmark") + parser.add_argument( + "--mode", + choices=["reference", "reward", "both", "paged-kv-compare"], + default="reference", + ) + parser.add_argument("--smoke", action="store_true") + parser.add_argument( + "--compare-paged-kv", + action="store_true", + help="Run the local generation-style paged-KV reservation baseline too.", + ) + parser.add_argument("--batch-size", type=int, default=8) + parser.add_argument("--seq-len", type=int, default=1024) + parser.add_argument("--completion-len", type=int, default=256) + parser.add_argument("--vocab-size", type=int, default=4096) + parser.add_argument("--hidden-dim", type=int, default=256) + parser.add_argument("--paged-kv-layers", type=int, default=4) + parser.add_argument("--paged-kv-heads", type=int, default=8) + parser.add_argument("--paged-kv-head-dim", type=int, default=32) + parser.add_argument("--paged-kv-block-size", type=int, default=16) + parser.add_argument("--paged-kv-blocks", type=int, default=None) + parser.add_argument("--paged-kv-dtype", choices=sorted(DTYPE_CHOICES), default="float16") + parser.add_argument("--valid-density", type=float, default=1.0) + parser.add_argument("--seed", type=int, default=47) + parser.add_argument("--device", default="cpu") + parser.add_argument( + "--output", + type=Path, + default=REPO_ROOT / "task-workspace/issues/issue_47/benchmark.csv", + ) + parser.add_argument("--json-output", type=Path, default=None) + args = parser.parse_args() + + if args.smoke: + args.batch_size = min(args.batch_size, 2) + args.seq_len = min(args.seq_len, 16) + args.completion_len = min(args.completion_len, 8) + args.vocab_size = min(args.vocab_size, 128) + args.hidden_dim = min(args.hidden_dim, 32) + args.paged_kv_layers = min(args.paged_kv_layers, 2) + args.paged_kv_heads = min(args.paged_kv_heads, 2) + args.paged_kv_head_dim = min(args.paged_kv_head_dim, 8) + args.paged_kv_block_size = min(args.paged_kv_block_size, 16) + if args.batch_size <= 0: + raise ValueError("batch-size must be greater than zero") + if args.seq_len < 2: + raise ValueError("seq-len must be at least 2") + if args.completion_len <= 0 or args.completion_len >= args.seq_len: + raise ValueError("completion-len must be greater than zero and less than seq-len") + if args.vocab_size <= 1: + raise ValueError("vocab-size must be greater than one") + if args.hidden_dim <= 0: + raise ValueError("hidden-dim must be greater than zero") + if args.paged_kv_layers <= 0: + raise ValueError("paged-kv-layers must be greater than zero") + if args.paged_kv_heads <= 0: + raise ValueError("paged-kv-heads must be greater than zero") + if args.paged_kv_head_dim <= 0: + raise ValueError("paged-kv-head-dim must be greater than zero") + if args.paged_kv_block_size <= 0: + raise ValueError("paged-kv-block-size must be greater than zero") + if args.paged_kv_blocks is not None and args.paged_kv_blocks <= 0: + raise ValueError("paged-kv-blocks must be greater than zero") + if not (0.0 < args.valid_density <= 1.0): + raise ValueError("valid-density must be in (0, 1]") + return args + + +def main() -> None: + args = parse_args() + row = _run(args) + _append_row(args.output, row) + _write_json(args.json_output, row) + print(json.dumps(row, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/rl_engine/executors/deepspeed_trainer.py b/rl_engine/executors/deepspeed_trainer.py index 44a1134ae..4a7acfb01 100644 --- a/rl_engine/executors/deepspeed_trainer.py +++ b/rl_engine/executors/deepspeed_trainer.py @@ -24,6 +24,7 @@ RolloutStageResult, TorchRLTrainingConfig, TrainingStageResult, + objective_reference_logps, ) from rl_engine.testing import ( compute_policy_ratio, @@ -113,7 +114,7 @@ def train(self, rollout: RolloutStageResult) -> TrainingStageResult: output_dtype=torch.float32, ) old_logps = current_logps.detach() - 0.01 - ref_logps = current_logps.detach() - 0.02 + ref_logps = objective_reference_logps(current_logps, batch) ratio = compute_policy_ratio(current_logps, old_logps, batch.completion_mask) unclipped = ratio * batch.advantages.float() clipped = torch.clamp(ratio, 0.8, 1.2) * batch.advantages.float() @@ -133,6 +134,7 @@ def train(self, rollout: RolloutStageResult) -> TrainingStageResult: finished_at = time.perf_counter() published = self._next_published_weight_version(rollout.weight_version) + active_advantages = batch.advantages.float()[batch.completion_mask] return TrainingStageResult( iteration=rollout.iteration, consumed_weight_version=rollout.weight_version, @@ -145,6 +147,16 @@ def train(self, rollout: RolloutStageResult) -> TrainingStageResult: "training_device": str(self.device), "deepspeed_engine": type(self.engine).__name__, "deepspeed_zero_stage": self.config.zero_stage, + "active_advantage_mean_global": ( + float(active_advantages.mean().detach().cpu().item()) + if active_advantages.numel() + else 0.0 + ), + "active_advantage_std_global": ( + float(active_advantages.std(unbiased=False).detach().cpu().item()) + if active_advantages.numel() + else 0.0 + ), **payload_metrics, }, started_at=started_at, diff --git a/rl_engine/executors/paged_kv_baseline.py b/rl_engine/executors/paged_kv_baseline.py new file mode 100644 index 000000000..0c4657200 --- /dev/null +++ b/rl_engine/executors/paged_kv_baseline.py @@ -0,0 +1,413 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import inspect +import time +from dataclasses import dataclass +from typing import Any, Mapping, Optional + +import torch + +from rl_engine.executors.stateless_executor import ( + RewardAdapter, + StatelessForwardInputs, + StatelessForwardMode, + StatelessForwardOutputs, + StatelessForwardResult, + TensorTreeSummary, + default_reward_adapter, + extract_kv_cache_outputs, + score_reference_logprobs, + score_rewards, + summarize_tensor_tree, +) + + +@dataclass(frozen=True) +class PagedKVScoringConfig: + """Configuration for a generation-style paged-KV scoring baseline.""" + + mode: StatelessForwardMode = "reference" + num_layers: int = 4 + num_kv_heads: int = 8 + head_dim: int = 32 + block_size: int = 16 + kv_cache_dtype: torch.dtype = torch.float16 + kv_cache_blocks: Optional[int] = None + use_cache: bool = True + detach_outputs: bool = True + return_token_scores: bool = False + temperature: float = 1.0 + output_dtype: torch.dtype = torch.float32 + + def __post_init__(self) -> None: + if self.mode not in {"reference", "reward", "both"}: + raise ValueError("mode must be 'reference', 'reward', or 'both'") + if self.num_layers <= 0: + raise ValueError("num_layers must be greater than zero") + if self.num_kv_heads <= 0: + raise ValueError("num_kv_heads must be greater than zero") + if self.head_dim <= 0: + raise ValueError("head_dim must be greater than zero") + if self.block_size <= 0: + raise ValueError("block_size must be greater than zero") + if self.kv_cache_blocks is not None and self.kv_cache_blocks <= 0: + raise ValueError("kv_cache_blocks must be greater than zero") + if not self.use_cache: + raise ValueError("PagedKVScoringConfig.use_cache must be True") + if self.temperature <= 0.0: + raise ValueError("temperature must be greater than zero") + + +@dataclass(frozen=True) +class PagedKVCacheReservation: + """Allocated paged-KV tensors plus the block table needed to address them.""" + + key_cache: torch.Tensor + value_cache: torch.Tensor + block_tables: torch.Tensor + sequence_lengths: torch.Tensor + blocks_per_sequence: torch.Tensor + block_size: int + required_blocks: int + reserved_blocks: int + cache_bytes: int + metadata_bytes: int + reserved_bytes: int + + @property + def reserved_mb(self) -> float: + return self.reserved_bytes / 1_048_576.0 + + +class PagedKVScoringBaseline: + """ + Correctness-first baseline for generation-engine-style scoring. + + This wrapper reserves a paged KV-cache and block table before running the + wrapped model with ``use_cache=True`` when the model accepts that keyword. + It does not implement a vLLM engine; it provides a local, reproducible + memory reservation baseline for comparing stateless scoring against a + generation path that keeps paged KV state alive while scoring a full + sequence. + """ + + def __init__( + self, + model: torch.nn.Module, + config: Optional[PagedKVScoringConfig] = None, + *, + reward_adapter: Optional[RewardAdapter] = None, + ): + self.model = model + self.config = config or PagedKVScoringConfig() + self.reward_adapter = reward_adapter or default_reward_adapter + + def score(self, inputs: StatelessForwardInputs) -> StatelessForwardResult: + _validate_inputs(inputs, self.config) + + device = inputs.input_ids.device + cuda_tracking = device.type == "cuda" and torch.cuda.is_available() + if cuda_tracking: + torch.cuda.reset_peak_memory_stats(device) + torch.cuda.synchronize(device) + + started_at = time.perf_counter() + reservation = reserve_paged_kv_cache(inputs, self.config) + with torch.no_grad(): + raw_outputs, use_cache_passed = _run_cache_forward( + self.model, + inputs, + use_cache=self.config.use_cache, + ) + kv_cache = extract_kv_cache_outputs(raw_outputs) + kv_cache_summary = summarize_tensor_tree(kv_cache) + outputs = StatelessForwardOutputs( + raw=raw_outputs, + logits=_extract_logits(raw_outputs), + kv_cache=kv_cache, + ) + + reference_logps: Optional[torch.Tensor] = None + rewards: Optional[torch.Tensor] = None + token_scores: Optional[torch.Tensor] = None + + if self.config.mode in {"reference", "both"}: + if outputs.logits is None: + raise ValueError("reference mode requires model outputs to expose logits") + reference_logps = score_reference_logprobs( + outputs.logits, + inputs, + temperature=self.config.temperature, + output_dtype=self.config.output_dtype, + ) + if self.config.return_token_scores: + token_scores = reference_logps + + if self.config.mode in {"reward", "both"}: + rewards = score_rewards( + outputs, + inputs, + reward_adapter=self.reward_adapter, + output_dtype=self.config.output_dtype, + ) + + if self.config.detach_outputs: + reference_logps = _detach_optional(reference_logps) + rewards = _detach_optional(rewards) + token_scores = _detach_optional(token_scores) + + if cuda_tracking: + torch.cuda.synchronize(device) + finished_at = time.perf_counter() + + metrics = collect_paged_kv_metrics( + inputs, + reservation, + config=self.config, + elapsed_seconds=finished_at - started_at, + use_cache_passed=use_cache_passed, + cuda_tracking=cuda_tracking, + model_kv_cache_summary=kv_cache_summary, + ) + return StatelessForwardResult( + reference_logps=reference_logps, + rewards=rewards, + token_scores=token_scores, + metrics=metrics, + ) + + +def reserve_paged_kv_cache( + inputs: StatelessForwardInputs, + config: PagedKVScoringConfig, +) -> PagedKVCacheReservation: + """Allocate paged KV tensors and a dense block table for the scoring batch.""" + + _validate_inputs(inputs, config) + + device = inputs.input_ids.device + sequence_lengths = _bool_mask(inputs.attention_mask, device=device).sum(dim=1).to(torch.int32) + if not bool((sequence_lengths > 0).all().item()): + raise ValueError("attention_mask must contain at least one token per sequence") + + blocks_per_sequence = torch.div( + sequence_lengths + config.block_size - 1, + config.block_size, + rounding_mode="floor", + ).to(torch.int32) + required_blocks = int(blocks_per_sequence.sum().item()) + reserved_blocks = int(config.kv_cache_blocks or required_blocks) + if reserved_blocks < required_blocks: + raise ValueError( + f"kv_cache_blocks={reserved_blocks} cannot fit required blocks={required_blocks}" + ) + + key_cache = torch.empty( + ( + config.num_layers, + reserved_blocks, + config.block_size, + config.num_kv_heads, + config.head_dim, + ), + device=device, + dtype=config.kv_cache_dtype, + ) + value_cache = torch.empty_like(key_cache) + + max_blocks_per_sequence = int(blocks_per_sequence.max().item()) + block_tables = torch.full( + (inputs.input_ids.shape[0], max_blocks_per_sequence), + -1, + device=device, + dtype=torch.int32, + ) + cursor = 0 + for row, block_count_tensor in enumerate(blocks_per_sequence): + block_count = int(block_count_tensor.item()) + block_tables[row, :block_count] = torch.arange( + cursor, + cursor + block_count, + device=device, + dtype=torch.int32, + ) + cursor += block_count + + cache_bytes = _tensor_nbytes(key_cache) + _tensor_nbytes(value_cache) + metadata_bytes = ( + _tensor_nbytes(block_tables) + + _tensor_nbytes(sequence_lengths) + + _tensor_nbytes(blocks_per_sequence) + ) + return PagedKVCacheReservation( + key_cache=key_cache, + value_cache=value_cache, + block_tables=block_tables, + sequence_lengths=sequence_lengths, + blocks_per_sequence=blocks_per_sequence, + block_size=config.block_size, + required_blocks=required_blocks, + reserved_blocks=reserved_blocks, + cache_bytes=cache_bytes, + metadata_bytes=metadata_bytes, + reserved_bytes=cache_bytes + metadata_bytes, + ) + + +def collect_paged_kv_metrics( + inputs: StatelessForwardInputs, + reservation: PagedKVCacheReservation, + *, + config: PagedKVScoringConfig, + elapsed_seconds: float, + use_cache_passed: bool, + cuda_tracking: bool, + model_kv_cache_summary: TensorTreeSummary, +) -> dict[str, float | int | str | bool]: + input_ids = inputs.input_ids + active_tokens = int(_bool_mask(inputs.completion_mask, device=input_ids.device).sum().item()) + total_kv_cache_bytes = reservation.reserved_bytes + model_kv_cache_summary.total_bytes + metrics: dict[str, float | int | str | bool] = { + "baseline_kind": "generation_engine_paged_kv_reservation", + "baseline_includes_model_kv_cache": model_kv_cache_summary.tensor_count > 0, + "mode": config.mode, + "batch_size": int(input_ids.shape[0]), + "sequence_len": int(input_ids.shape[1]), + "active_completion_tokens": active_tokens, + "attention_tokens": int(reservation.sequence_lengths.sum().item()), + "device": str(input_ids.device), + "dtype": str(input_ids.dtype).replace("torch.", ""), + "kv_cache_dtype": str(config.kv_cache_dtype).replace("torch.", ""), + "elapsed_ms": elapsed_seconds * 1000.0, + "use_cache": True, + "use_cache_passed": bool(use_cache_passed), + "detached_outputs": bool(config.detach_outputs), + "paged_kv_layers": int(config.num_layers), + "paged_kv_heads": int(config.num_kv_heads), + "paged_kv_head_dim": int(config.head_dim), + "paged_kv_block_size": int(config.block_size), + "paged_kv_required_blocks": int(reservation.required_blocks), + "paged_kv_blocks": int(reservation.reserved_blocks), + "paged_kv_max_blocks_per_sequence": int(reservation.block_tables.shape[1]), + "paged_kv_cache_reserved_mb": reservation.reserved_bytes / 1_048_576.0, + "paged_kv_cache_payload_mb": reservation.cache_bytes / 1_048_576.0, + "paged_kv_metadata_mb": reservation.metadata_bytes / 1_048_576.0, + "model_kv_cache_output_present": model_kv_cache_summary.tensor_count > 0, + "model_kv_cache_output_tensors": model_kv_cache_summary.tensor_count, + "model_kv_cache_output_bytes": model_kv_cache_summary.total_bytes, + "model_kv_cache_output_mb": model_kv_cache_summary.total_mb, + "total_kv_cache_bytes": total_kv_cache_bytes, + "total_kv_cache_mb": total_kv_cache_bytes / 1_048_576.0, + } + if cuda_tracking: + device = input_ids.device + metrics["peak_allocated_mb"] = torch.cuda.max_memory_allocated(device) / 1_048_576.0 + metrics["peak_reserved_mb"] = torch.cuda.max_memory_reserved(device) / 1_048_576.0 + return metrics + + +def _validate_inputs(inputs: StatelessForwardInputs, config: PagedKVScoringConfig) -> None: + input_ids = inputs.input_ids + if input_ids.ndim != 2: + raise ValueError(f"input_ids must have shape [B, S], got {tuple(input_ids.shape)}") + if inputs.attention_mask.shape != input_ids.shape: + raise ValueError("attention_mask shape must match input_ids shape") + if inputs.completion_mask.shape != input_ids.shape: + raise ValueError("completion_mask shape must match input_ids shape") + if inputs.labels is not None and inputs.labels.shape != input_ids.shape: + raise ValueError("labels shape must match input_ids shape") + if inputs.attention_mask.device != input_ids.device: + raise ValueError("attention_mask device must match input_ids device") + if inputs.completion_mask.device != input_ids.device: + raise ValueError("completion_mask device must match input_ids device") + if inputs.labels is not None and inputs.labels.device != input_ids.device: + raise ValueError("labels device must match input_ids device") + if config.mode in {"reference", "both"}: + if input_ids.shape[1] < 2: + raise ValueError("reference scoring requires sequence_len >= 2") + if _completion_mask_starts_at_position_zero( + inputs.completion_mask, + device=input_ids.device, + ): + raise ValueError( + "completion_mask[:, 0] must be False; completions cannot start at position 0" + ) + if not bool(_bool_mask(inputs.completion_mask, device=input_ids.device).any().item()): + raise ValueError("completion_mask must contain at least one active token") + + +def _run_cache_forward( + model: torch.nn.Module, + inputs: StatelessForwardInputs, + *, + use_cache: bool, +) -> tuple[Any, bool]: + kwargs: dict[str, Any] = { + "input_ids": inputs.input_ids, + "attention_mask": inputs.attention_mask, + } + if _call_accepts_keyword(model, "use_cache"): + kwargs["use_cache"] = use_cache + return model(**kwargs), True + + try: + kwargs["use_cache"] = use_cache + return model(**kwargs), True + except TypeError as exc: + if "use_cache" not in str(exc): + raise + kwargs.pop("use_cache", None) + return model(**kwargs), False + + +def _call_accepts_keyword(model: torch.nn.Module, keyword: str) -> bool: + try: + signature = inspect.signature(model.forward) + except (TypeError, ValueError): + return False + for parameter in signature.parameters.values(): + if parameter.kind == inspect.Parameter.VAR_KEYWORD: + return True + if parameter.name == keyword: + return True + return False + + +def _extract_logits(raw_outputs: Any) -> Optional[torch.Tensor]: + if isinstance(raw_outputs, torch.Tensor): + return raw_outputs + if isinstance(raw_outputs, Mapping): + value = raw_outputs.get("logits") + return value if isinstance(value, torch.Tensor) else None + logits = getattr(raw_outputs, "logits", None) + if isinstance(logits, torch.Tensor): + return logits + if isinstance(raw_outputs, (tuple, list)) and raw_outputs: + first = raw_outputs[0] + return first if isinstance(first, torch.Tensor) else None + return None + + +def _bool_mask(mask: torch.Tensor, *, device: torch.device) -> torch.Tensor: + return mask.to(device=device, dtype=torch.bool) + + +def _completion_mask_starts_at_position_zero( + completion_mask: torch.Tensor, + *, + device: torch.device, +) -> bool: + if completion_mask.shape[1] == 0: + return False + return bool(_bool_mask(completion_mask, device=device)[:, 0].any().item()) + + +def _detach_optional(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + return tensor.detach() if tensor is not None else None + + +def _tensor_nbytes(tensor: torch.Tensor) -> int: + return int(tensor.numel() * tensor.element_size()) diff --git a/rl_engine/executors/stateless_executor.py b/rl_engine/executors/stateless_executor.py new file mode 100644 index 000000000..2047218f5 --- /dev/null +++ b/rl_engine/executors/stateless_executor.py @@ -0,0 +1,636 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import inspect +import time +from contextlib import contextmanager +from dataclasses import dataclass, replace +from typing import Any, Callable, Iterator, Literal, Mapping, Optional + +import torch + +from rl_engine.testing.reference_ops import selected_logprobs_reference + +StatelessForwardMode = Literal["reference", "reward", "both"] +StatelessAttentionBackend = Literal["flash_attention_2", "sdpa", "eager", "model_default"] +RewardAdapter = Callable[["StatelessForwardOutputs", "StatelessForwardInputs"], torch.Tensor] +_MISSING = object() + + +@dataclass(frozen=True) +class StatelessForwardConfig: + """Configuration for no-cache reference/reward model scoring.""" + + mode: StatelessForwardMode = "both" + use_cache: bool = False + attention_backend: StatelessAttentionBackend = "flash_attention_2" + reject_kv_cache_outputs: bool = True + detach_outputs: bool = True + return_token_scores: bool = False + max_batch_size: Optional[int] = None + temperature: float = 1.0 + output_dtype: torch.dtype = torch.float32 + + def __post_init__(self) -> None: + if self.mode not in {"reference", "reward", "both"}: + raise ValueError("mode must be 'reference', 'reward', or 'both'") + if self.use_cache: + raise ValueError("StatelessForwardConfig.use_cache must be False") + if self.attention_backend not in {"flash_attention_2", "sdpa", "eager", "model_default"}: + raise ValueError( + "attention_backend must be 'flash_attention_2', 'sdpa', 'eager', " + "or 'model_default'" + ) + if self.max_batch_size is not None and self.max_batch_size <= 0: + raise ValueError("max_batch_size must be greater than zero") + if self.temperature <= 0.0: + raise ValueError("temperature must be greater than zero") + + +@dataclass(frozen=True) +class StatelessForwardInputs: + """Dense full-sequence batch for stateless scoring.""" + + input_ids: torch.Tensor + attention_mask: torch.Tensor + completion_mask: torch.Tensor + labels: Optional[torch.Tensor] = None + + +@dataclass(frozen=True) +class StatelessForwardOutputs: + """Normalized model outputs consumed by scoring adapters.""" + + raw: Any + logits: Optional[torch.Tensor] + kv_cache: Optional[Any] = None + + +@dataclass(frozen=True) +class StatelessForwardResult: + """Scoring tensors and scalar metrics produced by the stateless executor.""" + + reference_logps: Optional[torch.Tensor] + rewards: Optional[torch.Tensor] + token_scores: Optional[torch.Tensor] + metrics: Mapping[str, float | int | str | bool] + + +@dataclass(frozen=True) +class TensorTreeSummary: + """Count tensor payloads in nested optional runtime outputs.""" + + tensor_count: int + total_bytes: int + + @property + def total_mb(self) -> float: + return self.total_bytes / 1_048_576.0 + + +class StatelessForwardExecutor: + """ + Lightweight Reference/Reward scoring wrapper. + + The executor runs one full-sequence forward pass with ``use_cache=False`` + whenever the wrapped model accepts that argument. It never calls a + generation loop and it does not instantiate a paged-KV serving runtime. + """ + + def __init__( + self, + model: torch.nn.Module, + config: Optional[StatelessForwardConfig] = None, + *, + reward_adapter: Optional[RewardAdapter] = None, + ): + self.model = model + self.config = config or StatelessForwardConfig() + self.reward_adapter = reward_adapter or default_reward_adapter + + def score(self, inputs: StatelessForwardInputs) -> StatelessForwardResult: + _validate_inputs(inputs, self.config) + + device = inputs.input_ids.device + cuda_tracking = device.type == "cuda" and torch.cuda.is_available() + if cuda_tracking: + torch.cuda.reset_peak_memory_stats(device) + torch.cuda.synchronize(device) + + started_at = time.perf_counter() + with _temporarily_configure_stateless_model(self.model, self.config) as no_cache_policy: + with torch.no_grad(): + try: + raw_outputs, use_cache_passed = _run_no_cache_forward(self.model, inputs) + no_cache_policy["attention_backend_fallback"] = False + except Exception as exc: + if not _should_fallback_attention_backend(exc, self.config): + raise + fallback_config = replace(self.config, attention_backend="eager") + with _temporarily_configure_stateless_model( + self.model, + fallback_config, + ) as fallback_policy: + raw_outputs, use_cache_passed = _run_no_cache_forward(self.model, inputs) + no_cache_policy.update(fallback_policy) + no_cache_policy.update( + { + "attention_backend_requested": self.config.attention_backend, + "attention_backend_fallback": True, + "attention_backend_fallback_reason": ( + f"{type(exc).__name__}: {str(exc)[:200]}" + ), + } + ) + kv_cache = extract_kv_cache_outputs(raw_outputs) + kv_cache_summary = summarize_tensor_tree(kv_cache) + if self.config.reject_kv_cache_outputs and kv_cache is not None: + raise ValueError( + "stateless scoring received KV-cache outputs despite use_cache=False " + f"({kv_cache_summary.tensor_count} tensors, " + f"{kv_cache_summary.total_mb:.4f} MiB)." + ) + outputs = StatelessForwardOutputs( + raw=raw_outputs, + logits=_extract_logits(raw_outputs), + kv_cache=kv_cache, + ) + + reference_logps: Optional[torch.Tensor] = None + rewards: Optional[torch.Tensor] = None + token_scores: Optional[torch.Tensor] = None + + if self.config.mode in {"reference", "both"}: + if outputs.logits is None: + raise ValueError("reference mode requires model outputs to expose logits") + reference_logps = score_reference_logprobs( + outputs.logits, + inputs, + temperature=self.config.temperature, + output_dtype=self.config.output_dtype, + ) + if self.config.return_token_scores: + token_scores = reference_logps + + if self.config.mode in {"reward", "both"}: + rewards = score_rewards( + outputs, + inputs, + reward_adapter=self.reward_adapter, + output_dtype=self.config.output_dtype, + ) + + if self.config.detach_outputs: + reference_logps = _detach_optional(reference_logps) + rewards = _detach_optional(rewards) + token_scores = _detach_optional(token_scores) + + if cuda_tracking: + torch.cuda.synchronize(device) + finished_at = time.perf_counter() + + metrics = collect_stateless_metrics( + inputs, + mode=self.config.mode, + elapsed_seconds=finished_at - started_at, + use_cache_passed=use_cache_passed, + detached_outputs=self.config.detach_outputs, + cuda_tracking=cuda_tracking, + no_cache_policy=no_cache_policy, + kv_cache_summary=kv_cache_summary, + ) + return StatelessForwardResult( + reference_logps=reference_logps, + rewards=rewards, + token_scores=token_scores, + metrics=metrics, + ) + + +def score_reference_logprobs( + logits: torch.Tensor, + inputs: StatelessForwardInputs, + *, + temperature: float = 1.0, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Compute causal next-token selected logprobs aligned to ``[B, S]`` masks.""" + + if logits.ndim != 3: + raise ValueError(f"reference logits must have shape [B, S, V], got {tuple(logits.shape)}") + if logits.shape[:2] != inputs.input_ids.shape: + raise ValueError( + "reference logits leading shape must match input_ids shape, got " + f"{tuple(logits.shape[:2])} and {tuple(inputs.input_ids.shape)}" + ) + + labels = inputs.labels if inputs.labels is not None else inputs.input_ids + if labels.shape != inputs.input_ids.shape: + raise ValueError("labels shape must match input_ids shape") + if _completion_mask_starts_at_position_zero(inputs.completion_mask, device=logits.device): + raise ValueError( + "completion_mask[:, 0] must be False; completions cannot start at position 0" + ) + + shifted_logits = logits[:, :-1, :] + shifted_labels = labels[:, 1:] + shifted_mask = _bool_mask(inputs.completion_mask[:, 1:], device=logits.device) + shifted_logps = selected_logprobs_reference( + shifted_logits, + shifted_labels.to(device=logits.device), + mask=shifted_mask, + temperature=temperature, + output_dtype=output_dtype, + ) + result = torch.zeros( + inputs.input_ids.shape, + device=logits.device, + dtype=output_dtype, + ) + result[:, 1:] = shifted_logps + return result.masked_fill(~_bool_mask(inputs.completion_mask, device=result.device), 0.0) + + +def score_rewards( + outputs: StatelessForwardOutputs, + inputs: StatelessForwardInputs, + *, + reward_adapter: RewardAdapter, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Run a reward adapter and validate that it returns one scalar per sequence.""" + + rewards = reward_adapter(outputs, inputs) + if not isinstance(rewards, torch.Tensor): + raise TypeError("reward_adapter must return a torch.Tensor") + if rewards.ndim != 1: + raise ValueError(f"reward_adapter must return shape [B], got {tuple(rewards.shape)}") + if rewards.shape[0] != inputs.input_ids.shape[0]: + raise ValueError( + f"reward_adapter batch size {rewards.shape[0]} must match input batch size " + f"{inputs.input_ids.shape[0]}" + ) + return rewards.to(device=inputs.input_ids.device, dtype=output_dtype) + + +def default_reward_adapter( + outputs: StatelessForwardOutputs, + inputs: StatelessForwardInputs, +) -> torch.Tensor: + """Default adapter for common scalar reward model outputs.""" + + del inputs + reward = _extract_named_tensor(outputs.raw, ("rewards", "reward", "scores", "score")) + if reward is not None: + return _squeeze_reward_tensor(reward) + + logits = outputs.logits + if logits is None: + raise ValueError("reward mode requires scalar logits or a reward_adapter") + return _squeeze_reward_tensor(logits) + + +def collect_stateless_metrics( + inputs: StatelessForwardInputs, + *, + mode: StatelessForwardMode, + elapsed_seconds: float, + use_cache_passed: bool, + detached_outputs: bool, + cuda_tracking: bool, + no_cache_policy: Mapping[str, float | int | str | bool], + kv_cache_summary: TensorTreeSummary, +) -> dict[str, float | int | str | bool]: + """Collect common scoring metrics without importing optional runtimes.""" + + input_ids = inputs.input_ids + active_tokens = int(_bool_mask(inputs.completion_mask, device=input_ids.device).sum().item()) + metrics: dict[str, float | int | str | bool] = { + "mode": mode, + "batch_size": int(input_ids.shape[0]), + "sequence_len": int(input_ids.shape[1]), + "active_completion_tokens": active_tokens, + "device": str(input_ids.device), + "dtype": str(input_ids.dtype).replace("torch.", ""), + "elapsed_ms": elapsed_seconds * 1000.0, + "use_cache": False, + "use_cache_passed": bool(use_cache_passed), + "detached_outputs": bool(detached_outputs), + "zero_kv_cache": kv_cache_summary.tensor_count == 0, + "kv_cache_output_present": kv_cache_summary.tensor_count > 0, + "kv_cache_output_tensors": kv_cache_summary.tensor_count, + "kv_cache_output_bytes": kv_cache_summary.total_bytes, + "kv_cache_output_mb": kv_cache_summary.total_mb, + **dict(no_cache_policy), + } + if cuda_tracking: + device = input_ids.device + metrics["peak_allocated_mb"] = torch.cuda.max_memory_allocated(device) / 1_048_576.0 + metrics["peak_reserved_mb"] = torch.cuda.max_memory_reserved(device) / 1_048_576.0 + return metrics + + +def configure_stateless_model( + model: torch.nn.Module, + config: StatelessForwardConfig, +) -> dict[str, float | int | str | bool]: + """ + Apply best-effort no-cache scoring knobs without importing optional runtimes. + + Hugging Face-style models decide attention kernels from their config rather + than forward kwargs. Setting these attributes keeps the wrapper explicit + while still allowing plain PyTorch modules to run unchanged. + """ + + use_cache_targets = 0 + attention_targets = 0 + + for target in _model_config_targets(model): + if hasattr(target, "use_cache"): + target.use_cache = False + use_cache_targets += 1 + if config.attention_backend != "model_default": + target.attn_implementation = config.attention_backend + target._attn_implementation = config.attention_backend + attention_targets += 1 + + return { + "attention_backend": config.attention_backend, + "attention_backend_configured": attention_targets > 0, + "attention_backend_config_targets": attention_targets, + "model_config_use_cache_disabled": use_cache_targets > 0, + "model_config_use_cache_targets": use_cache_targets, + "reject_kv_cache_outputs": config.reject_kv_cache_outputs, + } + + +@contextmanager +def _temporarily_configure_stateless_model( + model: torch.nn.Module, + config: StatelessForwardConfig, +) -> Iterator[dict[str, float | int | str | bool]]: + saved = _model_config_snapshot(model, config) + policy = configure_stateless_model(model, config) + try: + yield policy + finally: + _restore_model_config_snapshot(saved) + + +def extract_kv_cache_outputs(raw_outputs: Any) -> Optional[Any]: + """Extract common cache-bearing output fields from model outputs.""" + + names = ( + "past_key_values", + "past_key_value", + "presents", + "present", + "kv_cache", + "cache", + ) + if isinstance(raw_outputs, Mapping): + for name in names: + if name in raw_outputs and raw_outputs[name] is not None: + return raw_outputs[name] + return None + + for name in names: + value = getattr(raw_outputs, name, None) + if value is not None: + return value + + if isinstance(raw_outputs, (tuple, list)) and len(raw_outputs) > 1: + candidate = raw_outputs[1] + if _looks_like_kv_cache(candidate): + return candidate + return None + + +def summarize_tensor_tree(value: Any) -> TensorTreeSummary: + """Count tensors and bytes in nested tuples/lists/mappings/dataclass-like outputs.""" + + seen: set[int] = set() + + def visit(node: Any) -> tuple[int, int]: + if node is None: + return 0, 0 + node_id = id(node) + if node_id in seen: + return 0, 0 + seen.add(node_id) + + if isinstance(node, torch.Tensor): + return 1, _tensor_nbytes(node) + if isinstance(node, Mapping): + counts = [visit(item) for item in node.values()] + elif isinstance(node, (tuple, list)): + counts = [visit(item) for item in node] + elif hasattr(node, "__dict__"): + counts = [visit(item) for item in vars(node).values()] + else: + return 0, 0 + return sum(count for count, _bytes in counts), sum(_bytes for _count, _bytes in counts) + + tensor_count, total_bytes = visit(value) + return TensorTreeSummary(tensor_count=tensor_count, total_bytes=total_bytes) + + +def _validate_inputs(inputs: StatelessForwardInputs, config: StatelessForwardConfig) -> None: + input_ids = inputs.input_ids + attention_mask = inputs.attention_mask + completion_mask = inputs.completion_mask + + if input_ids.ndim != 2: + raise ValueError(f"input_ids must have shape [B, S], got {tuple(input_ids.shape)}") + if attention_mask.shape != input_ids.shape: + raise ValueError("attention_mask shape must match input_ids shape") + if completion_mask.shape != input_ids.shape: + raise ValueError("completion_mask shape must match input_ids shape") + if inputs.labels is not None and inputs.labels.shape != input_ids.shape: + raise ValueError("labels shape must match input_ids shape") + if attention_mask.device != input_ids.device: + raise ValueError("attention_mask device must match input_ids device") + if completion_mask.device != input_ids.device: + raise ValueError("completion_mask device must match input_ids device") + if inputs.labels is not None and inputs.labels.device != input_ids.device: + raise ValueError("labels device must match input_ids device") + if config.max_batch_size is not None and input_ids.shape[0] > config.max_batch_size: + raise ValueError( + f"batch size {input_ids.shape[0]} exceeds max_batch_size {config.max_batch_size}" + ) + if config.mode in {"reference", "both"}: + if input_ids.shape[1] < 2: + raise ValueError("reference scoring requires sequence_len >= 2") + if _completion_mask_starts_at_position_zero(completion_mask, device=input_ids.device): + raise ValueError( + "completion_mask[:, 0] must be False; completions cannot start at position 0" + ) + if not bool(_bool_mask(completion_mask, device=input_ids.device).any().item()): + raise ValueError("completion_mask must contain at least one active token") + + +def _run_no_cache_forward( + model: torch.nn.Module, + inputs: StatelessForwardInputs, +) -> tuple[Any, bool]: + kwargs: dict[str, Any] = { + "input_ids": inputs.input_ids, + "attention_mask": inputs.attention_mask, + } + if _call_accepts_keyword(model, "use_cache"): + kwargs["use_cache"] = False + return model(**kwargs), True + + try: + kwargs["use_cache"] = False + return model(**kwargs), True + except TypeError as exc: + if "use_cache" not in str(exc): + raise + kwargs.pop("use_cache", None) + return model(**kwargs), False + + +def _should_fallback_attention_backend( + exc: Exception, + config: StatelessForwardConfig, +) -> bool: + if config.attention_backend != "flash_attention_2": + return False + if isinstance(exc, (ImportError, ModuleNotFoundError, StopIteration)): + return True + message = str(exc).lower() + markers = ( + "flash_attention", + "flash attention", + "flash-attn", + "flash_attn", + "flash attention kernels", + "attn_implementation", + ) + return any(marker in message for marker in markers) + + +def _call_accepts_keyword(model: torch.nn.Module, keyword: str) -> bool: + try: + signature = inspect.signature(model.forward) + except (TypeError, ValueError): + return False + for parameter in signature.parameters.values(): + if parameter.kind == inspect.Parameter.VAR_KEYWORD: + return True + if parameter.name == keyword: + return True + return False + + +def _extract_logits(raw_outputs: Any) -> Optional[torch.Tensor]: + if isinstance(raw_outputs, torch.Tensor): + return raw_outputs + if isinstance(raw_outputs, Mapping): + value = raw_outputs.get("logits") + return value if isinstance(value, torch.Tensor) else None + logits = getattr(raw_outputs, "logits", None) + if isinstance(logits, torch.Tensor): + return logits + if isinstance(raw_outputs, (tuple, list)) and raw_outputs: + first = raw_outputs[0] + return first if isinstance(first, torch.Tensor) else None + return None + + +def _extract_named_tensor(raw_outputs: Any, names: tuple[str, ...]) -> Optional[torch.Tensor]: + if isinstance(raw_outputs, Mapping): + for name in names: + value = raw_outputs.get(name) + if isinstance(value, torch.Tensor): + return value + for name in names: + value = getattr(raw_outputs, name, None) + if isinstance(value, torch.Tensor): + return value + return None + + +def _squeeze_reward_tensor(value: torch.Tensor) -> torch.Tensor: + if value.ndim == 1: + return value + if value.ndim == 2 and value.shape[1] == 1: + return value[:, 0] + raise ValueError(f"reward tensor must have shape [B] or [B, 1], got {tuple(value.shape)}") + + +def _bool_mask(mask: torch.Tensor, *, device: torch.device) -> torch.Tensor: + return mask.to(device=device, dtype=torch.bool) + + +def _completion_mask_starts_at_position_zero( + completion_mask: torch.Tensor, + *, + device: torch.device, +) -> bool: + if completion_mask.shape[1] == 0: + return False + return bool(_bool_mask(completion_mask, device=device)[:, 0].any().item()) + + +def _detach_optional(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + return tensor.detach() if tensor is not None else None + + +def _model_config_snapshot( + model: torch.nn.Module, + config: StatelessForwardConfig, +) -> list[tuple[Any, str, Any]]: + saved: list[tuple[Any, str, Any]] = [] + for target in _model_config_targets(model): + if hasattr(target, "use_cache"): + saved.append((target, "use_cache", target.use_cache)) + if config.attention_backend != "model_default": + saved.append( + ( + target, + "attn_implementation", + getattr(target, "attn_implementation", _MISSING), + ) + ) + saved.append( + ( + target, + "_attn_implementation", + getattr(target, "_attn_implementation", _MISSING), + ) + ) + return saved + + +def _restore_model_config_snapshot(saved: list[tuple[Any, str, Any]]) -> None: + for target, attribute, previous in reversed(saved): + if previous is _MISSING: + if hasattr(target, attribute): + delattr(target, attribute) + else: + setattr(target, attribute, previous) + + +def _model_config_targets(model: torch.nn.Module) -> list[Any]: + targets: list[Any] = [] + seen: set[int] = set() + for name in ("config", "generation_config"): + target = getattr(model, name, None) + if target is None or id(target) in seen: + continue + targets.append(target) + seen.add(id(target)) + return targets + + +def _looks_like_kv_cache(value: Any) -> bool: + summary = summarize_tensor_tree(value) + return summary.tensor_count > 0 + + +def _tensor_nbytes(tensor: torch.Tensor) -> int: + return int(tensor.numel() * tensor.element_size()) diff --git a/rl_engine/executors/training_contract.py b/rl_engine/executors/training_contract.py index 14019c9bc..bf732494e 100644 --- a/rl_engine/executors/training_contract.py +++ b/rl_engine/executors/training_contract.py @@ -5,10 +5,15 @@ import time from dataclasses import dataclass, field -from typing import Any, Mapping, Optional, Protocol, Sequence +from typing import Any, Callable, Mapping, Optional, Protocol, Sequence import torch +from rl_engine.executors.stateless_executor import ( + StatelessForwardExecutor, + StatelessForwardInputs, + StatelessForwardResult, +) from rl_engine.testing import SyntheticRLKernelBatch, make_synthetic_rl_kernel_batch @@ -64,6 +69,7 @@ class TorchRLTrainingConfig: dtype: torch.dtype = torch.float32 seed: int = 0 min_completion_len: int = 1 + advantage_eps: float = 1e-8 class RolloutBatchMixin: @@ -74,10 +80,34 @@ def _batch_from_rollout_or_synthetic( self, rollout: RolloutStageResult, ) -> tuple[SyntheticRLKernelBatch, dict[str, Any]]: - token_groups = extract_rollout_token_groups(rollout.payload) - if token_groups: - return self._batch_from_token_groups(token_groups, rollout), { + candidate_groups = extract_rollout_candidate_groups(rollout.payload) + if candidate_groups: + reward_groups = extract_rollout_reward_groups(rollout.payload) + reference_logp_groups = extract_rollout_reference_logp_groups(rollout.payload) + has_payload_rewards = _reward_groups_match(candidate_groups, reward_groups) + has_payload_reference_logps = _reference_logp_groups_match( + candidate_groups, + reference_logp_groups, + max_completion_len=self.config.completion_len, + ) + batch = self._batch_from_candidate_groups( + candidate_groups, + rollout, + reward_groups=reward_groups if has_payload_rewards else None, + reference_logp_groups=( + reference_logp_groups if has_payload_reference_logps else None + ), + ) + token_groups = [tokens for group in candidate_groups for tokens in group] + return batch, { "training_data_source": "rollout_payload", + "reward_source": "payload_rewards" if has_payload_rewards else "token_id_proxy", + "reference_logp_source": ( + "payload_reference_logps" + if has_payload_reference_logps + else "synthetic_current_offset" + ), + "rollout_prompt_groups": len(candidate_groups), "rollout_sequences": len(token_groups), "rollout_tokens": sum(len(group) for group in token_groups), } @@ -98,6 +128,8 @@ def _batch_from_rollout_or_synthetic( "training_data_source": "synthetic_fallback", "rollout_sequences": 0, "rollout_tokens": 0, + "reward_source": "synthetic_rewards", + "reference_logp_source": "synthetic_current_offset", } def _batch_from_token_groups( @@ -105,11 +137,27 @@ def _batch_from_token_groups( token_groups: Sequence[Sequence[int]], rollout: RolloutStageResult, ) -> SyntheticRLKernelBatch: + return self._batch_from_candidate_groups( + [[group] for group in token_groups], + rollout, + reward_groups=None, + reference_logp_groups=None, + ) + + def _batch_from_candidate_groups( + self, + candidate_groups: Sequence[Sequence[Sequence[int]]], + rollout: RolloutStageResult, + *, + reward_groups: Optional[Sequence[Sequence[float]]] = None, + reference_logp_groups: Optional[Sequence[Sequence[Sequence[float]]]] = None, + ) -> SyntheticRLKernelBatch: + flat_token_groups = [tokens for group in candidate_groups for tokens in group] completion_len = max( self.config.min_completion_len, - min(self.config.completion_len, max(len(group) for group in token_groups)), + min(self.config.completion_len, max(len(group) for group in flat_token_groups)), ) - batch_size = len(token_groups) + batch_size = len(flat_token_groups) token_ids = torch.zeros( (batch_size, completion_len), device=self.device, @@ -120,13 +168,43 @@ def _batch_from_token_groups( device=self.device, dtype=torch.bool, ) - for row, group in enumerate(token_groups): - clipped = [int(token) % self.config.vocab_size for token in group[:completion_len]] - if not clipped: - continue - values = torch.tensor(clipped, device=self.device, dtype=torch.long) - token_ids[row, : values.numel()] = values - completion_mask[row, : values.numel()] = True + flat_rewards: list[float] = [] + flat_group_ids: list[int] = [] + flat_reference_logps: list[list[float]] = [] + row = 0 + for group_index, group in enumerate(candidate_groups): + for candidate_index, candidate_tokens in enumerate(group): + clipped = [ + int(token) % self.config.vocab_size + for token in candidate_tokens[:completion_len] + ] + if clipped: + values = torch.tensor(clipped, device=self.device, dtype=torch.long) + token_ids[row, : values.numel()] = values + completion_mask[row, : values.numel()] = True + flat_rewards.append( + _candidate_reward_value( + candidate_tokens, + reward_groups, + group_index, + candidate_index, + vocab_size=self.config.vocab_size, + ) + ) + if reference_logp_groups is not None: + flat_reference_logps.append( + _candidate_reference_logps( + reference_logp_groups, + group_index, + candidate_index, + completion_len=completion_len, + ) + ) + flat_group_ids.append(group_index) + row += 1 + + if not bool(completion_mask.any().item()): + completion_mask[:, :1] = True prompt_tokens = torch.zeros( (batch_size, self.config.prompt_len), @@ -145,14 +223,21 @@ def _batch_from_token_groups( dim=1, ) - generator = torch.Generator(device=self.device) - generator.manual_seed(self.config.seed + int(rollout.iteration)) - advantages = torch.randn( - (batch_size, completion_len), + rewards = torch.tensor( + flat_rewards, device=self.device, - generator=generator, dtype=self.config.dtype, ) + group_ids = torch.tensor(flat_group_ids, device=self.device, dtype=torch.long) + sequence_advantages = _compute_group_relative_advantages( + rewards, + group_ids=group_ids, + num_groups=len(candidate_groups), + eps=self.config.advantage_eps, + ) + advantages = _broadcast_sequence_advantages(sequence_advantages, completion_mask).to( + dtype=self.config.dtype + ) old_logps = torch.zeros( (batch_size, completion_len), device=self.device, @@ -163,16 +248,19 @@ def _batch_from_token_groups( device=self.device, dtype=self.config.dtype, ) - rewards = torch.randn( - (batch_size,), - device=self.device, - generator=generator, - dtype=self.config.dtype, - ) + if reference_logp_groups is not None: + for row_index, reference_values in enumerate(flat_reference_logps): + if reference_values: + values = torch.tensor( + reference_values, + device=self.device, + dtype=self.config.dtype, + ) + ref_logps[row_index, : values.numel()] = values valid_indices = completion_mask.reshape(-1).nonzero(as_tuple=False).squeeze(-1) metadata: dict[str, Any] = { - "num_prompts": batch_size, - "samples_per_prompt": 1, + "num_prompts": len(candidate_groups), + "samples_per_prompt": max(len(group) for group in candidate_groups), "batch_size": batch_size, "prompt_len": self.config.prompt_len, "completion_len": completion_len, @@ -184,6 +272,13 @@ def _batch_from_token_groups( "device": str(self.device), "seed": self.config.seed + int(rollout.iteration), "source": "rollout_payload", + "group_ids": flat_group_ids, + "reward_source": "payload_rewards" if reward_groups is not None else "token_id_proxy", + "reference_logp_source": ( + "payload_reference_logps" + if reference_logp_groups is not None + else "synthetic_current_offset" + ), } return SyntheticRLKernelBatch( input_ids=input_ids, @@ -218,41 +313,473 @@ def make_rollout_result( ) +class StatelessScoringWorker: + """Attach no-cache reference/reward scores to a completed rollout payload.""" + + def __init__( + self, + executor: StatelessForwardExecutor, + collate_inputs: Callable[[RolloutStageResult], StatelessForwardInputs], + ): + self.executor = executor + self.collate_inputs = collate_inputs + + def score(self, rollout: RolloutStageResult) -> RolloutStageResult: + started_at = time.perf_counter() + inputs = self.collate_inputs(rollout) + result = self.executor.score(inputs) + finished_at = time.perf_counter() + return RolloutStageResult( + iteration=rollout.iteration, + weight_version=rollout.weight_version, + payload=attach_stateless_scores_to_payload(rollout.payload, result, inputs=inputs), + started_at=started_at, + finished_at=finished_at, + metrics={ + **dict(rollout.metrics), + "scoring_backend": "stateless", + "scoring_mode": result.metrics["mode"], + "scoring_elapsed_ms": result.metrics["elapsed_ms"], + "scoring_active_completion_tokens": result.metrics["active_completion_tokens"], + "scoring_zero_kv_cache": result.metrics["zero_kv_cache"], + "scoring_attention_backend": result.metrics["attention_backend"], + "scoring_kv_cache_output_mb": result.metrics["kv_cache_output_mb"], + }, + ) + + +def attach_stateless_scores_to_payload( + payload: Any, + result: StatelessForwardResult, + *, + inputs: Optional[StatelessForwardInputs] = None, +) -> dict[str, Any]: + """Return a rollout payload augmented with stateless scoring outputs.""" + + scored_payload = dict(payload) if isinstance(payload, Mapping) else {"raw_payload": payload} + scored_payload["stateless_scores"] = { + "reference_logps": result.reference_logps, + "rewards": result.rewards, + "token_scores": result.token_scores, + "metrics": dict(result.metrics), + } + if result.rewards is not None: + _attach_reward_tensor_to_grouped_candidates(scored_payload, result.rewards) + if result.reference_logps is not None and inputs is not None: + _attach_reference_logps_to_grouped_candidates( + scored_payload, + result.reference_logps, + inputs.completion_mask, + ) + return scored_payload + + +def build_stateless_inputs_from_rollout_payload( + payload: Any, + *, + prompt_len: Optional[int] = 1, + prompt_token_id: int = 0, + max_completion_len: Optional[int] = None, + device: torch.device | str = "cpu", +) -> StatelessForwardInputs: + """Build dense no-cache scoring inputs from grouped rollout token payloads.""" + + if prompt_len is not None and prompt_len < 0: + raise ValueError("prompt_len must be non-negative") + if max_completion_len is not None and max_completion_len <= 0: + raise ValueError("max_completion_len must be greater than zero") + + candidate_records = _candidate_records_from_payload(payload) + if not candidate_records: + raise ValueError("rollout payload does not contain candidate token ids") + + flat_token_groups = [tokens for _prompt_tokens, tokens in candidate_records] + completion_len = max(len(tokens) for tokens in flat_token_groups) + if max_completion_len is not None: + completion_len = min(completion_len, max_completion_len) + completion_len = max(1, completion_len) + resolved_prompt_len = ( + max((len(prompt_tokens) for prompt_tokens, _tokens in candidate_records), default=0) + if prompt_len is None + else prompt_len + ) + seq_len = resolved_prompt_len + completion_len + resolved_device = torch.device(device) + + input_ids = torch.full( + (len(flat_token_groups), seq_len), + int(prompt_token_id), + device=resolved_device, + dtype=torch.long, + ) + attention_mask = torch.zeros_like(input_ids, dtype=torch.bool) + completion_mask = torch.zeros_like(input_ids, dtype=torch.bool) + if resolved_prompt_len and prompt_len is not None: + attention_mask[:, :resolved_prompt_len] = True + + for row, (prompt_tokens, token_ids) in enumerate(candidate_records): + if resolved_prompt_len and prompt_tokens: + clipped_prompt = [int(token) for token in prompt_tokens[:resolved_prompt_len]] + values = torch.tensor(clipped_prompt, device=resolved_device, dtype=torch.long) + input_ids[row, : values.numel()] = values + if prompt_len is None: + attention_mask[row, : values.numel()] = True + clipped = [int(token) for token in token_ids[:completion_len]] + if clipped: + values = torch.tensor(clipped, device=resolved_device, dtype=torch.long) + start = resolved_prompt_len + end = resolved_prompt_len + values.numel() + input_ids[row, start:end] = values + attention_mask[row, start:end] = True + completion_mask[row, start:end] = True + + if not bool(completion_mask.any().item()): + raise ValueError("rollout payload does not contain active completion tokens") + return StatelessForwardInputs( + input_ids=input_ids, + attention_mask=attention_mask, + completion_mask=completion_mask, + ) + + def extract_rollout_token_groups(payload: Any) -> list[list[int]]: """Extract generated token ids from RL-Kernel/vLLM-style rollout payloads.""" - groups: list[list[int]] = [] + return [tokens for group in extract_rollout_candidate_groups(payload) for tokens in group] + + +def extract_rollout_candidate_groups(payload: Any) -> list[list[list[int]]]: + """Extract generated token ids while preserving prompt-level candidate groups.""" + if not isinstance(payload, Mapping): - return groups + return [] normalized_outputs = payload.get("normalized_outputs") if isinstance(normalized_outputs, Sequence) and not isinstance( normalized_outputs, (str, bytes) ): - for group in normalized_outputs: - if not isinstance(group, Sequence) or isinstance(group, (str, bytes)): - continue - for candidate in group: - token_ids = _candidate_token_ids(candidate) - if token_ids: - groups.append(token_ids) + grouped = _candidate_groups_from_grouped_outputs(normalized_outputs) + if grouped: + return grouped + + outputs = payload.get("outputs") + if isinstance(outputs, Sequence) and not isinstance(outputs, (str, bytes)): + return _candidate_groups_from_grouped_outputs(outputs) + return [] + + +def extract_rollout_reward_groups(payload: Any) -> list[list[float]]: + """Extract scalar reward groups from rollout payloads when they are present.""" + + if not isinstance(payload, Mapping): + return [] + + normalized_outputs = payload.get("normalized_outputs") + if isinstance(normalized_outputs, Sequence) and not isinstance( + normalized_outputs, (str, bytes) + ): + groups = _reward_groups_from_grouped_outputs(normalized_outputs) if groups: return groups outputs = payload.get("outputs") if isinstance(outputs, Sequence) and not isinstance(outputs, (str, bytes)): - for group in outputs: - if isinstance(group, Sequence) and not isinstance(group, (str, bytes)): - candidates = group - else: - candidates = [group] - for candidate in candidates: - token_ids = _candidate_token_ids(candidate) - if token_ids: - groups.append(token_ids) + return _reward_groups_from_grouped_outputs(outputs) + return [] + + +def extract_rollout_reference_logp_groups(payload: Any) -> list[list[list[float]]]: + """Extract per-candidate reference logprobs from rollout payloads when present.""" + + if not isinstance(payload, Mapping): + return [] + + normalized_outputs = payload.get("normalized_outputs") + if isinstance(normalized_outputs, Sequence) and not isinstance( + normalized_outputs, (str, bytes) + ): + groups = _reference_logp_groups_from_grouped_outputs(normalized_outputs) + if groups: + return groups + + outputs = payload.get("outputs") + if isinstance(outputs, Sequence) and not isinstance(outputs, (str, bytes)): + return _reference_logp_groups_from_grouped_outputs(outputs) + return [] + + +def _candidate_records_from_payload(payload: Any) -> list[tuple[list[int], list[int]]]: + if not isinstance(payload, Mapping): + return [] + + normalized_outputs = payload.get("normalized_outputs") + if isinstance(normalized_outputs, Sequence) and not isinstance( + normalized_outputs, + (str, bytes), + ): + records = _candidate_records_from_grouped_outputs(normalized_outputs) + if records: + return records + + outputs = payload.get("outputs") + if isinstance(outputs, Sequence) and not isinstance(outputs, (str, bytes)): + return _candidate_records_from_grouped_outputs(outputs) + return [] + + +def _candidate_records_from_grouped_outputs( + grouped_outputs: Sequence[Any], +) -> list[tuple[list[int], list[int]]]: + records: list[tuple[list[int], list[int]]] = [] + for group in grouped_outputs: + if isinstance(group, Sequence) and not isinstance(group, (str, bytes, Mapping)): + candidates = group + else: + candidates = [group] + for candidate in candidates: + token_ids = _candidate_token_ids(candidate) + if token_ids: + records.append((_candidate_prompt_token_ids(candidate), token_ids)) + return records + + +def _attach_reward_tensor_to_grouped_candidates( + payload: dict[str, Any], + rewards: torch.Tensor, +) -> None: + reward_values = [float(value) for value in rewards.detach().cpu().reshape(-1).tolist()] + if not reward_values: + return + + for key in ("normalized_outputs", "outputs"): + grouped_outputs = payload.get(key) + if not isinstance(grouped_outputs, Sequence) or isinstance(grouped_outputs, (str, bytes)): + continue + updated_outputs, used = _attach_rewards_to_grouped_outputs( + grouped_outputs, + reward_values, + ) + if used == 0: + continue + if used != len(reward_values): + raise ValueError( + "stateless reward count must match rollout candidate count, got " + f"{len(reward_values)} rewards for {used} candidates" + ) + payload[key] = updated_outputs + return + + +def _attach_reference_logps_to_grouped_candidates( + payload: dict[str, Any], + reference_logps: torch.Tensor, + completion_mask: torch.Tensor, +) -> None: + if reference_logps.shape != completion_mask.shape: + raise ValueError("reference_logps shape must match completion_mask shape") + + mask = completion_mask.detach().to(dtype=torch.bool, device=reference_logps.device) + logp_rows = [ + [float(value) for value in reference_logps[row][mask[row]].detach().cpu().tolist()] + for row in range(reference_logps.shape[0]) + ] + if not logp_rows: + return + + for key in ("normalized_outputs", "outputs"): + grouped_outputs = payload.get(key) + if not isinstance(grouped_outputs, Sequence) or isinstance(grouped_outputs, (str, bytes)): + continue + updated_outputs, used = _attach_reference_logps_to_grouped_outputs( + grouped_outputs, + logp_rows, + ) + if used == 0: + continue + if used != len(logp_rows): + raise ValueError( + "stateless reference_logps row count must match rollout candidate count, got " + f"{len(logp_rows)} rows for {used} candidates" + ) + payload[key] = updated_outputs + return + + +def _candidate_groups_from_grouped_outputs(grouped_outputs: Sequence[Any]) -> list[list[list[int]]]: + groups: list[list[list[int]]] = [] + for group in grouped_outputs: + if isinstance(group, Sequence) and not isinstance(group, (str, bytes, Mapping)): + candidates = group + else: + candidates = [group] + candidate_group = [] + for candidate in candidates: + token_ids = _candidate_token_ids(candidate) + if token_ids: + candidate_group.append(token_ids) + if candidate_group: + groups.append(candidate_group) + return groups + + +def _reward_groups_from_grouped_outputs(grouped_outputs: Sequence[Any]) -> list[list[float]]: + groups: list[list[float]] = [] + for group in grouped_outputs: + if isinstance(group, Sequence) and not isinstance(group, (str, bytes, Mapping)): + candidates = group + else: + candidates = [group] + reward_group = [] + for candidate in candidates: + reward = _candidate_reward(candidate) + if reward is not None: + reward_group.append(reward) + if reward_group: + groups.append(reward_group) + return groups + + +def _reference_logp_groups_from_grouped_outputs( + grouped_outputs: Sequence[Any], +) -> list[list[list[float]]]: + groups: list[list[list[float]]] = [] + for group in grouped_outputs: + if isinstance(group, Sequence) and not isinstance(group, (str, bytes, Mapping)): + candidates = group + else: + candidates = [group] + reference_group = [] + for candidate in candidates: + reference_logps = _candidate_reference_logp_values(candidate) + if reference_logps: + reference_group.append(reference_logps) + if reference_group: + groups.append(reference_group) return groups +def _attach_rewards_to_grouped_outputs( + grouped_outputs: Sequence[Any], + rewards: Sequence[float], +) -> tuple[list[Any], int]: + updated_groups: list[Any] = [] + reward_index = 0 + + for group in grouped_outputs: + if isinstance(group, Sequence) and not isinstance(group, (str, bytes, Mapping)): + updated_candidates = [] + for candidate in group: + updated_candidate, consumed = _attach_reward_to_candidate( + candidate, + rewards, + reward_index, + ) + reward_index += consumed + updated_candidates.append(updated_candidate) + updated_groups.append(updated_candidates) + else: + updated_candidate, consumed = _attach_reward_to_candidate( + group, + rewards, + reward_index, + ) + reward_index += consumed + updated_groups.append(updated_candidate) + + return updated_groups, reward_index + + +def _attach_reward_to_candidate( + candidate: Any, + rewards: Sequence[float], + reward_index: int, +) -> tuple[Any, int]: + token_ids = _candidate_token_ids(candidate) + if not token_ids: + return candidate, 0 + if reward_index >= len(rewards): + raise ValueError( + "stateless reward count must match rollout candidate count, got fewer rewards " + "than candidates" + ) + + reward = float(rewards[reward_index]) + if isinstance(candidate, Mapping): + updated = dict(candidate) + updated["reward"] = reward + updated["reward_source"] = "stateless_executor" + return updated, 1 + + try: + updated = dict(vars(candidate)) + except TypeError: + updated = {"candidate": candidate} + updated["reward"] = reward + updated["reward_source"] = "stateless_executor" + return updated, 1 + + +def _attach_reference_logps_to_grouped_outputs( + grouped_outputs: Sequence[Any], + logp_rows: Sequence[Sequence[float]], +) -> tuple[list[Any], int]: + updated_groups: list[Any] = [] + row_index = 0 + + for group in grouped_outputs: + if isinstance(group, Sequence) and not isinstance(group, (str, bytes, Mapping)): + updated_candidates = [] + for candidate in group: + updated_candidate, consumed = _attach_reference_logps_to_candidate( + candidate, + logp_rows, + row_index, + ) + row_index += consumed + updated_candidates.append(updated_candidate) + updated_groups.append(updated_candidates) + else: + updated_candidate, consumed = _attach_reference_logps_to_candidate( + group, + logp_rows, + row_index, + ) + row_index += consumed + updated_groups.append(updated_candidate) + + return updated_groups, row_index + + +def _attach_reference_logps_to_candidate( + candidate: Any, + logp_rows: Sequence[Sequence[float]], + row_index: int, +) -> tuple[Any, int]: + token_ids = _candidate_token_ids(candidate) + if not token_ids: + return candidate, 0 + if row_index >= len(logp_rows): + raise ValueError( + "stateless reference_logps row count must match rollout candidate count, got " + "fewer rows than candidates" + ) + + values = [float(value) for value in logp_rows[row_index]] + if isinstance(candidate, Mapping): + updated = dict(candidate) + updated["reference_logps"] = values + updated["reference_logp_source"] = "stateless_executor" + return updated, 1 + + try: + updated = dict(vars(candidate)) + except TypeError: + updated = {"candidate": candidate} + updated["reference_logps"] = values + updated["reference_logp_source"] = "stateless_executor" + return updated, 1 + + def _candidate_token_ids(candidate: Any) -> list[int]: if candidate is None: return [] @@ -278,9 +805,213 @@ def _candidate_token_ids(candidate: Any) -> list[int]: return [] +def _candidate_prompt_token_ids(candidate: Any) -> list[int]: + if candidate is None: + return [] + if isinstance(candidate, Mapping): + for key in ("prompt_token_ids", "prompt_ids", "input_token_ids", "input_ids"): + if key in candidate: + return _copy_int_list(candidate[key]) + return [] + + for attr in ("prompt_token_ids", "prompt_ids", "input_token_ids", "input_ids"): + value = getattr(candidate, attr, None) + if value is not None: + return _copy_int_list(value) + return [] + + +def _candidate_reward(candidate: Any) -> Optional[float]: + if candidate is None: + return None + if isinstance(candidate, Mapping): + nested_outputs = candidate.get("outputs") + if isinstance(nested_outputs, Sequence) and not isinstance(nested_outputs, (str, bytes)): + for nested in nested_outputs: + reward = _candidate_reward(nested) + if reward is not None: + return reward + for key in ("reward", "score", "scalar_reward", "reward_score"): + if key in candidate: + return _safe_float(candidate[key]) + return None + + nested_outputs = getattr(candidate, "outputs", None) + if isinstance(nested_outputs, Sequence) and not isinstance(nested_outputs, (str, bytes)): + for nested in nested_outputs: + reward = _candidate_reward(nested) + if reward is not None: + return reward + for attr in ("reward", "score", "scalar_reward", "reward_score"): + if hasattr(candidate, attr): + return _safe_float(getattr(candidate, attr)) + return None + + +def _candidate_reference_logp_values(candidate: Any) -> list[float]: + if candidate is None: + return [] + if isinstance(candidate, Mapping): + nested_outputs = candidate.get("outputs") + if isinstance(nested_outputs, Sequence) and not isinstance(nested_outputs, (str, bytes)): + for nested in nested_outputs: + reference_logps = _candidate_reference_logp_values(nested) + if reference_logps: + return reference_logps + for key in ("reference_logps", "ref_logps", "reference_logprobs", "ref_logprobs"): + if key in candidate: + return _copy_float_list(candidate[key]) + return [] + + nested_outputs = getattr(candidate, "outputs", None) + if isinstance(nested_outputs, Sequence) and not isinstance(nested_outputs, (str, bytes)): + for nested in nested_outputs: + reference_logps = _candidate_reference_logp_values(nested) + if reference_logps: + return reference_logps + for attr in ("reference_logps", "ref_logps", "reference_logprobs", "ref_logprobs"): + if hasattr(candidate, attr): + return _copy_float_list(getattr(candidate, attr)) + return [] + + +def _safe_float(value: Any) -> float: + if isinstance(value, torch.Tensor): + flat = value.detach().cpu().reshape(-1) + if flat.numel() != 1: + raise ValueError("rollout reward tensors must contain exactly one value") + return float(flat[0].item()) + return float(value) + + +def _reward_groups_match( + candidate_groups: Sequence[Sequence[Sequence[int]]], + reward_groups: Sequence[Sequence[float]], +) -> bool: + if len(candidate_groups) != len(reward_groups): + return False + return all( + len(candidates) == len(rewards) + for candidates, rewards in zip(candidate_groups, reward_groups, strict=False) + ) + + +def _reference_logp_groups_match( + candidate_groups: Sequence[Sequence[Sequence[int]]], + reference_logp_groups: Sequence[Sequence[Sequence[float]]], + *, + max_completion_len: int, +) -> bool: + if len(candidate_groups) != len(reference_logp_groups): + return False + for candidates, reference_group in zip( + candidate_groups, + reference_logp_groups, + strict=False, + ): + if len(candidates) != len(reference_group): + return False + for token_ids, reference_logps in zip(candidates, reference_group, strict=False): + required_len = min(len(token_ids), max_completion_len) + if len(reference_logps) < required_len: + return False + return True + + +def _candidate_reward_value( + token_ids: Sequence[int], + reward_groups: Optional[Sequence[Sequence[float]]], + group_index: int, + candidate_index: int, + *, + vocab_size: int, +) -> float: + if reward_groups is not None: + return float(reward_groups[group_index][candidate_index]) + if not token_ids: + return 0.0 + clipped = [int(token) % vocab_size for token in token_ids] + return float(sum(clipped)) / float(max(len(clipped), 1) * max(vocab_size - 1, 1)) + + +def _candidate_reference_logps( + reference_logp_groups: Sequence[Sequence[Sequence[float]]], + group_index: int, + candidate_index: int, + *, + completion_len: int, +) -> list[float]: + return [ + float(value) + for value in reference_logp_groups[group_index][candidate_index][:completion_len] + ] + + +def objective_reference_logps( + current_logps: torch.Tensor, + batch: SyntheticRLKernelBatch, +) -> torch.Tensor: + """Return payload reference logps when available, otherwise a synthetic offset.""" + + if batch.metadata.get("reference_logp_source") != "payload_reference_logps": + return current_logps.detach() - 0.02 + if batch.ref_logps.shape != current_logps.shape: + raise ValueError("payload reference logps shape must match current logps shape") + return batch.ref_logps.detach().to(device=current_logps.device, dtype=torch.float32) + + +def _compute_group_relative_advantages( + rewards: torch.Tensor, + *, + group_ids: torch.Tensor, + num_groups: int, + eps: float, +) -> torch.Tensor: + flat_rewards = rewards.reshape(-1).float() + flat_group_ids = group_ids.reshape(-1).to(device=flat_rewards.device, dtype=torch.long) + if flat_rewards.numel() != flat_group_ids.numel(): + raise ValueError("rewards and group_ids must contain the same number of elements") + + counts = flat_rewards.new_zeros(num_groups).index_add_( + 0, + flat_group_ids, + torch.ones_like(flat_rewards), + ) + sums = flat_rewards.new_zeros(num_groups).index_add_(0, flat_group_ids, flat_rewards) + sq_sums = flat_rewards.new_zeros(num_groups).index_add_( + 0, + flat_group_ids, + flat_rewards * flat_rewards, + ) + safe_counts = counts.clamp_min(1.0) + means = sums / safe_counts + variances = (sq_sums / safe_counts) - means * means + stds = torch.sqrt(torch.clamp(variances, min=0.0) + eps) + advantages = (flat_rewards - means[flat_group_ids]) / stds[flat_group_ids] + return advantages.to(dtype=rewards.dtype) + + +def _broadcast_sequence_advantages( + sequence_advantages: torch.Tensor, + completion_mask: torch.Tensor, +) -> torch.Tensor: + return sequence_advantages.reshape(-1, 1) * completion_mask.to( + device=sequence_advantages.device, + dtype=sequence_advantages.dtype, + ) + + def _copy_int_list(value: Any) -> list[int]: if isinstance(value, torch.Tensor): return [int(item) for item in value.detach().cpu().reshape(-1).tolist()] if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): return [int(item) for item in value] return [] + + +def _copy_float_list(value: Any) -> list[float]: + if isinstance(value, torch.Tensor): + return [float(item) for item in value.detach().cpu().reshape(-1).tolist()] + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return [float(item) for item in value] + return [] diff --git a/tests/test_deepspeed_training_worker.py b/tests/test_deepspeed_training_worker.py index 49b54e330..e09344477 100644 --- a/tests/test_deepspeed_training_worker.py +++ b/tests/test_deepspeed_training_worker.py @@ -222,6 +222,10 @@ def test_deepspeed_training_worker_uses_engine_backward_and_step(monkeypatch): assert result.metrics["rollout_sequences"] == 2 assert result.metrics["rollout_tokens"] == 6 assert math.isfinite(result.metrics["loss"]) + assert "advantage_mean" not in result.metrics + assert "advantage_std" not in result.metrics + assert math.isfinite(result.metrics["active_advantage_mean_global"]) + assert result.metrics["active_advantage_std_global"] >= 0.0 def test_deepspeed_training_worker_synthetic_fallback(monkeypatch): diff --git a/tests/test_paged_kv_baseline.py b/tests/test_paged_kv_baseline.py new file mode 100644 index 000000000..cf58871dc --- /dev/null +++ b/tests/test_paged_kv_baseline.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import importlib +import sys +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.executors.paged_kv_baseline import ( + PagedKVScoringBaseline, + PagedKVScoringConfig, + reserve_paged_kv_cache, +) +from rl_engine.executors.stateless_executor import StatelessForwardInputs, score_reference_logprobs + + +class FakeGenerationReferenceModel(torch.nn.Module): + def __init__(self, logits: torch.Tensor): + super().__init__() + self.register_buffer("fixed_logits", logits) + self.use_cache_calls: list[bool | None] = [] + + def forward(self, input_ids, attention_mask=None, use_cache=None): + del attention_mask + self.use_cache_calls.append(use_cache) + batch, seq_len = input_ids.shape + key = torch.empty(batch, 1, seq_len, 4, device=input_ids.device) + value = torch.empty_like(key) + return SimpleNamespace( + logits=self.fixed_logits[: input_ids.shape[0], : input_ids.shape[1]], + past_key_values=((key, value),) if use_cache else None, + ) + + +def _inputs() -> StatelessForwardInputs: + input_ids = torch.tensor( + [ + [0, 1, 2, 3, 0], + [0, 2, 1, 4, 5], + ], + dtype=torch.long, + ) + attention_mask = torch.tensor( + [ + [True, True, True, True, True], + [True, True, True, False, False], + ] + ) + completion_mask = torch.tensor( + [ + [False, False, True, True, False], + [False, True, True, False, False], + ] + ) + return StatelessForwardInputs( + input_ids=input_ids, + attention_mask=attention_mask, + completion_mask=completion_mask, + ) + + +def _logits_for(inputs: StatelessForwardInputs, vocab_size: int = 8) -> torch.Tensor: + logits = torch.full((*inputs.input_ids.shape, vocab_size), -3.0) + for batch_index in range(inputs.input_ids.shape[0]): + for pos in range(inputs.input_ids.shape[1] - 1): + token = int(inputs.input_ids[batch_index, pos + 1].item()) + logits[batch_index, pos, token] = 4.0 + batch_index + pos + return logits + + +def test_reserve_paged_kv_cache_builds_block_table_and_tracks_bytes(): + inputs = _inputs() + config = PagedKVScoringConfig( + num_layers=2, + num_kv_heads=2, + head_dim=4, + block_size=2, + kv_cache_dtype=torch.float32, + kv_cache_blocks=8, + ) + + reservation = reserve_paged_kv_cache(inputs, config) + + assert reservation.sequence_lengths.tolist() == [5, 3] + assert reservation.blocks_per_sequence.tolist() == [3, 2] + assert reservation.required_blocks == 5 + assert reservation.reserved_blocks == 8 + assert reservation.key_cache.shape == (2, 8, 2, 2, 4) + assert reservation.value_cache.shape == reservation.key_cache.shape + assert reservation.block_tables.tolist() == [[0, 1, 2], [3, 4, -1]] + + expected_cache_bytes = 2 * reservation.key_cache.numel() * reservation.key_cache.element_size() + assert reservation.cache_bytes == expected_cache_bytes + assert reservation.reserved_bytes == reservation.cache_bytes + reservation.metadata_bytes + assert reservation.reserved_mb > 0.0 + + +def test_paged_kv_baseline_scores_like_reference_path_and_enables_cache(): + inputs = _inputs() + logits = _logits_for(inputs).requires_grad_() + model = FakeGenerationReferenceModel(logits) + baseline = PagedKVScoringBaseline( + model, + PagedKVScoringConfig( + mode="reference", + num_layers=1, + num_kv_heads=1, + head_dim=4, + block_size=2, + kv_cache_dtype=torch.float32, + return_token_scores=True, + ), + ) + + result = baseline.score(inputs) + expected = score_reference_logprobs(logits, inputs) + + assert model.use_cache_calls == [True] + assert result.reference_logps is not None + assert torch.allclose(result.reference_logps, expected) + assert result.reference_logps.requires_grad is False + assert result.token_scores is not None + assert result.metrics["baseline_kind"] == "generation_engine_paged_kv_reservation" + assert result.metrics["use_cache"] is True + assert result.metrics["use_cache_passed"] is True + assert result.metrics["paged_kv_required_blocks"] == 5 + assert result.metrics["paged_kv_cache_reserved_mb"] > 0.0 + assert result.metrics["model_kv_cache_output_present"] is True + assert result.metrics["model_kv_cache_output_tensors"] == 2 + assert result.metrics["model_kv_cache_output_mb"] > 0.0 + assert result.metrics["total_kv_cache_mb"] > result.metrics["paged_kv_cache_reserved_mb"] + + +def test_paged_kv_reservation_rejects_capacity_smaller_than_batch_needs(): + inputs = _inputs() + config = PagedKVScoringConfig( + num_layers=1, + num_kv_heads=1, + head_dim=4, + block_size=2, + kv_cache_blocks=4, + ) + + with pytest.raises(ValueError, match="kv_cache_blocks=4"): + reserve_paged_kv_cache(inputs, config) + + +def test_importing_paged_kv_baseline_does_not_import_heavy_optional_runtimes(monkeypatch): + for name in ("vllm", "deepspeed", "ray", "flash_attn"): + monkeypatch.delitem(sys.modules, name, raising=False) + + module = importlib.import_module("rl_engine.executors.paged_kv_baseline") + + assert module.PagedKVScoringBaseline is not None + for name in ("vllm", "deepspeed", "ray", "flash_attn"): + assert name not in sys.modules diff --git a/tests/test_stateless_executor.py b/tests/test_stateless_executor.py new file mode 100644 index 000000000..fc6d9b3a4 --- /dev/null +++ b/tests/test_stateless_executor.py @@ -0,0 +1,362 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import importlib +import sys +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.executors.stateless_executor import ( + StatelessForwardConfig, + StatelessForwardExecutor, + StatelessForwardInputs, + StatelessForwardOutputs, + default_reward_adapter, + score_reference_logprobs, + score_rewards, +) +from rl_engine.testing.reference_ops import selected_logprobs_reference + + +class FakeReferenceModel(torch.nn.Module): + def __init__(self, logits: torch.Tensor): + super().__init__() + self.register_buffer("fixed_logits", logits) + self.config = SimpleNamespace(use_cache=True, _attn_implementation="eager") + self.generation_config = SimpleNamespace(use_cache=True) + self.use_cache_calls: list[bool | None] = [] + self.config_use_cache_calls: list[bool | None] = [] + self.generation_config_use_cache_calls: list[bool | None] = [] + self.attn_implementation_calls: list[str | None] = [] + + def forward(self, input_ids, attention_mask=None, use_cache=None): + del attention_mask + self.use_cache_calls.append(use_cache) + self.config_use_cache_calls.append(getattr(self.config, "use_cache", None)) + self.generation_config_use_cache_calls.append( + getattr(self.generation_config, "use_cache", None) + ) + self.attn_implementation_calls.append(getattr(self.config, "_attn_implementation", None)) + return SimpleNamespace( + logits=self.fixed_logits[: input_ids.shape[0], : input_ids.shape[1]], + past_key_values=None, + ) + + +class FakeRewardModel(torch.nn.Module): + def __init__(self, rewards: torch.Tensor): + super().__init__() + self.register_buffer("fixed_rewards", rewards) + self.use_cache_calls: list[bool | None] = [] + + def forward(self, input_ids, attention_mask=None, use_cache=None): + del attention_mask + self.use_cache_calls.append(use_cache) + return {"logits": self.fixed_rewards[: input_ids.shape[0]].unsqueeze(-1)} + + +class NoUseCacheModel(torch.nn.Module): + def forward(self, input_ids, attention_mask=None): + del attention_mask + batch, seq_len = input_ids.shape + vocab = 8 + return torch.zeros(batch, seq_len, vocab, device=input_ids.device) + + +class CacheReturningModel(torch.nn.Module): + def forward(self, input_ids, attention_mask=None, use_cache=None): + del attention_mask, use_cache + batch, seq_len = input_ids.shape + logits = torch.zeros(batch, seq_len, 8, device=input_ids.device) + key = torch.zeros(batch, 1, seq_len, 4, device=input_ids.device) + value = torch.zeros_like(key) + return {"logits": logits, "past_key_values": ((key, value),)} + + +def _inputs() -> StatelessForwardInputs: + input_ids = torch.tensor( + [ + [0, 1, 2, 3, 0], + [0, 2, 1, 4, 5], + ], + dtype=torch.long, + ) + attention_mask = torch.tensor( + [ + [1, 1, 1, 1, 0], + [1, 1, 1, 1, 1], + ], + dtype=torch.bool, + ) + completion_mask = torch.tensor( + [ + [False, False, True, True, False], + [False, False, True, True, True], + ] + ) + return StatelessForwardInputs( + input_ids=input_ids, + attention_mask=attention_mask, + completion_mask=completion_mask, + ) + + +def _logits_for(inputs: StatelessForwardInputs, vocab_size: int = 8) -> torch.Tensor: + logits = torch.full((*inputs.input_ids.shape, vocab_size), -3.0) + for batch_index in range(inputs.input_ids.shape[0]): + for pos in range(inputs.input_ids.shape[1] - 1): + token = int(inputs.input_ids[batch_index, pos + 1].item()) + logits[batch_index, pos, token] = 4.0 + batch_index + pos + return logits + + +def test_reference_scoring_matches_next_token_pytorch_reference_and_masks_prompt_tokens(): + inputs = _inputs() + logits = _logits_for(inputs) + + actual = score_reference_logprobs(logits, inputs) + shifted = selected_logprobs_reference( + logits[:, :-1, :], + inputs.input_ids[:, 1:], + mask=inputs.completion_mask[:, 1:], + ) + expected = torch.zeros_like(actual) + expected[:, 1:] = shifted + + assert torch.allclose(actual, expected) + assert torch.equal( + actual[~inputs.completion_mask], + torch.zeros_like(actual[~inputs.completion_mask]), + ) + assert actual[0, 2] != 0.0 + assert actual[0, 1] == 0.0 + + +def test_executor_runs_full_sequence_forward_with_use_cache_false_and_detaches_outputs(): + inputs = _inputs() + logits = _logits_for(inputs).requires_grad_() + model = FakeReferenceModel(logits) + executor = StatelessForwardExecutor( + model, + StatelessForwardConfig(mode="reference", return_token_scores=True), + ) + + result = executor.score(inputs) + + assert model.use_cache_calls == [False] + assert result.reference_logps is not None + assert result.token_scores is not None + assert result.reference_logps.requires_grad is False + assert result.token_scores.requires_grad is False + assert result.metrics["mode"] == "reference" + assert result.metrics["active_completion_tokens"] == int(inputs.completion_mask.sum().item()) + assert result.metrics["use_cache"] is False + assert result.metrics["use_cache_passed"] is True + assert result.metrics["detached_outputs"] is True + assert result.metrics["zero_kv_cache"] is True + assert result.metrics["kv_cache_output_tensors"] == 0 + assert result.metrics["attention_backend"] == "flash_attention_2" + assert result.metrics["attention_backend_configured"] is True + assert result.metrics["model_config_use_cache_disabled"] is True + assert model.config_use_cache_calls == [False] + assert model.generation_config_use_cache_calls == [False] + assert model.attn_implementation_calls == ["flash_attention_2"] + assert model.config.use_cache is True + assert model.generation_config.use_cache is True + assert model.config._attn_implementation == "eager" + assert not hasattr(model.config, "attn_implementation") + assert not hasattr(model.generation_config, "attn_implementation") + + +def test_executor_falls_back_for_models_without_use_cache_argument(): + inputs = _inputs() + executor = StatelessForwardExecutor( + NoUseCacheModel(), + StatelessForwardConfig(mode="reference"), + ) + + result = executor.score(inputs) + + assert result.reference_logps is not None + assert result.metrics["use_cache_passed"] is False + + +def test_executor_rejects_models_that_return_kv_cache_outputs(): + inputs = _inputs() + executor = StatelessForwardExecutor( + CacheReturningModel(), + StatelessForwardConfig(mode="reference"), + ) + + with pytest.raises(ValueError, match="KV-cache outputs"): + executor.score(inputs) + + +def test_reward_mode_returns_one_scalar_per_sequence_with_default_adapter(): + inputs = _inputs() + rewards = torch.tensor([1.5, -0.25]) + model = FakeRewardModel(rewards) + executor = StatelessForwardExecutor(model, StatelessForwardConfig(mode="reward")) + + result = executor.score(inputs) + + assert model.use_cache_calls == [False] + assert result.reference_logps is None + assert result.rewards is not None + assert torch.allclose(result.rewards, rewards) + assert result.rewards.requires_grad is False + assert result.metrics["mode"] == "reward" + + +def test_both_mode_runs_one_forward_and_returns_reference_logps_and_rewards(): + inputs = _inputs() + logits = _logits_for(inputs) + rewards = torch.tensor([0.2, 0.8]) + + class FakeBothModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.calls = 0 + + def forward(self, input_ids, attention_mask=None, use_cache=None): + del attention_mask + assert use_cache is False + self.calls += 1 + return { + "logits": logits[: input_ids.shape[0], : input_ids.shape[1]], + "rewards": rewards, + } + + model = FakeBothModel() + executor = StatelessForwardExecutor(model, StatelessForwardConfig(mode="both")) + + result = executor.score(inputs) + + assert model.calls == 1 + assert result.reference_logps is not None + assert result.rewards is not None + assert torch.allclose(result.rewards, rewards) + + +def test_custom_reward_adapter_and_shape_validation(): + inputs = _inputs() + outputs = StatelessForwardOutputs(raw={"hidden": torch.ones(2, 3)}, logits=None) + + def adapter(model_outputs, batch_inputs): + assert model_outputs is outputs + return batch_inputs.attention_mask.float().sum(dim=1) + + rewards = score_rewards(outputs, inputs, reward_adapter=adapter) + + assert torch.equal(rewards, torch.tensor([4.0, 5.0])) + + with pytest.raises(ValueError, match="shape \\[B\\]"): + score_rewards(outputs, inputs, reward_adapter=lambda _outputs, _inputs: torch.ones(2, 2)) + + with pytest.raises(ValueError, match="batch size"): + score_rewards(outputs, inputs, reward_adapter=lambda _outputs, _inputs: torch.ones(1)) + + +def test_default_reward_adapter_rejects_non_scalar_logits(): + inputs = _inputs() + outputs = StatelessForwardOutputs(raw={}, logits=torch.zeros(2, 3, 8)) + + with pytest.raises(ValueError, match="reward tensor"): + default_reward_adapter(outputs, inputs) + + +def test_executor_rejects_shape_mismatches_empty_masks_and_invalid_config(): + inputs = _inputs() + model = FakeReferenceModel(_logits_for(inputs)) + executor = StatelessForwardExecutor(model, StatelessForwardConfig(mode="reference")) + + with pytest.raises(ValueError, match="attention_mask shape"): + executor.score( + StatelessForwardInputs( + input_ids=inputs.input_ids, + attention_mask=torch.ones(2, 4, dtype=torch.bool), + completion_mask=inputs.completion_mask, + ) + ) + + invalid_start_mask = inputs.completion_mask.clone() + invalid_start_mask[0, 0] = True + with pytest.raises(ValueError, match=r"completion_mask\[:, 0\]"): + executor.score( + StatelessForwardInputs( + input_ids=inputs.input_ids, + attention_mask=inputs.attention_mask, + completion_mask=invalid_start_mask, + ) + ) + + with pytest.raises(ValueError, match="completion_mask must contain"): + executor.score( + StatelessForwardInputs( + input_ids=inputs.input_ids, + attention_mask=inputs.attention_mask, + completion_mask=torch.zeros_like(inputs.completion_mask), + ) + ) + + with pytest.raises(ValueError, match="use_cache"): + StatelessForwardConfig(use_cache=True) + + with pytest.raises(ValueError, match="attention_backend"): + StatelessForwardConfig(attention_backend="paged_attention") # type: ignore[arg-type] + + with pytest.raises(ValueError, match="mode"): + StatelessForwardConfig(mode="generate") # type: ignore[arg-type] + + +def test_importing_stateless_executor_does_not_import_heavy_optional_runtimes(monkeypatch): + for name in ("vllm", "deepspeed", "ray", "flash_attn"): + monkeypatch.delitem(sys.modules, name, raising=False) + + module = importlib.import_module("rl_engine.executors.stateless_executor") + + assert module.StatelessForwardExecutor is not None + for name in ("vllm", "deepspeed", "ray", "flash_attn"): + assert name not in sys.modules + + +def test_reference_scoring_rejects_completion_start_at_position_zero(): + inputs = _inputs() + invalid_start_mask = inputs.completion_mask.clone() + invalid_start_mask[0, 0] = True + invalid_inputs = StatelessForwardInputs( + input_ids=inputs.input_ids, + attention_mask=inputs.attention_mask, + completion_mask=invalid_start_mask, + ) + + with pytest.raises(ValueError, match=r"completion_mask\[:, 0\]"): + score_reference_logprobs(_logits_for(inputs), invalid_inputs) + + +def test_attention_backend_fallback_does_not_swallow_unrelated_kernel_errors(): + class KernelFailureModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(use_cache=True, _attn_implementation="eager") + self.calls = 0 + + def forward(self, input_ids, attention_mask=None, use_cache=None): + del input_ids, attention_mask, use_cache + self.calls += 1 + raise RuntimeError("CUDA kernels launch failed") + + model = KernelFailureModel() + executor = StatelessForwardExecutor(model, StatelessForwardConfig(mode="reference")) + + with pytest.raises(RuntimeError, match="CUDA kernels launch failed"): + executor.score(_inputs()) + + assert model.calls == 1 + assert model.config.use_cache is True + assert model.config._attn_implementation == "eager" + assert not hasattr(model.config, "attn_implementation") diff --git a/tests/test_stateless_hf_integration.py b/tests/test_stateless_hf_integration.py new file mode 100644 index 000000000..07ed3347b --- /dev/null +++ b/tests/test_stateless_hf_integration.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.executors.stateless_executor import ( + StatelessForwardConfig, + StatelessForwardExecutor, + StatelessForwardInputs, +) + +transformers = pytest.importorskip("transformers") +from transformers import ( # noqa: E402 + BertConfig, + BertForSequenceClassification, + GPT2Config, + GPT2LMHeadModel, +) + + +def _inputs(device: torch.device | str = "cpu") -> StatelessForwardInputs: + resolved_device = torch.device(device) + input_ids = torch.tensor( + [ + [0, 1, 2, 3, 4], + [0, 5, 6, 7, 0], + ], + dtype=torch.long, + device=resolved_device, + ) + attention_mask = torch.tensor( + [ + [True, True, True, True, True], + [True, True, True, True, False], + ], + device=resolved_device, + ) + completion_mask = torch.tensor( + [ + [False, False, True, True, True], + [False, False, True, True, False], + ], + device=resolved_device, + ) + return StatelessForwardInputs( + input_ids=input_ids, + attention_mask=attention_mask, + completion_mask=completion_mask, + ) + + +def _tiny_gpt2_reference_model() -> GPT2LMHeadModel: + config = GPT2Config( + vocab_size=32, + n_positions=16, + n_embd=16, + n_layer=1, + n_head=2, + bos_token_id=0, + eos_token_id=1, + use_cache=True, + ) + return GPT2LMHeadModel(config).eval() + + +def _tiny_bert_reward_model() -> BertForSequenceClassification: + config = BertConfig( + vocab_size=32, + hidden_size=16, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=32, + num_labels=1, + ) + return BertForSequenceClassification(config).eval() + + +def test_stateless_reference_scores_real_hf_causal_lm_without_kv_cache(): + torch.manual_seed(47) + model = _tiny_gpt2_reference_model() + executor = StatelessForwardExecutor(model, StatelessForwardConfig(mode="reference")) + + result = executor.score(_inputs()) + + assert result.reference_logps is not None + assert result.reference_logps.shape == (2, 5) + assert torch.equal( + result.reference_logps[~_inputs().completion_mask], + torch.zeros_like(result.reference_logps[~_inputs().completion_mask]), + ) + assert result.metrics["use_cache_passed"] is True + assert result.metrics["zero_kv_cache"] is True + assert result.metrics["kv_cache_output_tensors"] == 0 + assert result.metrics["attention_backend_fallback"] is True + assert result.metrics["attention_backend"] == "eager" + assert model.config.use_cache is True + + +def test_stateless_reward_scores_real_hf_sequence_classifier(): + torch.manual_seed(48) + model = _tiny_bert_reward_model() + executor = StatelessForwardExecutor(model, StatelessForwardConfig(mode="reward")) + + result = executor.score(_inputs()) + + assert result.rewards is not None + assert result.rewards.shape == (2,) + assert result.rewards.requires_grad is False + assert result.metrics["zero_kv_cache"] is True + assert result.metrics["attention_backend_fallback"] is True + assert result.metrics["attention_backend"] == "eager" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") +def test_stateless_reference_scores_real_hf_causal_lm_on_cuda_without_kv_cache(): + torch.manual_seed(49) + device = torch.device("cuda") + model = _tiny_gpt2_reference_model().to(device=device) + executor = StatelessForwardExecutor(model, StatelessForwardConfig(mode="reference")) + + result = executor.score(_inputs(device)) + + assert result.reference_logps is not None + assert result.reference_logps.is_cuda + assert result.metrics["zero_kv_cache"] is True + assert result.metrics["kv_cache_output_mb"] == 0.0 + assert result.metrics["attention_backend_fallback"] is True + assert result.metrics["attention_backend"] == "eager" + assert "peak_allocated_mb" in result.metrics diff --git a/tests/test_stateless_training_contract.py b/tests/test_stateless_training_contract.py new file mode 100644 index 000000000..f1b57d2da --- /dev/null +++ b/tests/test_stateless_training_contract.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import time + +import torch + +from rl_engine.executors.stateless_executor import StatelessForwardConfig, StatelessForwardExecutor +from rl_engine.executors.training_contract import ( + RolloutBatchMixin, + RolloutStageResult, + StatelessScoringWorker, + TorchRLTrainingConfig, + build_stateless_inputs_from_rollout_payload, + extract_rollout_reference_logp_groups, + extract_rollout_reward_groups, +) + + +class BatchOnlyWorker(RolloutBatchMixin): + def __init__(self, config: TorchRLTrainingConfig): + self.config = config + self.device = torch.device(config.device) + + +def _rollout(payload) -> RolloutStageResult: + return RolloutStageResult( + iteration=0, + weight_version=4, + payload=payload, + started_at=time.perf_counter(), + finished_at=time.perf_counter(), + ) + + +def test_build_stateless_inputs_from_rollout_payload_preserves_candidate_order_and_masks(): + payload = { + "normalized_outputs": [ + [{"token_ids": [1, 2, 3]}, {"token_ids": [4]}], + [{"token_ids": [5, 6]}], + ] + } + + inputs = build_stateless_inputs_from_rollout_payload( + payload, + prompt_len=2, + prompt_token_id=9, + max_completion_len=2, + ) + + assert inputs.input_ids.tolist() == [ + [9, 9, 1, 2], + [9, 9, 4, 9], + [9, 9, 5, 6], + ] + assert inputs.attention_mask.tolist() == [ + [True, True, True, True], + [True, True, True, False], + [True, True, True, True], + ] + assert inputs.completion_mask.tolist() == [ + [False, False, True, True], + [False, False, True, False], + [False, False, True, True], + ] + + +def test_build_stateless_inputs_preserves_payload_prompt_token_ids_when_available(): + payload = { + "normalized_outputs": [ + [{"prompt_token_ids": [7, 8, 9], "token_ids": [1, 2]}], + [{"prompt_token_ids": [5], "token_ids": [3]}], + ] + } + + inputs = build_stateless_inputs_from_rollout_payload( + payload, + prompt_len=None, + max_completion_len=2, + ) + + assert inputs.input_ids.tolist() == [ + [7, 8, 9, 1, 2], + [5, 0, 0, 3, 0], + ] + assert inputs.attention_mask.tolist() == [ + [True, True, True, True, True], + [True, False, False, True, False], + ] + assert inputs.completion_mask.tolist() == [ + [False, False, False, True, True], + [False, False, False, True, False], + ] + + +def test_stateless_scoring_worker_attaches_rewards_for_training_payload(): + class FakeRewardScorer(torch.nn.Module): + def __init__(self): + super().__init__() + self.use_cache_calls = [] + + def forward(self, input_ids, attention_mask=None, use_cache=None): + del attention_mask + self.use_cache_calls.append(use_cache) + rewards = torch.tensor([1.0, 3.0], device=input_ids.device) + return {"rewards": rewards} + + payload = { + "normalized_outputs": [ + [ + {"token_ids": [1, 2]}, + {"token_ids": [3, 4]}, + ] + ] + } + scorer_model = FakeRewardScorer() + scorer = StatelessScoringWorker( + StatelessForwardExecutor( + scorer_model, + StatelessForwardConfig(mode="reward"), + ), + lambda rollout: build_stateless_inputs_from_rollout_payload( + rollout.payload, + prompt_len=1, + device="cpu", + ), + ) + + scored = scorer.score(_rollout(payload)) + + assert scorer_model.use_cache_calls == [False] + assert scored.payload["stateless_scores"]["rewards"].tolist() == [1.0, 3.0] + assert extract_rollout_reward_groups(scored.payload) == [[1.0, 3.0]] + assert scored.metrics["scoring_backend"] == "stateless" + assert scored.metrics["scoring_mode"] == "reward" + assert scored.metrics["scoring_zero_kv_cache"] is True + assert scored.metrics["scoring_attention_backend"] == "flash_attention_2" + + worker = BatchOnlyWorker( + TorchRLTrainingConfig( + num_prompts=1, + samples_per_prompt=2, + prompt_len=1, + completion_len=2, + vocab_size=16, + valid_density=1.0, + seed=17, + ) + ) + batch, payload_metrics = worker._batch_from_rollout_or_synthetic(scored) + + assert payload_metrics["reward_source"] == "payload_rewards" + assert payload_metrics["rollout_prompt_groups"] == 1 + assert batch.rewards.tolist() == [1.0, 3.0] + assert torch.allclose( + batch.advantages, + torch.tensor([[-1.0, -1.0], [1.0, 1.0]], dtype=batch.advantages.dtype), + atol=1e-5, + ) + + +def test_stateless_scoring_worker_attaches_reference_logps_for_training_payload(): + class FakeReferenceRewardScorer(torch.nn.Module): + def __init__(self): + super().__init__() + self.use_cache_calls = [] + + def forward(self, input_ids, attention_mask=None, use_cache=None): + del attention_mask + self.use_cache_calls.append(use_cache) + logits = torch.full((*input_ids.shape, 16), -5.0, device=input_ids.device) + for row in range(input_ids.shape[0]): + for pos in range(input_ids.shape[1] - 1): + token = int(input_ids[row, pos + 1].item()) + logits[row, pos, token] = 5.0 + row + pos + return { + "logits": logits, + "rewards": torch.tensor([1.0, 3.0], device=input_ids.device), + } + + payload = { + "normalized_outputs": [ + [ + {"token_ids": [1, 2]}, + {"token_ids": [3, 4]}, + ] + ] + } + scorer_model = FakeReferenceRewardScorer() + scorer = StatelessScoringWorker( + StatelessForwardExecutor( + scorer_model, + StatelessForwardConfig(mode="both"), + ), + lambda rollout: build_stateless_inputs_from_rollout_payload( + rollout.payload, + prompt_len=1, + device="cpu", + ), + ) + + scored = scorer.score(_rollout(payload)) + + assert scorer_model.use_cache_calls == [False] + assert extract_rollout_reward_groups(scored.payload) == [[1.0, 3.0]] + assert scored.metrics["scoring_zero_kv_cache"] is True + reference_groups = extract_rollout_reference_logp_groups(scored.payload) + assert len(reference_groups) == 1 + assert [len(row) for row in reference_groups[0]] == [2, 2] + assert scored.payload["normalized_outputs"][0][0]["reference_logp_source"] == ( + "stateless_executor" + ) + + worker = BatchOnlyWorker( + TorchRLTrainingConfig( + num_prompts=1, + samples_per_prompt=2, + prompt_len=1, + completion_len=2, + vocab_size=16, + valid_density=1.0, + seed=19, + ) + ) + batch, payload_metrics = worker._batch_from_rollout_or_synthetic(scored) + + assert payload_metrics["reward_source"] == "payload_rewards" + assert payload_metrics["reference_logp_source"] == "payload_reference_logps" + assert torch.allclose( + batch.ref_logps, + torch.tensor(reference_groups[0], dtype=batch.ref_logps.dtype), + ) + + +def test_training_contract_accepts_reference_logps_for_truncated_completions(): + payload = { + "normalized_outputs": [ + [ + { + "token_ids": [1, 2, 3], + "reward": 1.0, + "reference_logps": [-0.1, -0.2], + } + ] + ] + } + worker = BatchOnlyWorker( + TorchRLTrainingConfig( + num_prompts=1, + samples_per_prompt=1, + prompt_len=1, + completion_len=2, + vocab_size=16, + valid_density=1.0, + ) + ) + + batch, payload_metrics = worker._batch_from_rollout_or_synthetic(_rollout(payload)) + + assert payload_metrics["reference_logp_source"] == "payload_reference_logps" + assert batch.token_ids.tolist() == [[1, 2]] + assert torch.allclose( + batch.ref_logps, + torch.tensor([[-0.1, -0.2]], dtype=batch.ref_logps.dtype), + )