From 7d140c894aa13e5ff29a167e6e621f6ebcbd5156 Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Tue, 11 Aug 2026 23:37:24 -0700 Subject: [PATCH 1/3] Add first-class Parquet prompt OPD runtime --- PROJECT_STATE.md | 23 ++ pyproject.toml | 1 + src/miniverl/cli.py | 37 +- src/miniverl/config/__init__.py | 8 + src/miniverl/config/models.py | 105 +++++- src/miniverl/data/__init__.py | 17 + src/miniverl/data/verl_parquet.py | 288 ++++++++++++++++ src/miniverl/evaluation/teacher_gate.py | 13 +- src/miniverl/models/adapter_io.py | 33 +- src/miniverl/models/base.py | 32 ++ src/miniverl/models/factory.py | 6 +- src/miniverl/models/hf.py | 61 +++- src/miniverl/models/sampling.py | 63 +++- src/miniverl/models/toy.py | 64 +++- src/miniverl/runtime/roles.py | 2 +- src/miniverl/runtime/rollout.py | 310 +++++++++++++++++ src/miniverl/training/trainer.py | 316 ++++++++++++++++-- tests/integration/test_prompt_opd_pipeline.py | 130 +++++++ tests/unit/test_packaging.py | 1 + tests/unit/test_prompt_rollout_runtime.py | 181 ++++++++++ tests/unit/test_prompt_source_config.py | 76 +++++ tests/unit/test_verl_parquet_source.py | 153 +++++++++ 22 files changed, 1857 insertions(+), 63 deletions(-) create mode 100644 src/miniverl/data/__init__.py create mode 100644 src/miniverl/data/verl_parquet.py create mode 100644 src/miniverl/runtime/rollout.py create mode 100644 tests/integration/test_prompt_opd_pipeline.py create mode 100644 tests/unit/test_prompt_rollout_runtime.py create mode 100644 tests/unit/test_prompt_source_config.py create mode 100644 tests/unit/test_verl_parquet_source.py diff --git a/PROJECT_STATE.md b/PROJECT_STATE.md index c0749b0..a7e4d2d 100644 --- a/PROJECT_STATE.md +++ b/PROJECT_STATE.md @@ -23,6 +23,29 @@ distributed dimensions fail closed. `bridge compile-opd` is deliberately an offline config compiler: it loads no weights and does not claim the prompt runtime or verl engine equivalence implemented by later PRs. +## v0.8.0 single-GPU verl OPD pivot — PR B + +The trainer now has a discriminated `environment` / `verl_parquet` source +contract. Parquet rows stream by record batch, reject every invalid row, retain +the supported verl metadata and record schema, content, row, tokenizer and +rendered-prompt digests. Pure prompt OPD needs neither a `ToolEnvironment` nor +a reward model; task rewards, oracle SFT and privileged teacher context fail +closed instead of being ignored. + +`PromptDatasetRolloutRuntime` applies the actor chat template once, performs +real masked padded greedy generation, restores logical row order and creates +trajectories whose selected model spans contain response tokens only. A +bounded padded-token budget is enforced before allocation. CUDA OOM retry may +split physical batches but never changes logical batch size, seeds or training +configuration. The existing multi-turn environment runner is retained behind +`ToolEnvironmentRolloutRuntime` and its prior integration suite remains green. + +Focused validation on 2026-08-11: 17 new config/data/runtime/end-to-end tests, +94 existing trainer/toy-pipeline tests, Ruff, and mypy over 126 source files. +The pure-OPD integration executes two Parquet prompts through rollout, teacher +scoring, a padded actor update, checkpoint and manifest without an environment +or reward. + ## v0.7.1 Product correction — RELEASE CANDIDATE Branch `v0.7.1-product-correction` starts from synchronized main diff --git a/pyproject.toml b/pyproject.toml index 64c0756..6beff07 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ train = [ "peft>=0.12", "accelerate>=0.33", "numpy>=1.24", + "pyarrow>=15", ] cuda = ["bitsandbytes>=0.43; platform_system != 'Darwin'"] dpo = [ diff --git a/src/miniverl/cli.py b/src/miniverl/cli.py index 8e18aa3..6503430 100644 --- a/src/miniverl/cli.py +++ b/src/miniverl/cli.py @@ -241,18 +241,20 @@ def validate( return warnings: list[str] = [] - try: - environment = make_environment(config.environment.name, **config.environment.params) - if config.models.teacher.mode.value == "privileged_context" and not hasattr( - environment, "privileged_context" - ): - warnings.append("environment provides no privileged context") - except MiniVerlError as exc: - if as_json: - _emit_json({"valid": False, "path": str(recipe), "errors": [exc.message]}) - raise typer.Exit(1) from None - _fail(exc) - return + environment_config = config.environment + if environment_config is not None: + try: + environment = make_environment(environment_config.name, **environment_config.params) + if config.models.teacher.mode.value == "privileged_context" and not hasattr( + environment, "privileged_context" + ): + warnings.append("environment provides no privileged context") + except MiniVerlError as exc: + if as_json: + _emit_json({"valid": False, "path": str(recipe), "errors": [exc.message]}) + raise typer.Exit(1) from None + _fail(exc) + return steps_per_cycle = max( 1, @@ -281,8 +283,9 @@ def validate( "backend": config.models.backend.value, "student": config.models.student.model_id, "teacher": config.models.teacher.model_id, - "environment": config.environment.name, - "difficulty": config.environment.difficulty, + "source_kind": config.source.kind.value, + "environment": environment_config.name if environment_config is not None else None, + "difficulty": environment_config.difficulty if environment_config is not None else None, "objective": ( "sft_cross_entropy" if config.run.mode.value == "sft" @@ -302,7 +305,9 @@ def validate( "optimizer_steps_per_cycle": steps_per_cycle, "planned_optimizer_steps": steps_per_cycle * (config.train.cycles + config.train.sft_warmup_cycles), - "eval_tasks": config.effective_eval_tasks, + "eval_tasks": ( + config.effective_eval_tasks if environment_config is not None else config.eval.tasks + ), "seed": config.run.seed, "warnings": warnings, } @@ -594,7 +599,7 @@ def qualify_teacher_command( try: config = RunConfig.from_yaml(recipe) - if config.environment.name != "sqlite_recovery": + if config.require_environment("qualify-teacher").name != "sqlite_recovery": raise ConfigError("qualify-teacher requires environment.name=sqlite_recovery") _require_training_stack("miniverl qualify-teacher") from miniverl.evaluation.teacher_gate import evaluate_teacher_candidate diff --git a/src/miniverl/config/__init__.py b/src/miniverl/config/__init__.py index fc05e49..72a5937 100644 --- a/src/miniverl/config/__init__.py +++ b/src/miniverl/config/__init__.py @@ -9,6 +9,7 @@ CacheConfig, Divergence, EnvironmentConfig, + EnvironmentSourceConfig, EvalConfig, GateConfig, GateSignal, @@ -25,6 +26,7 @@ OPDFreshness, OptimizerName, Precision, + PromptTruncation, Quantization, ReferenceModelConfig, ReportConfig, @@ -33,6 +35,7 @@ RunMeta, SelectionConfig, SelectorName, + SourceKind, StudentModelConfig, TeacherAdapterConfig, TeacherContextMode, @@ -41,6 +44,7 @@ ToyModelConfig, TrainConfig, TrainingMode, + VerlParquetSourceConfig, ) __all__ = [ @@ -50,6 +54,7 @@ "CacheConfig", "Divergence", "EnvironmentConfig", + "EnvironmentSourceConfig", "EvalConfig", "GateConfig", "GateSignal", @@ -66,6 +71,7 @@ "OPDFreshness", "OptimizerName", "Precision", + "PromptTruncation", "Quantization", "ReferenceModelConfig", "ReportConfig", @@ -75,6 +81,7 @@ "SelectionConfig", "SelectorName", "StudentModelConfig", + "SourceKind", "TeacherContextMode", "TeacherMode", "TeacherAdapterConfig", @@ -82,4 +89,5 @@ "ToyModelConfig", "TrainConfig", "TrainingMode", + "VerlParquetSourceConfig", ] diff --git a/src/miniverl/config/models.py b/src/miniverl/config/models.py index 9a15e6a..03c7740 100644 --- a/src/miniverl/config/models.py +++ b/src/miniverl/config/models.py @@ -14,7 +14,7 @@ from enum import Enum from pathlib import Path -from typing import Any, Literal +from typing import Annotated, Any, Literal import yaml from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator, model_validator @@ -60,6 +60,10 @@ "LossConfig", "SelectionConfig", "RolloutConfig", + "SourceKind", + "PromptTruncation", + "EnvironmentSourceConfig", + "VerlParquetSourceConfig", "EnvironmentConfig", "TrainConfig", "MemoryConfig", @@ -491,6 +495,45 @@ class RolloutConfig(_Base): # first parse error. max_parse_errors: int = Field(default=2, ge=0, le=32) max_repeated_calls: int = Field(default=2, ge=1, le=32) + prompt_batch_size: int = Field(default=1, ge=1, le=1024) + max_padded_tokens: int = Field(default=4096, ge=16, le=1048576) + + +class SourceKind(str, Enum): + """Where rollout inputs originate.""" + + ENVIRONMENT = "environment" + VERL_PARQUET = "verl_parquet" + + +class PromptTruncation(str, Enum): + """Explicit overlength handling for prompt datasets.""" + + ERROR = "error" + LEFT = "left" + RIGHT = "right" + + +class EnvironmentSourceConfig(_Base): + """Discriminator for the backward-compatible registered-environment path.""" + + kind: Literal[SourceKind.ENVIRONMENT] = SourceKind.ENVIRONMENT + + +class VerlParquetSourceConfig(_Base): + """Bounded, deterministic verl-style Parquet prompt input.""" + + kind: Literal[SourceKind.VERL_PARQUET] = SourceKind.VERL_PARQUET + train_files: list[str] = Field(min_length=1) + val_files: list[str] = Field(default_factory=list) + prompt_key: str = Field(default="prompt", min_length=1) + allow_plain_string_prompts: bool = False + use_task_rewards: bool = False + max_prompt_length: int = Field(default=512, ge=1, le=131072) + truncation: PromptTruncation = PromptTruncation.ERROR + shuffle: bool = True + seed: int = Field(default=0, ge=0) + row_batch_size: int = Field(default=256, ge=1, le=65536) class EnvironmentConfig(_Base): @@ -681,7 +724,11 @@ class RunConfig(_Base): schema_version: int = CONFIG_SCHEMA_VERSION run: RunMeta = Field(default_factory=RunMeta) models: ModelsConfig - environment: EnvironmentConfig + source: Annotated[ + EnvironmentSourceConfig | VerlParquetSourceConfig, + Field(discriminator="kind"), + ] = Field(default_factory=EnvironmentSourceConfig) + environment: EnvironmentConfig | None = None rollout: RolloutConfig = Field(default_factory=RolloutConfig) selection: SelectionConfig = Field(default_factory=SelectionConfig) loss: LossConfig = Field(default_factory=LossConfig) @@ -725,6 +772,30 @@ def _validate_combination(self) -> RunConfig: ) mode = self.run.mode + if self.source.kind is SourceKind.ENVIRONMENT and self.environment is None: + raise ValueError("source.kind=environment requires an environment configuration") + if self.source.kind is SourceKind.VERL_PARQUET and self.environment is not None: + raise ValueError("source.kind=verl_parquet must not define environment") + if self.source.kind is SourceKind.VERL_PARQUET: + if mode is not TrainingMode.OPD: + raise ValueError( + "source.kind=verl_parquet supports pure OPD in v0.8; it has no oracle " + "labels for sft or offline_kd" + ) + if self.train.sft_warmup_cycles: + raise ValueError( + "source.kind=verl_parquet cannot run sft_warmup_cycles without oracle labels" + ) + if self.models.teacher.mode is not TeacherContextMode.STANDARD: + raise ValueError( + "source.kind=verl_parquet requires the actor and teacher to score the " + "same rendered prompt; privileged_context is outside the v0.8 profile" + ) + if self.source.use_task_rewards: + raise ValueError( + "source.use_task_rewards=true is not implemented for the pure OPD prompt " + "runtime; refusing to ignore row reward_model metadata" + ) if mode is TrainingMode.OPD and self.cache.reuse_across_policy_versions: raise ValueError( "cache.reuse_across_policy_versions=true contradicts run.mode=opd: " @@ -829,7 +900,7 @@ def _validate_combination(self) -> RunConfig: alignment = self.alignment - if self.eval.enabled: + if self.eval.enabled and self.environment is not None: split_sizes = { "train": self.environment.train_tasks, "eval": self.environment.eval_tasks, @@ -917,8 +988,22 @@ def effective_eval_tasks(self) -> int: """Number of evaluation tasks after applying the eval override.""" if self.eval.tasks is not None: return self.eval.tasks + if self.environment is None: + raise ConfigError( + "eval.tasks is required when source.kind=verl_parquet", + hint="set eval.tasks to a bounded number no larger than the validation rows", + ) return self.environment.eval_tasks + def require_environment(self, operation: str = "this operation") -> EnvironmentConfig: + """Return the registered environment or fail with a source-aware message.""" + if self.environment is None: + raise ConfigError( + f"{operation} requires source.kind=environment", + hint="the verl_parquet source uses prompt rollout and has no ToolEnvironment", + ) + return self.environment + @classmethod def from_yaml(cls, path: str | Path) -> RunConfig: """Load and validate a recipe from a YAML file.""" @@ -988,6 +1073,20 @@ def resolved_for_runtime(self) -> RunConfig: source = private.get("_source_path") if isinstance(private, dict) else None base = source.parent if isinstance(source, Path) else Path.cwd() runtime.offline_kd.dataset_path = str((base / path).resolve()) + if isinstance(runtime.source, VerlParquetSourceConfig): + private = getattr(self, "__pydantic_private__", None) + source_path = private.get("_source_path") if isinstance(private, dict) else None + base = source_path.parent if isinstance(source_path, Path) else Path.cwd() + + def resolve_files(paths: list[str]) -> list[str]: + return [ + str(path if path.is_absolute() else (base / path).resolve()) + for raw in paths + for path in [Path(raw)] + ] + + runtime.source.train_files = resolve_files(runtime.source.train_files) + runtime.source.val_files = resolve_files(runtime.source.val_files) return runtime def to_yaml(self) -> str: diff --git a/src/miniverl/data/__init__.py b/src/miniverl/data/__init__.py new file mode 100644 index 0000000..cae73a0 --- /dev/null +++ b/src/miniverl/data/__init__.py @@ -0,0 +1,17 @@ +"""Typed, source-agnostic training data inputs.""" + +from miniverl.data.verl_parquet import ( + PromptDatasetManifest, + PromptRecord, + RenderedPrompt, + VerlParquetDataset, + render_prompt, +) + +__all__ = [ + "PromptDatasetManifest", + "PromptRecord", + "RenderedPrompt", + "VerlParquetDataset", + "render_prompt", +] diff --git a/src/miniverl/data/verl_parquet.py b/src/miniverl/data/verl_parquet.py new file mode 100644 index 0000000..47f2992 --- /dev/null +++ b/src/miniverl/data/verl_parquet.py @@ -0,0 +1,288 @@ +"""Bounded reader and prompt renderer for the supported verl Parquet subset.""" + +from __future__ import annotations + +import hashlib +import json +import random +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +from miniverl.config.models import PromptTruncation, VerlParquetSourceConfig +from miniverl.errors import ConfigError, MissingDependencyError + +SplitName = Literal["train", "val"] +_PRESERVED_FIELDS = ("data_source", "ability", "reward_model", "extra_info") + + +def _canonical(value: Any) -> str: + try: + return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + except (TypeError, ValueError) as exc: + raise ConfigError( + f"Parquet row contains metadata that is not canonical JSON: {exc}", + hint="use JSON-compatible scalars, lists and mappings in preserved metadata fields", + ) from exc + + +def _digest(value: Any) -> str: + return hashlib.sha256(_canonical(value).encode("utf-8")).hexdigest() + + +def _parquet() -> Any: + try: + import pyarrow.parquet as pq + except ImportError as exc: # pragma: no cover - dependency gate + raise MissingDependencyError("pyarrow", "train", "verl Parquet prompt loading") from exc + return pq + + +@dataclass(frozen=True) +class PromptRecord: + """One validated prompt row with its complete supported provenance.""" + + prompt: str | list[dict[str, str]] + data_source: Any + ability: Any + reward_model: Any + extra_info: Any + source_file: str + source_row_index: int + row_digest: str + canonical_payload: str + + +@dataclass(frozen=True) +class PromptDatasetManifest: + """Content and schema identity obtained through a bounded scan.""" + + rows: dict[str, int] + schema_digest: str + content_digest: str + files: tuple[str, ...] + + +@dataclass(frozen=True) +class RenderedPrompt: + """The single actor-rendered prompt shared by actor and teacher scoring.""" + + record: PromptRecord + text: str + token_ids: tuple[int, ...] + tokenizer_identity: dict[str, Any] + rendered_prompt_digest: str + prompt_token_count: int + truncation_decision: str + original_prompt_token_count: int + + +class VerlParquetDataset: + """Read prompt rows by record batch; never load a complete table.""" + + def __init__(self, config: VerlParquetSourceConfig) -> None: + self.config = config + + def _files(self, split: SplitName) -> list[Path]: + raw = self.config.train_files if split == "train" else self.config.val_files + paths = [Path(item).resolve() for item in raw] + missing = [str(path) for path in paths if not path.is_file()] + if missing: + raise ConfigError( + f"{split} Parquet file not found: {missing[0]}", + hint="resolve recipe-relative paths before constructing the dataset", + ) + return paths + + def _records(self, split: SplitName) -> Iterator[PromptRecord]: + pq = _parquet() + prompt_key = self.config.prompt_key + for path in self._files(split): + try: + parquet = pq.ParquetFile(path) + except Exception as exc: + raise ConfigError(f"cannot open Parquet file {path}: {exc}") from exc + if prompt_key not in parquet.schema_arrow.names: + raise ConfigError( + f"Parquet file {path} has no prompt field {prompt_key!r}", + hint=f"available columns: {', '.join(parquet.schema_arrow.names)}", + ) + source_index = 0 + columns = [ + prompt_key, + *[name for name in _PRESERVED_FIELDS if name in parquet.schema_arrow.names], + ] + try: + batches = parquet.iter_batches( + batch_size=self.config.row_batch_size, columns=columns + ) + for batch in batches: + for row in batch.to_pylist(): + yield self._validate_row(row, path=path, source_index=source_index) + source_index += 1 + except ConfigError: + raise + except Exception as exc: + raise ConfigError( + f"failed reading {path} near row {source_index}: {exc}", + hint="no rows were silently skipped", + ) from exc + + def _validate_row(self, row: dict[str, Any], *, path: Path, source_index: int) -> PromptRecord: + prompt = row.get(self.config.prompt_key) + location = f"{path} row {source_index}" + if isinstance(prompt, str): + if not self.config.allow_plain_string_prompts: + raise ConfigError( + f"{location} contains a plain-string prompt without explicit opt-in", + hint="set source.allow_plain_string_prompts=true or store a chat message list", + ) + if not prompt: + raise ConfigError(f"{location} contains an empty prompt") + validated: str | list[dict[str, str]] = prompt + elif isinstance(prompt, list) and prompt: + messages: list[dict[str, str]] = [] + for index, message in enumerate(prompt): + if not isinstance(message, dict): + raise ConfigError(f"{location} prompt message {index} is not a mapping") + role = message.get("role") + content = message.get("content") + if not isinstance(role, str) or not role or not isinstance(content, str): + raise ConfigError( + f"{location} prompt message {index} requires string role and content" + ) + messages.append({"role": role, "content": content}) + validated = messages + else: + raise ConfigError( + f"{location} field {self.config.prompt_key!r} must be a non-empty message list" + " or an explicitly enabled plain string" + ) + preserved = {name: row.get(name) for name in _PRESERVED_FIELDS} + if self.config.use_task_rewards and preserved["reward_model"] is None: + raise ConfigError( + f"{location} has no reward_model but source.use_task_rewards=true", + hint="provide reward_model per row or disable task rewards for pure OPD", + ) + payload = {"prompt": validated, **preserved} + canonical = _canonical(payload) + return PromptRecord( + prompt=validated, + source_file=str(path), + source_row_index=source_index, + row_digest=hashlib.sha256(canonical.encode("utf-8")).hexdigest(), + canonical_payload=canonical, + **preserved, + ) + + def iter_split(self, split: SplitName, *, epoch: int = 0) -> Iterator[PromptRecord]: + """Yield all rows, optionally with a deterministic bounded-buffer shuffle.""" + records = self._records(split) + if not self.config.shuffle: + yield from records + return + rng = random.Random(self.config.seed ^ epoch ^ (0x56414C if split == "val" else 0)) + buffer: list[PromptRecord] = [] + for record in records: + if len(buffer) < self.config.row_batch_size: + buffer.append(record) + continue + index = rng.randrange(len(buffer)) + yield buffer[index] + buffer[index] = record + rng.shuffle(buffer) + yield from buffer + + def inspect(self) -> PromptDatasetManifest: + """Scan schemas and canonical row digests without retaining row payloads.""" + pq = _parquet() + schema_items: list[dict[str, str]] = [] + content = hashlib.sha256() + rows: dict[str, int] = {"train": 0, "val": 0} + files: list[str] = [] + for split in ("train", "val"): + typed_split: SplitName = split + for path in self._files(typed_split): + files.append(str(path)) + parquet = pq.ParquetFile(path) + schema_items.append({"split": split, "schema": str(parquet.schema_arrow)}) + for record in self._records(typed_split): + rows[split] += 1 + content.update(split.encode("ascii")) + content.update(b"\0") + content.update(record.row_digest.encode("ascii")) + content.update(b"\n") + return PromptDatasetManifest( + rows=rows, + schema_digest=_digest(schema_items), + content_digest=content.hexdigest(), + files=tuple(files), + ) + + +def _apply_chat_template(tokenizer: Any, messages: list[dict[str, str]]) -> str: + target = ( + tokenizer + if callable(getattr(tokenizer, "apply_chat_template", None)) + else getattr(tokenizer, "_tok", None) + ) + method = getattr(target, "apply_chat_template", None) + if not callable(method): + raise ConfigError( + "the actor tokenizer has no chat template for message-list prompts", + hint="use a tokenizer with apply_chat_template or explicitly supply plain strings", + ) + try: + rendered = method(messages, tokenize=False, add_generation_prompt=True) + except TypeError: + rendered = method(messages) + if not isinstance(rendered, str) or not rendered: + raise ConfigError("the actor tokenizer chat template produced no text") + return rendered + + +def render_prompt( + record: PromptRecord, + tokenizer: Any, + config: VerlParquetSourceConfig, +) -> RenderedPrompt: + """Render exactly once, then enforce the configured token bound exactly.""" + text = ( + _apply_chat_template(tokenizer, record.prompt) + if isinstance(record.prompt, list) + else record.prompt + ) + token_ids = list(tokenizer.encode(text)) + original_count = len(token_ids) + decision = "not_needed" + if original_count > config.max_prompt_length: + if config.truncation is PromptTruncation.ERROR: + raise ConfigError( + f"prompt {record.row_digest[:12]} has {original_count} tokens; " + f"max_prompt_length={config.max_prompt_length}", + hint="raise the bound or explicitly select source.truncation=left/right", + ) + if config.truncation is PromptTruncation.LEFT: + token_ids = token_ids[-config.max_prompt_length :] + else: + token_ids = token_ids[: config.max_prompt_length] + text = tokenizer.decode(token_ids) + decision = f"truncated_{config.truncation.value}" + identity = dict( + getattr( + tokenizer, + "identity", + {"behavioral_fingerprint_v1": getattr(tokenizer, "fingerprint", "unknown")}, + ) + ) + return RenderedPrompt( + record=record, + text=text, + token_ids=tuple(int(token) for token in token_ids), + tokenizer_identity=identity, + rendered_prompt_digest=hashlib.sha256(text.encode("utf-8")).hexdigest(), + prompt_token_count=len(token_ids), + truncation_decision=decision, + original_prompt_token_count=original_count, + ) diff --git a/src/miniverl/evaluation/teacher_gate.py b/src/miniverl/evaluation/teacher_gate.py index 4c2acae..473f71a 100644 --- a/src/miniverl/evaluation/teacher_gate.py +++ b/src/miniverl/evaluation/teacher_gate.py @@ -72,25 +72,26 @@ def evaluate_teacher_candidate( ) -> dict[str, Any]: """Evaluate one configured frozen teacher as a tool policy on eval only.""" validate_gate_split(split) - count = tasks or config.environment.eval_tasks - if count < 1 or count > config.environment.eval_tasks: + environment_config = config.require_environment("teacher qualification") + count = tasks or environment_config.eval_tasks + if count < 1 or count > environment_config.eval_tasks: raise ConfigError( - f"teacher qualification tasks must be within 1..{config.environment.eval_tasks}" + f"teacher qualification tasks must be within 1..{environment_config.eval_tasks}" ) destination = Path(out).resolve() destination.mkdir(parents=True, exist_ok=False) trajectories_path = destination / "trajectories.jsonl" task_results_path = destination / "task-results.jsonl" - environment = make_environment(config.environment.name, **config.environment.params) + environment = make_environment(environment_config.name, **environment_config.params) teacher = None started = time.perf_counter() try: task_splits = make_splits( environment, counts={"train": 0, "eval": count, "test": 0}, - seed=config.environment.split_seed, - difficulty=config.environment.difficulty, + seed=environment_config.split_seed, + difficulty=environment_config.difficulty, ) from miniverl.models.factory import build_teacher, build_tokenizer, resolve_device diff --git a/src/miniverl/models/adapter_io.py b/src/miniverl/models/adapter_io.py index e5e270b..1469201 100644 --- a/src/miniverl/models/adapter_io.py +++ b/src/miniverl/models/adapter_io.py @@ -474,7 +474,12 @@ def export_adapter( } env = collect_environment() policy_evaluation = None - protocol_version = str(config.environment.params.get("protocol_version", "v1")) + environment_config = config.environment + protocol_version = str( + (environment_config.params if environment_config is not None else {}).get( + "protocol_version", "v1" + ) + ) if paths.eval_json.is_file(): summary = read_json(paths.eval_json) final_eval = summary.get("eval") if isinstance(summary, dict) else None @@ -493,13 +498,25 @@ def export_adapter( "source_checkpoint_digest": digest_tree(checkpoint_path), "lora": config.models.student.lora.model_dump(mode="json"), "training_environment": env, - "training_task": { - "environment": config.environment.name, - "difficulty": config.environment.difficulty, - "protocol": f"miniverl_tool_protocol_{protocol_version}", - "protocol_version": protocol_version, - "mode": config.run.mode.value, - }, + "training_task": ( + { + "source_kind": "environment", + "environment": environment_config.name, + "difficulty": environment_config.difficulty, + "protocol": f"miniverl_tool_protocol_{protocol_version}", + "protocol_version": protocol_version, + "mode": config.run.mode.value, + } + if environment_config is not None + else { + "source_kind": "verl_parquet", + "environment": None, + "difficulty": None, + "protocol": None, + "protocol_version": None, + "mode": config.run.mode.value, + } + ), "policy_evaluation": policy_evaluation, "checksums": checksums, } diff --git a/src/miniverl/models/base.py b/src/miniverl/models/base.py index dc05e05..8678195 100644 --- a/src/miniverl/models/base.py +++ b/src/miniverl/models/base.py @@ -102,6 +102,38 @@ def generate( """Sample a continuation, stopping on EOS, a stop string, or the budget.""" ... + def generate_batch( + self, + prefix_token_ids: Sequence[Sequence[int]], + *, + max_new_tokens: int, + stop_sequences: Sequence[str] = (), + temperature: float = 0.0, + top_p: float = 1.0, + top_k: int = 0, + seeds: Sequence[int | None] | None = None, + ) -> list[GenerationOutput]: + """Compatibility batch API; concrete local backends use one padded forward. + + Stochastic sampling deliberately falls back to the established per-example + generator so seed semantics remain unchanged. + """ + chosen_seeds = list(seeds or [None] * len(prefix_token_ids)) + if len(chosen_seeds) != len(prefix_token_ids): + raise ValueError("generate_batch needs exactly one seed per prompt") + return [ + self.generate( + prefix, + max_new_tokens=max_new_tokens, + stop_sequences=stop_sequences, + temperature=temperature, + top_p=top_p, + top_k=top_k, + seed=seed, + ) + for prefix, seed in zip(prefix_token_ids, chosen_seeds, strict=True) + ] + # -- scoring -------------------------------------------------------- @abstractmethod diff --git a/src/miniverl/models/factory.py b/src/miniverl/models/factory.py index b90f84c..a7f611f 100644 --- a/src/miniverl/models/factory.py +++ b/src/miniverl/models/factory.py @@ -182,7 +182,11 @@ def build_teacher( tokenizer=tokenizer, trainable=False, local_files_only=local_files_only, - protocol_version=str(config.environment.params.get("protocol_version", "v1")), + protocol_version=str( + (config.environment.params if config.environment is not None else {}).get( + "protocol_version", "v1" + ) + ), ) diff --git a/src/miniverl/models/hf.py b/src/miniverl/models/hf.py index 1b528bb..82fb106 100644 --- a/src/miniverl/models/hf.py +++ b/src/miniverl/models/hf.py @@ -29,7 +29,7 @@ from miniverl.errors import BackendError, MissingDependencyError from miniverl.models.adapters import ArchitectureAdapter from miniverl.models.base import BackendCapabilities, CausalLMBackend, GenerationOutput -from miniverl.models.sampling import run_generation +from miniverl.models.sampling import run_generation, run_greedy_padded_generation from miniverl.utils.lazy import have_module, require_peft, require_transformers __all__ = ["HFBackend", "resolve_dtype", "supports_bfloat16"] @@ -417,6 +417,65 @@ def step(new_ids: list[int], state: Any) -> tuple[torch.Tensor, Any]: if was_training: self.model.train() + def generate_batch( + self, + prefix_token_ids: Sequence[Sequence[int]], + *, + max_new_tokens: int, + stop_sequences: Sequence[str] = (), + temperature: float = 0.0, + top_p: float = 1.0, + top_k: int = 0, + seeds: Sequence[int | None] | None = None, + ) -> list[GenerationOutput]: + """Use one masked padded forward per greedy decoding step.""" + if temperature != 0.0: + return super().generate_batch( + prefix_token_ids, + max_new_tokens=max_new_tokens, + stop_sequences=stop_sequences, + temperature=temperature, + top_p=top_p, + top_k=top_k, + seeds=seeds, + ) + was_training = self.model.training + self.model.eval() + + def step(rows: list[list[int]]) -> torch.Tensor: + lengths = [len(row) for row in rows] + width = max(lengths) + ids = torch.full( + (len(rows), width), + self.tokenizer.pad_token_id, + dtype=torch.long, + device=self._device, + ) + mask = torch.zeros((len(rows), width), dtype=torch.bool, device=self._device) + for index, row in enumerate(rows): + ids[index, : len(row)] = torch.tensor(row, dtype=torch.long, device=self._device) + mask[index, : len(row)] = True + with torch.no_grad(): + hidden, _ = self._backbone_forward(ids, attention_mask=mask, use_cache=False) + positions = torch.tensor( + [length - 1 for length in lengths], dtype=torch.long, device=self._device + ) + selected = hidden[torch.arange(len(rows), device=self._device), positions] + return self.adapter.lm_head(selected) + + try: + return run_greedy_padded_generation( + step=step, + prefix_token_ids=prefix_token_ids, + decode=self.tokenizer.decode, + eos_token_id=self.tokenizer.eos_token_id, + max_new_tokens=max_new_tokens, + stop_sequences=stop_sequences, + ) + finally: + if was_training: + self.model.train() + # -- scoring ------------------------------------------------------------ def hidden_states_at( diff --git a/src/miniverl/models/sampling.py b/src/miniverl/models/sampling.py index e6d3351..5de0e69 100644 --- a/src/miniverl/models/sampling.py +++ b/src/miniverl/models/sampling.py @@ -14,7 +14,7 @@ from miniverl.models.base import GenerationOutput -__all__ = ["sample_from_logits", "run_generation", "StepFn"] +__all__ = ["sample_from_logits", "run_generation", "run_greedy_padded_generation", "StepFn"] #: ``step(new_token_ids, state) -> (next_token_logits [V], new_state)`` StepFn = Callable[[list[int], Any], "tuple[torch.Tensor, Any]"] @@ -125,3 +125,64 @@ def run_generation( matched_stop=matched_stop, logprobs=logprobs, ) + + +def run_greedy_padded_generation( + *, + step: Callable[[list[list[int]]], torch.Tensor], + prefix_token_ids: Sequence[Sequence[int]], + decode: Callable[[Sequence[int]], str], + eos_token_id: int, + max_new_tokens: int, + stop_sequences: Sequence[str] = (), +) -> list[GenerationOutput]: + """Greedy decode a real padded batch while retaining per-row stop state. + + ``step`` receives compact unpadded rows and performs the one padded model + forward. Re-padding on every step keeps each row's next-token state at its + true sequence end; padding can never become context or a selected token. + """ + if not prefix_token_ids: + return [] + if any(not row for row in prefix_token_ids): + raise ValueError("padded generation cannot contain an empty prompt") + if max_new_tokens < 1: + raise ValueError(f"max_new_tokens must be >= 1, got {max_new_tokens}") + sequences = [list(row) for row in prefix_token_ids] + generated: list[list[int]] = [[] for _ in sequences] + reasons = ["max_new_tokens"] * len(sequences) + matched: list[str | None] = [None] * len(sequences) + active = [True] * len(sequences) + for _ in range(max_new_tokens): + logits = step(sequences) + if tuple(logits.shape[:1]) != (len(sequences),): + raise ValueError( + f"padded generation step returned {tuple(logits.shape)}, expected [batch, vocab]" + ) + for index in range(len(sequences)): + if not active[index]: + continue + token = int(torch.argmax(logits[index].detach().to(torch.float32)).item()) + generated[index].append(token) + sequences[index].append(token) + if token == eos_token_id: + reasons[index] = "eos" + active[index] = False + continue + text = decode(generated[index]) + hit = next((value for value in stop_sequences if value and value in text), None) + if hit is not None: + reasons[index] = "stop_sequence" + matched[index] = hit + active[index] = False + if not any(active): + break + return [ + GenerationOutput( + token_ids=row, + text=decode(row), + stop_reason=reasons[index], + matched_stop=matched[index], + ) + for index, row in enumerate(generated) + ] diff --git a/src/miniverl/models/toy.py b/src/miniverl/models/toy.py index 885e4a4..4e6ad0c 100644 --- a/src/miniverl/models/toy.py +++ b/src/miniverl/models/toy.py @@ -23,7 +23,7 @@ from miniverl.errors import BackendError from miniverl.models.base import BackendCapabilities, CausalLMBackend, GenerationOutput -from miniverl.models.sampling import run_generation +from miniverl.models.sampling import run_generation, run_greedy_padded_generation from miniverl.models.tokenizers import ToyTokenizer __all__ = ["ToyCausalLM", "ToyBackend", "fit_toy_model"] @@ -333,6 +333,68 @@ def step(new_ids: list[int], state: Any) -> tuple[torch.Tensor, Any]: if was_training: self.model.train() + def generate_batch( + self, + prefix_token_ids: Sequence[Sequence[int]], + *, + max_new_tokens: int, + stop_sequences: Sequence[str] = (), + temperature: float = 0.0, + top_p: float = 1.0, + top_k: int = 0, + seeds: Sequence[int | None] | None = None, + ) -> list[GenerationOutput]: + """Use one masked padded forward per greedy decoding step.""" + if temperature != 0.0: + return super().generate_batch( + prefix_token_ids, + max_new_tokens=max_new_tokens, + stop_sequences=stop_sequences, + temperature=temperature, + top_p=top_p, + top_k=top_k, + seeds=seeds, + ) + was_training = self.model.training + self.model.eval() + + def step(rows: list[list[int]]) -> torch.Tensor: + lengths = [len(row) for row in rows] + width = max(lengths) + ids = torch.full( + (len(rows), width), + self.tokenizer.pad_token_id, + dtype=torch.long, + device=self._device, + ) + mask = torch.zeros((len(rows), width), dtype=torch.bool, device=self._device) + for index, row in enumerate(rows): + ids[index, : len(row)] = torch.tensor(row, dtype=torch.long, device=self._device) + mask[index, : len(row)] = True + with torch.no_grad(): + hidden, _ = self.model( + ids, past_key_values=None, use_cache=False, attention_mask=mask + ) + positions = torch.tensor( + [length - 1 for length in lengths], dtype=torch.long, device=self._device + ) + return self.model.lm_head( + hidden[torch.arange(len(rows), device=self._device), positions] + ) + + try: + return run_greedy_padded_generation( + step=step, + prefix_token_ids=prefix_token_ids, + decode=self.tokenizer.decode, + eos_token_id=self.tokenizer.eos_token_id, + max_new_tokens=max_new_tokens, + stop_sequences=stop_sequences, + ) + finally: + if was_training: + self.model.train() + # -- scoring -------------------------------------------------------- def hidden_states_at( diff --git a/src/miniverl/runtime/roles.py b/src/miniverl/runtime/roles.py index 898ce61..4a2a2ad 100644 --- a/src/miniverl/runtime/roles.py +++ b/src/miniverl/runtime/roles.py @@ -98,7 +98,7 @@ class LocalRoleGraph: rollout_runtime: RolloutRuntime teacher_policy: TeacherPolicy | None reference_policy: ReferencePolicy | None - reward_or_verifier: RewardOrVerifier + reward_or_verifier: RewardOrVerifier | None target_builder: TargetBuilder | None update_runtime: UpdateRuntime evaluation_runtime: EvaluationRuntime diff --git a/src/miniverl/runtime/rollout.py b/src/miniverl/runtime/rollout.py new file mode 100644 index 0000000..5028750 --- /dev/null +++ b/src/miniverl/runtime/rollout.py @@ -0,0 +1,310 @@ +"""Source-agnostic rollout runtimes for tool episodes and prompt datasets.""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass +from typing import Any, Protocol, runtime_checkable + +from miniverl.agent.loop import RolloutRunner +from miniverl.config.models import RolloutConfig, VerlParquetSourceConfig +from miniverl.data.verl_parquet import RenderedPrompt +from miniverl.errors import GpuMemoryError +from miniverl.models.base import GenerationOutput +from miniverl.schemas.trajectory import Span, SpanType, TerminationReason, Trajectory + +__all__ = [ + "GeneratedPromptBatch", + "PreparedPromptBatch", + "PromptDatasetRolloutRuntime", + "RolloutRuntime", + "ToolEnvironmentRolloutRuntime", +] + + +@runtime_checkable +class RolloutRuntime(Protocol): + """Minimal lifecycle used by the source-agnostic trainer path.""" + + def prepare_batch(self, inputs: Any) -> Any: + """Validate and physically batch logical rollout inputs.""" + ... + + def generate(self, batch: Any, *, policy_version: int, seed: int) -> Any: + """Generate under one explicitly identified actor policy.""" + ... + + def to_trajectories( + self, batch: Any, generated: Any, *, policy_version: int + ) -> list[Trajectory]: + """Freeze exact token provenance for training and teacher scoring.""" + ... + + def close(self) -> None: + """Release runtime-owned resources.""" + ... + + +@dataclass(frozen=True) +class PreparedPromptBatch: + """One logical prompt batch and its deterministic physical partition.""" + + prompts: tuple[RenderedPrompt, ...] + physical_batches: tuple[tuple[int, ...], ...] + + +@dataclass(frozen=True) +class GeneratedPromptBatch: + """Outputs restored to logical order plus the actual physical batch sizes.""" + + outputs: tuple[GenerationOutput, ...] + policy_version: int + physical_batch_sizes: tuple[int, ...] + oom_downshifts: int = 0 + + +class PromptDatasetRolloutRuntime: + """Single-turn padded actor generation over already-rendered prompts.""" + + def __init__( + self, + *, + backend: Any, + source_config: VerlParquetSourceConfig, + rollout_config: RolloutConfig, + ) -> None: + self.backend = backend + self.source_config = source_config + self.config = rollout_config + self._closed = False + + def prepare_batch(self, inputs: list[RenderedPrompt]) -> PreparedPromptBatch: + if self._closed: + raise RuntimeError("prompt rollout runtime is closed") + if not inputs: + raise ValueError("prompt rollout batch cannot be empty") + max_new = self.config.max_new_tokens_per_turn + for prompt in inputs: + if len(prompt.token_ids) + max_new > self.config.max_total_tokens: + raise ValueError( + f"prompt {prompt.record.row_digest[:12]} plus {max_new} response tokens " + f"exceeds rollout.max_total_tokens={self.config.max_total_tokens}" + ) + groups: list[tuple[int, ...]] = [] + current: list[int] = [] + current_width = 0 + for index, prompt in enumerate(inputs): + candidate_width = max(current_width, len(prompt.token_ids)) + candidate_size = len(current) + 1 + padded_tokens = candidate_size * (candidate_width + max_new) + if current and ( + candidate_size > self.config.prompt_batch_size + or padded_tokens > self.config.max_padded_tokens + ): + groups.append(tuple(current)) + current = [] + current_width = 0 + current.append(index) + current_width = max(current_width, len(prompt.token_ids)) + if current_width + max_new > self.config.max_padded_tokens: + raise ValueError( + f"one prompt needs {current_width + max_new} padded tokens, above " + f"rollout.max_padded_tokens={self.config.max_padded_tokens}" + ) + if current: + groups.append(tuple(current)) + return PreparedPromptBatch(prompts=tuple(inputs), physical_batches=tuple(groups)) + + def _generate_group( + self, + batch: PreparedPromptBatch, + indices: tuple[int, ...], + *, + base_seed: int, + ) -> tuple[list[tuple[int, GenerationOutput]], list[int], int]: + try: + output = self.backend.generate_batch( + [batch.prompts[index].token_ids for index in indices], + max_new_tokens=self.config.max_new_tokens_per_turn, + temperature=self.config.temperature, + top_p=self.config.top_p, + top_k=self.config.top_k, + seeds=[base_seed * 1_000_003 + index for index in indices], + ) + except BaseException as exc: + message = str(exc).lower() + is_oom = type(exc).__name__ in {"OutOfMemoryError", "CudaOutOfMemoryError"} or ( + "out of memory" in message and "cuda" in message + ) + if not is_oom: + raise + if len(indices) == 1: + raise GpuMemoryError( + "prompt generation ran out of GPU memory at physical batch size 1", + hint="lower max_prompt_length/max_response_length or choose a smaller model", + ) from exc + midpoint = len(indices) // 2 + left, left_sizes, left_down = self._generate_group( + batch, indices[:midpoint], base_seed=base_seed + ) + right, right_sizes, right_down = self._generate_group( + batch, indices[midpoint:], base_seed=base_seed + ) + return left + right, left_sizes + right_sizes, left_down + right_down + 1 + if len(output) != len(indices): + raise RuntimeError( + f"backend returned {len(output)} generations for physical batch {len(indices)}" + ) + return list(zip(indices, output, strict=True)), [len(indices)], 0 + + def generate( + self, + batch: PreparedPromptBatch, + *, + policy_version: int, + seed: int, + ) -> GeneratedPromptBatch: + if self._closed: + raise RuntimeError("prompt rollout runtime is closed") + indexed: list[tuple[int, GenerationOutput]] = [] + physical_sizes: list[int] = [] + downshifts = 0 + for group in batch.physical_batches: + rows, sizes, count = self._generate_group(batch, group, base_seed=seed) + indexed.extend(rows) + physical_sizes.extend(sizes) + downshifts += count + indexed.sort(key=lambda item: item[0]) + return GeneratedPromptBatch( + outputs=tuple(output for _, output in indexed), + policy_version=policy_version, + physical_batch_sizes=tuple(physical_sizes), + oom_downshifts=downshifts, + ) + + def to_trajectories( + self, + batch: PreparedPromptBatch, + generated: GeneratedPromptBatch, + *, + policy_version: int, + ) -> list[Trajectory]: + if generated.policy_version != policy_version: + raise ValueError( + f"generated policy version {generated.policy_version} does not match {policy_version}" + ) + if len(generated.outputs) != len(batch.prompts): + raise ValueError("generated prompt count does not match the prepared batch") + trajectories: list[Trajectory] = [] + model_id = str(getattr(self.backend, "model_id", self.backend.capabilities.name)) + revision = getattr(self.backend, "model_revision", None) + for prompt, output in zip(batch.prompts, generated.outputs, strict=True): + prompt_ids = list(prompt.token_ids) + response_ids = list(output.token_ids) + if not response_ids: + raise ValueError( + f"actor produced an empty response for prompt {prompt.record.row_digest[:12]}" + ) + boundary = len(prompt_ids) + token_ids = [*prompt_ids, *response_ids] + trajectory_id = hashlib.sha256( + f"{prompt.record.row_digest}:{policy_version}".encode("ascii") + ).hexdigest() + metadata = { + "source_kind": "verl_parquet", + "data_source": prompt.record.data_source, + "ability": prompt.record.ability, + "reward_model": prompt.record.reward_model, + "extra_info": prompt.record.extra_info, + "source_file": prompt.record.source_file, + "source_row_index": prompt.record.source_row_index, + "row_digest": prompt.record.row_digest, + "rendered_prompt_digest": prompt.rendered_prompt_digest, + "tokenizer_identity": prompt.tokenizer_identity, + "prompt_token_count": prompt.prompt_token_count, + "original_prompt_token_count": prompt.original_prompt_token_count, + "truncation_decision": prompt.truncation_decision, + "response_token_count": len(response_ids), + "generation_stop_reason": output.stop_reason, + } + trajectories.append( + Trajectory( + trajectory_id=trajectory_id, + task_id=prompt.record.row_digest, + environment="verl_parquet", + token_ids=token_ids, + attention_mask=[1] * len(token_ids), + model_generated_mask=[False] * boundary + [True] * len(response_ids), + critical_mask=[False] * len(token_ids), + spans=[ + Span( + span_type=SpanType.USER, + start=0, + end=boundary, + turn_id=0, + text=prompt.text, + metadata={"rendered_prompt_digest": prompt.rendered_prompt_digest}, + ), + Span( + span_type=SpanType.ASSISTANT_TEXT, + start=boundary, + end=len(token_ids), + turn_id=0, + text=output.text, + ), + ], + policy_version=policy_version, + tokenizer_fingerprint=str(self.backend.tokenizer.fingerprint), + model_id=model_id, + model_revision=revision, + termination_reason=( + TerminationReason.EOS_WITHOUT_FINAL + if output.stop_reason == "eos" + else TerminationReason.MAX_TOKENS + ), + generated_token_count=len(response_ids), + assistant_turns=1, + metadata=metadata, + ) + ) + return trajectories + + def close(self) -> None: + self._closed = True + + +class ToolEnvironmentRolloutRuntime: + """Adapter preserving the existing multi-turn tool runner byte semantics.""" + + def __init__(self, runner: RolloutRunner) -> None: + self.runner = runner + + def prepare_batch(self, inputs: Any) -> tuple[Any, ...]: + return tuple(inputs) + + def generate( + self, batch: tuple[Any, ...], *, policy_version: int, seed: int + ) -> list[Trajectory]: + return [ + self.runner.rollout( + task, + policy_version=policy_version, + seed=seed * 1_000_003 + index, + ) + for index, task in enumerate(batch) + ] + + def to_trajectories( + self, + batch: tuple[Any, ...], + generated: list[Trajectory], + *, + policy_version: int, + ) -> list[Trajectory]: + del batch + if any(item.policy_version != policy_version for item in generated): + raise ValueError("tool rollout policy version changed inside one batch") + return generated + + def close(self) -> None: + """The trainer owns and closes the backward-compatible environment.""" diff --git a/src/miniverl/training/trainer.py b/src/miniverl/training/trainer.py index ca87b49..4ddba57 100644 --- a/src/miniverl/training/trainer.py +++ b/src/miniverl/training/trainer.py @@ -45,8 +45,10 @@ OPDFreshness, Quantization, RunConfig, + SourceKind, TeacherContextMode, TrainingMode, + VerlParquetSourceConfig, ) from miniverl.environments.base import Task, ToolEnvironment, make_splits from miniverl.environments.registry import make_environment @@ -181,8 +183,11 @@ def __init__( validated_config: RunConfig, paths: RunPaths, run_id: str, - environment: ToolEnvironment, + environment: ToolEnvironment | None, splits: dict[str, list[Task]], + rollout_runtime: Any, + prompt_dataset: Any | None, + prompt_dataset_manifest: Any | None, tokenizer: Any, student: Any, teacher: Any | None, @@ -203,6 +208,11 @@ def __init__( self.run_id = run_id self.environment = environment self.splits = splits + self.rollout_runtime = rollout_runtime + self.prompt_dataset = prompt_dataset + self.prompt_dataset_manifest = prompt_dataset_manifest + self._prompt_train_iterator: Any | None = None + self._prompt_train_epoch = 0 self.tokenizer = tokenizer self.student = student self.teacher = teacher @@ -220,7 +230,7 @@ def __init__( self.metrics_log = JsonlWriter(paths.metrics) self.events = EventLog(JsonlWriter(paths.events)) - self.runner = RolloutRunner(backend=student, environment=environment, config=config.rollout) + self.runner = getattr(rollout_runtime, "runner", None) # Imported here rather than at module scope: the scorer pulls in torch, and # `import miniverl.trainer` must stay readable on a bare install so the CLI # can raise MissingDependencyError instead of ModuleNotFoundError. @@ -244,7 +254,7 @@ def __init__( self.artifact_bridge = LocalArtifactBridge(paths.root) self.role_graph = LocalRoleGraph( actor_policy=self.student, - rollout_runtime=self.runner, + rollout_runtime=self.rollout_runtime, teacher_policy=self.teacher, reference_policy=self.reference, reward_or_verifier=self.environment, @@ -463,6 +473,9 @@ def from_config( raise environment: ToolEnvironment | None = None + rollout_runtime: Any | None = None + prompt_dataset: Any | None = None + prompt_dataset_manifest: Any | None = None student: Any | None = None teacher: Any | None = None reference: Any | None = None @@ -474,18 +487,44 @@ def from_config( write_bytes(paths.config_submitted, validated_config.submitted_bytes) write_text(paths.config_validated, validated_config.to_yaml()) write_text(paths.config_original, config.to_yaml()) - environment = make_environment(config.environment.name, **config.environment.params) - splits = make_splits( - environment, - counts={ - "train": config.environment.train_tasks, - "eval": config.environment.eval_tasks, - "test": config.environment.test_tasks, - }, - seed=config.environment.split_seed, - difficulty=config.environment.difficulty, - ) + splits: dict[str, list[Task]] = {"train": [], "eval": [], "test": []} + if config.source.kind is SourceKind.ENVIRONMENT: + assert config.environment is not None + environment = make_environment(config.environment.name, **config.environment.params) + splits = make_splits( + environment, + counts={ + "train": config.environment.train_tasks, + "eval": config.environment.eval_tasks, + "test": config.environment.test_tasks, + }, + seed=config.environment.split_seed, + difficulty=config.environment.difficulty, + ) + else: + from miniverl.data.verl_parquet import VerlParquetDataset + + assert isinstance(config.source, VerlParquetSourceConfig) + prompt_dataset = VerlParquetDataset(config.source) + prompt_dataset_manifest = prompt_dataset.inspect() + if prompt_dataset_manifest.rows["train"] == 0: + raise ConfigError("the Parquet training source contains zero prompt rows") + if config.eval.enabled and prompt_dataset_manifest.rows["val"] == 0: + raise ConfigError( + "evaluation is enabled but source.val_files contains zero prompt rows", + hint="provide validation prompts or set eval.enabled=false", + ) + if ( + config.eval.enabled + and config.eval.tasks is not None + and config.eval.task_offset + config.eval.tasks + > prompt_dataset_manifest.rows["val"] + ): + raise ConfigError( + "eval task range exceeds the number of Parquet validation rows" + ) if config.models.teacher.mode is TeacherContextMode.PRIVILEGED_CONTEXT: + assert environment is not None probe = environment.privileged_context(splits["train"][0]) if not probe: raise ConfigError( @@ -564,6 +603,27 @@ def teacher_fits() -> bool: hint="miniVERL requires a shared output vocabulary for distillation", ) + if config.source.kind is SourceKind.ENVIRONMENT: + assert environment is not None + from miniverl.runtime.rollout import ToolEnvironmentRolloutRuntime + + rollout_runtime = ToolEnvironmentRolloutRuntime( + RolloutRunner( + backend=student, + environment=environment, + config=config.rollout, + ) + ) + else: + from miniverl.runtime.rollout import PromptDatasetRolloutRuntime + + assert isinstance(config.source, VerlParquetSourceConfig) + rollout_runtime = PromptDatasetRolloutRuntime( + backend=student, + source_config=config.source, + rollout_config=config.rollout, + ) + trainer = cls( config=config, validated_config=validated_config, @@ -571,6 +631,9 @@ def teacher_fits() -> bool: run_id=resolved_id, environment=environment, splits=splits, + rollout_runtime=rollout_runtime, + prompt_dataset=prompt_dataset, + prompt_dataset_manifest=prompt_dataset_manifest, tokenizer=tokenizer, student=student, teacher=teacher, @@ -810,6 +873,32 @@ def build_manifest(self) -> dict[str, Any]: else None ), } + if self.environment is not None and config.environment is not None: + environment_info: dict[str, Any] | None = { + **self.environment.describe(), + "difficulty": config.environment.difficulty, + "split_seed": config.environment.split_seed, + "split_sizes": {k: len(v) for k, v in self.splits.items()}, + } + source_info: dict[str, Any] = { + "kind": "environment", + "environment": environment_info, + } + else: + assert isinstance(config.source, VerlParquetSourceConfig) + assert self.prompt_dataset_manifest is not None + environment_info = None + source_info = { + "kind": "verl_parquet", + "prompt_key": config.source.prompt_key, + "use_task_rewards": config.source.use_task_rewards, + "rows": self.prompt_dataset_manifest.rows, + "schema_digest": self.prompt_dataset_manifest.schema_digest, + "content_digest": self.prompt_dataset_manifest.content_digest, + "files": list(self.prompt_dataset_manifest.files), + "shuffle": config.source.shuffle, + "seed": config.source.seed, + } return { "miniverl_version": __version__, "run_id": self.run_id, @@ -829,12 +918,8 @@ def build_manifest(self) -> dict[str, Any]: ), "seed": config.run.seed, "deterministic": config.run.deterministic, - "environment": { - **self.environment.describe(), - "difficulty": config.environment.difficulty, - "split_seed": config.environment.split_seed, - "split_sizes": {k: len(v) for k, v in self.splits.items()}, - }, + "source": source_info, + "environment": environment_info, "models": { "backend": config.models.backend.value, "runtime": config.models.runtime.value, @@ -1031,11 +1116,32 @@ def _finalize_manifest( # -- task sampling -------------------------------------------------------- def _build_task_order(self) -> list[int]: + if self.prompt_dataset is not None: + return [] order = list(range(len(self.splits["train"]))) random.Random(self.config.run.seed ^ 0x5EED).shuffle(order) return order - def _next_tasks(self, count: int) -> list[Task]: + def _next_tasks(self, count: int) -> list[Any]: + if self.prompt_dataset is not None: + from miniverl.data.verl_parquet import render_prompt + + assert isinstance(self.config.source, VerlParquetSourceConfig) + output: list[Any] = [] + while len(output) < count: + if self._prompt_train_iterator is None: + self._prompt_train_iterator = iter( + self.prompt_dataset.iter_split("train", epoch=self._prompt_train_epoch) + ) + try: + record = next(self._prompt_train_iterator) + except StopIteration: + self._prompt_train_epoch += 1 + self._prompt_train_iterator = None + continue + output.append(render_prompt(record, self.tokenizer, self.config.source)) + self.task_cursor += 1 + return output train = self.splits["train"] out: list[Task] = [] for _ in range(count): @@ -1053,6 +1159,7 @@ def _prepare_toy_teacher(self) -> None: config.models.backend is not ModelBackend.TOY or self.teacher is None or config.models.teacher.toy_pretrain_steps <= 0 + or self.environment is None ): return from miniverl.models.toy import ToyBackend, fit_toy_model @@ -1102,13 +1209,38 @@ def _prepare_toy_teacher(self) -> None: def _collect( self, - tasks: list[Task], + tasks: list[Any], *, oracle: bool, rollout_seed_base: int | None = None, ) -> tuple[list[Trajectory], RolloutStats]: stats = RolloutStats() trajectories: list[Trajectory] = [] + if self.prompt_dataset is not None: + if oracle: + raise ConfigError("Parquet prompt rollouts have no oracle trajectory source") + seed = ( + rollout_seed_base + if rollout_seed_base is not None + else self.config.run.seed + self.global_step * 1013 + ) + prepared = self.rollout_runtime.prepare_batch(tasks) + generated = self.rollout_runtime.generate( + prepared, + policy_version=self.policy_version, + seed=seed, + ) + trajectories = self.rollout_runtime.to_trajectories( + prepared, + generated, + policy_version=self.policy_version, + ) + for offset, trajectory in enumerate(trajectories): + trajectory.metadata["generation_seed"] = seed * 1_000_003 + offset + stats.observe(trajectory) + append_trajectories(self.paths.trajectories, trajectories) + return trajectories, stats + assert self.runner is not None for offset, task in enumerate(tasks): if oracle: traj = self.runner.oracle_rollout( @@ -1356,6 +1488,7 @@ def _load_offline_dataset(self, *, expected_digest: str) -> None: ) teacher_view = None if privileged: + assert self.runner is not None task = task_by_id.get(trajectory.task_id) if task is None: raise CheckpointError( @@ -1474,6 +1607,7 @@ def _build_samples(self, trajectories: list[Trajectory]) -> list[TrainSample]: continue teacher_view: Trajectory | None = None if privileged and self.scorer is not None: + assert self.runner is not None task = task_by_id.get(traj.task_id) if task is not None: teacher_view = self.runner.privileged_render(traj, task) @@ -1491,6 +1625,7 @@ def _build_samples(self, trajectories: list[Trajectory]) -> list[TrainSample]: for sample in samples: teacher_view = None if privileged: + assert self.runner is not None task = task_by_id.get(sample.trajectory.task_id) if task is not None: teacher_view = self.runner.privileged_render(sample.trajectory, task) @@ -2189,7 +2324,14 @@ def _run_cycle(self) -> list[dict[str, Any]]: parameter_version=self.parameter_version, rollout_policy_version=rollout_policy_version, trajectories=stats.rollouts, - success_rate=round(stats.to_dict()["success_rate"], 4), + success_rate=( + None + if self.prompt_dataset is not None + else round(stats.to_dict()["success_rate"], 4) + ), + reward_status=( + "not_applicable_pure_opd" if self.prompt_dataset is not None else "measured" + ), generated_tokens=stats.generated_tokens, rollout_tokens_per_second=round(stats.generated_tokens / rollout_seconds, 2), ) @@ -2381,7 +2523,7 @@ def evaluate( self, *, split: str | None = None, - tasks: list[Task] | None = None, + tasks: list[Any] | None = None, tag: str = "eval", write: bool = True, ) -> dict[str, Any]: @@ -2412,13 +2554,20 @@ def _evaluate_impl( self, *, split: str | None = None, - tasks: list[Task] | None = None, + tasks: list[Any] | None = None, tag: str = "eval", write: bool = True, ) -> dict[str, Any]: """Deterministic greedy evaluation owned by the current operation.""" config = self.config chosen_split = split or config.eval.split + if self.prompt_dataset is not None: + return self._evaluate_prompt_impl( + chosen_split=chosen_split, + prompts=tasks, + tag=tag, + write=write, + ) pool = tasks if tasks is not None else self.splits.get(chosen_split, []) limit = config.effective_eval_tasks if tasks is None else len(pool) task_offset = config.eval.task_offset if tasks is None else 0 @@ -2451,6 +2600,7 @@ def _evaluate_impl( stats = RolloutStats() trajectories: list[Trajectory] = [] by_difficulty: dict[str, list[int]] = {} + assert self.runner is not None for offset, task in enumerate(pool): traj = self.runner.rollout( task, @@ -2528,6 +2678,116 @@ def _evaluate_impl( finally: self.student.set_train(was_training) + def _evaluate_prompt_impl( + self, + *, + chosen_split: str, + prompts: list[Any] | None, + tag: str, + write: bool, + ) -> dict[str, Any]: + """Generate validation responses without inventing a reward or success metric.""" + from itertools import islice + + from miniverl.data.verl_parquet import render_prompt + from miniverl.runtime.rollout import PromptDatasetRolloutRuntime + + config = self.config + assert isinstance(config.source, VerlParquetSourceConfig) + assert self.prompt_dataset is not None + assert self.prompt_dataset_manifest is not None + if chosen_split not in {"eval", "val"}: + raise ConfigError( + "Parquet prompt evaluation uses the validation files", + hint="set eval.split=eval (the compatibility name for source.val_files)", + ) + if prompts is None: + available = self.prompt_dataset_manifest.rows["val"] + limit = config.eval.tasks if config.eval.tasks is not None else available + records = islice( + self.prompt_dataset.iter_split("val", epoch=0), + config.eval.task_offset, + config.eval.task_offset + limit, + ) + prompts = [render_prompt(record, self.tokenizer, config.source) for record in records] + if not prompts: + return { + "tag": tag, + "split": "val", + "tasks": 0, + "success_rate": None, + "reward_status": "not_applicable_pure_opd", + "note": "no validation prompts", + } + model = getattr(self.student, "model", None) + was_training = getattr(model, "training", None) + if not isinstance(was_training, bool): + raise BackendError("the student backend does not expose train/eval mode") + self.student.set_train(False) + runtime = PromptDatasetRolloutRuntime( + backend=self.student, + source_config=config.source, + rollout_config=config.rollout.model_copy( + update={"temperature": config.eval.temperature} + ), + ) + try: + gpu.reset_peak_stats() + started = time.perf_counter() + prepared = runtime.prepare_batch(prompts) + generated = runtime.generate( + prepared, + policy_version=self.policy_version, + seed=config.eval.seed + config.eval.task_offset, + ) + trajectories = runtime.to_trajectories( + prepared, + generated, + policy_version=self.policy_version, + ) + elapsed = max(time.perf_counter() - started, 1e-9) + append_trajectories(self.paths.eval_trajectories, trajectories) + generated_tokens = sum(item.generated_token_count for item in trajectories) + payload = { + "tag": tag, + "split": "val", + "tasks": len(trajectories), + "policy_version": self.policy_version, + "parameter_version": self.parameter_version, + "global_step": self.global_step, + "global_optimizer_step": self.global_step, + "rollout_iteration": self._cycles_completed, + "rollout_policy_version": self._last_rollout_policy_version, + "temperature": config.eval.temperature, + "task_offset": config.eval.task_offset, + "seconds": round(elapsed, 3), + "generated_tokens": generated_tokens, + "rollout_tokens_per_second": round(generated_tokens / elapsed, 2), + "physical_batch_sizes": list(generated.physical_batch_sizes), + "oom_downshifts": generated.oom_downshifts, + "success_rate": None, + "reward_status": "not_applicable_pure_opd", + "measurement_status": { + "response_generation": "measured", + "task_reward": "not_configured", + "task_success": "not_measured", + }, + "memory": gpu.snapshot().to_dict(), + } + if write: + self.metrics_log.write({"phase": "eval", **payload, "ts": utc_now()}) + self.events.emit( + "eval", + tag=tag, + tasks=len(trajectories), + success_rate=None, + reward_status="not_applicable_pure_opd", + ) + return payload + finally: + runtime.close() + self.student.set_train(was_training) + # -- checkpointing ----------------------------------------------------------- def _config_digest(self) -> str: @@ -2785,6 +3045,12 @@ def clear_optimizer() -> None: cleanup("student release", student.release) student = None + rollout_runtime = self.rollout_runtime + self.rollout_runtime = None # type: ignore[assignment] # destructive close + if rollout_runtime is not None: + cleanup("rollout runtime close", rollout_runtime.close) + del rollout_runtime + environment = self.environment self.environment = None # type: ignore[assignment] # destructive close if environment is not None: diff --git a/tests/integration/test_prompt_opd_pipeline.py b/tests/integration/test_prompt_opd_pipeline.py new file mode 100644 index 0000000..6941dc7 --- /dev/null +++ b/tests/integration/test_prompt_opd_pipeline.py @@ -0,0 +1,130 @@ +"""End-to-end pure OPD over first-class verl-style Parquet prompts.""" + +from __future__ import annotations + +import json + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from tests.conftest import requires_torch + +pytestmark = [requires_torch, pytest.mark.torch] + + +def test_prompt_opd_trains_without_an_environment_or_reward(tmp_path) -> None: + from miniverl.config import RunConfig + from miniverl.trainer import OPDTrainer + + train = tmp_path / "train.parquet" + pq.write_table( + pa.Table.from_pylist( + [ + { + "prompt": "Compute 1 + 1.", + "data_source": "unit", + "ability": "arithmetic", + "extra_info": {"row": 0}, + }, + { + "prompt": "Compute 2 + 2.", + "data_source": "unit", + "ability": "arithmetic", + "extra_info": {"row": 1}, + }, + ] + ), + train, + row_group_size=1, + ) + config = RunConfig.model_validate( + { + "run": { + "name": "prompt-opd", + "mode": "opd", + "seed": 9, + "output_dir": str(tmp_path / "runs"), + }, + "models": { + "backend": "toy", + "device": "cpu", + "student": { + "model_id": "toy-student", + "lora": {"enabled": False}, + "toy": { + "hidden_size": 16, + "num_layers": 1, + "num_heads": 2, + "intermediate_size": 32, + "max_position_embeddings": 128, + }, + }, + "teacher": { + "model_id": "toy-teacher", + "toy_pretrain_steps": 0, + "toy": { + "hidden_size": 16, + "num_layers": 1, + "num_heads": 2, + "intermediate_size": 32, + "max_position_embeddings": 128, + }, + }, + }, + "source": { + "kind": "verl_parquet", + "train_files": [str(train)], + "allow_plain_string_prompts": True, + "max_prompt_length": 32, + "shuffle": False, + }, + "rollout": { + "max_turns": 1, + "max_new_tokens_per_turn": 3, + "max_total_tokens": 64, + "temperature": 0.0, + "prompt_batch_size": 2, + "max_padded_tokens": 128, + }, + "selection": {"selector": "all_model_tokens"}, + "loss": { + "mode": "bucketed_topk_tail", + "divergence": "forward_kl", + "top_k": 4, + "chunk_size": 16, + }, + "train": { + "cycles": 1, + "rollouts_per_cycle": 2, + "gradient_accumulation_steps": 2, + "learning_rate": 0.001, + }, + "memory": {"strategy": "resident"}, + "cache": {"entries_per_shard": 2, "dtype": "float32"}, + "eval": {"enabled": False}, + "report": {"enabled": False}, + } + ) + + trainer = OPDTrainer.from_config(config, run_id="prompt-opd-test") + try: + result = trainer.train() + finally: + trainer.close() + + assert result.global_step == 1 + assert result.policy_version == 1 + manifest = json.loads((result.run_dir / "manifest.json").read_text(encoding="utf-8")) + assert manifest["source"]["kind"] == "verl_parquet" + assert manifest["source"]["rows"] == {"train": 2, "val": 0} + rows = [ + json.loads(line) + for line in (result.run_dir / "trajectories.jsonl").read_text(encoding="utf-8").splitlines() + ] + assert len(rows) == 2 + assert all( + not any(row["model_generated_mask"][: row["metadata"]["prompt_token_count"]]) + for row in rows + ) + assert all(row["metadata"]["reward_model"] is None for row in rows) diff --git a/tests/unit/test_packaging.py b/tests/unit/test_packaging.py index 251fa88..e34e845 100644 --- a/tests/unit/test_packaging.py +++ b/tests/unit/test_packaging.py @@ -39,6 +39,7 @@ "cache", "community", "config", + "data", "environments", "evidence", "evaluation", diff --git a/tests/unit/test_prompt_rollout_runtime.py b/tests/unit/test_prompt_rollout_runtime.py new file mode 100644 index 0000000..9d29e63 --- /dev/null +++ b/tests/unit/test_prompt_rollout_runtime.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +from miniverl.config.models import RolloutConfig, VerlParquetSourceConfig +from miniverl.data.verl_parquet import PromptRecord, RenderedPrompt +from miniverl.models.tokenizers import ToyTokenizer +from miniverl.models.toy import ToyBackend +from miniverl.runtime.rollout import PromptDatasetRolloutRuntime + + +def _record(index: int) -> PromptRecord: + return PromptRecord( + prompt=f"prompt {index}", + data_source="unit", + ability=None, + reward_model=None, + extra_info={"index": index}, + source_file="train.parquet", + source_row_index=index, + row_digest=f"{index:064x}", + canonical_payload="{}", + ) + + +def test_real_padded_greedy_generation_matches_sequential() -> None: + tokenizer = ToyTokenizer() + backend = ToyBackend(tokenizer=tokenizer, model_id="toy", seed=19, trainable=False) + prompts = [tokenizer.encode("short"), tokenizer.encode("a longer prompt")] + + expected = [ + backend.generate(row, max_new_tokens=7, temperature=0.0, seed=100 + index) + for index, row in enumerate(prompts) + ] + actual = backend.generate_batch( + prompts, + max_new_tokens=7, + temperature=0.0, + seeds=[100, 101], + ) + + assert [row.token_ids for row in actual] == [row.token_ids for row in expected] + assert [row.text for row in actual] == [row.text for row in expected] + assert [len(row.token_ids) for row in actual] == [len(row.token_ids) for row in expected] + + single = backend.generate_batch(prompts[:1], max_new_tokens=7, temperature=0.0, seeds=[100]) + assert single[0].token_ids == expected[0].token_ids + + +def test_prompt_trajectories_never_select_prompt_or_padding_tokens() -> None: + tokenizer = ToyTokenizer() + backend = ToyBackend(tokenizer=tokenizer, model_id="toy", seed=7, trainable=False) + runtime = PromptDatasetRolloutRuntime( + backend=backend, + source_config=VerlParquetSourceConfig( + train_files=["unused.parquet"], allow_plain_string_prompts=True + ), + rollout_config=RolloutConfig( + max_new_tokens_per_turn=5, + max_total_tokens=32, + temperature=0.0, + prompt_batch_size=2, + max_padded_tokens=64, + ), + ) + rendered = [ + RenderedPrompt( + record=_record(index), + text=text, + token_ids=tuple(tokenizer.encode(text)), + tokenizer_identity=tokenizer.identity, + rendered_prompt_digest=f"{index + 10:064x}", + prompt_token_count=len(tokenizer.encode(text)), + truncation_decision="not_needed", + original_prompt_token_count=len(tokenizer.encode(text)), + ) + for index, text in enumerate(("short", "a longer prompt")) + ] + + batch = runtime.prepare_batch(rendered) + generated = runtime.generate(batch, policy_version=4, seed=90) + trajectories = runtime.to_trajectories(batch, generated, policy_version=4) + + assert [trajectory.task_id for trajectory in trajectories] == [ + row.record.row_digest for row in rendered + ] + for prompt, output, trajectory in zip(rendered, generated.outputs, trajectories, strict=True): + assert trajectory.policy_version == 4 + assert trajectory.token_ids == [*prompt.token_ids, *output.token_ids] + assert not any(trajectory.model_generated_mask[: prompt.prompt_token_count]) + assert all(trajectory.model_generated_mask[prompt.prompt_token_count :]) + assert trajectory.model_token_positions() == list( + range(prompt.prompt_token_count, len(trajectory.token_ids)) + ) + assert trajectory.metadata["response_token_count"] == len(output.token_ids) + + +def test_physical_batches_respect_padded_token_budget_without_reordering() -> None: + tokenizer = ToyTokenizer() + backend = ToyBackend(tokenizer=tokenizer, model_id="toy", seed=7, trainable=False) + runtime = PromptDatasetRolloutRuntime( + backend=backend, + source_config=VerlParquetSourceConfig( + train_files=["unused.parquet"], allow_plain_string_prompts=True + ), + rollout_config=RolloutConfig( + max_new_tokens_per_turn=4, + max_total_tokens=32, + temperature=0.0, + prompt_batch_size=8, + max_padded_tokens=24, + ), + ) + rendered = [ + RenderedPrompt( + record=_record(index), + text="x" * length, + token_ids=tuple(tokenizer.encode("x" * length)), + tokenizer_identity=tokenizer.identity, + rendered_prompt_digest=f"{index + 20:064x}", + prompt_token_count=length, + truncation_decision="not_needed", + original_prompt_token_count=length, + ) + for index, length in enumerate((2, 5, 3, 6)) + ] + + batch = runtime.prepare_batch(rendered) + assert [item.record.source_row_index for item in batch.prompts] == [0, 1, 2, 3] + assert all( + len(group) * (max(len(batch.prompts[index].token_ids) for index in group) + 4) <= 24 + for group in batch.physical_batches + ) + generated = runtime.generate(batch, policy_version=0, seed=1) + assert generated.physical_batch_sizes == tuple(len(group) for group in batch.physical_batches) + assert len(generated.outputs) == 4 + + +def test_oom_downshift_changes_only_physical_batching(monkeypatch) -> None: + tokenizer = ToyTokenizer() + backend = ToyBackend(tokenizer=tokenizer, model_id="toy", seed=7, trainable=False) + original = backend.generate_batch + + def fail_multi(prefixes, **kwargs): # type: ignore[no-untyped-def] + if len(prefixes) > 1: + raise RuntimeError("CUDA out of memory: injected") + return original(prefixes, **kwargs) + + monkeypatch.setattr(backend, "generate_batch", fail_multi) + rollout_config = RolloutConfig( + max_new_tokens_per_turn=2, + max_total_tokens=32, + temperature=0.0, + prompt_batch_size=2, + max_padded_tokens=64, + ) + runtime = PromptDatasetRolloutRuntime( + backend=backend, + source_config=VerlParquetSourceConfig( + train_files=["unused.parquet"], allow_plain_string_prompts=True + ), + rollout_config=rollout_config, + ) + rendered = [ + RenderedPrompt( + record=_record(index), + text="x", + token_ids=tuple(tokenizer.encode("x")), + tokenizer_identity=tokenizer.identity, + rendered_prompt_digest=f"{index + 30:064x}", + prompt_token_count=1, + truncation_decision="not_needed", + original_prompt_token_count=1, + ) + for index in range(2) + ] + + generated = runtime.generate(runtime.prepare_batch(rendered), policy_version=3, seed=4) + + assert generated.physical_batch_sizes == (1, 1) + assert generated.oom_downshifts == 1 + assert rollout_config.prompt_batch_size == 2 + assert generated.policy_version == 3 diff --git a/tests/unit/test_prompt_source_config.py b/tests/unit/test_prompt_source_config.py new file mode 100644 index 0000000..f28759d --- /dev/null +++ b/tests/unit/test_prompt_source_config.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from miniverl.config.models import RunConfig, SourceKind + + +def _models() -> dict[str, object]: + return { + "backend": "toy", + "student": {"model_id": "toy-student"}, + "teacher": {"model_id": "toy-teacher"}, + } + + +def test_verl_parquet_source_does_not_require_an_environment() -> None: + config = RunConfig.model_validate( + { + "models": _models(), + "source": { + "kind": "verl_parquet", + "train_files": ["train.parquet"], + "val_files": ["val.parquet"], + "prompt_key": "prompt", + }, + "eval": {"tasks": 4}, + } + ) + + assert config.environment is None + assert config.source.kind is SourceKind.VERL_PARQUET + assert config.effective_eval_tasks == 4 + + +def test_legacy_environment_recipe_remains_backward_compatible() -> None: + config = RunConfig.model_validate( + { + "models": _models(), + "environment": {"name": "calculator"}, + } + ) + + assert config.source.kind is SourceKind.ENVIRONMENT + assert config.environment is not None + + +def test_environment_and_parquet_source_cannot_be_mixed() -> None: + with pytest.raises(ValidationError, match="must not define environment"): + RunConfig.model_validate( + { + "models": _models(), + "environment": {"name": "calculator"}, + "source": { + "kind": "verl_parquet", + "train_files": ["train.parquet"], + }, + } + ) + + +def test_plain_string_prompts_require_an_explicit_opt_in() -> None: + config = RunConfig.model_validate( + { + "models": _models(), + "source": { + "kind": "verl_parquet", + "train_files": ["train.parquet"], + "allow_plain_string_prompts": True, + "truncation": "left", + }, + } + ) + + assert config.source.allow_plain_string_prompts is True + assert config.source.truncation.value == "left" diff --git a/tests/unit/test_verl_parquet_source.py b/tests/unit/test_verl_parquet_source.py new file mode 100644 index 0000000..4c78bda --- /dev/null +++ b/tests/unit/test_verl_parquet_source.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import json +from typing import ClassVar + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from miniverl.config.models import PromptTruncation, VerlParquetSourceConfig +from miniverl.data.verl_parquet import VerlParquetDataset, render_prompt +from miniverl.errors import ConfigError + + +class TinyTokenizer: + fingerprint = "f" * 64 + identity: ClassVar[dict[str, str]] = { + "behavioral_fingerprint_v1": fingerprint, + "structural_digest_v2": "s" * 64, + } + eos_token_id = 0 + pad_token_id = 0 + vocab_size = 256 + + def encode(self, text: str) -> list[int]: + return list(text.encode("utf-8")) + + def decode(self, token_ids: list[int]) -> str: + return bytes(token_ids).decode("utf-8") + + def apply_chat_template(self, messages: list[dict[str, str]]) -> str: + return ( + "".join(f"<{m['role']}>{m['content']}" for m in messages) + "" + ) + + +def _write(path, rows: list[dict[str, object]], *, row_group_size: int = 2): + pq.write_table(pa.Table.from_pylist(rows), path, row_group_size=row_group_size) + return path + + +def _row(index: int) -> dict[str, object]: + return { + "prompt": [{"role": "user", "content": f"question {index}"}], + "data_source": "unit", + "ability": "reasoning", + "reward_model": {"style": "rule"}, + "extra_info": {"index": index}, + } + + +def test_streaming_loader_preserves_metadata_and_is_deterministic(tmp_path) -> None: + train = _write(tmp_path / "train.parquet", [_row(i) for i in range(9)]) + source = VerlParquetSourceConfig( + train_files=[str(train)], row_batch_size=2, shuffle=True, seed=71 + ) + dataset = VerlParquetDataset(source) + + first = list(dataset.iter_split("train", epoch=3)) + second = list(dataset.iter_split("train", epoch=3)) + + assert [row.row_digest for row in first] == [row.row_digest for row in second] + assert sorted(row.source_row_index for row in first) == list(range(9)) + assert first[0].data_source == "unit" + assert first[0].ability == "reasoning" + assert first[0].reward_model == {"style": "rule"} + assert isinstance(first[0].extra_info, dict) + assert first[0].source_file == str(train.resolve()) + manifest = dataset.inspect() + assert manifest.rows == {"train": 9, "val": 0} + assert len(manifest.content_digest) == len(manifest.schema_digest) == 64 + + +def test_rejects_every_invalid_row_instead_of_silently_filtering(tmp_path) -> None: + train = _write(tmp_path / "train.parquet", [{"prompt": "plain text"}]) + dataset = VerlParquetDataset(VerlParquetSourceConfig(train_files=[str(train)])) + + with pytest.raises(ConfigError, match="plain-string prompt"): + list(dataset.iter_split("train")) + + +def test_task_rewards_require_reward_model_but_pure_opd_does_not(tmp_path) -> None: + train = _write( + tmp_path / "train.parquet", + [{"prompt": [{"role": "user", "content": "hello"}]}], + ) + pure = VerlParquetDataset(VerlParquetSourceConfig(train_files=[str(train)])) + assert next(iter(pure.iter_split("train"))).reward_model is None + + rewarded = VerlParquetDataset( + VerlParquetSourceConfig(train_files=[str(train)], use_task_rewards=True) + ) + with pytest.raises(ConfigError, match="reward_model"): + list(rewarded.iter_split("train")) + + +def test_chat_template_is_applied_once_and_records_provenance(tmp_path) -> None: + train = _write(tmp_path / "train.parquet", [_row(1)]) + source = VerlParquetSourceConfig(train_files=[str(train)], max_prompt_length=200) + record = next(VerlParquetDataset(source).iter_split("train")) + rendered = render_prompt(record, TinyTokenizer(), source) + + assert rendered.text == "question 1" + assert rendered.prompt_token_count == len(rendered.token_ids) + assert rendered.truncation_decision == "not_needed" + assert rendered.tokenizer_identity["structural_digest_v2"] == "s" * 64 + assert len(rendered.rendered_prompt_digest) == 64 + + +@pytest.mark.parametrize( + ("mode", "expected"), + [(PromptTruncation.LEFT, "def"), (PromptTruncation.RIGHT, "abc")], +) +def test_truncation_is_exact(mode, expected, tmp_path) -> None: + train = _write(tmp_path / "train.parquet", [{"prompt": "abcdef"}]) + source = VerlParquetSourceConfig( + train_files=[str(train)], + allow_plain_string_prompts=True, + max_prompt_length=3, + truncation=mode, + ) + record = next(VerlParquetDataset(source).iter_split("train")) + rendered = render_prompt(record, TinyTokenizer(), source) + + assert rendered.text == expected + assert rendered.truncation_decision == f"truncated_{mode.value}" + + +def test_overlength_error_is_actionable(tmp_path) -> None: + train = _write(tmp_path / "train.parquet", [{"prompt": "abcdef"}]) + source = VerlParquetSourceConfig( + train_files=[str(train)], allow_plain_string_prompts=True, max_prompt_length=3 + ) + record = next(VerlParquetDataset(source).iter_split("train")) + + with pytest.raises(ConfigError, match=r"6 tokens.*max_prompt_length=3"): + render_prompt(record, TinyTokenizer(), source) + + +def test_row_digest_binds_preserved_payload(tmp_path) -> None: + a = _write(tmp_path / "a.parquet", [_row(1)]) + changed = _row(1) + changed["extra_info"] = {"index": 2} + b = _write(tmp_path / "b.parquet", [changed]) + first = next( + VerlParquetDataset(VerlParquetSourceConfig(train_files=[str(a)])).iter_split("train") + ) + second = next( + VerlParquetDataset(VerlParquetSourceConfig(train_files=[str(b)])).iter_split("train") + ) + + assert first.row_digest != second.row_digest + assert json.loads(first.canonical_payload)["extra_info"] == {"index": 1} From c9aa397f5c53b8b79260ee1493d40f946b41b1f2 Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Tue, 11 Aug 2026 23:39:19 -0700 Subject: [PATCH 2/3] Mark prompt rollout tests as torch-dependent --- tests/unit/test_prompt_rollout_runtime.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/unit/test_prompt_rollout_runtime.py b/tests/unit/test_prompt_rollout_runtime.py index 9d29e63..f644ea2 100644 --- a/tests/unit/test_prompt_rollout_runtime.py +++ b/tests/unit/test_prompt_rollout_runtime.py @@ -1,5 +1,10 @@ from __future__ import annotations +import pytest + +pytest.importorskip("torch") +pytestmark = pytest.mark.torch + from miniverl.config.models import RolloutConfig, VerlParquetSourceConfig from miniverl.data.verl_parquet import PromptRecord, RenderedPrompt from miniverl.models.tokenizers import ToyTokenizer From db2d8433af6d09ee8e9cd058132f9b0d9476aa0b Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Tue, 11 Aug 2026 23:39:47 -0700 Subject: [PATCH 3/3] Keep torch marker collection-safe --- tests/unit/test_prompt_rollout_runtime.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/test_prompt_rollout_runtime.py b/tests/unit/test_prompt_rollout_runtime.py index f644ea2..997b3f3 100644 --- a/tests/unit/test_prompt_rollout_runtime.py +++ b/tests/unit/test_prompt_rollout_runtime.py @@ -1,3 +1,5 @@ +# ruff: noqa: E402 + from __future__ import annotations import pytest