From 6fe3083bd9bc260e0a37804471605d00c9dbb396 Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Thu, 13 Aug 2026 23:44:37 -0700 Subject: [PATCH 1/7] Add pinned sampled-k1 OPD profile --- .github/workflows/verl-bridge.yml | 4 +- README.md | 19 +- README.zh-CN.md | 17 +- docs/adr/0010-verl-v0.8-pg-k1-contract.md | 35 ++ docs/for-verl-users.md | 16 +- docs/profiles/index.md | 17 +- docs/verl-opd-runtime.md | 8 +- docs/verl-opd-scaleout.md | 5 +- examples/verl-opd-v0.8-single-gpu-pg-k1.yaml | 90 +++++ mkdocs.yml | 1 + scripts/run_verl_pg_k1_workload.py | 219 ++++++++++ src/miniverl/bridge/doctor.py | 39 +- src/miniverl/bridge/export.py | 69 +++- src/miniverl/bridge/materialize.py | 44 +- src/miniverl/bridge/opd_pg_v08.py | 380 ++++++++++++++++++ src/miniverl/bridge/opd_runtime.py | 90 +++-- src/miniverl/bridge/opd_v08.py | 27 +- src/miniverl/bridge/profiles.py | 190 ++++++++- src/miniverl/cache/store.py | 124 +++++- src/miniverl/cli.py | 34 +- src/miniverl/config/models.py | 63 +++ src/miniverl/losses/chunked.py | 53 +++ src/miniverl/losses/verl_pg.py | 138 +++++++ src/miniverl/models/base.py | 3 + src/miniverl/models/hf.py | 5 + src/miniverl/models/sampling.py | 6 + src/miniverl/models/shared.py | 2 + src/miniverl/models/toy.py | 5 + src/miniverl/runtime/rollout.py | 9 + src/miniverl/schemas/cache.py | 31 +- src/miniverl/teachers/local.py | 115 +++++- src/miniverl/training/trainer.py | 109 ++++- tests/cli/test_opd_product_cli.py | 43 ++ tests/conformance/test_verl_v08_pg_k1.py | 193 +++++++++ tests/integration/test_prompt_opd_pipeline.py | 139 +++++++ tests/unit/test_cache.py | 54 +++ tests/unit/test_profile_registry.py | 3 +- tests/unit/test_prompt_rollout_runtime.py | 42 ++ tests/unit/test_verl_opd_export.py | 56 ++- tests/unit/test_verl_opd_materialize.py | 45 ++- tests/unit/test_verl_pg_k1.py | 103 +++++ tests/unit/test_verl_pg_profile.py | 99 +++++ 42 files changed, 2582 insertions(+), 162 deletions(-) create mode 100644 docs/adr/0010-verl-v0.8-pg-k1-contract.md create mode 100644 examples/verl-opd-v0.8-single-gpu-pg-k1.yaml create mode 100644 scripts/run_verl_pg_k1_workload.py create mode 100644 src/miniverl/bridge/opd_pg_v08.py create mode 100644 src/miniverl/losses/verl_pg.py create mode 100644 tests/conformance/test_verl_v08_pg_k1.py create mode 100644 tests/unit/test_verl_pg_k1.py create mode 100644 tests/unit/test_verl_pg_profile.py diff --git a/.github/workflows/verl-bridge.yml b/.github/workflows/verl-bridge.yml index a61ac31..1e36e07 100644 --- a/.github/workflows/verl-bridge.yml +++ b/.github/workflows/verl-bridge.yml @@ -47,7 +47,7 @@ jobs: changed=$(git diff --name-only "origin/$BASE_REF"...HEAD) echo "changed files:" echo "$changed" - if echo "$changed" | grep -Eq '^(\.github/workflows/verl-bridge\.yml|pyproject\.toml|scripts/[^/]*verl_bridge[^/]*|src/miniverl/bridge/.*|src/miniverl/losses/verl_topk\.py|tests/.*verl_(bridge|opd).*|tests/conformance/test_verl_v08_loss\.py)$'; then + if echo "$changed" | grep -Eq '^(\.github/workflows/verl-bridge\.yml|pyproject\.toml|scripts/[^/]*verl_bridge[^/]*|src/miniverl/bridge/.*|src/miniverl/losses/verl_(topk|pg)\.py|tests/.*verl_(bridge|opd|pg).*|tests/conformance/test_verl_v08_(loss|pg_k1)\.py)$'; then echo "bridge=true" >> "$GITHUB_OUTPUT" else echo "bridge=false" >> "$GITHUB_OUTPUT" @@ -75,6 +75,8 @@ jobs: "git+https://github.com/verl-project/verl.git@7aed6b230776f963fa09509c10d9c3a767d1102c" - name: Compare forward_kl_topk values, diagnostics and gradients run: python -m pytest -q tests/conformance/test_verl_v08_loss.py -m verl_conformance + - name: Compare sampled k1 policy-loss values, gradients and optimizer step + run: python -m pytest -q tests/conformance/test_verl_v08_pg_k1.py -m verl_conformance - name: Materialize and recheck a tiny launchable pure-OPD bundle run: >- python -m pytest -q tests/conformance/test_verl_v08_materialize.py diff --git a/README.md b/README.md index 371c019..beb641a 100644 --- a/README.md +++ b/README.md @@ -130,19 +130,24 @@ explain/check` to inspect the closed built-in registry and distinguish an accepted field from one with a demonstrated native effect. New plans, caches, checkpoints and exports bind the complete independent profile identity. -`verl-opd-v0.8-single-gpu-v1` pins official verl `v0.8.0` at commit -`7aed6b230776f963fa09509c10d9c3a767d1102c`. Its executable path is intentionally -narrow: +Both profiles pin official verl `v0.8.0` at commit +`7aed6b230776f963fa09509c10d9c3a767d1102c`: + +| Profile | Objective | Teacher target | Status | +| --- | --- | --- | --- | +| direct GKD | `forward_kl_topk` | top-k IDs/log-probabilities | measured | +| PG OPD | sampled `k1` + vanilla policy loss | sampled-token teacher log-probability | measured | + +Their shared executable boundary is intentionally narrow: - one trainable actor and one teacher; - one generated response per prompt (`n=1`); -- reward-free generalized knowledge distillation; -- `forward_kl_topk` with teacher top-k IDs/log-probabilities, mass/overlap - diagnostics and token-mean aggregation; +- reward-free distillation with token-mean aggregation; - LoRA or QLoRA adapter updates on one CUDA GPU; - immutable model revisions and verl-style structured prompt Parquet. -Policy-gradient OPD, task rewards, KL penalties, multi-teacher routing, +The PG path is not PPO: it uses a detached distillation-derived advantage with +the pinned vanilla policy-loss form. Task rewards, KL penalties, multi-teacher routing, multimodal inputs, PPO, GRPO, critics, Ray, FSDP, Megatron, multi-GPU and multi-node execution are unsupported. Known unsupported values receive a machine-readable classification; unknown fields and unresolved `${...}` diff --git a/README.zh-CN.md b/README.zh-CN.md index 94e20d9..b63d7ad 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -109,17 +109,24 @@ explain/check` 查看封闭的内置 registry，并区分“字段可接受” 生效”。新生成的 plan、cache、checkpoint 与 export 都绑定完整且独立版本化的 profile identity。 -`verl-opd-v0.8-single-gpu-v1` 固定官方 verl `v0.8.0` 与 commit -`7aed6b230776f963fa09509c10d9c3a767d1102c`。可执行边界有意保持狭窄： +两个 profile 都固定官方 verl `v0.8.0` 与 commit +`7aed6b230776f963fa09509c10d9c3a767d1102c`： + +| Profile | 目标 | Teacher target | 状态 | +| --- | --- | --- | --- | +| direct GKD | `forward_kl_topk` | top-k ID/log-probability | 已实测 | +| PG OPD | sampled `k1` + vanilla policy loss | sampled-token teacher log-probability | 已实测 | + +两者的可执行边界有意保持狭窄： - 一个可训练 actor 与一个 teacher； - 每个 prompt 只生成一个 response（`n=1`）； -- 无 reward 的 generalized knowledge distillation； -- `forward_kl_topk` 使用 teacher top-k token ID/log-probability、top-k mass/overlap 诊断与 token-mean 聚合； +- 无 reward 的 distillation 与 token-mean 聚合； - 一张 CUDA GPU 上的 LoRA/QLoRA adapter 更新； - 不可变模型 revision 与 verl 风格结构化 prompt Parquet。 -PG OPD、task reward、KL penalty、多 teacher、多模态、PPO、GRPO、critic、Ray、FSDP、 +PG 路径不是 PPO：它用 detach 的 distillation advantage 和固定 vanilla policy-loss 形式。 +Task reward、KL penalty、多 teacher、多模态、PPO、GRPO、critic、Ray、FSDP、 Megatron、多 GPU 与多节点均不支持。已知不支持值会得到机器可读分类；未知字段和未解析 的 `${...}` 会被拒绝。输入必须是 resolved profile 子集，而不是任意启动脚本。 diff --git a/docs/adr/0010-verl-v0.8-pg-k1-contract.md b/docs/adr/0010-verl-v0.8-pg-k1-contract.md new file mode 100644 index 0000000..4616cda --- /dev/null +++ b/docs/adr/0010-verl-v0.8-pg-k1-contract.md @@ -0,0 +1,35 @@ +# ADR 0010: pinned verl v0.8 sampled-k1 policy-gradient contract + +Status: accepted for `verl-opd-v0.8-single-gpu-pg-k1-v1`. + +## Decision + +miniVERL implements one closed, reward-free policy-gradient OPD profile pinned +to verl `v0.8.0` commit `7aed6b230776f963fa09509c10d9c3a767d1102c`. +The contract is the composition of upstream `kl_penalty_forward(..., "k1")` +and `compute_policy_loss_vanilla`: rollout-time actor log probability minus +teacher log probability forms the sampled k1 estimator, its detached negative +is the token advantage, and the current/old actor probability ratio is clipped +with the pinned vanilla settings before token-mean aggregation. + +The cache binds sampled token IDs, old actor and teacher log probabilities, +response mask, actor policy version, teacher/tokenizer identity, profile +identity and estimator implementation version. A target cannot cross policy +versions. Each accepted group is generated by the current policy and produces +exactly one optimizer update. + +## Boundary + +This is not PPO. It has no task reward, critic, externally supplied advantage, +reference KL, replay, multiple generations, multi-teacher routing, multimodal +input or distributed execution. Only `loss_mode: k1`, +`use_policy_gradient: true`, `policy_loss_mode: vanilla`, `n: 1`, and the +conformance-tested clipping values are accepted. Other upstream options fail +closed instead of being approximated. + +## Verification + +The optional conformance suite loads the exact pinned upstream function bodies +and compares estimator tensors, detached advantages, scalar loss, metrics, +actor-logit gradients and a deterministic optimizer step. Runtime evidence is +systems/conformance evidence only; it makes no task-quality or alignment claim. diff --git a/docs/for-verl-users.md b/docs/for-verl-users.md index c79e9c3..6c7cd6d 100644 --- a/docs/for-verl-users.md +++ b/docs/for-verl-users.md @@ -94,10 +94,13 @@ bytes and schema are copied into export provenance. ## Actor, teacher and reference roles -The actor is the trainable PEFT policy. The teacher produces top-k token IDs -and log-probabilities on actor-generated tokens; top-k mass and overlap are -diagnostics, not an explicit tail bucket. Pure GKD in this profile has no reference -policy, critic, task reward or policy-gradient term. Local runtime strategies +The actor is the trainable PEFT policy. In the direct-GKD profile the teacher +produces top-k token IDs and log-probabilities on actor-generated tokens; +top-k mass and overlap are diagnostics, not an explicit tail bucket. In the PG +k1 profile the target is only the sampled token's teacher log-probability, +bound to the rollout actor log-probability and policy version; current actor +log-probability is recomputed at update time. Neither profile has a reference +policy, critic or task reward. Local runtime strategies may keep quantized roles resident, swap movable unquantized roles, or share a compatible backbone, but role identities never collapse in provenance. @@ -121,8 +124,9 @@ tested. - **Unknown field:** capture a resolved config and remove fields outside the documented profile; inspect-only compilation can still explain known unsupported values. -- **Algorithm field rejected:** PG OPD, rewards, KL penalties, `n>1`, multiple - teachers and distributed counts are intentionally unsupported. +- **Algorithm field rejected:** select the explicit PG-k1 profile for its + narrow sampled policy-loss path. Rewards, critics, external advantages, KL + penalties, `n>1`, multiple teachers and distributed counts remain unsupported. - **Interpolation rejected:** resolve Hydra/OmegaConf in your trusted verl environment first. miniVERL will not execute `${...}`. - **High-risk reinterpretation not accepted:** inspect `miniverl plan`, then diff --git a/docs/profiles/index.md b/docs/profiles/index.md index c2274ba..0ba0718 100644 --- a/docs/profiles/index.md +++ b/docs/profiles/index.md @@ -39,6 +39,17 @@ scale-out reports carry the complete profile identity. Changing any version axis produces a new digest. Existing profile names cannot silently acquire new semantics; a later algorithm path is registered under a different name. -Currently measured: `verl-opd-v0.8-single-gpu-v1`, the reward-free direct-GKD -`forward_kl_topk` path. See [For verl users](../for-verl-users.md) for the -algorithm and execution boundary. +## Which profile should I use? + +| Profile | Objective | Teacher target | Trade-off | Status | +| --- | --- | --- | --- | --- | +| `verl-opd-v0.8-single-gpu-v1` | direct GKD `forward_kl_topk` | top-k token IDs and log-probabilities | fuller distributional signal; larger target artifact | measured | +| `verl-opd-v0.8-single-gpu-pg-k1-v1` | sampled `k1` + vanilla policy loss | sampled-token teacher log-probability | sampled-token signal; smaller target artifact | measured | + +Both are reward-free, strict current-policy, one-actor/one-teacher profiles on +one CUDA GPU. The PG profile is not PPO: it uses a policy-loss form with a +detached distillation-derived advantage and has no task reward, critic, +reference KL or replay. Neither profile is claimed to be more accurate or +better aligned. See [For verl users](../for-verl-users.md), the +[PG-k1 ADR](../adr/0010-verl-v0.8-pg-k1-contract.md), and the runtime evidence +for the measured boundary. diff --git a/docs/verl-opd-runtime.md b/docs/verl-opd-runtime.md index feea46f..7ad7592 100644 --- a/docs/verl-opd-runtime.md +++ b/docs/verl-opd-runtime.md @@ -1,6 +1,6 @@ # Current verl-style OPD runtime -This is miniVERL's current executable path: one documented, resolved subset of +This is miniVERL's current executable path: two documented, resolved subsets of official verl `v0.8.0`, pinned at `7aed6b23`, compiled into local phases on one NVIDIA CUDA GPU. It is not a distributed verl runtime and does not accept arbitrary Hydra YAML. @@ -38,6 +38,12 @@ reinterpretations; the built-in profile has a value-bound approval manifest. reports top-k mass/overlap diagnostics. It does not create the explicit K+1 tail bucket used by miniVERL's separate native `bucketed_topk_tail` objective. +The second profile, `verl-opd-v0.8-single-gpu-pg-k1-v1`, records the sampled +token, rollout-time actor log-probability and teacher log-probability, then +recomputes the current actor log-probability for pinned `k1` + vanilla +policy-loss semantics. It exports no top-k requirement. See [Which profile +should I use?](profiles/index.md) for the neutral choice boundary. + ## Placement is fail-closed - NF4/int8 roles use resident local phases; bitsandbytes parameters cannot swap. diff --git a/docs/verl-opd-scaleout.md b/docs/verl-opd-scaleout.md index 38984e3..aa7cbe9 100644 --- a/docs/verl-opd-scaleout.md +++ b/docs/verl-opd-scaleout.md @@ -1,6 +1,6 @@ # Current OPD scale-out contract -A completed `verl-opd-v0.8-single-gpu-v1` run can export portable PEFT, +A completed direct-GKD or sampled-k1 PG profile run can export portable PEFT, Parquet, config and provenance artifacts for the exact pinned upstream profile. This is an artifact handoff, not proof that a distributed job ran or that miniVERL and verl have algorithmic parity. @@ -13,7 +13,8 @@ miniverl bridge doctor scaleout --require-verl The export preserves source-run student and teacher identities, model revisions, prompt/response bounds, learning rate, schedule, Parquet bytes and -pure `forward_kl_topk` overrides. Any inserted placeholder is listed; missing +the selected profile's exact `forward_kl_topk` or sampled-k1 policy-loss +overrides. The PG bundle has no top-k requirement. Any inserted placeholder is listed; missing validation data remains missing rather than being replaced with training data. ## Readiness states stay separate diff --git a/examples/verl-opd-v0.8-single-gpu-pg-k1.yaml b/examples/verl-opd-v0.8-single-gpu-pg-k1.yaml new file mode 100644 index 0000000..edc5947 --- /dev/null +++ b/examples/verl-opd-v0.8-single-gpu-pg-k1.yaml @@ -0,0 +1,90 @@ +# Resolved fixture for verl-opd-v0.8-single-gpu-pg-k1-v1. +data: + train_files: [data/train.parquet] + val_files: [data/val.parquet] + prompt_key: prompt + train_batch_size: 8 + max_prompt_length: 256 + max_response_length: 64 + filter_overlong_prompts: false + truncation: error + shuffle: true + seed: 17 +actor_rollout_ref: + model: + path: Qwen/Qwen3-0.6B + enable_gradient_checkpointing: true + lora_rank: 16 + lora_alpha: 32 + target_modules: [q_proj, v_proj] + lora_adapter_path: null + actor: + optim: {lr: 1e-5, weight_decay: 0.01, lr_warmup_steps: 0} + loss_agg_mode: token-mean + use_kl_loss: false + ppo_mini_batch_size: 8 + ppo_max_token_len_per_gpu: 2048 + use_dynamic_bsz: true + rollout: + name: vllm + n: 1 + temperature: 1.0 + top_p: 0.95 + tensor_model_parallel_size: 1 + gpu_memory_utilization: 0.5 + max_model_len: 320 + max_num_batched_tokens: 2048 + max_num_seqs: 8 +algorithm: + adv_estimator: null + use_kl_in_reward: false +distillation: + enabled: true + teacher_key: data_source + n_gpus_per_node: 1 + nnodes: 1 + teacher_models: + teacher_model: + model_path: Qwen/Qwen3-1.7B + num_replicas: 1 + inference: + name: vllm + dtype: bfloat16 + tensor_model_parallel_size: 1 + data_parallel_size: 1 + pipeline_model_parallel_size: 1 + gpu_memory_utilization: 0.5 + max_model_len: 321 + distillation_loss: + loss_mode: k1 + topk: null + use_task_rewards: false + distillation_loss_coef: 1.0 + loss_max_clamp: 10.0 + log_prob_min_clamp: null + use_policy_gradient: true + policy_loss_mode: vanilla + clip_ratio: 0.2 + clip_ratio_low: 0.2 + clip_ratio_high: 0.2 +trainer: + project_name: mini-verl + experiment_name: pg-k1-smoke + save_freq: 10 + test_freq: 10 + total_epochs: 1 + total_training_steps: 4 + n_gpus_per_node: 1 + nnodes: 1 +miniverl: + student_revision: c1899de289a04d12100db370d81485cdf75e47ca + teacher_revision: 70d244cc86ccca08cf5af4e1e306ecf908b1ad5e + actor_runtime: {dtype: auto, quantization: nf4, attn_implementation: sdpa} + teacher_runtime: {quantization: nf4, attn_implementation: sdpa} + runtime: {mode: auto} + memory: {vram_limit_gib: 16, headroom_gib: 1.5} + batching: + rollout_batch_size: 2 + teacher_score_batch_size: 2 + update_trajectory_batch_size: 2 + teacher_adapter: {path: null, revision: null} diff --git a/mkdocs.yml b/mkdocs.yml index ee2be8f..2f5b6f4 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -43,6 +43,7 @@ nav: - Home: index.md - For verl users: for-verl-users.md - Compatibility profiles: profiles/index.md + - PG-k1 semantic ADR: adr/0010-verl-v0.8-pg-k1-contract.md - Config overrides: config-overrides.md - Immutable plans: immutable-plans.md - Hardware planning: hardware-planning.md diff --git a/scripts/run_verl_pg_k1_workload.py b/scripts/run_verl_pg_k1_workload.py new file mode 100644 index 0000000..0d1c7b3 --- /dev/null +++ b/scripts/run_verl_pg_k1_workload.py @@ -0,0 +1,219 @@ +"""Run the bounded RTX 4080 sampled-k1 policy-gradient OPD workload. + +This records systems and semantic-conformance evidence only. It does not score +task quality or compare algorithms. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.metadata +import json +import subprocess +import time +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from run_verl_opd_reference_workload import ( + INTERRUPT_AFTER, + LOGICAL_BATCH, + PROMPT_LIMIT, + RESPONSE_LIMIT, + UPDATES, + _equivalence, + _train_resumed, + _train_uninterrupted, + _tree_bytes, + summarize_run, + write_dataset, +) + +from miniverl import __version__ +from miniverl.bridge.opd_pg_v08 import ( + VERL_OPD_PG_K1_V08_PROFILE, + load_verl_pg_k1_v08_source, +) +from miniverl.bridge.opd_plan import build_immutable_opd_plan, write_immutable_opd_plan +from miniverl.bridge.opd_runtime import build_system_plan +from miniverl.utils.runs import write_json_atomic + +EXAMPLE = Path("examples/verl-opd-v0.8-single-gpu-pg-k1.yaml") + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _package_version(name: str) -> str | None: + try: + return importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + return None + + +def _driver() -> str: + return ( + subprocess.run( + ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"], + check=True, + capture_output=True, + text=True, + ) + .stdout.splitlines()[0] + .strip() + ) + + +def _overrides(dataset: Path) -> list[str]: + return [ + f'data.train_files=["{dataset.resolve().as_posix()}"]', + "data.val_files=[]", + f"data.train_batch_size={LOGICAL_BATCH}", + f"data.max_prompt_length={PROMPT_LIMIT}", + f"data.max_response_length={RESPONSE_LIMIT}", + f"actor_rollout_ref.actor.ppo_mini_batch_size={LOGICAL_BATCH}", + "actor_rollout_ref.actor.ppo_max_token_len_per_gpu=768", + "actor_rollout_ref.rollout.max_model_len=192", + "actor_rollout_ref.rollout.max_num_batched_tokens=768", + "actor_rollout_ref.rollout.max_num_seqs=4", + "trainer.experiment_name=qwen3-pg-k1-systems-workload", + "trainer.save_freq=4", + f"trainer.total_training_steps={UPDATES}", + "miniverl.batching.rollout_batch_size=4", + "miniverl.batching.teacher_score_batch_size=4", + "miniverl.batching.update_trajectory_batch_size=1", + ] + + +def build_plan(dataset: Path, path: Path) -> tuple[Any, Any]: + compiled = load_verl_pg_k1_v08_source( + EXAMPLE, + overrides=_overrides(dataset), + accept_local_reinterpretations=True, + ) + system = build_system_plan(compiled) + plan = build_immutable_opd_plan(compiled, source=str(EXAMPLE), system_plan=system) + write_immutable_opd_plan(path, plan) + from miniverl.config import RunConfig + + return plan, RunConfig.model_validate(plan.resolved_native_config) + + +def run_workload(out: Path, result_path: Path, *, offline: bool) -> dict[str, Any]: + if out.exists(): + raise FileExistsError(f"output directory already exists: {out}") + out.mkdir(parents=True) + started = time.perf_counter() + dataset = out / "data/reference-workload.parquet" + dataset_identity = write_dataset(dataset) + plan, native = build_plan(dataset, out / "plan.json") + reference_run = _train_uninterrupted(native, plan, out / "runs", offline=offline) + reference = summarize_run( + reference_run["run"], construction_seconds=reference_run["construction_seconds"] + ) + resumed_run = _train_resumed(native, plan, out / "runs", offline=offline) + resumed = summarize_run(resumed_run["run"], construction_seconds=0.0) + equivalence = _equivalence(reference, resumed) + + import torch + + device = torch.cuda.get_device_properties(torch.cuda.current_device()) + payload: dict[str, Any] = { + "schema_version": 1, + "kind": "single_gpu_verl_pg_k1_systems_workload", + "status": "measured", + "measured_at": datetime.now(timezone.utc) + .replace(microsecond=0) + .isoformat() + .replace("+00:00", "Z"), + "miniverl_version": __version__, + "source_commit": subprocess.run( + ["git", "rev-parse", "HEAD"], check=True, capture_output=True, text=True + ).stdout.strip(), + "workload_driver_sha256": _sha256(Path(__file__)), + "profile": VERL_OPD_PG_K1_V08_PROFILE, + "profile_identity": plan.profile_identity, + "verl": {**plan.pinned_verl, "distributed_execution_tested": False}, + "hardware": { + "gpu": device.name, + "gpu_count": 1, + "vram_gib": round(int(device.total_memory) / 1024**3, 3), + "driver": _driver(), + "torch": torch.__version__, + "cuda_runtime": str(torch.version.cuda), + "transformers": _package_version("transformers"), + "peft": _package_version("peft"), + "bitsandbytes": _package_version("bitsandbytes"), + }, + "models": { + "student": native.models.student.model_dump(mode="json"), + "teacher": native.models.teacher.model_dump(mode="json"), + }, + "recipe": { + "prompts_available": dataset_identity["distinct_prompts"], + "prompts_consumed": reference["distinct_prompts_consumed"], + "prompt_limit": PROMPT_LIMIT, + "response_limit": RESPONSE_LIMIT, + "logical_batch": LOGICAL_BATCH, + "optimizer_updates": UPDATES, + "interrupt_after_updates": INTERRUPT_AFTER, + "teacher_target": "sampled_token_log_probability", + "estimator": "k1", + "policy_loss": "vanilla", + "compiled_plan_sha256": plan.plan_digest, + "input_parquet_sha256": dataset_identity["sha256"], + }, + "measurements": { + key: value + for key, value in reference.items() + if key not in {"checkpoint_hashes", "checkpoint_state", "trajectory_sha256"} + }, + "resume": { + "resume_load_seconds": round(resumed_run["resume_load_seconds"], 4), + "interrupt_checkpoint_bytes": resumed_run["interrupt_checkpoint_bytes"], + **equivalence, + }, + "artifacts": { + "checkpoint_hashes": reference["checkpoint_hashes"], + "trajectory_sha256": reference["trajectory_sha256"], + "standard_peft_adapter_sha256": reference_run["adapter_manifest"]["checksums"][ + "adapter_model.safetensors" + ], + "total_workload_bytes": _tree_bytes(out), + }, + "resource_contract": { + "peak_reserved_limit_gib": 14.5, + "peak_reserved_within_limit": reference["peak_reserved_gib"] <= 14.5, + "total_gpu_workload_seconds": round(time.perf_counter() - started, 2), + "gpu_hour_limit": 3.0, + }, + "scientific_scope": { + "runtime_correctness_only": True, + "task_quality_evaluated": False, + "algorithm_comparison": False, + }, + } + if not payload["resource_contract"]["peak_reserved_within_limit"]: + raise RuntimeError("PG-k1 workload exceeded the 14.5 GiB reserved-memory limit") + result_path.parent.mkdir(parents=True, exist_ok=True) + write_json_atomic(result_path, payload) + return payload + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", required=True, type=Path) + parser.add_argument("--result", required=True, type=Path) + parser.add_argument("--offline", action="store_true") + args = parser.parse_args() + print(json.dumps(run_workload(args.out, args.result, offline=args.offline), indent=2)) + + +if __name__ == "__main__": + main() diff --git a/src/miniverl/bridge/doctor.py b/src/miniverl/bridge/doctor.py index d59b569..48446d6 100644 --- a/src/miniverl/bridge/doctor.py +++ b/src/miniverl/bridge/doctor.py @@ -18,6 +18,8 @@ VERL_REPOSITORY, VERL_TAG, ) +from miniverl.bridge.opd_pg_v08 import VERL_OPD_PG_K1_V08_PROFILE +from miniverl.bridge.opd_v08 import VERL_OPD_V08_PROFILE from miniverl.bridge.preflight import preflight_bundle_tree from miniverl.bridge.reward_static import REWARD_LEVELS, inspect_reward_scaffold from miniverl.bridge.safetensors_check import SAFETENSORS_LEVELS, inspect_safetensors @@ -388,11 +390,15 @@ def _check_config(root: Path) -> dict[str, Any]: teacher_identity = json.loads( (root / "teacher" / "teacher-model.json").read_text(encoding="utf-8") ) + compatibility = json.loads( + (root / "provenance" / "compatibility-report.json").read_text(encoding="utf-8") + ) except (OSError, json.JSONDecodeError, yaml.YAMLError) as exc: return {"status": "fail", "detail": str(exc)} required_roots = {"data", "actor_rollout_ref", "algorithm", "distillation", "trainer"} actual = set(payload) if isinstance(payload, dict) else set() opd_problems: list[str] = [] + profile = compatibility.get("profile") if not required_roots.issubset(actual): opd_problems.append("missing required pure-OPD root") try: @@ -408,11 +414,24 @@ def _check_config(root: Path) -> dict[str, Any]: if model.get(field) != value: opd_problems.append(f"actor_rollout_ref.model.{field}") loss = payload["distillation"]["distillation_loss"] - semantic_expected = { - "loss_mode": "forward_kl_topk", - "use_task_rewards": False, - "use_policy_gradient": False, - } + semantic_expected = ( + { + "loss_mode": "k1", + "use_task_rewards": False, + "use_policy_gradient": True, + "policy_loss_mode": "vanilla", + } + if profile == VERL_OPD_PG_K1_V08_PROFILE + else { + "loss_mode": "forward_kl_topk", + "use_task_rewards": False, + "use_policy_gradient": False, + } + ) + if profile not in {VERL_OPD_V08_PROFILE, VERL_OPD_PG_K1_V08_PROFILE}: + opd_problems.append("unregistered pure-OPD profile") + if profile == VERL_OPD_PG_K1_V08_PROFILE and "topk" in loss: + opd_problems.append("distillation.distillation_loss.topk") for field, value in semantic_expected.items(): if loss.get(field) != value: opd_problems.append(f"distillation.distillation_loss.{field}") @@ -435,7 +454,8 @@ def _check_config(root: Path) -> dict[str, Any]: opd_problems.append("invalid pure-OPD handoff structure") return { "status": "ok" if not opd_problems else "fail", - "profile": "verl-opd-v0.8-single-gpu-v1", + "profile": profile, + "target_representation": compatibility.get("target_representation"), "roots": sorted(actual), "model_handoff_problems": opd_problems, } @@ -1229,7 +1249,10 @@ def inspect_bridge_bundle( snapshot_checks = (student_snapshot, teacher_snapshot) if snapshot_required else () checks = (target, model, tokenizer, parquet, config, reward, hashes, privacy, *snapshot_checks) artifact_failed = any(check.get("status") != "ok" for check in checks) - opd_profile = config.get("profile") == "verl-opd-v0.8-single-gpu-v1" + opd_profile = config.get("profile") in { + VERL_OPD_V08_PROFILE, + VERL_OPD_PG_K1_V08_PROFILE, + } materialized_complete = not opd_profile or ( student_snapshot.get("status") == "ok" and teacher_snapshot.get("status") == "ok" ) @@ -1298,7 +1321,7 @@ def inspect_bridge_bundle( "teacher_artifact_loadable": teacher_snapshot.get("status") == "ok", "dataset_loadable": parquet.get("status") == "ok", "reward_implementation_complete": bool(reward.get("implementation_complete", False)), - "reward_required": config.get("profile") != "verl-opd-v0.8-single-gpu-v1", + "reward_required": not opd_profile, "launchable": bool( require_verl and upstream["status"] == "passed" diff --git a/src/miniverl/bridge/export.py b/src/miniverl/bridge/export.py index 2b73d6b..6deea8c 100644 --- a/src/miniverl/bridge/export.py +++ b/src/miniverl/bridge/export.py @@ -25,6 +25,8 @@ required_verl_text, validate_target_verl, ) +from miniverl.bridge.opd_pg_v08 import VERL_OPD_PG_K1_V08_PROFILE +from miniverl.bridge.opd_v08 import VERL_OPD_V08_PROFILE from miniverl.errors import ConfigError from miniverl.utils.privacy import portable_payload from miniverl.utils.runs import read_json, write_json, write_text @@ -426,7 +428,7 @@ def _write_hashes(root: Path) -> None: write_text(checksum, "\n".join(lines) + "\n") -_OPD_PROFILE = "verl-opd-v0.8-single-gpu-v1" +_OPD_PROFILES = frozenset({VERL_OPD_V08_PROFILE, VERL_OPD_PG_K1_V08_PROFILE}) def _opd_run_contract(run: Path) -> tuple[dict[str, Any], dict[str, Any]] | None: @@ -440,17 +442,28 @@ def _opd_run_contract(run: Path) -> tuple[dict[str, Any], dict[str, Any]] | None source = read_json(source_path) if source_path.is_file() else report.get("source") except (OSError, json.JSONDecodeError) as exc: raise ConfigError(f"cannot read source-run OPD compatibility artifacts: {exc}") from exc - if report.get("profile") != _OPD_PROFILE: + profile = report.get("profile") + if profile not in _OPD_PROFILES: return None if report.get("executable") is not True or not isinstance(source, dict): raise ConfigError("source run does not carry an executable verl OPD compatibility plan") - required = { - "distillation.distillation_loss.loss_mode": "forward_kl_topk", + required: dict[str, Any] = { "distillation.distillation_loss.use_task_rewards": False, - "distillation.distillation_loss.use_policy_gradient": False, "actor_rollout_ref.actor.use_kl_loss": False, "algorithm.use_kl_in_reward": False, } + required.update( + { + "distillation.distillation_loss.loss_mode": "k1", + "distillation.distillation_loss.use_policy_gradient": True, + "distillation.distillation_loss.policy_loss_mode": "vanilla", + } + if profile == VERL_OPD_PG_K1_V08_PROFILE + else { + "distillation.distillation_loss.loss_mode": "forward_kl_topk", + "distillation.distillation_loss.use_policy_gradient": False, + } + ) for field, expected in required.items(): if _get(source, field) != expected: raise ConfigError( @@ -508,7 +521,11 @@ def _copy_opd_data( def _opd_overrides( - source: dict[str, Any], adapter: dict[str, Any], data_paths: dict[str, list[str]] + source: dict[str, Any], + adapter: dict[str, Any], + data_paths: dict[str, list[str]], + *, + profile: str, ) -> dict[str, Any]: """Preserve the supported source profile while rebasing portable paths.""" overrides: dict[str, Any] = { @@ -532,6 +549,8 @@ def _opd_overrides( # no such teacher inference key. It therefore belongs in provenance, not # in an override file that promises to parse upstream. teacher.get("inference", {}).pop("pipeline_model_parallel_size", None) + if profile == VERL_OPD_PG_K1_V08_PROFILE: + overrides["distillation"]["distillation_loss"].pop("topk", None) return overrides @@ -559,11 +578,16 @@ def _opd_launch_script(adapter: dict[str, Any], teacher: dict[str, Any]) -> str: """ -def _opd_bundle_readme() -> str: +def _opd_bundle_readme(profile: str) -> str: + semantics = ( + "sampled k1 teacher log-probabilities and pinned vanilla policy-gradient loss" + if profile == VERL_OPD_PG_K1_V08_PROFILE + else "pure GKD `forward_kl_topk` targets" + ) return f"""# miniVERL pure-OPD scale-out bundle -This checksummed bundle preserves a local `{_OPD_PROFILE}` run as standard -PEFT, tokenizer and Parquet artifacts plus pure GKD `forward_kl_topk` +This checksummed bundle preserves a local `{profile}` run as standard +PEFT, tokenizer and Parquet artifacts plus {semantics} overrides for [`verl {VERL_TAG}`]({VERL_REPOSITORY}/tree/{VERL_TAG}) at `{VERL_COMMIT}`. It contains no reward scaffold because task rewards are disabled by contract. @@ -585,6 +609,7 @@ def _export_opd_bundle( compatibility: dict[str, Any], destination: Path, ) -> dict[str, Any]: + profile = str(compatibility["profile"]) student_id = _get(source, "actor_rollout_ref.model.path") student_revision = _get(source, "miniverl.student_revision") if student_id != adapter["base_model"] or student_revision != adapter["revision"]: @@ -640,7 +665,7 @@ def _export_opd_bundle( if teacher_adapter_path is not None and teacher_adapter_source is not None: _copy_adapter_payload(teacher_adapter_source, teacher_dir / "adapter") data_paths, data_evidence = _copy_opd_data(run, source, temporary / "data") - overrides = _opd_overrides(source, adapter, data_paths) + overrides = _opd_overrides(source, adapter, data_paths, profile=profile) recipe = temporary / "recipe" recipe.mkdir() write_text( @@ -662,7 +687,7 @@ def _export_opd_bundle( ) report: dict[str, Any] = { "schema_version": 2, - "profile": _OPD_PROFILE, + "profile": profile, "profile_identity": read_json(manifest_path).get("profile_identity") or {}, "target_verl": { "repository": VERL_REPOSITORY, @@ -670,7 +695,16 @@ def _export_opd_bundle( "commit": VERL_COMMIT, }, "miniverl_version": __version__, - "target_semantics": "pure GKD forward_kl_topk OPD", + "target_semantics": ( + "pure OPD sampled k1 policy gradient" + if profile == VERL_OPD_PG_K1_V08_PROFILE + else "pure GKD forward_kl_topk OPD" + ), + "target_representation": ( + "sampled_token_log_probability" + if profile == VERL_OPD_PG_K1_V08_PROFILE + else "topk_distribution" + ), "artifact_complete": False, "artifact_bundle_complete": True, "config_semantics_supported": True, @@ -687,11 +721,18 @@ def _export_opd_bundle( "reward_required": False, "launch_blockers": blockers, "data_round_trip": data_evidence, - "unsupported_semantics": list(_UNSUPPORTED), + "unsupported_semantics": [ + item + for item in _UNSUPPORTED + if not ( + profile == VERL_OPD_PG_K1_V08_PROFILE + and item == "PPO advantage or clipping semantics" + ) + ], "distributed_execution_status": "not tested", } write_json(provenance / "compatibility-report.json", report) - write_text(temporary / "README.md", _opd_bundle_readme()) + write_text(temporary / "README.md", _opd_bundle_readme(profile)) _write_hashes(temporary) temporary.replace(destination) except BaseException: diff --git a/src/miniverl/bridge/materialize.py b/src/miniverl/bridge/materialize.py index eb6d7e7..1762b4e 100644 --- a/src/miniverl/bridge/materialize.py +++ b/src/miniverl/bridge/materialize.py @@ -21,6 +21,8 @@ from miniverl import __version__ from miniverl.bridge.contract import VERL_COMMIT +from miniverl.bridge.opd_pg_v08 import VERL_OPD_PG_K1_V08_PROFILE +from miniverl.bridge.opd_v08 import VERL_OPD_V08_PROFILE from miniverl.bridge.preflight import preflight_bundle_tree from miniverl.bridge.safetensors_check import inspect_safetensors from miniverl.errors import ConfigError @@ -450,6 +452,8 @@ def _load_local_transformers_metadata( def _validate_upstream_bundle(root: Path) -> dict[str, Any]: + report = read_json(root / "provenance/compatibility-report.json") + profile = report.get("profile") try: distribution = importlib.metadata.distribution("verl") except importlib.metadata.PackageNotFoundError as exc: @@ -499,12 +503,34 @@ def _validate_upstream_bundle(root: Path) -> dict[str, Any]: ) teacher = _load_local_transformers_metadata(root / "teacher/base", role="teacher") recipe = yaml.safe_load((root / "recipe/verl-opd-overrides.yaml").read_text(encoding="utf-8")) - topk = recipe["distillation"]["distillation_loss"].get("topk") - if isinstance(topk, bool) or not isinstance(topk, int) or topk < 1: - raise ConfigError("pure-OPD top-k must be a positive integer") - for role, metadata in (("student", student), ("teacher", teacher)): - if topk > metadata["vocab_size"]: - raise ConfigError(f"top-k {topk} exceeds the {role} tokenizer vocabulary") + loss = recipe["distillation"]["distillation_loss"] + semantic_checks: dict[str, Any] + if profile == VERL_OPD_PG_K1_V08_PROFILE: + semantic_checks = { + "loss_mode": "k1", + "use_policy_gradient": True, + "policy_loss_mode": "vanilla", + "use_task_rewards": False, + } + if "topk" in loss: + raise ConfigError("sampled-k1 PG materialization forbids top-k targets") + target_check = "sampled_k1_policy_gradient_contract" + else: + semantic_checks = { + "loss_mode": "forward_kl_topk", + "use_policy_gradient": False, + "use_task_rewards": False, + } + topk = loss.get("topk") + if isinstance(topk, bool) or not isinstance(topk, int) or topk < 1: + raise ConfigError("pure-OPD top-k must be a positive integer") + for role, metadata in (("student", student), ("teacher", teacher)): + if topk > metadata["vocab_size"]: + raise ConfigError(f"top-k {topk} exceeds the {role} tokenizer vocabulary") + target_check = "topk_contract" + for field, expected in semantic_checks.items(): + if loss.get(field) != expected: + raise ConfigError(f"profile loss contract requires {field}={expected!r}") if student["tokenizer_structural_digest_v2"] != teacher["tokenizer_structural_digest_v2"]: raise ConfigError("student and teacher tokenizer structural identities differ") return { @@ -521,7 +547,7 @@ def _validate_upstream_bundle(root: Path) -> dict[str, Any]: "student_peft": "passed", "student_snapshot": student, "teacher_snapshot": teacher, - "topk_contract": "passed", + target_check: "passed", }, } @@ -580,8 +606,8 @@ def materialize_verl_bundle( teacher_identity = read_json(root / "teacher/teacher-model.json") except (OSError, json.JSONDecodeError) as exc: raise ConfigError(f"cannot read scale-out bundle identities: {exc}") from exc - if report.get("profile") != "verl-opd-v0.8-single-gpu-v1": - raise ConfigError("materialize supports only the pure-OPD v0.8 single-GPU profile") + if report.get("profile") not in {VERL_OPD_V08_PROFILE, VERL_OPD_PG_K1_V08_PROFILE}: + raise ConfigError("materialize supports only a registered pure-OPD v0.8 profile") student_id = str(student_identity.get("model_id") or "") teacher_id = str(teacher_identity.get("model_id") or "") student_revision = _immutable_revision(student_identity.get("revision"), role="student") diff --git a/src/miniverl/bridge/opd_pg_v08.py b/src/miniverl/bridge/opd_pg_v08.py new file mode 100644 index 0000000..e38f60c --- /dev/null +++ b/src/miniverl/bridge/opd_pg_v08.py @@ -0,0 +1,380 @@ +"""Typed compiler for the pinned verl v0.8 sampled-k1 PG subset.""" + +from __future__ import annotations + +import copy +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any, Literal + +import yaml +from pydantic import BaseModel, ValidationError + +from miniverl.bridge.contract import VERL_COMMIT, VERL_REPOSITORY, VERL_TAG +from miniverl.bridge.interpolation import reject_interpolation +from miniverl.bridge.opd_v08 import ( + _FIELD_RULES, + CompatibilityEntry, + CompiledLocalExecutionPlan, + FieldClassification, + VerlOPDV08Profile, + _canonical_digest, + _flatten, + _reject_non_finite_source_numbers, + _resolve_overrides, + _Rule, + _rule, +) +from miniverl.errors import ConfigError + +__all__ = [ + "VERL_OPD_PG_K1_V08_PROFILE", + "compile_verl_pg_k1_v08", + "load_verl_pg_k1_v08", + "load_verl_pg_k1_v08_source", + "pg_compatibility_rule", + "pg_field_rules_digest", +] + +VERL_OPD_PG_K1_V08_PROFILE = "verl-opd-v0.8-single-gpu-pg-k1-v1" + +_PG_FIELD_RULES: dict[str, _Rule] = dict(_FIELD_RULES) +_PG_FIELD_RULES.update( + { + "distillation.distillation_loss.loss_mode": _rule( + "loss.mode", "semantically_conformant", "dedicated pinned sampled k1 estimator" + ), + "distillation.distillation_loss.topk": _rule( + None, + "informational_only", + "not applicable because k1 stores only the sampled-token teacher log-probability", + ), + "distillation.distillation_loss.use_policy_gradient": _rule( + "loss.use_policy_gradient", "exact", "must remain enabled for the PG-k1 profile" + ), + "distillation.distillation_loss.policy_loss_mode": _rule( + "loss.policy_loss_mode", "exact", "pinned verl vanilla policy loss" + ), + "distillation.distillation_loss.clip_ratio": _rule( + "loss.clip_ratio", "exact", "pinned vanilla policy-loss clip ratio" + ), + "distillation.distillation_loss.clip_ratio_low": _rule( + "loss.clip_ratio_low", "exact", "pinned vanilla policy-loss lower clip ratio" + ), + "distillation.distillation_loss.clip_ratio_high": _rule( + "loss.clip_ratio_high", "exact", "pinned vanilla policy-loss upper clip ratio" + ), + "distillation.distillation_loss.log_prob_min_clamp": _rule( + None, + "informational_only", + "not applied by the sampled k1 estimator; the supported value is null", + ), + } +) + + +def pg_compatibility_rule(field: str) -> dict[str, Any]: + """Return one immutable PG-profile field rule.""" + try: + rule = _PG_FIELD_RULES[field] + except KeyError as exc: + raise ConfigError( + f"field {field!r} is not part of profile {VERL_OPD_PG_K1_V08_PROFILE!r}" + ) from exc + return { + "upstream_field": field, + "local_target": rule.target, + "classification": rule.classification, + "reason": rule.reason, + "semantic_risk": rule.risk, + "user_confirmation_required": rule.confirmation, + } + + +def pg_field_rules_digest() -> str: + """Bind the PG profile identity to its complete field-rule table.""" + return _canonical_digest([pg_compatibility_rule(field) for field in sorted(_PG_FIELD_RULES)]) + + +def _pg_semantic_blocker(path: str, value: Any) -> str | None: + required = { + "distillation.enabled": (True, "distillation must be enabled"), + "distillation.distillation_loss.loss_mode": ( + "k1", + "only the pinned sampled k1 estimator is supported", + ), + "distillation.distillation_loss.use_policy_gradient": ( + True, + "policy-gradient optimization must be enabled", + ), + "distillation.distillation_loss.policy_loss_mode": ( + "vanilla", + "only the pinned vanilla policy loss is supported", + ), + "distillation.distillation_loss.use_task_rewards": ( + False, + "task-reward mixtures are unsupported", + ), + "distillation.distillation_loss.topk": ( + None, + "top-k targets are not part of sampled k1", + ), + "distillation.distillation_loss.log_prob_min_clamp": ( + None, + "log-probability clamping is not part of the pinned sampled k1 path", + ), + "distillation.distillation_loss.clip_ratio": ( + 0.2, + "only the conformance-tested clip_ratio=0.2 is supported", + ), + "distillation.distillation_loss.clip_ratio_low": ( + 0.2, + "only the conformance-tested clip_ratio_low=0.2 is supported", + ), + "distillation.distillation_loss.clip_ratio_high": ( + 0.2, + "only the conformance-tested clip_ratio_high=0.2 is supported", + ), + "actor_rollout_ref.actor.use_kl_loss": ( + False, + "actor/reference KL loss is outside pure OPD", + ), + "algorithm.use_kl_in_reward": (False, "KL reward is outside pure OPD"), + "actor_rollout_ref.actor.loss_agg_mode": ( + "token-mean", + "only token-mean aggregation is supported", + ), + "actor_rollout_ref.rollout.n": (1, "one generation per prompt is required"), + "actor_rollout_ref.rollout.tensor_model_parallel_size": ( + 1, + "tensor parallelism greater than one is unsupported", + ), + "distillation.teacher_models.teacher_model.inference.tensor_model_parallel_size": ( + 1, + "teacher tensor parallelism greater than one is unsupported", + ), + "distillation.teacher_models.teacher_model.inference.data_parallel_size": ( + 1, + "teacher data parallelism greater than one is unsupported", + ), + "distillation.teacher_models.teacher_model.inference.pipeline_model_parallel_size": ( + 1, + "teacher pipeline parallelism greater than one is unsupported", + ), + "trainer.nnodes": (1, "multi-node execution is unsupported"), + } + expected = required.get(path) + if expected is not None and value != expected[0]: + return expected[1] + if path == "algorithm.adv_estimator" and value is not None: + return "external PPO/GRPO advantages are unsupported" + if path == "trainer.n_gpus_per_node" and value != 1: + return "more than one local training GPU is unsupported" + if path == "distillation.nnodes" and value not in {0, 1}: + return "multi-node teacher execution is unsupported" + if path == "data.filter_overlong_prompts" and value: + return "filter_overlong_prompts=true is unsupported; use explicit truncation" + if ( + path + in { + "actor_rollout_ref.actor.fsdp_config.param_offload", + "actor_rollout_ref.actor.fsdp_config.optimizer_offload", + } + and value + ): + return "FSDP offload semantics are unsupported by the one-GPU local runtime" + return None + + +def compile_verl_pg_k1_v08( + payload: Mapping[str, Any], + *, + override_files: Sequence[Path] = (), + overrides: Sequence[str] = (), + trailing_overrides: Sequence[str] = (), + reinterpretations_accepted: bool = True, + acceptance_source: str = "library_call", + require_executable: bool = True, +) -> CompiledLocalExecutionPlan: + """Compile the closed sampled-k1 profile into a deterministic local plan.""" + merged = copy.deepcopy(dict(payload)) + reject_interpolation(merged, label="verl PG-k1 OPD config") + records = _resolve_overrides( + merged, + override_files=override_files, + overrides=overrides, + trailing_overrides=trailing_overrides, + ) + reject_interpolation(merged, label="resolved verl PG-k1 OPD config") + + teachers = ( + merged.get("distillation", {}).get("teacher_models", {}) + if isinstance(merged.get("distillation"), Mapping) + else {} + ) + if isinstance(teachers, Mapping) and set(teachers) != {"teacher_model"}: + raise ConfigError( + "verl PG-k1 OPD config is not executable: multi-teacher routing is unsupported; " + f"found {sorted(map(str, teachers))}" + ) + + flat = _flatten(merged) + _reject_non_finite_source_numbers(flat) + unknown = sorted(path for path in flat if path not in _PG_FIELD_RULES) + if unknown: + raise ConfigError( + f"verl PG-k1 OPD config is not executable: unsupported field {unknown[0]!r}", + hint="the profile accepts only its documented resolved verl v0.8 subset", + ) + try: + source = VerlOPDV08Profile.model_validate(merged) + except ValidationError as exc: + finite = any("finite" in error.get("msg", "").lower() for error in exc.errors()) + message = "verl PG-k1 numeric fields must be finite" if finite else "invalid PG-k1 profile" + raise ConfigError(message, hint=str(exc)) from exc + + compatibility: list[CompatibilityEntry] = [] + blockers: list[str] = [] + for path, value in sorted(flat.items()): + rule = _PG_FIELD_RULES[path] + blocker = _pg_semantic_blocker(path, value) + if blocker: + blockers.append(path) + classification: FieldClassification = "unsupported" + reason = blocker + risk: Literal["none", "low", "medium", "high"] = "high" + else: + classification = rule.classification + reason = rule.reason + risk = rule.risk + if path == "trainer.total_epochs" and source.trainer.total_training_steps is not None: + classification = "informational_only" + reason = "recorded but superseded by the explicit total_training_steps cap" + risk = "none" + if ( + path + in { + "actor_rollout_ref.model.lora_adapter_path", + "miniverl.teacher_adapter.path", + "miniverl.teacher_adapter.revision", + } + and value is None + ): + classification = "informational_only" + reason = "null records that no existing adapter input is configured" + risk = "none" + confirmation_required = (rule.confirmation and classification != "informational_only") or ( + classification == "locally_reinterpreted" and risk == "high" + ) + compatibility.append( + CompatibilityEntry( + upstream_field=path, + source_value=value, + local_target=rule.target, + classification=classification, + reason=reason, + semantic_risk=risk, + user_confirmation_required=confirmation_required, + executable=blocker is None, + ) + ) + + executable = not blockers + local_execution = { + "device_count": 1, + "distributed_execution": False, + "compiler_scope": "config_semantics_only", + "model_weights_loaded": False, + "roles": ["actor_rollout", "teacher_scoring", "actor_update", "artifact"], + "runtime_mode": source.miniverl.runtime.mode, + "student_model": source.actor_rollout_ref.model.path, + "teacher_model": source.distillation.teacher_models.teacher_model.model_path, + "loss_mode": "k1", + "loss_reduction": "token-mean", + "task_rewards": False, + "policy_gradient": True, + "policy_loss_mode": "vanilla", + "teacher_target": "sampled_token_log_probability", + } + high_risk = [item.upstream_field for item in compatibility if item.user_confirmation_required] + common = { + "profile": VERL_OPD_PG_K1_V08_PROFILE, + "upstream": {"repository": VERL_REPOSITORY, "tag": VERL_TAG, "commit": VERL_COMMIT}, + "source_digest": _canonical_digest(merged), + "source_leaf_fields": sorted(flat), + "source": source, + "overrides": records, + "compatibility": compatibility, + "reinterpretation_acceptance": { + "required_fields": high_risk, + "accepted": reinterpretations_accepted or not high_risk, + "source": acceptance_source if high_risk else "not_required", + }, + "local_execution": local_execution, + "executable": executable, + } + digest_payload = { + key: value.model_dump(mode="json") if isinstance(value, BaseModel) else value + for key, value in common.items() + } + digest_payload["overrides"] = [item.model_dump(mode="json") for item in records] + digest_payload["compatibility"] = [item.model_dump(mode="json") for item in compatibility] + plan = CompiledLocalExecutionPlan( + **common, + compiled_digest=_canonical_digest(digest_payload), + ) + if require_executable and not executable: + raise ConfigError( + "verl PG-k1 OPD config is not executable: unsupported semantics in " + + ", ".join(blockers) + ) + return plan + + +def load_verl_pg_k1_v08( + path: Path, + *, + override_files: Sequence[Path] = (), + overrides: Sequence[str] = (), + trailing_overrides: Sequence[str] = (), + accept_local_reinterpretations: bool = False, + require_executable: bool = True, +) -> CompiledLocalExecutionPlan: + """Load one resolved YAML document without executing interpolation.""" + if path.suffix.lower() in {".sh", ".bash", ".ps1", ".cmd", ".bat"}: + raise ConfigError("verl PG-k1 input must be resolved YAML, not a shell script") + try: + payload = yaml.safe_load(path.read_text(encoding="utf-8")) + except (OSError, UnicodeError, yaml.YAMLError) as exc: + raise ConfigError(f"cannot read resolved verl PG-k1 YAML {path}: {exc}") from exc + if not isinstance(payload, Mapping): + raise ConfigError("verl PG-k1 YAML must contain one mapping") + return compile_verl_pg_k1_v08( + payload, + override_files=override_files, + overrides=overrides, + trailing_overrides=trailing_overrides, + reinterpretations_accepted=accept_local_reinterpretations, + acceptance_source="cli_flag" if accept_local_reinterpretations else "not_accepted", + require_executable=require_executable, + ) + + +def load_verl_pg_k1_v08_source( + source: str | Path, + *, + override_files: Sequence[Path] = (), + overrides: Sequence[str] = (), + trailing_overrides: Sequence[str] = (), + accept_local_reinterpretations: bool = False, + require_executable: bool = True, +) -> CompiledLocalExecutionPlan: + """Load the PG profile from an explicit resolved YAML path.""" + return load_verl_pg_k1_v08( + Path(source), + override_files=override_files, + overrides=overrides, + trailing_overrides=trailing_overrides, + accept_local_reinterpretations=accept_local_reinterpretations, + require_executable=require_executable, + ) diff --git a/src/miniverl/bridge/opd_runtime.py b/src/miniverl/bridge/opd_runtime.py index 5a7d247..e02c7c0 100644 --- a/src/miniverl/bridge/opd_runtime.py +++ b/src/miniverl/bridge/opd_runtime.py @@ -9,6 +9,7 @@ from pydantic import BaseModel, ConfigDict from miniverl.bridge.opd_capabilities import decide_placement +from miniverl.bridge.opd_pg_v08 import VERL_OPD_PG_K1_V08_PROFILE from miniverl.bridge.opd_v08 import CompiledLocalExecutionPlan from miniverl.config.models import RunConfig from miniverl.errors import ConfigError @@ -96,12 +97,15 @@ def build_system_plan(compiled: CompiledLocalExecutionPlan) -> OPDSystemPlan: max_tokens = source.data.max_prompt_length + source.data.max_response_length rollout_batch = source.miniverl.batching.rollout_batch_size token_buffers = max_tokens * rollout_batch * 64 * 1024 - target_bytes = ( - source.data.train_batch_size - * source.data.max_response_length - * source.distillation.distillation_loss.topk - * 12 - ) + pg_k1 = compiled.profile == VERL_OPD_PG_K1_V08_PROFILE + if pg_k1: + # token id + position + actor/teacher log-probs + response-mask weight + target_bytes = source.data.train_batch_size * source.data.max_response_length * 28 + else: + top_k = source.distillation.distillation_loss.topk + if top_k is None: + raise ConfigError("direct forward_kl_topk profile requires a positive topk") + target_bytes = source.data.train_batch_size * source.data.max_response_length * top_k * 12 known_static = sum( value for value in (student_static, teacher_static, trainable_optimizer) @@ -174,12 +178,23 @@ def build_system_plan(compiled: CompiledLocalExecutionPlan) -> OPDSystemPlan: "prompt_tokens_max": source.data.max_prompt_length, "response_tokens_max": source.data.max_response_length, }, - loss={ - "mode": source.distillation.distillation_loss.loss_mode, - "aggregation": source.actor_rollout_ref.actor.loss_agg_mode, - "top_k": source.distillation.distillation_loss.topk, - "task_rewards": False, - }, + loss=( + { + "mode": "k1", + "aggregation": "token-mean", + "teacher_target": "sampled_token_log_probability", + "policy_gradient": True, + "policy_loss_mode": "vanilla", + "task_rewards": False, + } + if pg_k1 + else { + "mode": source.distillation.distillation_loss.loss_mode, + "aggregation": source.actor_rollout_ref.actor.loss_agg_mode, + "top_k": source.distillation.distillation_loss.topk, + "task_rewards": False, + } + ), local_execution={ "strategy": decision.strategy, "requested_strategy": source.miniverl.runtime.mode, @@ -294,6 +309,37 @@ def compile_native_run_config( from miniverl.bridge.profiles import get_profile profile_identity = get_profile(compiled.profile).identity.model_dump(mode="json") + pg_k1 = compiled.profile == VERL_OPD_PG_K1_V08_PROFILE + loss_payload = ( + { + "mode": "verl_pg_k1", + "aggregation": "token-mean", + "divergence": "reverse_kl", + "temperature": 1.0, + "scale_by_temperature_squared": False, + "top_k": 1, + "loss_max_clamp": source.distillation.distillation_loss.loss_max_clamp, + "sampled_token_nll_weight": 0.0, + "policy_loss_mode": "vanilla", + "clip_ratio": 0.2, + "clip_ratio_low": 0.2, + "clip_ratio_high": 0.2, + "clip_ratio_c": 3.0, + "estimator_implementation_version": "verl-v0.8-pg-k1-v1", + } + if pg_k1 + else { + "mode": "forward_kl_topk", + "aggregation": "token-mean", + "divergence": "forward_kl", + "temperature": 1.0, + "scale_by_temperature_squared": False, + "top_k": source.distillation.distillation_loss.topk, + "log_prob_min_clamp": source.distillation.distillation_loss.log_prob_min_clamp, + "loss_max_clamp": source.distillation.distillation_loss.loss_max_clamp, + "sampled_token_nll_weight": 0.0, + } + ) payload = { "schema_version": 1, "run": { @@ -302,7 +348,12 @@ def compile_native_run_config( "seed": source.data.seed or 0, "output_dir": "runs", "deterministic": True, - "tags": [compiled.profile, "verl-v0.8.0", "pure-opd"], + "tags": [ + compiled.profile, + "verl-v0.8.0", + "pure-opd", + "pg-k1" if pg_k1 else "direct-gkd", + ], "profile_identity": profile_identity, }, "models": { @@ -357,6 +408,7 @@ def compile_native_run_config( "max_total_tokens": source.data.max_prompt_length + source.data.max_response_length, "temperature": source.actor_rollout_ref.rollout.temperature, "top_p": source.actor_rollout_ref.rollout.top_p, + "record_logprobs": pg_k1, "prompt_batch_size": min( source.miniverl.batching.rollout_batch_size, source.actor_rollout_ref.rollout.max_num_seqs @@ -367,17 +419,7 @@ def compile_native_run_config( * source.miniverl.batching.rollout_batch_size, }, "selection": {"selector": "all_model_tokens"}, - "loss": { - "mode": "forward_kl_topk", - "aggregation": "token-mean", - "divergence": "forward_kl", - "temperature": 1.0, - "scale_by_temperature_squared": False, - "top_k": source.distillation.distillation_loss.topk, - "log_prob_min_clamp": source.distillation.distillation_loss.log_prob_min_clamp, - "loss_max_clamp": source.distillation.distillation_loss.loss_max_clamp, - "sampled_token_nll_weight": 0.0, - }, + "loss": loss_payload, "train": { "cycles": total_steps, "rollouts_per_cycle": logical_batch, diff --git a/src/miniverl/bridge/opd_v08.py b/src/miniverl/bridge/opd_v08.py index e96866d..01bdc9e 100644 --- a/src/miniverl/bridge/opd_v08.py +++ b/src/miniverl/bridge/opd_v08.py @@ -188,12 +188,16 @@ class VerlOPDTeacherModelsConfig(_StrictModel): class VerlOPDLossConfig(_StrictModel): loss_mode: str - topk: int = Field(ge=1) + topk: int | None = Field(default=None, ge=1) use_task_rewards: bool = False distillation_loss_coef: float = Field(default=1.0, ge=0) loss_max_clamp: float | None = Field(default=None, ge=0) log_prob_min_clamp: float | None = None use_policy_gradient: bool = False + policy_loss_mode: str = "vanilla" + clip_ratio: float = Field(default=0.2, ge=0, lt=1) + clip_ratio_low: float | None = Field(default=0.2, ge=0, lt=1) + clip_ratio_high: float | None = Field(default=0.2, ge=0, lt=1) @model_validator(mode="after") def _numbers_are_finite(self) -> VerlOPDLossConfig: @@ -202,6 +206,11 @@ def _numbers_are_finite(self) -> VerlOPDLossConfig: _finite(self.loss_max_clamp, "loss_max_clamp") if self.log_prob_min_clamp is not None: _finite(self.log_prob_min_clamp, "log_prob_min_clamp") + _finite(self.clip_ratio, "clip_ratio") + if self.clip_ratio_low is not None: + _finite(self.clip_ratio_low, "clip_ratio_low") + if self.clip_ratio_high is not None: + _finite(self.clip_ratio_high, "clip_ratio_high") return self @@ -380,7 +389,7 @@ class CompatibilityEntry(_StrictModel): class CompiledLocalExecutionPlan(_StrictModel): schema_version: Literal[2] = 2 - profile: Literal["verl-opd-v0.8-single-gpu-v1"] = "verl-opd-v0.8-single-gpu-v1" + profile: str = VERL_OPD_V08_PROFILE upstream: dict[str, str] source_digest: str compiled_digest: str @@ -624,6 +633,18 @@ def _rule( "distillation.distillation_loss.use_policy_gradient": _rule( "loss.use_policy_gradient", "exact", "must remain disabled for direct GKD OPD" ), + "distillation.distillation_loss.policy_loss_mode": _rule( + None, "informational_only", "inactive because direct GKD does not use a policy loss" + ), + "distillation.distillation_loss.clip_ratio": _rule( + None, "informational_only", "inactive because direct GKD does not use a policy loss" + ), + "distillation.distillation_loss.clip_ratio_low": _rule( + None, "informational_only", "inactive because direct GKD does not use a policy loss" + ), + "distillation.distillation_loss.clip_ratio_high": _rule( + None, "informational_only", "inactive because direct GKD does not use a policy loss" + ), "trainer.project_name": _rule("run.project", "exact", "same provenance label"), "trainer.experiment_name": _rule("run.name", "exact", "same provenance label"), "trainer.save_freq": _rule( @@ -1034,6 +1055,8 @@ def _semantic_blocker(path: str, value: Any) -> str | None: return expected[1] if path == "trainer.n_gpus_per_node" and value != 1: return "more than one local training GPU is unsupported" + if path == "distillation.distillation_loss.topk" and value is None: + return "direct forward_kl_topk requires a positive topk" if path == "distillation.nnodes" and value not in {0, 1}: return "multi-node teacher execution is unsupported" if path == "algorithm.adv_estimator" and value is not None: diff --git a/src/miniverl/bridge/profiles.py b/src/miniverl/bridge/profiles.py index 29535ac..efc4555 100644 --- a/src/miniverl/bridge/profiles.py +++ b/src/miniverl/bridge/profiles.py @@ -8,6 +8,7 @@ import hashlib from abc import ABC, abstractmethod +from collections.abc import Sequence from dataclasses import dataclass from pathlib import Path from typing import Any, Literal @@ -15,6 +16,12 @@ from pydantic import BaseModel, ConfigDict, Field from miniverl.bridge.contract import VERL_COMMIT, VERL_REPOSITORY, VERL_TAG +from miniverl.bridge.opd_pg_v08 import ( + VERL_OPD_PG_K1_V08_PROFILE, + load_verl_pg_k1_v08_source, + pg_compatibility_rule, + pg_field_rules_digest, +) from miniverl.bridge.opd_v08 import ( VERL_OPD_V08_PROFILE, VerlOPDV08Profile, @@ -34,6 +41,7 @@ "check_profile", "get_profile", "list_profiles", + "load_profile_source", ] @@ -253,13 +261,160 @@ def show(self) -> dict[str, Any]: } +@dataclass(frozen=True) +class _PGK1Profile(CompatibilityProfile): + @property + def identity(self) -> ProfileIdentity: + return ProfileIdentity.create( + profile_name=VERL_OPD_PG_K1_V08_PROFILE, + profile_schema_version=1, + upstream_repository=VERL_REPOSITORY, + upstream_tag=VERL_TAG, + upstream_commit=VERL_COMMIT, + field_rule_digest=pg_field_rules_digest(), + native_compiler_version="pg-k1-native-v1", + loss_conformance_version="pg-k1-verl-v0.8-v1", + export_version="verl-opd-pg-k1-export-v1", + ) + + @property + def summary(self) -> ProfileSummary: + return ProfileSummary( + name=VERL_OPD_PG_K1_V08_PROFILE, + objective="k1 + vanilla policy loss", + teacher_target="sampled-token teacher log-probability", + status="conformance_only", + identity=self.identity, + ) + + def config_schema(self) -> dict[str, Any]: + return VerlOPDV08Profile.model_json_schema() + + def explain(self, field: str) -> CompatibilityExplanation: + rule = pg_compatibility_rule(field) + classification = str(rule["classification"]) + unsupported = classification == "unsupported" + informational = classification == "informational_only" + return CompatibilityExplanation( + profile=VERL_OPD_PG_K1_V08_PROFILE, + upstream_field=field, + classification=classification, + local_target=rule["local_target"], + reason=str(rule["reason"]), + semantic_risk=str(rule["semantic_risk"]), + user_confirmation_required=bool(rule["user_confirmation_required"]), + supported_algorithm=True, + field_accepted=not unsupported, + field_effective=( + not unsupported and not informational and rule["local_target"] is not None + ), + field_locally_reinterpreted=classification == "locally_reinterpreted", + field_informational=informational, + field_unsupported=unsupported, + profile_applicable=True, + ) + + def check( + self, source: str | Path, *, accept_local_reinterpretations: bool = False + ) -> CompatibilityCheck: + compiled = load_verl_pg_k1_v08_source( + source, + accept_local_reinterpretations=accept_local_reinterpretations, + require_executable=False, + ) + fields = [item.model_dump(mode="json") for item in compiled.compatibility] + unsupported = sum(not item.executable for item in compiled.compatibility) + informational = sum( + item.classification == "informational_only" for item in compiled.compatibility + ) + reinterpreted = sum( + item.classification == "locally_reinterpreted" for item in compiled.compatibility + ) + effective = sum( + item.executable + and item.classification != "informational_only" + and item.local_target is not None + for item in compiled.compatibility + ) + accepted = bool(compiled.reinterpretation_acceptance["accepted"]) + status: Literal["compatible", "needs_user_confirmation", "unsupported"] + if unsupported: + status = "unsupported" + elif not accepted: + status = "needs_user_confirmation" + else: + status = "compatible" + return CompatibilityCheck( + status=status, + executable=compiled.executable and accepted, + profile_identity=self.identity, + source_digest=compiled.source_digest, + compiled_digest=compiled.compiled_digest, + summary={ + "supported_algorithm": True, + "field_accepted": len(fields) - unsupported, + "field_effective": effective, + "field_locally_reinterpreted": reinterpreted, + "field_informational": informational, + "field_unsupported": unsupported, + "profile_not_applicable": False, + }, + fields=fields, + ) + + def show(self) -> dict[str, Any]: + from importlib.resources import files + + import yaml + + payload = yaml.safe_load( + files("miniverl") + .joinpath("resources/qwen3_0_6b_1_7b_opd.yaml") + .read_text(encoding="utf-8") + ) + payload["distillation"]["distillation_loss"].update( + { + "loss_mode": "k1", + "topk": None, + "use_policy_gradient": True, + "policy_loss_mode": "vanilla", + "clip_ratio": 0.2, + "clip_ratio_low": 0.2, + "clip_ratio_high": 0.2, + "loss_max_clamp": 10.0, + "log_prob_min_clamp": None, + } + ) + payload["algorithm"]["adv_estimator"] = None + return { + **self.summary.model_dump(mode="json"), + "identity": self.identity.model_dump(mode="json"), + "algorithm_contract": { + "actor_count": 1, + "teacher_count": 1, + "generations_per_prompt": 1, + "objective": "sampled k1 through vanilla policy loss", + "task_rewards": False, + "distributed_execution": False, + "device": "one local CUDA GPU", + }, + "minimal_yaml": yaml.safe_dump(payload, sort_keys=False, allow_unicode=True), + "override_invocation": ( + "miniverl plan --profile verl-opd-v0.8-single-gpu-pg-k1-v1 " + "--config pg-k1.yaml --set data.train_batch_size=4 --json" + ), + } + + _REGISTRY: dict[str, CompatibilityProfile] = { VERL_OPD_V08_PROFILE: _DirectGKDProfile(), + VERL_OPD_PG_K1_V08_PROFILE: _PGK1Profile(), } def list_profiles() -> tuple[ProfileSummary, ...]: - return tuple(_REGISTRY[name].summary for name in sorted(_REGISTRY)) + """Return profiles in stable product order: direct GKD, then sampled PG.""" + return tuple(profile.summary for profile in _REGISTRY.values()) def get_profile(name: str) -> CompatibilityProfile: @@ -281,3 +436,36 @@ def check_profile( return get_profile(name).check( source, accept_local_reinterpretations=accept_local_reinterpretations ) + + +def load_profile_source( + name: str, + source: str | Path, + *, + override_files: Sequence[Path] = (), + overrides: Sequence[str] = (), + trailing_overrides: Sequence[str] = (), + accept_local_reinterpretations: bool = False, + require_executable: bool = True, +) -> Any: + """Compile one source through the selected closed built-in profile.""" + get_profile(name) # fail closed before dispatch; never load third-party code + if name == VERL_OPD_V08_PROFILE: + return load_verl_opd_v08_source( + source, + override_files=override_files, + overrides=overrides, + trailing_overrides=trailing_overrides, + accept_local_reinterpretations=accept_local_reinterpretations, + require_executable=require_executable, + ) + if name == VERL_OPD_PG_K1_V08_PROFILE: + return load_verl_pg_k1_v08_source( + source, + override_files=override_files, + overrides=overrides, + trailing_overrides=trailing_overrides, + accept_local_reinterpretations=accept_local_reinterpretations, + require_executable=require_executable, + ) + raise AssertionError(f"registered profile {name!r} has no compiler dispatch") diff --git a/src/miniverl/cache/store.py b/src/miniverl/cache/store.py index 2e7f2f2..4341bf9 100644 --- a/src/miniverl/cache/store.py +++ b/src/miniverl/cache/store.py @@ -52,7 +52,7 @@ _INDEX_NAME = "index.json" _SHARD_NAME = re.compile(r"^shard-(\d+)\.safetensors$") -_TENSOR_FIELDS = ( +_DISTRIBUTION_TENSOR_FIELDS = ( "positions", "topk_indices", "topk_log_probs", @@ -60,6 +60,13 @@ "target_token_ids", "weights", ) +_PG_TENSOR_FIELDS = ( + "positions", + "target_token_ids", + "weights", + "old_actor_log_probs", + "teacher_sampled_token_log_probs", +) logger = get_logger("cache") @@ -178,7 +185,9 @@ def create( top_k: int, temperature: float, loss_mode: str, + target_representation: str = "topk_distribution", score_implementation_version: str | None = None, + estimator_implementation_version: str | None = None, execution_plan_digest: str | None = None, profile_identity: dict[str, Any] | None = None, dtype: str = "float32", @@ -220,7 +229,9 @@ def create( top_k=top_k, temperature=temperature, loss_mode=loss_mode, + target_representation=target_representation, score_implementation_version=score_implementation_version, + estimator_implementation_version=estimator_implementation_version, execution_plan_digest=execution_plan_digest, profile_identity=dict(profile_identity or {}), dtype=dtype, @@ -274,7 +285,9 @@ def assert_compatible( top_k: int, temperature: float, loss_mode: str, + target_representation: str = "topk_distribution", score_implementation_version: str | None = None, + estimator_implementation_version: str | None = None, execution_plan_digest: str | None = None, profile_identity: dict[str, Any] | None = None, dtype: str, @@ -304,10 +317,13 @@ def assert_compatible( "top_k": top_k, "temperature": temperature, "loss_mode": loss_mode, + "target_representation": target_representation, "dtype": dtype, } if score_implementation_version is not None: expected["score_implementation_version"] = score_implementation_version + if estimator_implementation_version is not None: + expected["estimator_implementation_version"] = estimator_implementation_version if self.index.schema_version >= 2: expected["tokenizer_identity"] = dict(tokenizer_identity or {}) expected["teacher_adapter_provenance"] = ( @@ -359,20 +375,60 @@ def write( if batch.trajectory_id in self.index.entries or batch.trajectory_id in self._pending: raise CacheError(f"trajectory {batch.trajectory_id!r} is already in the cache") + if batch.target_representation != self.index.target_representation: + raise CacheError( + f"entry target representation {batch.target_representation!r} does not match " + f"cache {self.index.target_representation!r}" + ) + if batch.estimator_implementation_version != self.index.estimator_implementation_version: + raise CacheError("entry estimator implementation version does not match cache identity") float_dtype = self._dtype() - tensors = { + common_tensors = { "positions": batch.positions.to(torch.int64).cpu().contiguous(), - "topk_indices": batch.topk_indices.to(torch.int32).cpu().contiguous(), - "topk_log_probs": batch.topk_log_probs.to(float_dtype).cpu().contiguous(), - "tail_log_prob": _finite_tail(batch.tail_log_prob, torch) - .to(float_dtype) - .cpu() - .contiguous(), "target_token_ids": batch.target_token_ids.to(torch.int64).cpu().contiguous(), "weights": batch.weights.to(torch.float32).cpu().contiguous(), } + tensor_fields: tuple[str, ...] + if self.index.target_representation == "sampled_token_log_probs": + if batch.old_actor_log_probs is None or batch.teacher_sampled_token_log_probs is None: + raise CacheError( + "sampled-token targets require actor and teacher log-probabilities" + ) + tensors = { + **common_tensors, + "old_actor_log_probs": batch.old_actor_log_probs.to(torch.float32) + .cpu() + .contiguous(), + "teacher_sampled_token_log_probs": batch.teacher_sampled_token_log_probs.to( + torch.float32 + ) + .cpu() + .contiguous(), + } + tensor_fields = _PG_TENSOR_FIELDS + elif self.index.target_representation == "topk_distribution": + if ( + batch.topk_indices is None + or batch.topk_log_probs is None + or batch.tail_log_prob is None + ): + raise CacheError("top-k distribution targets require top-k and tail tensors") + tensors = { + **common_tensors, + "topk_indices": batch.topk_indices.to(torch.int32).cpu().contiguous(), + "topk_log_probs": batch.topk_log_probs.to(float_dtype).cpu().contiguous(), + "tail_log_prob": _finite_tail(batch.tail_log_prob, torch) + .to(float_dtype) + .cpu() + .contiguous(), + } + tensor_fields = _DISTRIBUTION_TENSOR_FIELDS + else: + raise CacheError( + f"unsupported teacher target representation {self.index.target_representation!r}" + ) digest = hashlib.sha256() - for key in _TENSOR_FIELDS: + for key in tensor_fields: digest.update(key.encode("utf-8")) digest.update(tensors[key].numpy().tobytes()) span_counts: dict[str, int] = {} @@ -398,6 +454,7 @@ def write( "prompt_row_digest": batch.prompt_row_digest, "actor_response_token_ids": batch.actor_response_token_ids, "binding_checksum": binding_checksum, + "tensor_fields": list(tensor_fields), }, } self._pending_order.append(batch.trajectory_id) @@ -414,7 +471,7 @@ def write( loss_mode=self.index.loss_mode, temperature=self.index.temperature, created_at=_utc_now(), - tensor_keys=list(_TENSOR_FIELDS), + tensor_keys=list(tensor_fields), checksum=digest.hexdigest(), selected_span_types=span_counts, ordered_span_types=list(batch.span_types), @@ -472,7 +529,7 @@ def flush(self) -> None: loss_mode=next_index.loss_mode, temperature=next_index.temperature, created_at=created, - tensor_keys=list(_TENSOR_FIELDS), + tensor_keys=list(meta["tensor_fields"]), checksum=meta["checksum"], selected_span_types=meta["selected_span_types"], ordered_span_types=meta["ordered_span_types"], @@ -556,7 +613,16 @@ def read( tensors = load_file(str(shard_path), device="cpu") loaded = {} digest = hashlib.sha256() - for key in _TENSOR_FIELDS: + expected_fields = ( + _PG_TENSOR_FIELDS + if self.index.target_representation == "sampled_token_log_probs" + else _DISTRIBUTION_TENSOR_FIELDS + ) + if tuple(entry.tensor_keys) != tuple(expected_fields): + raise CacheCorruptionError( + f"entry {trajectory_id!r} tensor schema does not match target representation" + ) + for key in expected_fields: full_key = f"{trajectory_id}|{key}" if full_key not in tensors: raise CacheCorruptionError(f"shard {entry.shard} is missing tensor {full_key!r}") @@ -580,15 +646,25 @@ def read( f"v1 cache entry {trajectory_id!r} does not preserve ordered span types", hint="re-score this legacy cache before using per-span metrics", ) - tail_log_prob = loaded["tail_log_prob"].to(device=device, dtype=torch.float32) - if entry.tail_is_exact_zero: - tail_log_prob = torch.full_like(tail_log_prob, float("-inf")) + tail_log_prob = None + if "tail_log_prob" in loaded: + tail_log_prob = loaded["tail_log_prob"].to(device=device, dtype=torch.float32) + if entry.tail_is_exact_zero: + tail_log_prob = torch.full_like(tail_log_prob, float("-inf")) return TeacherTargetBatch( trajectory_id=trajectory_id, policy_version=entry.policy_version, positions=loaded["positions"].to(device), - topk_indices=loaded["topk_indices"].to(device=device, dtype=torch.long), - topk_log_probs=loaded["topk_log_probs"].to(device=device, dtype=torch.float32), + topk_indices=( + loaded["topk_indices"].to(device=device, dtype=torch.long) + if "topk_indices" in loaded + else None + ), + topk_log_probs=( + loaded["topk_log_probs"].to(device=device, dtype=torch.float32) + if "topk_log_probs" in loaded + else None + ), tail_log_prob=tail_log_prob, target_token_ids=loaded["target_token_ids"].to(device), weights=loaded["weights"].to(device), @@ -597,6 +673,18 @@ def read( span_types=span_types, prompt_row_digest=entry.prompt_row_digest, actor_response_token_ids=entry.actor_response_token_ids, + old_actor_log_probs=( + loaded["old_actor_log_probs"].to(device=device, dtype=torch.float32) + if "old_actor_log_probs" in loaded + else None + ), + teacher_sampled_token_log_probs=( + loaded["teacher_sampled_token_log_probs"].to(device=device, dtype=torch.float32) + if "teacher_sampled_token_log_probs" in loaded + else None + ), + target_representation=self.index.target_representation, + estimator_implementation_version=self.index.estimator_implementation_version, ) def __contains__(self, trajectory_id: object) -> bool: @@ -625,7 +713,7 @@ def validate(self, *, verify_checksums: bool = True) -> list[str]: f"{traj_id}|{field}" for traj_id, entry in self.index.entries.items() if entry.shard == name - for field in _TENSOR_FIELDS + for field in entry.tensor_keys } missing = expected - tensor_names if missing: diff --git a/src/miniverl/cli.py b/src/miniverl/cli.py index 3fc51f4..e0da093 100644 --- a/src/miniverl/cli.py +++ b/src/miniverl/cli.py @@ -955,16 +955,10 @@ def plan_command( """Plan pinned single-GPU verl-style OPD without loading model weights.""" try: from miniverl.bridge.opd_runtime import build_system_plan - from miniverl.bridge.opd_v08 import ( - VERL_OPD_V08_PROFILE, - load_verl_opd_v08_source, - ) + from miniverl.bridge.profiles import load_profile_source - if profile != VERL_OPD_V08_PROFILE: - raise ConfigError( - f"unsupported OPD profile {profile!r}", hint=f"use --profile {VERL_OPD_V08_PROFILE}" - ) - compiled = load_verl_opd_v08_source( + compiled = load_profile_source( + profile, config, override_files=override_files, overrides=overrides, @@ -1074,15 +1068,8 @@ def verl_run_command( """Execute the pinned verl v0.8 pure-OPD subset on one local CUDA GPU.""" try: from miniverl.bridge.opd_runtime import build_system_plan, compile_native_run_config - from miniverl.bridge.opd_v08 import ( - VERL_OPD_V08_PROFILE, - load_verl_opd_v08_source, - ) + from miniverl.bridge.profiles import load_profile_source - if profile != VERL_OPD_V08_PROFILE: - raise ConfigError( - f"unsupported OPD profile {profile!r}", hint=f"use --profile {VERL_OPD_V08_PROFILE}" - ) if plan_path is not None and config is not None: raise ConfigError("--plan and --config are mutually exclusive") if plan_path is not None and ( @@ -1108,7 +1095,8 @@ def verl_run_command( else: if config is None: raise ConfigError("one of --config or --plan is required") - compiled = load_verl_opd_v08_source( + compiled = load_profile_source( + profile, config, override_files=override_files, overrides=overrides, @@ -1912,14 +1900,10 @@ def bridge_compile_opd_command( ) -> None: """Compile the pinned single-GPU verl v0.8 OPD config subset offline.""" try: - from miniverl.bridge.opd_v08 import VERL_OPD_V08_PROFILE, load_verl_opd_v08_source + from miniverl.bridge.profiles import load_profile_source - if profile != VERL_OPD_V08_PROFILE: - raise ConfigError( - f"unsupported OPD profile {profile!r}", - hint=f"use --profile {VERL_OPD_V08_PROFILE}", - ) - plan = load_verl_opd_v08_source( + plan = load_profile_source( + profile, config, override_files=override_files, overrides=overrides, diff --git a/src/miniverl/config/models.py b/src/miniverl/config/models.py index b2ef7be..867bb99 100644 --- a/src/miniverl/config/models.py +++ b/src/miniverl/config/models.py @@ -155,6 +155,7 @@ class LossMode(str, Enum): EXACT_FULL_VOCAB = "exact_full_vocab" BUCKETED_TOPK_TAIL = "bucketed_topk_tail" VERL_FORWARD_KL_TOPK = "forward_kl_topk" + VERL_PG_K1 = "verl_pg_k1" class LossAggregation(str, Enum): @@ -475,6 +476,14 @@ class LossConfig(_Base): ge=0.0, le=1.0, ) + #: Closed verl v0.8 policy-gradient subset. These values are inactive for + #: every other loss mode and are rejected if changed outside that profile. + policy_loss_mode: Literal["vanilla"] = "vanilla" + clip_ratio: float = Field(default=0.2, ge=0.0, lt=1.0) + clip_ratio_low: float = Field(default=0.2, ge=0.0, lt=1.0) + clip_ratio_high: float = Field(default=0.2, ge=0.0, lt=1.0) + clip_ratio_c: float = Field(default=3.0, gt=1.0) + estimator_implementation_version: str | None = None @model_validator(mode="before") @classmethod @@ -523,6 +532,9 @@ class RolloutConfig(_Base): temperature: float = Field(default=1.0, ge=0.0, le=5.0) top_p: float = Field(default=1.0, gt=0.0, le=1.0) top_k: int = Field(default=0, ge=0) + #: Capture the exact sampled-token policy log-probability at rollout time. + #: Required by the pinned PG-k1 profile; otherwise disabled by default. + record_logprobs: bool = False # Maximum parse errors tolerated in an episode. The rollout terminates as # soon as the count reaches this limit; zero therefore terminates on the # first parse error. @@ -905,6 +917,57 @@ def _validate_combination(self) -> RunConfig: "supported verl v0.8 profile" ) + if self.loss.mode is LossMode.VERL_PG_K1: + from miniverl.losses.verl_pg import VERL_PG_K1_IMPLEMENTATION_VERSION + + if mode is not TrainingMode.OPD: + raise ValueError("loss.mode=verl_pg_k1 requires run.mode=opd") + if self.source.kind is not SourceKind.VERL_PARQUET: + raise ValueError("loss.mode=verl_pg_k1 requires source.kind=verl_parquet") + if self.models.teacher.mode is not TeacherContextMode.STANDARD: + raise ValueError("loss.mode=verl_pg_k1 requires a standard same-prompt teacher") + if self.loss.divergence is not Divergence.REVERSE_KL: + raise ValueError("loss.mode=verl_pg_k1 requires divergence=reverse_kl") + if self.loss.aggregation is not LossAggregation.TOKEN_MEAN: + raise ValueError("loss.mode=verl_pg_k1 requires loss.aggregation=token-mean") + if self.loss.temperature != 1.0 or self.loss.scale_by_temperature_squared: + raise ValueError( + "loss.mode=verl_pg_k1 uses sampled log-probabilities directly; set " + "temperature=1.0 and scale_by_temperature_squared=false" + ) + if self.loss.sampled_token_nll_weight != 0.0: + raise ValueError("loss.mode=verl_pg_k1 cannot mix sampled-token NLL") + if self.loss.top_k != 1: + raise ValueError("loss.mode=verl_pg_k1 has no top-k target; set loss.top_k=1") + if self.loss.policy_loss_mode != "vanilla" or ( + self.loss.clip_ratio, + self.loss.clip_ratio_low, + self.loss.clip_ratio_high, + self.loss.clip_ratio_c, + ) != (0.2, 0.2, 0.2, 3.0): + raise ValueError( + "loss.mode=verl_pg_k1 supports only the pinned vanilla 0.2/0.2/0.2, " + "dual-clip 3.0 policy-loss contract" + ) + if self.loss.estimator_implementation_version != VERL_PG_K1_IMPLEMENTATION_VERSION: + raise ValueError( + "loss.mode=verl_pg_k1 requires estimator_implementation_version=" + f"{VERL_PG_K1_IMPLEMENTATION_VERSION}" + ) + if not self.rollout.record_logprobs: + raise ValueError("loss.mode=verl_pg_k1 requires rollout.record_logprobs=true") + if self.selection.selector is not SelectorName.ALL_MODEL_TOKENS: + raise ValueError( + "loss.mode=verl_pg_k1 requires selection.selector=all_model_tokens" + ) + if ( + self.selection.max_positions_per_trajectory is not None + or self.selection.critical_weight != 1.0 + or self.selection.other_weight != 1.0 + ): + raise ValueError( + "loss.mode=verl_pg_k1 requires every response token with unit response-mask weight" + ) if mode is TrainingMode.SFT and self.loss.sampled_token_nll_weight not in (0.0, 1.0): raise ValueError( "run.mode=sft trains with oracle cross-entropy only; " diff --git a/src/miniverl/losses/chunked.py b/src/miniverl/losses/chunked.py index b1a6aee..8a394e7 100644 --- a/src/miniverl/losses/chunked.py +++ b/src/miniverl/losses/chunked.py @@ -46,6 +46,7 @@ "ExactTargetProvider", "BucketedTargetProvider", "VerlTopKTargetProvider", + "VerlPGK1TargetProvider", "LossOutput", "chunked_selected_position_loss", ] @@ -174,6 +175,58 @@ def teacher_entropy(self, start: int, end: int) -> torch.Tensor: return torch.full((end - start,), float("nan"), dtype=torch.float32) +@dataclass +class VerlPGK1TargetProvider: + """Sampled-token teacher and rollout-policy targets for pinned PG-k1.""" + + target_token_ids: torch.Tensor + old_actor_log_probs: torch.Tensor + teacher_sampled_token_log_probs: torch.Tensor + clip_ratio: float = 0.2 + clip_ratio_low: float = 0.2 + clip_ratio_high: float = 0.2 + clip_ratio_c: float = 3.0 + loss_max_clamp: float | None = None + kind: str = "verl_pg_k1" + diagnostics: list[dict[str, torch.Tensor | float]] = field(default_factory=list) + + def divergence(self, start: int, end: int, student_logits: torch.Tensor) -> torch.Tensor: + from miniverl.losses.verl_pg import verl_pg_k1_loss + + targets = self.target_token_ids[start:end].to(student_logits.device) + current = ( + torch.log_softmax(student_logits.to(torch.float32), dim=-1) + .gather(-1, targets.unsqueeze(-1)) + .squeeze(-1) + ) + output = verl_pg_k1_loss( + current_log_probs=current, + old_log_probs=self.old_actor_log_probs[start:end].to(student_logits.device), + teacher_log_probs=self.teacher_sampled_token_log_probs[start:end].to( + student_logits.device + ), + weights=torch.ones(end - start, device=student_logits.device), + clip_ratio=self.clip_ratio, + clip_ratio_low=self.clip_ratio_low, + clip_ratio_high=self.clip_ratio_high, + clip_ratio_c=self.clip_ratio_c, + loss_max_clamp=self.loss_max_clamp, + ) + self.diagnostics.append( + { + "estimator": output.estimator.detach().cpu(), + "advantages": output.advantages.detach().cpu(), + "ratio": output.ratio.detach().cpu(), + **output.metrics, + } + ) + return output.per_token_loss + + def teacher_entropy(self, start: int, end: int) -> torch.Tensor: + """Entropy is not available from one sampled-token log-probability.""" + return torch.full((end - start,), float("nan"), dtype=torch.float32) + + @dataclass class LossOutput: """Result of one chunked objective evaluation.""" diff --git a/src/miniverl/losses/verl_pg.py b/src/miniverl/losses/verl_pg.py new file mode 100644 index 0000000..3f4b8b6 --- /dev/null +++ b/src/miniverl/losses/verl_pg.py @@ -0,0 +1,138 @@ +"""Pinned verl v0.8 sampled-k1 policy-gradient objective. + +This module implements only the closed +``verl-opd-v0.8-single-gpu-pg-k1-v1`` contract. It is deliberately not a +general PPO implementation: task rewards, external advantages, rollout +importance weights, critics, and alternative policy losses are outside the +profile. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +__all__ = [ + "VERL_PG_K1_IMPLEMENTATION_VERSION", + "VerlPGK1Output", + "verl_k1_estimator", + "verl_pg_k1_loss", +] + +VERL_PG_K1_IMPLEMENTATION_VERSION = "verl-v0.8-pg-k1-v1" + + +@dataclass(frozen=True) +class VerlPGK1Output: + """Exact tensors and scalar produced by the supported pinned path.""" + + loss: torch.Tensor + estimator: torch.Tensor + advantages: torch.Tensor + per_token_loss: torch.Tensor + ratio: torch.Tensor + metrics: dict[str, float] + + +def _require_vector(name: str, value: torch.Tensor, length: int | None = None) -> None: + if value.ndim != 1: + raise ValueError(f"{name} must be a one-dimensional token vector") + if length is not None and value.numel() != length: + raise ValueError(f"{name} has {value.numel()} tokens, expected {length}") + if not torch.isfinite(value).all(): + raise ValueError(f"{name} must contain only finite values") + + +def verl_k1_estimator( + old_actor_log_probs: torch.Tensor, + teacher_log_probs: torch.Tensor, +) -> torch.Tensor: + """Return the pinned sampled k1 estimate: ``old_actor - teacher``.""" + _require_vector("old_actor_log_probs", old_actor_log_probs) + _require_vector("teacher_log_probs", teacher_log_probs, old_actor_log_probs.numel()) + return old_actor_log_probs.to(torch.float32) - teacher_log_probs.to(torch.float32) + + +def _masked_mean(value: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: + total = weights.sum() + if not bool(total > 0): + raise ValueError("weights must select at least one response token") + return (value * weights).sum() / total + + +def verl_pg_k1_loss( + *, + current_log_probs: torch.Tensor, + old_log_probs: torch.Tensor, + teacher_log_probs: torch.Tensor, + weights: torch.Tensor, + clip_ratio: float = 0.2, + clip_ratio_low: float | None = None, + clip_ratio_high: float | None = None, + clip_ratio_c: float = 3.0, + loss_max_clamp: float | None = None, +) -> VerlPGK1Output: + """Apply verl v0.8 k1 -> detached advantage -> vanilla policy loss. + + ``weights`` is miniVERL's flattened response mask. The supported profile + requires unit weights, but accepting a numeric mask here keeps the loss + primitive independently testable and exactly reproduces token-mean + aggregation. + """ + n = current_log_probs.numel() + _require_vector("current_log_probs", current_log_probs) + _require_vector("old_log_probs", old_log_probs, n) + _require_vector("teacher_log_probs", teacher_log_probs, n) + _require_vector("weights", weights, n) + if bool((weights < 0).any()): + raise ValueError("weights must be non-negative") + if not 0.0 <= clip_ratio < 1.0: + raise ValueError("clip_ratio must be in [0, 1)") + low = clip_ratio if clip_ratio_low is None else clip_ratio_low + high = clip_ratio if clip_ratio_high is None else clip_ratio_high + if not 0.0 <= low < 1.0 or not 0.0 <= high < 1.0: + raise ValueError("clip ratios must be in [0, 1)") + if clip_ratio_c <= 1.0: + raise ValueError("clip_ratio_c must be greater than 1") + if loss_max_clamp is not None and loss_max_clamp <= 0.0: + raise ValueError("loss_max_clamp must be positive") + + work_weights = weights.to(device=current_log_probs.device, dtype=torch.float32) + estimator = verl_k1_estimator(old_log_probs, teacher_log_probs).to(current_log_probs.device) + if loss_max_clamp is not None: + estimator = estimator.clamp(min=-loss_max_clamp, max=loss_max_clamp) + advantages = -estimator.detach() + + negative_approx_kl = ( + current_log_probs.to(torch.float32) - old_log_probs.to(torch.float32) + ).clamp(min=-20.0, max=20.0) + ratio = negative_approx_kl.exp() + losses1 = -advantages * ratio + losses2 = -advantages * ratio.clamp(1.0 - low, 1.0 + high) + clipped = torch.maximum(losses1, losses2) + dual_clipped = torch.minimum(-advantages * clip_ratio_c, clipped) + per_token_loss = torch.where(advantages < 0, dual_clipped, clipped) + loss = _masked_mean(per_token_loss, work_weights) + + metrics = { + "pg_clipfrac": float( + _masked_mean((losses2 > losses1).to(torch.float32), work_weights).detach() + ), + "ppo_kl": float(_masked_mean(-negative_approx_kl, work_weights).detach()), + "pg_clipfrac_lower": float( + _masked_mean( + ((clipped > -advantages * clip_ratio_c) & (advantages < 0)).to(torch.float32), + work_weights, + ).detach() + ), + "abs_k1": float(_masked_mean(estimator.abs(), work_weights).detach()), + } + return VerlPGK1Output( + loss=loss, + estimator=estimator, + advantages=advantages, + per_token_loss=per_token_loss, + ratio=ratio, + metrics=metrics, + ) diff --git a/src/miniverl/models/base.py b/src/miniverl/models/base.py index 8678195..fb43559 100644 --- a/src/miniverl/models/base.py +++ b/src/miniverl/models/base.py @@ -98,6 +98,7 @@ def generate( top_p: float = 1.0, top_k: int = 0, seed: int | None = None, + record_logprobs: bool = False, ) -> GenerationOutput: """Sample a continuation, stopping on EOS, a stop string, or the budget.""" ... @@ -112,6 +113,7 @@ def generate_batch( top_p: float = 1.0, top_k: int = 0, seeds: Sequence[int | None] | None = None, + record_logprobs: bool = False, ) -> list[GenerationOutput]: """Compatibility batch API; concrete local backends use one padded forward. @@ -130,6 +132,7 @@ def generate_batch( top_p=top_p, top_k=top_k, seed=seed, + record_logprobs=record_logprobs, ) for prefix, seed in zip(prefix_token_ids, chosen_seeds, strict=True) ] diff --git a/src/miniverl/models/hf.py b/src/miniverl/models/hf.py index 82d7d30..38c2346 100644 --- a/src/miniverl/models/hf.py +++ b/src/miniverl/models/hf.py @@ -417,6 +417,7 @@ def generate( top_p: float = 1.0, top_k: int = 0, seed: int | None = None, + record_logprobs: bool = False, ) -> GenerationOutput: """Sample a continuation, projecting one position per step.""" was_training = self.model.training @@ -445,6 +446,7 @@ def step(new_ids: list[int], state: Any) -> tuple[torch.Tensor, Any]: top_p=top_p, top_k=top_k, generator=generator, + record_logprobs=record_logprobs, ) finally: self.model.config.use_cache = cache_flag @@ -461,6 +463,7 @@ def generate_batch( top_p: float = 1.0, top_k: int = 0, seeds: Sequence[int | None] | None = None, + record_logprobs: bool = False, ) -> list[GenerationOutput]: """Use one masked padded forward per greedy decoding step.""" if temperature != 0.0: @@ -472,6 +475,7 @@ def generate_batch( top_p=top_p, top_k=top_k, seeds=seeds, + record_logprobs=record_logprobs, ) was_training = self.model.training self.model.eval() @@ -505,6 +509,7 @@ def step(rows: list[list[int]]) -> torch.Tensor: eos_token_id=self.tokenizer.eos_token_id, max_new_tokens=max_new_tokens, stop_sequences=stop_sequences, + record_logprobs=record_logprobs, ) finally: if was_training: diff --git a/src/miniverl/models/sampling.py b/src/miniverl/models/sampling.py index 5de0e69..22f19e9 100644 --- a/src/miniverl/models/sampling.py +++ b/src/miniverl/models/sampling.py @@ -135,6 +135,7 @@ def run_greedy_padded_generation( eos_token_id: int, max_new_tokens: int, stop_sequences: Sequence[str] = (), + record_logprobs: bool = False, ) -> list[GenerationOutput]: """Greedy decode a real padded batch while retaining per-row stop state. @@ -152,6 +153,7 @@ def run_greedy_padded_generation( generated: list[list[int]] = [[] for _ in sequences] reasons = ["max_new_tokens"] * len(sequences) matched: list[str | None] = [None] * len(sequences) + logprobs: list[list[float]] = [[] for _ in sequences] active = [True] * len(sequences) for _ in range(max_new_tokens): logits = step(sequences) @@ -163,6 +165,9 @@ def run_greedy_padded_generation( if not active[index]: continue token = int(torch.argmax(logits[index].detach().to(torch.float32)).item()) + if record_logprobs: + row_log_probs = torch.log_softmax(logits[index].detach().to(torch.float32), dim=-1) + logprobs[index].append(float(row_log_probs[token].item())) generated[index].append(token) sequences[index].append(token) if token == eos_token_id: @@ -183,6 +188,7 @@ def run_greedy_padded_generation( text=decode(row), stop_reason=reasons[index], matched_stop=matched[index], + logprobs=logprobs[index], ) for index, row in enumerate(generated) ] diff --git a/src/miniverl/models/shared.py b/src/miniverl/models/shared.py index 1952703..e4dd42e 100644 --- a/src/miniverl/models/shared.py +++ b/src/miniverl/models/shared.py @@ -194,6 +194,7 @@ def generate( top_p: float = 1.0, top_k: int = 0, seed: int | None = None, + record_logprobs: bool = False, ) -> GenerationOutput: with self.activated(): return self.owner.generate( @@ -204,6 +205,7 @@ def generate( top_p=top_p, top_k=top_k, seed=seed, + record_logprobs=record_logprobs, ) def hidden_states_at( diff --git a/src/miniverl/models/toy.py b/src/miniverl/models/toy.py index 4e6ad0c..1fcd826 100644 --- a/src/miniverl/models/toy.py +++ b/src/miniverl/models/toy.py @@ -302,6 +302,7 @@ def generate( top_p: float = 1.0, top_k: int = 0, seed: int | None = None, + record_logprobs: bool = False, ) -> GenerationOutput: """Sample a continuation with a KV cache.""" was_training = self.model.training @@ -328,6 +329,7 @@ def step(new_ids: list[int], state: Any) -> tuple[torch.Tensor, Any]: top_p=top_p, top_k=top_k, generator=generator, + record_logprobs=record_logprobs, ) finally: if was_training: @@ -343,6 +345,7 @@ def generate_batch( top_p: float = 1.0, top_k: int = 0, seeds: Sequence[int | None] | None = None, + record_logprobs: bool = False, ) -> list[GenerationOutput]: """Use one masked padded forward per greedy decoding step.""" if temperature != 0.0: @@ -354,6 +357,7 @@ def generate_batch( top_p=top_p, top_k=top_k, seeds=seeds, + record_logprobs=record_logprobs, ) was_training = self.model.training self.model.eval() @@ -390,6 +394,7 @@ def step(rows: list[list[int]]) -> torch.Tensor: eos_token_id=self.tokenizer.eos_token_id, max_new_tokens=max_new_tokens, stop_sequences=stop_sequences, + record_logprobs=record_logprobs, ) finally: if was_training: diff --git a/src/miniverl/runtime/rollout.py b/src/miniverl/runtime/rollout.py index 5ef6b5d..2499559 100644 --- a/src/miniverl/runtime/rollout.py +++ b/src/miniverl/runtime/rollout.py @@ -135,6 +135,7 @@ def _generate_group( top_p=self.config.top_p, top_k=self.config.top_k, seeds=[base_seed * 1_000_003 + index for index in indices], + record_logprobs=self.config.record_logprobs, ) except BaseException as exc: message = str(exc).lower() @@ -232,6 +233,14 @@ def to_trajectories( "response_token_count": len(response_ids), "generation_stop_reason": output.stop_reason, } + if self.config.record_logprobs: + if len(output.logprobs) != len(response_ids): + raise ValueError( + "PG rollout requested sampled-token log-probabilities but the backend " + "did not return exactly one value per response token" + ) + metadata["actor_rollout_log_probs"] = list(output.logprobs) + metadata["actor_rollout_policy_version"] = policy_version trajectories.append( Trajectory( trajectory_id=trajectory_id, diff --git a/src/miniverl/schemas/cache.py b/src/miniverl/schemas/cache.py index 44e436e..4665299 100644 --- a/src/miniverl/schemas/cache.py +++ b/src/miniverl/schemas/cache.py @@ -8,7 +8,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any +from typing import Any, Literal from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -22,8 +22,8 @@ "CacheCompressionStats", ] -CACHE_SCHEMA_VERSION = 2 -READABLE_CACHE_SCHEMA_VERSIONS = frozenset({1, CACHE_SCHEMA_VERSION}) +CACHE_SCHEMA_VERSION = 3 +READABLE_CACHE_SCHEMA_VERSIONS = frozenset({1, 2, CACHE_SCHEMA_VERSION}) class CacheEntryMeta(BaseModel): @@ -77,7 +77,11 @@ class CacheIndex(BaseModel): top_k: int = Field(ge=1) temperature: float = Field(gt=0.0) loss_mode: str + target_representation: Literal["topk_distribution", "sampled_token_log_probs"] = ( + "topk_distribution" + ) score_implementation_version: str | None = None + estimator_implementation_version: str | None = None execution_plan_digest: str | None = None profile_identity: dict[str, Any] = Field(default_factory=dict) dtype: str = "float32" @@ -91,6 +95,13 @@ def _validate(self) -> CacheIndex: raise ValueError(f"unsupported cache schema_version {self.schema_version}") if self.top_k > self.vocab_size: raise ValueError(f"top_k={self.top_k} exceeds vocab_size={self.vocab_size}") + if ( + self.target_representation == "sampled_token_log_probs" + and not self.estimator_implementation_version + ): + raise ValueError( + "sampled-token teacher targets require an estimator implementation version" + ) for traj_id, entry in self.entries.items(): if entry.trajectory_id != traj_id: raise ValueError( @@ -120,21 +131,29 @@ class TeacherTargetBatch: ``tail_log_prob`` ``[N]`` float32 -- ``log(1 - sum_k p_teacher)`` ``target_token_ids`` ``[N]`` int64 ``weights`` ``[N]`` float32 + + For ``sampled_token_log_probs`` targets the top-k/tail tensors are absent + and two ``[N]`` float32 vectors are present instead: rollout-time actor + log-probabilities and teacher log-probabilities for the exact sampled IDs. """ trajectory_id: str policy_version: int positions: Any - topk_indices: Any - topk_log_probs: Any - tail_log_prob: Any target_token_ids: Any weights: Any + topk_indices: Any = None + topk_log_probs: Any = None + tail_log_prob: Any = None + old_actor_log_probs: Any = None + teacher_sampled_token_log_probs: Any = None temperature: float = 1.0 top_k: int = 0 span_types: list[str] = field(default_factory=list) prompt_row_digest: str | None = None actor_response_token_ids: list[int] | None = None + target_representation: str = "topk_distribution" + estimator_implementation_version: str | None = None class CacheCompressionStats(BaseModel): diff --git a/src/miniverl/teachers/local.py b/src/miniverl/teachers/local.py index 36a5327..01e93dc 100644 --- a/src/miniverl/teachers/local.py +++ b/src/miniverl/teachers/local.py @@ -13,6 +13,7 @@ from miniverl.losses.chunked import ( BucketedTargetProvider, ExactTargetProvider, + VerlPGK1TargetProvider, VerlTopKTargetProvider, ) from miniverl.losses.exact import exact_teacher_entropy @@ -101,7 +102,14 @@ def score( policy_version=student.policy_version, shape="bucketed", provider=( - VerlTopKTargetProvider( + VerlPGK1TargetProvider( + target_token_ids=torch.zeros(0, dtype=torch.long), + old_actor_log_probs=empty, + teacher_sampled_token_log_probs=empty, + loss_max_clamp=self.loss.loss_max_clamp, + ) + if self.loss.mode is LossMode.VERL_PG_K1 + else VerlTopKTargetProvider( topk_indices=torch.zeros(0, 1, dtype=torch.long), topk_log_probs=torch.zeros(0, 1), log_prob_min_clamp=self.loss.log_prob_min_clamp, @@ -135,6 +143,95 @@ def score( with role_context, torch.no_grad(): hidden = self.backend.hidden_states_at(source.token_ids, positions, with_grad=False) + if self.loss.mode is LossMode.VERL_PG_K1: + if teacher_view is not None: + raise AlignmentError("PG-k1 supports only a standard same-prompt teacher") + response_ids = [ + token_id + for token_id, generated in zip( + student.token_ids, student.model_generated_mask, strict=True + ) + if generated + ] + if alignment.target_token_ids != response_ids: + raise AlignmentError( + "PG-k1 requires every sampled response token in its exact original order" + ) + raw_old = student.metadata.get("actor_rollout_log_probs") + rollout_version = student.metadata.get("actor_rollout_policy_version") + if rollout_version != student.policy_version: + raise AlignmentError( + "PG-k1 actor log-probabilities are not bound to the trajectory policy version" + ) + if not isinstance(raw_old, list) or len(raw_old) != n: + raise AlignmentError( + "PG-k1 requires one rollout-time actor log-probability per response token" + ) + old_actor_log_probs = torch.tensor(raw_old, dtype=torch.float32) + if not torch.isfinite(old_actor_log_probs).all(): + raise AlignmentError("PG-k1 actor rollout log-probabilities must be finite") + teacher_parts = [] + for start in range(0, n, self.loss.chunk_size): + logits = self.backend.project(hidden[start : start + self.loss.chunk_size]).to( + torch.float32 + ) + chunk_targets = target_ids[start : start + self.loss.chunk_size].to( + logits.device + ) + teacher_parts.append( + torch.log_softmax(logits, dim=-1) + .gather(-1, chunk_targets.unsqueeze(-1)) + .squeeze(-1) + .cpu() + ) + del logits + teacher_log_probs = torch.cat(teacher_parts) + del hidden, teacher_parts + pg_provider = VerlPGK1TargetProvider( + target_token_ids=target_ids, + old_actor_log_probs=old_actor_log_probs, + teacher_sampled_token_log_probs=teacher_log_probs, + clip_ratio=self.loss.clip_ratio, + clip_ratio_low=self.loss.clip_ratio_low, + clip_ratio_high=self.loss.clip_ratio_high, + clip_ratio_c=self.loss.clip_ratio_c, + loss_max_clamp=self.loss.loss_max_clamp, + ) + from miniverl.losses.verl_pg import VERL_PG_K1_IMPLEMENTATION_VERSION + + batch = TeacherTargetBatch( + trajectory_id=student.trajectory_id, + policy_version=student.policy_version, + positions=torch.tensor(positions, dtype=torch.long), + target_token_ids=target_ids, + weights=weights, + old_actor_log_probs=old_actor_log_probs, + teacher_sampled_token_log_probs=teacher_log_probs, + temperature=1.0, + top_k=1, + span_types=list(alignment.span_types), + prompt_row_digest=student.metadata.get("row_digest"), + actor_response_token_ids=response_ids, + target_representation="sampled_token_log_probs", + estimator_implementation_version=VERL_PG_K1_IMPLEMENTATION_VERSION, + ) + return TeacherScoreResult( + trajectory_id=student.trajectory_id, + policy_version=student.policy_version, + shape="sampled_token_log_probs", + provider=pg_provider, + target_token_ids=target_ids, + weights=weights, + span_types=list(alignment.span_types), + teacher_entropy=torch.zeros(0, dtype=torch.float32), + num_positions=n, + cacheable=batch, + metrics={ + "selected_positions": float(n), + "teacher_sampled_token_positions": float(n), + }, + ) + if exact_resident: # Keep only [N, H]; the [chunk, V] distribution is rebuilt on demand. # Bound to its own name because the bucketed branch below `del`s @@ -148,7 +245,7 @@ def score( exact_teacher_entropy(chunk, temperature=self.loss.temperature) ) del chunk - provider: Any = ExactTargetProvider( + exact_provider: Any = ExactTargetProvider( teacher_logits_fn=lambda a, b: project(teacher_hidden[a:b]), divergence_name=self.loss.divergence.value, temperature=self.loss.temperature, @@ -159,7 +256,7 @@ def score( trajectory_id=student.trajectory_id, policy_version=student.policy_version, shape="exact_hidden", - provider=provider, + provider=exact_provider, target_token_ids=target_ids, weights=weights, span_types=list(alignment.span_types), @@ -220,7 +317,7 @@ def score( if generated ], ) - provider = ( + distribution_provider = ( VerlTopKTargetProvider( topk_indices=topk_indices, topk_log_probs=topk_log_probs, @@ -244,7 +341,7 @@ def score( trajectory_id=student.trajectory_id, policy_version=student.policy_version, shape="bucketed", - provider=provider, + provider=distribution_provider, target_token_ids=target_ids, weights=weights, span_types=list(alignment.span_types), @@ -270,8 +367,12 @@ def describe(self) -> dict[str, Any]: "score_implementation_version": ( "verl-v0.8.0-forward-kl-topk-v1" if self.loss.mode is LossMode.VERL_FORWARD_KL_TOPK - else "miniverl-native-v1" + else ( + "verl-v0.8-pg-k1-v1" + if self.loss.mode is LossMode.VERL_PG_K1 + else "miniverl-native-v1" + ) ), - "top_k": self._effective_top_k(), + "top_k": None if self.loss.mode is LossMode.VERL_PG_K1 else self._effective_top_k(), "temperature": self.loss.temperature, } diff --git a/src/miniverl/training/trainer.py b/src/miniverl/training/trainer.py index 85d5beb..47ba742 100644 --- a/src/miniverl/training/trainer.py +++ b/src/miniverl/training/trainer.py @@ -1288,13 +1288,19 @@ def _collect( def _open_cache(self) -> TeacherCache: if self._cache is None: + from miniverl.losses.verl_pg import VERL_PG_K1_IMPLEMENTATION_VERSION from miniverl.losses.verl_topk import VERL_TOPK_SCORE_IMPLEMENTATION assert self.teacher is not None + pg_k1 = self.config.loss.mode is LossMode.VERL_PG_K1 top_k = ( - self.student.vocab_size - if self.config.loss.mode is LossMode.EXACT_FULL_VOCAB - else min(self.config.loss.top_k, self.student.vocab_size) + 1 + if pg_k1 + else ( + self.student.vocab_size + if self.config.loss.mode is LossMode.EXACT_FULL_VOCAB + else min(self.config.loss.top_k, self.student.vocab_size) + ) ) path = Path(self.config.cache.dir or self.paths.teacher_cache) identity = { @@ -1307,10 +1313,20 @@ def _open_cache(self) -> TeacherCache: "top_k": top_k, "temperature": self.config.loss.temperature, "loss_mode": self.config.loss.mode.value, + "target_representation": ( + "sampled_token_log_probs" if pg_k1 else "topk_distribution" + ), "score_implementation_version": ( - VERL_TOPK_SCORE_IMPLEMENTATION - if self.config.loss.mode is LossMode.VERL_FORWARD_KL_TOPK - else "miniverl-native-v1" + VERL_PG_K1_IMPLEMENTATION_VERSION + if pg_k1 + else ( + VERL_TOPK_SCORE_IMPLEMENTATION + if self.config.loss.mode is LossMode.VERL_FORWARD_KL_TOPK + else "miniverl-native-v1" + ) + ), + "estimator_implementation_version": ( + VERL_PG_K1_IMPLEMENTATION_VERSION if pg_k1 else None ), "execution_plan_digest": self.config.run.execution_plan_digest, "profile_identity": self.config.run.profile_identity, @@ -1741,6 +1757,10 @@ def _compute_group_gradients( verl_teacher_mass: list[Any] = [] verl_overlap_count: list[Any] = [] verl_overlap_advantage: list[Any] = [] + verl_pg_estimators: list[Any] = [] + verl_pg_advantages: list[Any] = [] + verl_pg_ratios: list[Any] = [] + verl_pg_metric_rows: list[dict[str, float]] = [] group_scale = 1.0 / max(len(group), 1) token_mean = config.loss.aggregation is LossAggregation.TOKEN_MEAN group_weight_total = sum(sum(sample.alignment.token_weights) for sample in group) @@ -1820,10 +1840,22 @@ def _compute_group_gradients( if not isinstance(diagnostics, list): continue for values in diagnostics: - verl_student_mass.append(values["student_mass"]) - verl_teacher_mass.append(values["teacher_mass"]) - verl_overlap_count.append(values["overlap_count"]) - verl_overlap_advantage.append(values["overlap_token_advantage"]) + if "student_mass" in values: + verl_student_mass.append(values["student_mass"]) + verl_teacher_mass.append(values["teacher_mass"]) + verl_overlap_count.append(values["overlap_count"]) + verl_overlap_advantage.append(values["overlap_token_advantage"]) + elif "estimator" in values: + verl_pg_estimators.append(values["estimator"]) + verl_pg_advantages.append(values["advantages"]) + verl_pg_ratios.append(values["ratio"]) + verl_pg_metric_rows.append( + { + key: float(value) + for key, value in values.items() + if isinstance(value, (int, float)) + } + ) diagnostics.clear() loss_total += float(output.loss) * microbatch_scale positions_total += output.num_positions @@ -1902,6 +1934,22 @@ def _compute_group_gradients( else 0.0 ), } + if verl_pg_estimators: + estimators = torch.cat(verl_pg_estimators).float() + advantages = torch.cat(verl_pg_advantages).float() + ratios = torch.cat(verl_pg_ratios).float() + metric_names = sorted({key for row in verl_pg_metric_rows for key in row}) + result["verl_pg_k1"] = { + "estimator_mean": float(estimators.mean()), + "estimator_abs_mean": float(estimators.abs().mean()), + "advantage_mean": float(advantages.mean()), + "ratio_mean": float(ratios.mean()), + **{ + key: sum(row.get(key, 0.0) for row in verl_pg_metric_rows) + / len(verl_pg_metric_rows) + for key in metric_names + }, + } return result def _commit_update(self) -> dict[str, float]: @@ -1957,6 +2005,16 @@ def _optimize(self, samples: list[TrainSample], *, phase: str) -> list[dict[str, rollout_policy_version = ( next(iter(rollout_versions)) if rollout_versions else self.parameter_version ) + if ( + self.config.run.mode is TrainingMode.OPD + and self.config.train.opd_freshness is OPDFreshness.STRICT + and rollout_policy_version != self.parameter_version + ): + raise LifecycleError( + "strict OPD requires rollout policy version to equal the current parameter " + f"version before update (rollout={rollout_policy_version}, " + f"parameters={self.parameter_version})" + ) for start in range(0, len(samples), accum): group = samples[start : start + accum] if not group: @@ -2593,7 +2651,11 @@ def _write_token_analysis(self, samples: list[TrainSample]) -> int: def _reload_targets_from_cache(self, samples: list[TrainSample]) -> list[TrainSample]: """Re-attach providers from the on-disk cache after the teacher is gone.""" - from miniverl.losses.chunked import BucketedTargetProvider, VerlTopKTargetProvider + from miniverl.losses.chunked import ( + BucketedTargetProvider, + VerlPGK1TargetProvider, + VerlTopKTargetProvider, + ) cache = self._open_cache() config = self.config @@ -2617,15 +2679,31 @@ def _reload_targets_from_cache(self, samples: list[TrainSample]) -> list[TrainSa ], device=self.plan.device, ) - sample.teacher.provider = ( - VerlTopKTargetProvider( + if config.loss.mode is LossMode.VERL_PG_K1: + if ( + batch.old_actor_log_probs is None + or batch.teacher_sampled_token_log_probs is None + ): + raise CheckpointError("PG-k1 cache is missing sampled-token log-probabilities") + sample.teacher.provider = VerlPGK1TargetProvider( + target_token_ids=batch.target_token_ids, + old_actor_log_probs=batch.old_actor_log_probs, + teacher_sampled_token_log_probs=batch.teacher_sampled_token_log_probs, + clip_ratio=config.loss.clip_ratio, + clip_ratio_low=config.loss.clip_ratio_low, + clip_ratio_high=config.loss.clip_ratio_high, + clip_ratio_c=config.loss.clip_ratio_c, + loss_max_clamp=config.loss.loss_max_clamp, + ) + elif config.loss.mode is LossMode.VERL_FORWARD_KL_TOPK: + sample.teacher.provider = VerlTopKTargetProvider( topk_indices=batch.topk_indices, topk_log_probs=batch.topk_log_probs, log_prob_min_clamp=config.loss.log_prob_min_clamp, loss_max_clamp=config.loss.loss_max_clamp, ) - if config.loss.mode is LossMode.VERL_FORWARD_KL_TOPK - else BucketedTargetProvider( + else: + sample.teacher.provider = BucketedTargetProvider( topk_indices=batch.topk_indices, topk_log_probs=batch.topk_log_probs, tail_log_prob=batch.tail_log_prob, @@ -2635,7 +2713,6 @@ def _reload_targets_from_cache(self, samples: list[TrainSample]) -> list[TrainSa jsd_beta=config.loss.jsd_beta, tail_epsilon=config.loss.tail_epsilon, ) - ) return samples def evaluate( diff --git a/tests/cli/test_opd_product_cli.py b/tests/cli/test_opd_product_cli.py index 83820e3..dd523a7 100644 --- a/tests/cli/test_opd_product_cli.py +++ b/tests/cli/test_opd_product_cli.py @@ -50,6 +50,49 @@ def test_run_dry_compiles_a_valid_native_recipe() -> None: assert payload["resolved_native_config"]["loss"]["mode"] == "forward_kl_topk" +def test_pg_profile_plans_and_dry_runs_through_public_cli(tmp_path: Path) -> None: + from miniverl.bridge.profiles import get_profile + + profile = "verl-opd-v0.8-single-gpu-pg-k1-v1" + config = tmp_path / "pg-k1.yaml" + config.write_text(get_profile(profile).show()["minimal_yaml"], encoding="utf-8") + + planned = runner.invoke( + app, + ["plan", "--profile", profile, "--config", str(config), "--json"], + ) + assert planned.exit_code == 0, planned.output + plan = json.loads(planned.stdout) + assert plan["profile"] == profile + assert plan["loss"] == { + "mode": "k1", + "aggregation": "token-mean", + "teacher_target": "sampled_token_log_probability", + "policy_gradient": True, + "policy_loss_mode": "vanilla", + "task_rewards": False, + } + + dry = runner.invoke( + app, + [ + "run", + "--profile", + profile, + "--config", + str(config), + "--accept-local-reinterpretations", + "--dry-run", + "--offline", + "--json", + ], + ) + assert dry.exit_code == 0, dry.output + native = json.loads(dry.stdout)["resolved_native_config"] + assert native["loss"]["mode"] == "verl_pg_k1" + assert native["rollout"]["record_logprobs"] is True + + def test_plan_accepts_trailing_hydra_style_overrides_and_override_files(tmp_path) -> None: override_file = tmp_path / "plan.overrides" override_file.write_text( diff --git a/tests/conformance/test_verl_v08_pg_k1.py b/tests/conformance/test_verl_v08_pg_k1.py new file mode 100644 index 0000000..d78e884 --- /dev/null +++ b/tests/conformance/test_verl_v08_pg_k1.py @@ -0,0 +1,193 @@ +# ruff: noqa: E402 - verify the optional official checkout before torch-backed imports + +"""Scalar, metric, gradient and optimizer-step conformance for pinned PG-k1.""" + +from __future__ import annotations + +import ast +import importlib.metadata +import json +import subprocess +from pathlib import Path +from types import SimpleNamespace +from typing import ClassVar +from urllib.parse import unquote, urlparse + +import pytest + +torch = pytest.importorskip("torch") + +from miniverl.bridge.opd_v08 import VERL_COMMIT +from miniverl.losses.verl_pg import verl_pg_k1_loss + +pytestmark = [pytest.mark.torch, pytest.mark.verl_conformance] + + +class _VerlFunctional: + @staticmethod + def masked_sum(value, mask): # type: ignore[no-untyped-def] + return (value * mask).sum() + + @staticmethod + def masked_mean(value, mask): # type: ignore[no-untyped-def] + return (value * mask).sum() / mask.sum() + + +def _official_functions(): # type: ignore[no-untyped-def] + try: + distribution = importlib.metadata.distribution("verl") + except importlib.metadata.PackageNotFoundError: + pytest.skip("official verl v0.8.0 is not installed") + direct_url = json.loads(distribution.read_text("direct_url.json") or "{}") + vcs_info = direct_url.get("vcs_info") + if vcs_info is not None: + assert vcs_info["commit_id"] == VERL_COMMIT + else: + checkout = Path(unquote(urlparse(direct_url["url"]).path.lstrip("/"))) + completed = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=checkout, + check=True, + capture_output=True, + text=True, + ) + assert completed.stdout.strip() == VERL_COMMIT + source = Path(distribution.locate_file("")) / "verl" / "trainer" / "ppo" / "core_algos.py" + tree = ast.parse(source.read_text(encoding="utf-8"), filename=str(source)) + selected = { + node.name: node + for node in tree.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name in {"agg_loss", "compute_policy_loss_vanilla", "kl_penalty_forward"} + } + assert set(selected) == {"agg_loss", "compute_policy_loss_vanilla", "kl_penalty_forward"} + module = ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + selected["agg_loss"], + selected["compute_policy_loss_vanilla"], + selected["kl_penalty_forward"], + ], + type_ignores=[], + ) + ast.fix_missing_locations(module) + namespace = { + "torch": torch, + "verl_F": _VerlFunctional, + "AlgoConfig": type("AlgoConfig", (), {}), + "register_policy_loss": lambda _name: lambda function: function, + } + exec(compile(module, str(source), "exec"), namespace) + return SimpleNamespace(**namespace) + + +class _PolicyConfig: + clip_ratio = 0.2 + clip_ratio_low = 0.2 + clip_ratio_high = 0.2 + global_batch_info: ClassVar[dict[str, object]] = {} + + def get(self, name: str, default): # type: ignore[no-untyped-def] + return getattr(self, name, default) + + +def test_pg_k1_scalar_metrics_gradient_and_step_match_exact_pinned_functions() -> None: + official = _official_functions() + logits_data = torch.tensor( + [[0.4, -0.2, 0.1, 0.8], [-0.3, 1.0, 0.2, -0.1], [0.7, 0.0, -0.4, 0.3]], + dtype=torch.float32, + ) + sampled = torch.tensor([3, 1, 0]) + old = torch.tensor([-0.9, -0.7, -1.4]) + teacher = torch.tensor([-0.4, -1.1, -0.8]) + mask = torch.tensor([1.0, 1.0, 1.0]) + + official_logits = logits_data.clone().requires_grad_(True) + official_current = ( + torch.log_softmax(official_logits, dim=-1).gather(-1, sampled[:, None]).squeeze(-1) + ) + estimator = official.kl_penalty_forward(old, teacher, "k1") + official_loss, official_metrics = official.compute_policy_loss_vanilla( + old_log_prob=old, + log_prob=official_current, + advantages=-estimator.detach(), + response_mask=mask, + loss_agg_mode="token-mean", + config=_PolicyConfig(), + ) + official_loss.backward() + + local_logits = logits_data.clone().requires_grad_(True) + local_current = torch.log_softmax(local_logits, dim=-1).gather(-1, sampled[:, None]).squeeze(-1) + local = verl_pg_k1_loss( + current_log_probs=local_current, + old_log_probs=old, + teacher_log_probs=teacher, + weights=mask, + ) + local.loss.backward() + + torch.testing.assert_close(local.estimator, estimator, rtol=0.0, atol=0.0) + torch.testing.assert_close(local.loss, official_loss, rtol=1e-6, atol=1e-7) + torch.testing.assert_close(local_logits.grad, official_logits.grad, rtol=1e-6, atol=1e-7) + assert local.metrics["pg_clipfrac"] == pytest.approx(official_metrics["actor/pg_clipfrac"]) + assert local.metrics["ppo_kl"] == pytest.approx(official_metrics["actor/ppo_kl"]) + assert local.metrics["pg_clipfrac_lower"] == pytest.approx( + official_metrics["actor/pg_clipfrac_lower"] + ) + + official_parameter = torch.nn.Parameter(logits_data.clone()) + local_parameter = torch.nn.Parameter(logits_data.clone()) + official_optimizer = torch.optim.SGD([official_parameter], lr=0.03) + local_optimizer = torch.optim.SGD([local_parameter], lr=0.03) + official_current = ( + torch.log_softmax(official_parameter, dim=-1).gather(-1, sampled[:, None]).squeeze(-1) + ) + official_step_loss, _ = official.compute_policy_loss_vanilla( + old_log_prob=old, + log_prob=official_current, + advantages=-estimator.detach(), + response_mask=mask, + loss_agg_mode="token-mean", + config=_PolicyConfig(), + ) + official_step_loss.backward() + official_optimizer.step() + local_current = ( + torch.log_softmax(local_parameter, dim=-1).gather(-1, sampled[:, None]).squeeze(-1) + ) + local_step = verl_pg_k1_loss( + current_log_probs=local_current, + old_log_probs=old, + teacher_log_probs=teacher, + weights=mask, + ) + local_step.loss.backward() + local_optimizer.step() + torch.testing.assert_close(local_parameter, official_parameter, rtol=1e-6, atol=1e-7) + + +def test_teacher_logprob_perturbation_reverses_sampled_action_gradient() -> None: + logits = torch.tensor([[0.1, 0.2, -0.3]], requires_grad=True) + sampled = torch.tensor([1]) + current = torch.log_softmax(logits, dim=-1).gather(-1, sampled[:, None]).squeeze(-1) + old = current.detach() + favored = verl_pg_k1_loss( + current_log_probs=current, + old_log_probs=old, + teacher_log_probs=old + 0.75, + weights=torch.ones(1), + ) + favored.loss.backward() + favored_gradient = logits.grad.detach().clone() + logits.grad = None + current = torch.log_softmax(logits, dim=-1).gather(-1, sampled[:, None]).squeeze(-1) + disfavored = verl_pg_k1_loss( + current_log_probs=current, + old_log_probs=old, + teacher_log_probs=old - 0.75, + weights=torch.ones(1), + ) + disfavored.loss.backward() + assert logits.grad is not None + assert favored_gradient[0, 1] < 0 < logits.grad[0, 1] diff --git a/tests/integration/test_prompt_opd_pipeline.py b/tests/integration/test_prompt_opd_pipeline.py index 4841cc5..8d44176 100644 --- a/tests/integration/test_prompt_opd_pipeline.py +++ b/tests/integration/test_prompt_opd_pipeline.py @@ -191,3 +191,142 @@ def test_prompt_opd_trains_without_an_environment_or_reward(tmp_path) -> None: "overlap_ratio", "overlap_token_advantage", } + + +def test_prompt_pg_k1_uses_rollout_actor_and_sampled_teacher_logprobs(tmp_path) -> None: + from miniverl.config import RunConfig + from miniverl.trainer import OPDTrainer + + train = tmp_path / "pg-train.parquet" + pq.write_table( + pa.Table.from_pylist( + [ + {"prompt": "alpha", "data_source": "unit", "extra_info": {"row": 0}}, + {"prompt": "beta", "data_source": "unit", "extra_info": {"row": 1}}, + ] + ), + train, + row_group_size=1, + ) + config = RunConfig.model_validate( + { + "run": { + "name": "prompt-pg-k1", + "mode": "opd", + "seed": 13, + "output_dir": str(tmp_path / "runs"), + "execution_plan_digest": "b" * 64, + "profile_identity": { + "profile_name": "verl-opd-v0.8-single-gpu-pg-k1-v1", + "digest": "c" * 64, + }, + }, + "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, + "record_logprobs": True, + }, + "selection": {"selector": "all_model_tokens"}, + "loss": { + "mode": "verl_pg_k1", + "divergence": "reverse_kl", + "aggregation": "token-mean", + "temperature": 1.0, + "scale_by_temperature_squared": False, + "top_k": 1, + "chunk_size": 16, + "loss_max_clamp": 10.0, + "estimator_implementation_version": "verl-v0.8-pg-k1-v1", + }, + "train": { + "cycles": 2, + "rollouts_per_cycle": 2, + "gradient_accumulation_steps": 2, + "learning_rate": 0.001, + }, + "memory": {"strategy": "resident"}, + "cache": {"entries_per_shard": 2, "dtype": "float32", "keep_cycles": 1}, + "eval": {"enabled": False}, + "report": {"enabled": False}, + } + ) + + trainer = OPDTrainer.from_config(config, run_id="prompt-pg-k1-test") + try: + result = trainer.train() + finally: + trainer.close() + + assert result.global_step == 2 + assert result.parameter_version == 2 + metrics = [ + json.loads(line) + for line in (result.run_dir / "metrics.jsonl").read_text(encoding="utf-8").splitlines() + ] + updates = [row for row in metrics if row.get("phase") == "opd"] + assert [row["rollout_policy_version"] for row in updates] == [0, 1] + assert [row["parameter_version"] for row in updates] == [1, 2] + assert all(row["verl_pg_k1"]["ratio_mean"] == pytest.approx(1.0) for row in updates) + trajectories = [ + json.loads(line) + for line in (result.run_dir / "trajectories.jsonl").read_text(encoding="utf-8").splitlines() + ] + assert {row["metadata"]["actor_rollout_policy_version"] for row in trajectories} == {0, 1} + assert all( + len(row["metadata"]["actor_rollout_log_probs"]) == row["metadata"]["response_token_count"] + for row in trajectories + ) + cache_index = json.loads( + (result.run_dir / "teacher-cache" / "index.json").read_text(encoding="utf-8") + ) + assert cache_index["target_representation"] == "sampled_token_log_probs" + assert cache_index["estimator_implementation_version"] == "verl-v0.8-pg-k1-v1" + assert {entry["policy_version"] for entry in cache_index["entries"].values()} == {1} + assert all( + entry["tensor_keys"] + == [ + "positions", + "target_token_ids", + "weights", + "old_actor_log_probs", + "teacher_sampled_token_log_probs", + ] + for entry in cache_index["entries"].values() + ) diff --git a/tests/unit/test_cache.py b/tests/unit/test_cache.py index 1b25c76..e6bfa66 100644 --- a/tests/unit/test_cache.py +++ b/tests/unit/test_cache.py @@ -124,6 +124,60 @@ def test_float32_round_trip_is_exact(tmp_path: Path): assert loaded.span_types == batch.span_types +def test_pg_sampled_token_target_round_trip_is_version_bound(tmp_path: Path) -> None: + cache = TeacherCache.create( + tmp_path / "pg-cache", + miniverl_version="0.10.0.dev0", + teacher_model_id="toy-teacher", + teacher_model_revision="rev-pg", + tokenizer_fingerprint="fp-pg", + vocab_size=VOCAB, + top_k=1, + temperature=1.0, + loss_mode="verl_pg_k1", + target_representation="sampled_token_log_probs", + estimator_implementation_version="verl-v0.8-pg-k1-v1", + score_implementation_version="verl-v0.8-pg-k1-v1", + ) + batch = TeacherTargetBatch( + trajectory_id="pg:0", + policy_version=9, + positions=torch.tensor([3, 4, 5]), + target_token_ids=torch.tensor([11, 12, 13]), + weights=torch.ones(3), + old_actor_log_probs=torch.tensor([-2.0, -1.5, -3.0]), + teacher_sampled_token_log_probs=torch.tensor([-1.0, -2.0, -2.5]), + temperature=1.0, + top_k=1, + span_types=["assistant_text"] * 3, + prompt_row_digest="a" * 64, + actor_response_token_ids=[11, 12, 13], + target_representation="sampled_token_log_probs", + estimator_implementation_version="verl-v0.8-pg-k1-v1", + ) + cache.write(batch, selector="all_model_tokens") + cache.flush() + + loaded = cache.read( + "pg:0", + expect_policy_version=9, + expect_prompt_row_digest="a" * 64, + expect_actor_response_token_ids=[11, 12, 13], + ) + assert loaded.topk_indices is None + assert loaded.topk_log_probs is None + assert loaded.tail_log_prob is None + assert loaded.target_representation == "sampled_token_log_probs" + assert loaded.estimator_implementation_version == "verl-v0.8-pg-k1-v1" + assert torch.equal(loaded.old_actor_log_probs, batch.old_actor_log_probs) + assert torch.equal( + loaded.teacher_sampled_token_log_probs, + batch.teacher_sampled_token_log_probs, + ) + with pytest.raises(StaleCacheError, match="off-policy"): + cache.read("pg:0", expect_policy_version=10) + + def test_prompt_target_binding_rejects_a_different_actor_response(tmp_path: Path): cache = TeacherCache.create( tmp_path / "tc", diff --git a/tests/unit/test_profile_registry.py b/tests/unit/test_profile_registry.py index 85f1b0c..a639ab7 100644 --- a/tests/unit/test_profile_registry.py +++ b/tests/unit/test_profile_registry.py @@ -7,13 +7,14 @@ from miniverl.errors import ConfigError DIRECT_PROFILE = "verl-opd-v0.8-single-gpu-v1" +PG_PROFILE = "verl-opd-v0.8-single-gpu-pg-k1-v1" def test_registry_exposes_immutable_direct_gkd_identity() -> None: from miniverl.bridge.profiles import get_profile, list_profiles profiles = list_profiles() - assert [item.name for item in profiles] == [DIRECT_PROFILE] + assert [item.name for item in profiles] == [DIRECT_PROFILE, PG_PROFILE] profile = get_profile(DIRECT_PROFILE) identity = profile.identity diff --git a/tests/unit/test_prompt_rollout_runtime.py b/tests/unit/test_prompt_rollout_runtime.py index 45bd9db..d59e86e 100644 --- a/tests/unit/test_prompt_rollout_runtime.py +++ b/tests/unit/test_prompt_rollout_runtime.py @@ -52,6 +52,48 @@ def test_real_padded_greedy_generation_matches_sequential() -> None: assert single[0].token_ids == expected[0].token_ids +def test_prompt_runtime_records_exact_rollout_logprobs_for_pg_targets() -> None: + tokenizer = ToyTokenizer() + backend = ToyBackend(tokenizer=tokenizer, model_id="toy", seed=23, 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=2, + max_padded_tokens=64, + record_logprobs=True, + ), + ) + rendered = [ + RenderedPrompt( + record=_record(index), + text=text, + token_ids=tuple(tokenizer.encode(text)), + tokenizer_identity=tokenizer.identity, + rendered_prompt_digest=f"{index + 40: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", "longer prompt")) + ] + + batch = runtime.prepare_batch(rendered) + generated = runtime.generate(batch, policy_version=7, seed=11) + trajectories = runtime.to_trajectories(batch, generated, policy_version=7) + + for output, trajectory in zip(generated.outputs, trajectories, strict=True): + assert len(output.logprobs) == len(output.token_ids) + assert all(value <= 0.0 for value in output.logprobs) + assert trajectory.metadata["actor_rollout_log_probs"] == output.logprobs + assert trajectory.metadata["actor_rollout_policy_version"] == 7 + + def test_prompt_trajectories_never_select_prompt_or_padding_tokens() -> None: tokenizer = ToyTokenizer() backend = ToyBackend(tokenizer=tokenizer, model_id="toy", seed=7, trainable=False) diff --git a/tests/unit/test_verl_opd_export.py b/tests/unit/test_verl_opd_export.py index 389c231..0b82e46 100644 --- a/tests/unit/test_verl_opd_export.py +++ b/tests/unit/test_verl_opd_export.py @@ -21,13 +21,18 @@ def _safetensors_bytes() -> bytes: return struct.pack(" tuple[Path, Path, Path]: +def _opd_run( + tmp_path: Path, + *, + teacher_adapter: bool = False, + profile: str = "verl-opd-v0.8-single-gpu-v1", +) -> tuple[Path, Path, Path]: from miniverl.bridge.profiles import get_profile run = tmp_path / "run" model = run / "final-peft-adapter" model.mkdir(parents=True) - profile_identity = get_profile("verl-opd-v0.8-single-gpu-v1").identity.model_dump(mode="json") + profile_identity = get_profile(profile).identity.model_dump(mode="json") (run / "manifest.json").write_text( json.dumps({"status": "complete", "profile_identity": profile_identity}), encoding="utf-8", @@ -161,11 +166,24 @@ def _opd_run(tmp_path: Path, *, teacher_adapter: bool = False) -> tuple[Path, Pa }, }, } + if profile == "verl-opd-v0.8-single-gpu-pg-k1-v1": + source["distillation"]["distillation_loss"] = { + "loss_mode": "k1", + "topk": None, + "use_task_rewards": False, + "distillation_loss_coef": 1.0, + "log_prob_min_clamp": None, + "use_policy_gradient": True, + "policy_loss_mode": "vanilla", + "clip_ratio": 0.2, + "clip_ratio_low": 0.2, + "clip_ratio_high": 0.2, + } (run / "verl-source-config.json").write_text(json.dumps(source), encoding="utf-8") (run / "verl-compatibility-report.json").write_text( json.dumps( { - "profile": "verl-opd-v0.8-single-gpu-v1", + "profile": profile, "executable": True, "compiled_digest": "b" * 64, "source": source, @@ -174,12 +192,42 @@ def _opd_run(tmp_path: Path, *, teacher_adapter: bool = False) -> tuple[Path, Pa encoding="utf-8", ) (run / "local-execution-plan.json").write_text( - json.dumps({"profile": "verl-opd-v0.8-single-gpu-v1", "compiled_digest": "b" * 64}), + json.dumps({"profile": profile, "compiled_digest": "b" * 64}), encoding="utf-8", ) return run, train, val +def test_pg_k1_export_preserves_sampled_logprob_semantics_without_topk(tmp_path: Path) -> None: + from miniverl.bridge.contract import VERL_TAG + from miniverl.bridge.doctor import inspect_bridge_bundle + from miniverl.bridge.export import export_verl_bundle + + profile = "verl-opd-v0.8-single-gpu-pg-k1-v1" + run, _, _ = _opd_run(tmp_path, profile=profile) + out = tmp_path / "pg-bundle" + + report = export_verl_bundle(run, target_verl=VERL_TAG, out=out) + overrides = yaml.safe_load((out / "recipe/verl-opd-overrides.yaml").read_text()) + loss = overrides["distillation"]["distillation_loss"] + + assert report["profile"] == profile + assert report["target_semantics"] == "pure OPD sampled k1 policy gradient" + assert report["reward_required"] is False + assert loss["loss_mode"] == "k1" + assert loss["use_policy_gradient"] is True + assert loss["policy_loss_mode"] == "vanilla" + assert loss["use_task_rewards"] is False + assert "topk" not in loss + assert overrides["actor_rollout_ref"]["actor"]["use_kl_loss"] is False + assert overrides["algorithm"]["use_kl_in_reward"] is False + + diagnosis = inspect_bridge_bundle(out) + assert diagnosis["verdict"] == "ok" + assert diagnosis["config_profile"]["profile"] == profile + assert diagnosis["config_profile"]["target_representation"] == ("sampled_token_log_probability") + + def test_pure_opd_export_is_reward_free_and_preserves_data_bytes(tmp_path: Path) -> None: from miniverl.bridge.contract import VERL_TAG from miniverl.bridge.doctor import inspect_bridge_bundle diff --git a/tests/unit/test_verl_opd_materialize.py b/tests/unit/test_verl_opd_materialize.py index 67833c7..3d6ae28 100644 --- a/tests/unit/test_verl_opd_materialize.py +++ b/tests/unit/test_verl_opd_materialize.py @@ -8,6 +8,7 @@ from typing import Any import pytest +import yaml def _safetensors_bytes() -> bytes: @@ -49,12 +50,17 @@ def _snapshot(root: Path, revision: str) -> Path: return snapshot -def _bundle(tmp_path: Path, *, teacher_adapter: bool = False) -> tuple[Path, str, str]: +def _bundle( + tmp_path: Path, + *, + teacher_adapter: bool = False, + profile: str = "verl-opd-v0.8-single-gpu-v1", +) -> tuple[Path, str, str]: from miniverl.bridge.contract import VERL_TAG from miniverl.bridge.export import export_verl_bundle from tests.unit.test_verl_opd_export import _opd_run - run, _, _ = _opd_run(tmp_path, teacher_adapter=teacher_adapter) + run, _, _ = _opd_run(tmp_path, teacher_adapter=teacher_adapter, profile=profile) bundle = tmp_path / "bundle" export_verl_bundle(run, target_verl=VERL_TAG, out=bundle) student_revision = "c1899de289a04d12100db370d81485cdf75e47ca" @@ -62,6 +68,41 @@ def _bundle(tmp_path: Path, *, teacher_adapter: bool = False) -> tuple[Path, str return bundle, student_revision, teacher_revision +def test_pg_k1_bundle_materializes_without_a_topk_contract( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from miniverl.bridge import materialize + + profile = "verl-opd-v0.8-single-gpu-pg-k1-v1" + bundle, student_revision, teacher_revision = _bundle(tmp_path, profile=profile) + student = _snapshot(tmp_path / "student-cache", student_revision) + teacher = _snapshot(tmp_path / "teacher-cache", teacher_revision) + + def passed_pg(root: Path) -> dict[str, Any]: + recipe = yaml.safe_load( + (root / "recipe/verl-opd-overrides.yaml").read_text(encoding="utf-8") + ) + loss = recipe["distillation"]["distillation_loss"] + assert "topk" not in loss + assert loss["loss_mode"] == "k1" + assert loss["use_policy_gradient"] is True + return {"status": "passed", "checks": {"sampled_k1_policy_gradient_contract": "passed"}} + + monkeypatch.setattr(materialize, "_validate_upstream_bundle", passed_pg) + report = materialize.materialize_verl_bundle( + bundle, + student_snapshot=student, + teacher_snapshot=teacher, + offline=True, + ) + + assert report["profile"] == profile + assert report["launchable"] is True + assert report["materialization"]["upstream_validation"]["checks"] == { + "sampled_k1_policy_gradient_contract": "passed" + } + + def _tree_hashes(root: Path) -> dict[str, str]: return { path.relative_to(root).as_posix(): hashlib.sha256(path.read_bytes()).hexdigest() diff --git a/tests/unit/test_verl_pg_k1.py b/tests/unit/test_verl_pg_k1.py new file mode 100644 index 0000000..fe214c0 --- /dev/null +++ b/tests/unit/test_verl_pg_k1.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import pytest +import torch + +from miniverl.losses.verl_pg import ( + VERL_PG_K1_IMPLEMENTATION_VERSION, + verl_k1_estimator, + verl_pg_k1_loss, +) + + +def test_k1_estimator_and_detached_advantage_match_pinned_verl() -> None: + old = torch.tensor([-2.0, -0.5, -4.0], requires_grad=True) + teacher = torch.tensor([-1.0, -1.5, -2.5], requires_grad=True) + + estimator = verl_k1_estimator(old, teacher) + output = verl_pg_k1_loss( + current_log_probs=old, + old_log_probs=old, + teacher_log_probs=teacher, + weights=torch.ones(3), + ) + + torch.testing.assert_close(estimator, old - teacher) + torch.testing.assert_close(output.advantages, (teacher - old).detach()) + assert output.advantages.requires_grad is False + assert VERL_PG_K1_IMPLEMENTATION_VERSION == "verl-v0.8-pg-k1-v1" + + +def test_vanilla_policy_loss_matches_pinned_verl_equations_and_metrics() -> None: + current = torch.tensor([-1.8, -0.9, -4.4], requires_grad=True) + old = torch.tensor([-2.0, -0.5, -4.0]) + teacher = torch.tensor([-1.0, -1.5, -2.5]) + weights = torch.tensor([1.0, 0.0, 1.0]) + + output = verl_pg_k1_loss( + current_log_probs=current, + old_log_probs=old, + teacher_log_probs=teacher, + weights=weights, + clip_ratio=0.2, + clip_ratio_low=0.2, + clip_ratio_high=0.2, + clip_ratio_c=3.0, + ) + + advantage = (teacher - old).detach() + negative_approx_kl = (current - old).clamp(-20.0, 20.0) + ratio = negative_approx_kl.exp() + loss1 = -advantage * ratio + loss2 = -advantage * ratio.clamp(0.8, 1.2) + clipped = torch.maximum(loss1, loss2) + expected_tokens = torch.where(advantage < 0, torch.minimum(-advantage * 3.0, clipped), clipped) + expected = (expected_tokens * weights).sum() / weights.sum() + + torch.testing.assert_close(output.per_token_loss, expected_tokens) + torch.testing.assert_close(output.loss, expected) + assert output.metrics["pg_clipfrac"] == pytest.approx(0.5) + assert output.metrics["ppo_kl"] == pytest.approx(0.1) + assert output.metrics["pg_clipfrac_lower"] == pytest.approx(0.0) + + +def test_teacher_preference_changes_actor_gradient_in_expected_direction() -> None: + logits = torch.tensor([[0.2, -0.1, 0.0]], requires_grad=True) + sampled = torch.tensor([1]) + current = torch.log_softmax(logits, dim=-1).gather(-1, sampled[:, None]).squeeze(-1) + old = current.detach() + + favored = verl_pg_k1_loss( + current_log_probs=current, + old_log_probs=old, + teacher_log_probs=old + 1.0, + weights=torch.ones(1), + ) + favored.loss.backward() + favored_grad = logits.grad.detach().clone() + + logits.grad = None + current = torch.log_softmax(logits, dim=-1).gather(-1, sampled[:, None]).squeeze(-1) + disfavored = verl_pg_k1_loss( + current_log_probs=current, + old_log_probs=old, + teacher_log_probs=old - 1.0, + weights=torch.ones(1), + ) + disfavored.loss.backward() + + assert favored_grad[0, 1] < 0.0 + assert logits.grad is not None + assert logits.grad[0, 1] > 0.0 + torch.testing.assert_close(favored_grad, -logits.grad) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) +def test_pg_loss_rejects_nonfinite_inputs(value: float) -> None: + with pytest.raises(ValueError, match="finite"): + verl_pg_k1_loss( + current_log_probs=torch.tensor([value]), + old_log_probs=torch.tensor([-1.0]), + teacher_log_probs=torch.tensor([-1.0]), + weights=torch.ones(1), + ) diff --git a/tests/unit/test_verl_pg_profile.py b/tests/unit/test_verl_pg_profile.py new file mode 100644 index 0000000..6318c86 --- /dev/null +++ b/tests/unit/test_verl_pg_profile.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +import yaml + +from miniverl.errors import ConfigError + +PG_PROFILE = "verl-opd-v0.8-single-gpu-pg-k1-v1" + + +def _payload() -> dict[str, object]: + source = Path(__file__).resolve().parents[2] / "examples/verl-opd-v0.8-single-gpu.yaml" + payload = yaml.safe_load(source.read_text(encoding="utf-8")) + loss = payload["distillation"]["distillation_loss"] + loss.update( + { + "loss_mode": "k1", + "topk": None, + "use_policy_gradient": True, + "policy_loss_mode": "vanilla", + "clip_ratio": 0.2, + "clip_ratio_low": 0.2, + "clip_ratio_high": 0.2, + "loss_max_clamp": 10.0, + "log_prob_min_clamp": None, + } + ) + payload["algorithm"]["adv_estimator"] = None + return payload + + +def _set(payload: dict[str, object], path: str, value: object) -> None: + current = payload + parts = path.split(".") + for part in parts[:-1]: + current = current[part] # type: ignore[assignment] + current[parts[-1]] = value + + +def test_pg_profile_compiles_as_a_separate_executable_contract() -> None: + from miniverl.bridge.opd_pg_v08 import compile_verl_pg_k1_v08 + from miniverl.bridge.opd_runtime import build_system_plan, compile_native_run_config + + plan = compile_verl_pg_k1_v08(_payload()) + assert plan.profile == PG_PROFILE + assert plan.executable is True + assert plan.local_execution["loss_mode"] == "k1" + assert plan.local_execution["policy_gradient"] is True + assert plan.local_execution["policy_loss_mode"] == "vanilla" + assert plan.local_execution["task_rewards"] is False + assert plan.source.distillation.distillation_loss.topk is None + native = compile_native_run_config(plan, system_plan=build_system_plan(plan)) + assert native.loss.mode.value == "verl_pg_k1" + assert native.loss.policy_loss_mode == "vanilla" + assert native.loss.clip_ratio == pytest.approx(0.2) + assert native.rollout.record_logprobs is True + assert native.train.opd_freshness.value == "strict" + assert native.train.gradient_accumulation_steps == native.train.rollouts_per_cycle + + +@pytest.mark.parametrize( + ("path", "value"), + [ + ("distillation.distillation_loss.loss_mode", "k3"), + ("distillation.distillation_loss.use_policy_gradient", False), + ("distillation.distillation_loss.policy_loss_mode", "dppo_tv"), + ("distillation.distillation_loss.use_task_rewards", True), + ("distillation.distillation_loss.topk", 32), + ("distillation.distillation_loss.clip_ratio", 0.1), + ("actor_rollout_ref.actor.use_kl_loss", True), + ("algorithm.use_kl_in_reward", True), + ("algorithm.adv_estimator", "grpo"), + ("actor_rollout_ref.rollout.n", 2), + ("trainer.n_gpus_per_node", 2), + ], +) +def test_pg_profile_rejects_out_of_contract_semantics(path: str, value: object) -> None: + from miniverl.bridge.opd_pg_v08 import compile_verl_pg_k1_v08 + + payload = _payload() + _set(payload, path, value) + with pytest.raises(ConfigError, match="not executable") as exc_info: + compile_verl_pg_k1_v08(payload) + assert path in str(exc_info.value) + + +def test_registry_lists_pg_profile_with_independent_identity() -> None: + from miniverl.bridge.profiles import get_profile, list_profiles + + names = [profile.name for profile in list_profiles()] + assert names == ["verl-opd-v0.8-single-gpu-v1", PG_PROFILE] + direct = get_profile(names[0]) + pg = get_profile(PG_PROFILE) + assert pg.identity.digest != direct.identity.digest + assert pg.identity.loss_conformance_version == "pg-k1-verl-v0.8-v1" + assert pg.summary.objective == "k1 + vanilla policy loss" + assert pg.summary.teacher_target == "sampled-token teacher log-probability" From 5c2482c34244508c4581056e09f46f6eab0948fb Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Thu, 13 Aug 2026 23:45:46 -0700 Subject: [PATCH 2/7] Handle sampled targets in token analysis --- src/miniverl/training/trainer.py | 14 +++++++++++++- tests/integration/test_prompt_opd_pipeline.py | 15 ++++++++++++++- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/src/miniverl/training/trainer.py b/src/miniverl/training/trainer.py index 47ba742..7036eda 100644 --- a/src/miniverl/training/trainer.py +++ b/src/miniverl/training/trainer.py @@ -2618,8 +2618,15 @@ def _write_token_analysis(self, samples: list[TrainSample]) -> int: else [None] * len(per_token) ) teacher_top: list[int | None] = [None] * len(per_token) + teacher_sampled_log_prob: list[float | None] = [None] * len(per_token) if sample.teacher is not None and sample.teacher.cacheable is not None: - teacher_top = sample.teacher.cacheable.topk_indices[:, 0].tolist() + cacheable = sample.teacher.cacheable + if cacheable.topk_indices is not None: + teacher_top = cacheable.topk_indices[:, 0].tolist() + if cacheable.teacher_sampled_token_log_probs is not None: + teacher_sampled_log_prob = ( + cacheable.teacher_sampled_token_log_probs.tolist() + ) records = [] max_tokens = self.config.report.max_tokens_per_trajectory for i, target_position in enumerate(alignment.student_prediction_positions): @@ -2639,6 +2646,11 @@ def _write_token_analysis(self, samples: list[TrainSample]) -> int: "token_loss": per_token[i] if i < len(per_token) else None, "teacher_entropy": entropy[i] if i < len(entropy) else None, "teacher_top_token": teacher_top[i] if i < len(teacher_top) else None, + "teacher_sampled_token_log_prob": ( + teacher_sampled_log_prob[i] + if i < len(teacher_sampled_log_prob) + else None + ), "student_top_token": student_top[i] if i < len(student_top) else None, } ) diff --git a/tests/integration/test_prompt_opd_pipeline.py b/tests/integration/test_prompt_opd_pipeline.py index 8d44176..2de9623 100644 --- a/tests/integration/test_prompt_opd_pipeline.py +++ b/tests/integration/test_prompt_opd_pipeline.py @@ -284,7 +284,11 @@ def test_prompt_pg_k1_uses_rollout_actor_and_sampled_teacher_logprobs(tmp_path) "memory": {"strategy": "resident"}, "cache": {"entries_per_shard": 2, "dtype": "float32", "keep_cycles": 1}, "eval": {"enabled": False}, - "report": {"enabled": False}, + "report": { + "enabled": True, + "max_trajectories": 2, + "max_tokens_per_trajectory": 8, + }, } ) @@ -330,3 +334,12 @@ def test_prompt_pg_k1_uses_rollout_actor_and_sampled_teacher_logprobs(tmp_path) ] for entry in cache_index["entries"].values() ) + token_rows = [ + json.loads(line) + for line in (result.run_dir / "token_analysis.jsonl") + .read_text(encoding="utf-8") + .splitlines() + ] + assert token_rows + assert all(row["teacher_top_token"] is None for row in token_rows) + assert all(row["teacher_sampled_token_log_prob"] is not None for row in token_rows) From c12756b8f64c256a989e3295fe0e89e6fc2cd976 Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Thu, 13 Aug 2026 23:53:26 -0700 Subject: [PATCH 3/7] Bundle tokenizer metadata with PEFT exports --- src/miniverl/models/adapter_io.py | 4 ++++ tests/integration/test_hf_backend_offline.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/src/miniverl/models/adapter_io.py b/src/miniverl/models/adapter_io.py index ca0a711..f63a1e2 100644 --- a/src/miniverl/models/adapter_io.py +++ b/src/miniverl/models/adapter_io.py @@ -535,6 +535,10 @@ def export_adapter( safe_serialization=True, **({"selected_adapters": [selected_adapter]} if selected_adapter else {}), ) + save_tokenizer = getattr(getattr(tokenizer, "raw", None), "save_pretrained", None) + if not callable(save_tokenizer): + raise BackendError("loaded HF tokenizer cannot be serialized with the PEFT adapter") + save_tokenizer(target) if selected_adapter is not None: nested = target / selected_adapter if nested.is_dir(): diff --git a/tests/integration/test_hf_backend_offline.py b/tests/integration/test_hf_backend_offline.py index 2bb292c..c71ebba 100644 --- a/tests/integration/test_hf_backend_offline.py +++ b/tests/integration/test_hf_backend_offline.py @@ -1237,6 +1237,8 @@ def test_miniverl_checkpoint_exports_and_reloads_as_standard_peft( "adapter_model.safetensors", } assert (exported / "miniverl_adapter_manifest.json").is_file() + assert (exported / "tokenizer_config.json").is_file() + assert (exported / "tokenizer.json").is_file() assert manifest["policy_evaluation"] is None opd_mapping = config.model_dump(mode="json") From 9b502e05937a84b1851985fda8cde4e3bc256a50 Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Thu, 13 Aug 2026 23:59:41 -0700 Subject: [PATCH 4/7] Record measured PG profile evidence --- docs/evidence/rtx4080-verl-pg-k1-v1.json | 171 ++++++++++++++++++ docs/verl-pg-k1-workload.md | 38 ++++ mkdocs.yml | 1 + src/miniverl/bridge/doctor.py | 29 ++- src/miniverl/bridge/materialize.py | 45 ++++- src/miniverl/bridge/profiles.py | 2 +- .../test_verl_bridge_parquet_streaming.py | 11 ++ tests/unit/test_verl_opd_materialize.py | 8 + 8 files changed, 296 insertions(+), 9 deletions(-) create mode 100644 docs/evidence/rtx4080-verl-pg-k1-v1.json create mode 100644 docs/verl-pg-k1-workload.md diff --git a/docs/evidence/rtx4080-verl-pg-k1-v1.json b/docs/evidence/rtx4080-verl-pg-k1-v1.json new file mode 100644 index 0000000..d234fb6 --- /dev/null +++ b/docs/evidence/rtx4080-verl-pg-k1-v1.json @@ -0,0 +1,171 @@ +{ + "artifacts": { + "checkpoint_hashes": { + "adapter.safetensors": "315673b72ab3c871b6225e32981bce21877fca9d63d1c2e8e34d0b4b9295b8e6", + "optimizer.safetensors": "972cfa720e394da0435eeffde2a6602d764dfc774060b082e6dbb754e8842ffc", + "state.json": "e2b5135ba3bb78c09cdd1d946cd5b2a22456cbec4ff8f4b5bfe364da7e550e61" + }, + "standard_peft_adapter_sha256": "769a2be71627d1e0d78e35aa7131d360d0750e71c8dce89c0b7c354c2f5dee01", + "total_workload_bytes": 130048818, + "trajectory_sha256": "765d76ed2f16da571d35cd71176f08ffe447b78e7fb94a13e73474ed3b7cb893" + }, + "hardware": { + "bitsandbytes": "0.50.0", + "cuda_runtime": "13.0", + "driver": "596.49", + "gpu": "NVIDIA GeForce RTX 4080", + "gpu_count": 1, + "peft": "0.18.0", + "torch": "2.13.0+cu130", + "transformers": "5.14.1", + "vram_gib": 15.992 + }, + "kind": "single_gpu_verl_pg_k1_systems_workload", + "measured_at": "2026-08-14T06:51:51Z", + "measurements": { + "adapter_bytes": 9198518, + "batch_downshifts": { + "rollout_oom": 0, + "update_chunk_oom": 0 + }, + "cache_bytes": 26382, + "checkpoint_bytes": 27600769, + "cold_startup_seconds": 8.1302, + "distinct_prompts_consumed": 32, + "observed_rollout_physical_batch_sizes": [ + 4 + ], + "peak_allocated_gib": 2.3032, + "peak_reserved_gib": 3.1914, + "prompts_consumed": 32, + "steady_state_median_seconds": { + "actor_update": 1.1579, + "rollout": 19.0062, + "teacher_scoring": 0.2864 + }, + "steady_state_median_throughput": { + "rollout_tokens_per_second": 13.4693, + "teacher_scored_positions_per_second": 893.86, + "update_positions_per_second": 221.0899 + }, + "time_to_first_rollout_seconds": 28.4604, + "time_to_first_teacher_targets_seconds": 28.7706, + "time_to_first_update_seconds": 30.0902, + "total_run_bytes": 64966843 + }, + "miniverl_version": "0.10.0.dev0", + "models": { + "student": { + "adapter": null, + "attn_implementation": "sdpa", + "dtype": "auto", + "gradient_checkpointing": true, + "lora": { + "alpha": 32, + "bias": "none", + "dropout": 0.0, + "enabled": true, + "r": 16, + "target_modules": [ + "q_proj", + "v_proj" + ] + }, + "model_id": "Qwen/Qwen3-0.6B", + "prepare_kbit_training": true, + "quantization": "nf4", + "revision": "c1899de289a04d12100db370d81485cdf75e47ca", + "tokenizer_id": null, + "tokenizer_revision": null, + "toy": { + "hidden_size": 64, + "intermediate_size": 128, + "max_position_embeddings": 1024, + "num_heads": 4, + "num_layers": 2 + }, + "trust_remote_code": false + }, + "teacher": { + "adapter": null, + "attn_implementation": "sdpa", + "dtype": "bfloat16", + "mode": "standard", + "model_id": "Qwen/Qwen3-1.7B", + "quantization": "nf4", + "revision": "70d244cc86ccca08cf5af4e1e306ecf908b1ad5e", + "tokenizer_id": null, + "tokenizer_revision": null, + "toy": { + "hidden_size": 64, + "intermediate_size": 128, + "max_position_embeddings": 1024, + "num_heads": 4, + "num_layers": 2 + }, + "toy_pretrain_lr": 0.003, + "toy_pretrain_steps": 0, + "toy_teacher_seed": 99, + "trust_remote_code": false + } + }, + "profile": "verl-opd-v0.8-single-gpu-pg-k1-v1", + "profile_identity": { + "digest": "8bbff2027f44d9601d3ed3143621a60280a1bacafd1b08d2ad95bf4d9caac3ca", + "export_version": "verl-opd-pg-k1-export-v1", + "field_rule_digest": "d53a91b0094a12428006d973683ac85e91fc6b55f47dc33636ce1d2e661b7e57", + "loss_conformance_version": "pg-k1-verl-v0.8-v1", + "native_compiler_version": "pg-k1-native-v1", + "profile_name": "verl-opd-v0.8-single-gpu-pg-k1-v1", + "profile_schema_version": 1, + "upstream_commit": "7aed6b230776f963fa09509c10d9c3a767d1102c", + "upstream_repository": "https://github.com/verl-project/verl", + "upstream_tag": "v0.8.0" + }, + "recipe": { + "compiled_plan_sha256": "47f7a2218c823abf8e021d6569805db7c4d2805f33493a920d627322533cea5e", + "estimator": "k1", + "input_parquet_sha256": "baa99ba8c59d4c30b2e3081346860e7e8b2665dd70fdbcd7952935545a0cd770", + "interrupt_after_updates": 4, + "logical_batch": 4, + "optimizer_updates": 8, + "policy_loss": "vanilla", + "prompt_limit": 128, + "prompts_available": 64, + "prompts_consumed": 32, + "response_limit": 64, + "teacher_target": "sampled_token_log_probability" + }, + "resource_contract": { + "gpu_hour_limit": 3.0, + "peak_reserved_limit_gib": 14.5, + "peak_reserved_within_limit": true, + "total_gpu_workload_seconds": 359.7 + }, + "resume": { + "adapter_and_optimizer_byte_identical": true, + "excluded_run_identity_field": "resolved_config_digest", + "global_optimizer_steps": 8, + "interrupt_checkpoint_bytes": 27600769, + "resume_load_seconds": 5.5309, + "status": "exact_match", + "task_cursor": 32, + "training_state_fields_identical": true, + "trajectories_byte_identical": true + }, + "schema_version": 1, + "scientific_scope": { + "algorithm_comparison": false, + "runtime_correctness_only": true, + "task_quality_evaluated": false + }, + "source_commit": "5c2482c34244508c4581056e09f46f6eab0948fb", + "status": "measured", + "verl": { + "commit": "7aed6b230776f963fa09509c10d9c3a767d1102c", + "distributed_execution_tested": false, + "repository": "https://github.com/verl-project/verl", + "tag": "v0.8.0" + }, + "workload_driver_sha256": "5b712a40ef2932800ce43072fa88f175a2454febe7a60b0642356bc5be4dfdd3" +} diff --git a/docs/verl-pg-k1-workload.md b/docs/verl-pg-k1-workload.md new file mode 100644 index 0000000..a163f85 --- /dev/null +++ b/docs/verl-pg-k1-workload.md @@ -0,0 +1,38 @@ +# RTX 4080 sampled-k1 PG systems workload + +This is runtime and semantic-conformance evidence for +`verl-opd-v0.8-single-gpu-pg-k1-v1`, not a task-quality or algorithm +comparison result. The source profile pins verl `v0.8.0` at +`7aed6b230776f963fa09509c10d9c3a767d1102c`. + +| Item | Measured value | +| --- | --- | +| Hardware | one NVIDIA RTX 4080, 15.992 GiB | +| Models | Qwen3-0.6B actor / Qwen3-1.7B teacher, pinned NF4 snapshots | +| Workload | 32 distinct prompts, response bound 64, 8 strict updates | +| Target | sampled-token teacher log-probability; `k1` + vanilla policy loss | +| Peak CUDA | 2.3032 GiB allocated / 3.1914 GiB reserved | +| Time to first update | 30.0902 s | +| Steady rollout | 13.4693 generated tokens/s | +| Steady teacher scoring | 893.86 sampled positions/s | +| Steady actor update | 221.0899 positions/s | + +The matched run stopped after update 4, reconstructed the trainer from its +checkpoint, and finished at update 8 with byte-identical trajectories, adapter +and optimizer tensors and identical training-state fields (apart from the +deliberately run-specific resolved-config digest). Every update recorded +`rollout policy version == parameter version before update`; each successful +update incremented the parameter version once. + +The standard PEFT adapter exported with tokenizer metadata. A separate +checksummed scale-out bundle preserved the sampled-k1 fields without a top-k +requirement; exact student and teacher snapshots were materialized +transactionally, the pinned upstream config merge and sequential CPU +model/data load smoke passed, and the local bundle became `launchable: true`. +No distributed verl job ran, and algorithm-wide parity is not claimed. + +The complete machine-readable record is +[`evidence/rtx4080-verl-pg-k1-v1.json`](evidence/rtx4080-verl-pg-k1-v1.json) +(SHA-256 `b549905b2b62b5e51d721ba2932befb7b910284ff1b4016095d341ec232a725f`). +The estimator/scalar/gradient contract is documented in +[ADR 0010](adr/0010-verl-v0.8-pg-k1-contract.md). diff --git a/mkdocs.yml b/mkdocs.yml index 2f5b6f4..cd72f22 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -50,6 +50,7 @@ nav: - Current verl-style OPD runtime: verl-opd-runtime.md - Run verl-style OPD: opd-quickstart.md - RTX 4080 developer workload: verl-opd-reference-workload.md + - RTX 4080 PG-k1 workload: verl-pg-k1-workload.md - Start on one GPU: single-gpu-guide.md - Align: - "External Alignment Gate v1: early stop": alignment-external/alignment-external-v1.md diff --git a/src/miniverl/bridge/doctor.py b/src/miniverl/bridge/doctor.py index 48446d6..32ed0fb 100644 --- a/src/miniverl/bridge/doctor.py +++ b/src/miniverl/bridge/doctor.py @@ -6,9 +6,11 @@ import importlib.metadata import json import re +import subprocess from collections.abc import Mapping from pathlib import Path from typing import Any +from urllib.parse import unquote, urlparse import yaml @@ -336,7 +338,7 @@ def _finalize_tokenizer(check: dict[str, Any], *, require_load: bool) -> dict[st return check -def _check_parquet(root: Path) -> dict[str, Any]: +def _check_parquet(root: Path, *, require_reward_model: bool = True) -> dict[str, Any]: """Validate the exchange schema from Parquet metadata alone. Column names, row counts and row-group counts all live in the footer, so a @@ -346,7 +348,9 @@ def _check_parquet(root: Path) -> dict[str, Any]: import pyarrow.parquet as pq except ImportError: return {"status": "fail", "detail": "pyarrow is not installed"} - required = {"data_source", "prompt", "ability", "reward_model", "extra_info"} + required = {"data_source", "prompt", "ability", "extra_info"} + if require_reward_model: + required.add("reward_model") schemas: dict[str, list[str]] = {} rows: dict[str, int] = {} files = sorted((root / "data").glob("*.parquet")) @@ -1003,6 +1007,21 @@ def _installed_verl() -> dict[str, Any]: except json.JSONDecodeError: direct_url = {"invalid": True} commit = ((direct_url or {}).get("vcs_info") or {}).get("commit_id") + if commit is None and isinstance(direct_url, dict): + raw_url = direct_url.get("url") + if isinstance(raw_url, str) and urlparse(raw_url).scheme == "file": + checkout = Path(unquote(urlparse(raw_url).path.lstrip("/"))) + try: + commit = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=checkout, + check=True, + capture_output=True, + text=True, + timeout=10, + ).stdout.strip() + except (OSError, subprocess.SubprocessError): + commit = None return { "status": "ok" if commit == VERL_COMMIT else "unverified", "version": version, @@ -1225,8 +1244,12 @@ def inspect_bridge_bundle( student_snapshot = _check_materialized_snapshot(bundle, role="student") teacher_snapshot = _check_materialized_snapshot(bundle, role="teacher") tokenizer = _check_tokenizer(bundle, require_load=require_tokenizer_load) - parquet = _check_parquet(bundle) config = _check_config(bundle) + parquet = _check_parquet( + bundle, + require_reward_model=config.get("profile") + not in {VERL_OPD_V08_PROFILE, VERL_OPD_PG_K1_V08_PROFILE}, + ) reward = _check_reward(bundle, trust_and_import=trust_and_import_reward_code) hashes = _check_hashes(bundle) privacy = _check_privacy( diff --git a/src/miniverl/bridge/materialize.py b/src/miniverl/bridge/materialize.py index 1762b4e..0aca06b 100644 --- a/src/miniverl/bridge/materialize.py +++ b/src/miniverl/bridge/materialize.py @@ -13,9 +13,11 @@ import importlib.metadata import json import shutil +import subprocess import uuid from pathlib import Path from typing import Any +from urllib.parse import unquote, urlparse import yaml @@ -36,6 +38,12 @@ _IGNORED_TOP_LEVEL = {".cache", ".git"} +def _validate_vocab_domain(*, model_vocab_size: int, tokenizer_max_token_id: int) -> None: + """Accept padded model vocabularies while rejecting unreachable tokenizer IDs.""" + if model_vocab_size <= tokenizer_max_token_id: + raise ValueError("causal-LM logits do not cover the tokenizer ID domain") + + def _sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: @@ -433,14 +441,21 @@ def _load_local_transformers_metadata( encoded = tokenizer("materialization smoke", return_tensors="pt") with torch.no_grad(): output = model(**encoded, use_cache=False) - if output.logits.ndim != 3 or output.logits.shape[-1] != len(tokenizer): - raise ValueError("causal-LM logits do not match the tokenizer vocabulary") + model_vocab_size = int(output.logits.shape[-1]) if output.logits.ndim == 3 else 0 + tokenizer_max_id = max(int(value) for value in tokenizer.get_vocab().values()) + if output.logits.ndim != 3: + raise ValueError("causal-LM logits must have batch, sequence and vocabulary axes") + _validate_vocab_domain( + model_vocab_size=model_vocab_size, tokenizer_max_token_id=tokenizer_max_id + ) except Exception as exc: raise ConfigError(f"{role} local model/tokenizer smoke failed: {exc}") from exc metadata = { "config_class": type(config).__name__, "tokenizer_class": type(tokenizer).__name__, "vocab_size": len(tokenizer), + "model_vocab_size": model_vocab_size, + "tokenizer_max_token_id": tokenizer_max_id, "tokenizer_structural_digest_v2": tokenizer_digest, "parameters": sum(parameter.numel() for parameter in model.parameters()), "peft_adapter_loaded": adapter is not None, @@ -467,6 +482,21 @@ def _validate_upstream_bundle(root: Path) -> dict[str, Any]: except json.JSONDecodeError as exc: raise ConfigError("installed verl has invalid direct_url.json") from exc actual_commit = ((direct or {}).get("vcs_info") or {}).get("commit_id") + if actual_commit is None: + raw_url = (direct or {}).get("url") + if isinstance(raw_url, str) and urlparse(raw_url).scheme == "file": + checkout = Path(unquote(urlparse(raw_url).path.lstrip("/"))) + try: + actual_commit = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=checkout, + check=True, + capture_output=True, + text=True, + timeout=10, + ).stdout.strip() + except (OSError, subprocess.SubprocessError): + actual_commit = None if actual_commit != VERL_COMMIT: raise ConfigError( f"installed verl commit is {actual_commit or 'unverified'}, expected {VERL_COMMIT}" @@ -493,7 +523,7 @@ def _validate_upstream_bundle(root: Path) -> dict[str, Any]: from miniverl.bridge.doctor import _check_model, _check_parquet adapter = _check_model(root, require_payload=True) - parquet = _check_parquet(root) + parquet = _check_parquet(root, require_reward_model=False) if adapter["status"] != "ok": raise ConfigError(f"student PEFT validation failed: {adapter.get('detail')}") if parquet["status"] != "ok": @@ -538,8 +568,13 @@ def _validate_upstream_bundle(root: Path) -> dict[str, Any]: "resolved_config": "recipe/verl-opd-resolved.yaml", "scope": ( "exact verl config merge, Parquet footer, PEFT payload, local model config and " - "tokenizer loads, sequential CPU model loads and tiny forwards, and top-k bounds; " - "no distributed job was run" + "tokenizer loads, sequential CPU model loads and tiny forwards, and " + + ( + "the sampled-k1 policy-gradient contract; " + if profile == VERL_OPD_PG_K1_V08_PROFILE + else "top-k bounds; " + ) + + "no distributed job was run" ), "checks": { "config_parse": "passed", diff --git a/src/miniverl/bridge/profiles.py b/src/miniverl/bridge/profiles.py index efc4555..232faf7 100644 --- a/src/miniverl/bridge/profiles.py +++ b/src/miniverl/bridge/profiles.py @@ -283,7 +283,7 @@ def summary(self) -> ProfileSummary: name=VERL_OPD_PG_K1_V08_PROFILE, objective="k1 + vanilla policy loss", teacher_target="sampled-token teacher log-probability", - status="conformance_only", + status="measured", identity=self.identity, ) diff --git a/tests/unit/test_verl_bridge_parquet_streaming.py b/tests/unit/test_verl_bridge_parquet_streaming.py index 62ef968..700e56f 100644 --- a/tests/unit/test_verl_bridge_parquet_streaming.py +++ b/tests/unit/test_verl_bridge_parquet_streaming.py @@ -99,6 +99,17 @@ def _forbidden(*args: Any, **kwargs: Any) -> Any: assert "footer metadata only" in check["read_scope"] +def test_reward_free_opd_schema_does_not_require_reward_model(tmp_path: Path) -> None: + data = tmp_path / "data" + data.mkdir() + row = _row(0) + row.pop("reward_model") + pq.write_table(pa.Table.from_pylist([row]), data / "train.parquet") + + assert _check_parquet(tmp_path)["status"] == "fail" + assert _check_parquet(tmp_path, require_reward_model=False)["status"] == "ok" + + # --------------------------------------------------------- streaming bounds diff --git a/tests/unit/test_verl_opd_materialize.py b/tests/unit/test_verl_opd_materialize.py index 3d6ae28..d7709f6 100644 --- a/tests/unit/test_verl_opd_materialize.py +++ b/tests/unit/test_verl_opd_materialize.py @@ -128,6 +128,14 @@ def _passed_upstream(root: Path) -> dict[str, Any]: } +def test_materialization_accepts_a_padded_model_vocabulary() -> None: + from miniverl.bridge.materialize import _validate_vocab_domain + + _validate_vocab_domain(model_vocab_size=151_936, tokenizer_max_token_id=151_668) + with pytest.raises(ValueError, match="tokenizer ID domain"): + _validate_vocab_domain(model_vocab_size=151_668, tokenizer_max_token_id=151_668) + + def test_materialize_publishes_complete_launchable_bundle_transactionally( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: From 9c0ecde776333790571f383cfce5a80260bdd00c Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Fri, 14 Aug 2026 00:06:18 -0700 Subject: [PATCH 5/7] Keep direct profile schema immutable --- PYPI.md | 19 ++++++++++------ scripts/__init__.py | 1 + src/miniverl/bridge/opd_pg_v08.py | 35 +++++++++++++++++++++++++++-- src/miniverl/bridge/opd_v08.py | 37 +++++++++---------------------- src/miniverl/bridge/profiles.py | 3 ++- 5 files changed, 59 insertions(+), 36 deletions(-) create mode 100644 scripts/__init__.py diff --git a/PYPI.md b/PYPI.md index 8dbe0bd..5260151 100644 --- a/PYPI.md +++ b/PYPI.md @@ -130,19 +130,24 @@ explain/check` to inspect the closed built-in registry and distinguish an accepted field from one with a demonstrated native effect. New plans, caches, checkpoints and exports bind the complete independent profile identity. -`verl-opd-v0.8-single-gpu-v1` pins official verl `v0.8.0` at commit -`7aed6b230776f963fa09509c10d9c3a767d1102c`. Its executable path is intentionally -narrow: +Both profiles pin official verl `v0.8.0` at commit +`7aed6b230776f963fa09509c10d9c3a767d1102c`: + +| Profile | Objective | Teacher target | Status | +| --- | --- | --- | --- | +| direct GKD | `forward_kl_topk` | top-k IDs/log-probabilities | measured | +| PG OPD | sampled `k1` + vanilla policy loss | sampled-token teacher log-probability | measured | + +Their shared executable boundary is intentionally narrow: - one trainable actor and one teacher; - one generated response per prompt (`n=1`); -- reward-free generalized knowledge distillation; -- `forward_kl_topk` with teacher top-k IDs/log-probabilities, mass/overlap - diagnostics and token-mean aggregation; +- reward-free distillation with token-mean aggregation; - LoRA or QLoRA adapter updates on one CUDA GPU; - immutable model revisions and verl-style structured prompt Parquet. -Policy-gradient OPD, task rewards, KL penalties, multi-teacher routing, +The PG path is not PPO: it uses a detached distillation-derived advantage with +the pinned vanilla policy-loss form. Task rewards, KL penalties, multi-teacher routing, multimodal inputs, PPO, GRPO, critics, Ray, FSDP, Megatron, multi-GPU and multi-node execution are unsupported. Known unsupported values receive a machine-readable classification; unknown fields and unresolved `${...}` diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100644 index 0000000..604f8e3 --- /dev/null +++ b/scripts/__init__.py @@ -0,0 +1 @@ +"""Repository maintenance scripts; never installed as the miniVERL runtime.""" diff --git a/src/miniverl/bridge/opd_pg_v08.py b/src/miniverl/bridge/opd_pg_v08.py index e38f60c..5d85d37 100644 --- a/src/miniverl/bridge/opd_pg_v08.py +++ b/src/miniverl/bridge/opd_pg_v08.py @@ -8,7 +8,7 @@ from typing import Any, Literal import yaml -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, Field, ValidationError, model_validator from miniverl.bridge.contract import VERL_COMMIT, VERL_REPOSITORY, VERL_TAG from miniverl.bridge.interpolation import reject_interpolation @@ -17,8 +17,11 @@ CompatibilityEntry, CompiledLocalExecutionPlan, FieldClassification, + VerlOPDDistillationConfig, + VerlOPDLossConfig, VerlOPDV08Profile, _canonical_digest, + _finite, _flatten, _reject_non_finite_source_numbers, _resolve_overrides, @@ -38,6 +41,34 @@ VERL_OPD_PG_K1_V08_PROFILE = "verl-opd-v0.8-single-gpu-pg-k1-v1" + +class VerlPGK1LossConfig(VerlOPDLossConfig): + """PG-only loss fields, isolated from the published direct profile schema.""" + + topk: int | None = Field(default=None, ge=1) # type: ignore[assignment] + policy_loss_mode: str = "vanilla" + clip_ratio: float = Field(default=0.2, ge=0, lt=1) + clip_ratio_low: float | None = Field(default=0.2, ge=0, lt=1) + clip_ratio_high: float | None = Field(default=0.2, ge=0, lt=1) + + @model_validator(mode="after") + def _pg_numbers_are_finite(self) -> VerlPGK1LossConfig: + _finite(self.clip_ratio, "clip_ratio") + if self.clip_ratio_low is not None: + _finite(self.clip_ratio_low, "clip_ratio_low") + if self.clip_ratio_high is not None: + _finite(self.clip_ratio_high, "clip_ratio_high") + return self + + +class VerlPGK1DistillationConfig(VerlOPDDistillationConfig): + distillation_loss: VerlPGK1LossConfig + + +class VerlPGK1V08Profile(VerlOPDV08Profile): + distillation: VerlPGK1DistillationConfig + + _PG_FIELD_RULES: dict[str, _Rule] = dict(_FIELD_RULES) _PG_FIELD_RULES.update( { @@ -227,7 +258,7 @@ def compile_verl_pg_k1_v08( hint="the profile accepts only its documented resolved verl v0.8 subset", ) try: - source = VerlOPDV08Profile.model_validate(merged) + source = VerlPGK1V08Profile.model_validate(merged) except ValidationError as exc: finite = any("finite" in error.get("msg", "").lower() for error in exc.errors()) message = "verl PG-k1 numeric fields must be finite" if finite else "invalid PG-k1 profile" diff --git a/src/miniverl/bridge/opd_v08.py b/src/miniverl/bridge/opd_v08.py index 01bdc9e..4cd6f54 100644 --- a/src/miniverl/bridge/opd_v08.py +++ b/src/miniverl/bridge/opd_v08.py @@ -16,7 +16,15 @@ from typing import Any, Literal import yaml -from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator, model_validator +from pydantic import ( + BaseModel, + ConfigDict, + Field, + SerializeAsAny, + ValidationError, + field_validator, + model_validator, +) from miniverl.bridge.contract import VERL_COMMIT, VERL_REPOSITORY, VERL_TAG from miniverl.bridge.interpolation import reject_interpolation @@ -188,16 +196,12 @@ class VerlOPDTeacherModelsConfig(_StrictModel): class VerlOPDLossConfig(_StrictModel): loss_mode: str - topk: int | None = Field(default=None, ge=1) + topk: int = Field(ge=1) use_task_rewards: bool = False distillation_loss_coef: float = Field(default=1.0, ge=0) loss_max_clamp: float | None = Field(default=None, ge=0) log_prob_min_clamp: float | None = None use_policy_gradient: bool = False - policy_loss_mode: str = "vanilla" - clip_ratio: float = Field(default=0.2, ge=0, lt=1) - clip_ratio_low: float | None = Field(default=0.2, ge=0, lt=1) - clip_ratio_high: float | None = Field(default=0.2, ge=0, lt=1) @model_validator(mode="after") def _numbers_are_finite(self) -> VerlOPDLossConfig: @@ -206,11 +210,6 @@ def _numbers_are_finite(self) -> VerlOPDLossConfig: _finite(self.loss_max_clamp, "loss_max_clamp") if self.log_prob_min_clamp is not None: _finite(self.log_prob_min_clamp, "log_prob_min_clamp") - _finite(self.clip_ratio, "clip_ratio") - if self.clip_ratio_low is not None: - _finite(self.clip_ratio_low, "clip_ratio_low") - if self.clip_ratio_high is not None: - _finite(self.clip_ratio_high, "clip_ratio_high") return self @@ -394,7 +393,7 @@ class CompiledLocalExecutionPlan(_StrictModel): source_digest: str compiled_digest: str source_leaf_fields: list[str] - source: VerlOPDV08Profile + source: SerializeAsAny[VerlOPDV08Profile] overrides: list[OverrideRecord] compatibility: list[CompatibilityEntry] reinterpretation_acceptance: dict[str, Any] @@ -633,18 +632,6 @@ def _rule( "distillation.distillation_loss.use_policy_gradient": _rule( "loss.use_policy_gradient", "exact", "must remain disabled for direct GKD OPD" ), - "distillation.distillation_loss.policy_loss_mode": _rule( - None, "informational_only", "inactive because direct GKD does not use a policy loss" - ), - "distillation.distillation_loss.clip_ratio": _rule( - None, "informational_only", "inactive because direct GKD does not use a policy loss" - ), - "distillation.distillation_loss.clip_ratio_low": _rule( - None, "informational_only", "inactive because direct GKD does not use a policy loss" - ), - "distillation.distillation_loss.clip_ratio_high": _rule( - None, "informational_only", "inactive because direct GKD does not use a policy loss" - ), "trainer.project_name": _rule("run.project", "exact", "same provenance label"), "trainer.experiment_name": _rule("run.name", "exact", "same provenance label"), "trainer.save_freq": _rule( @@ -1055,8 +1042,6 @@ def _semantic_blocker(path: str, value: Any) -> str | None: return expected[1] if path == "trainer.n_gpus_per_node" and value != 1: return "more than one local training GPU is unsupported" - if path == "distillation.distillation_loss.topk" and value is None: - return "direct forward_kl_topk requires a positive topk" if path == "distillation.nnodes" and value not in {0, 1}: return "multi-node teacher execution is unsupported" if path == "algorithm.adv_estimator" and value is not None: diff --git a/src/miniverl/bridge/profiles.py b/src/miniverl/bridge/profiles.py index 232faf7..d0c3d07 100644 --- a/src/miniverl/bridge/profiles.py +++ b/src/miniverl/bridge/profiles.py @@ -18,6 +18,7 @@ from miniverl.bridge.contract import VERL_COMMIT, VERL_REPOSITORY, VERL_TAG from miniverl.bridge.opd_pg_v08 import ( VERL_OPD_PG_K1_V08_PROFILE, + VerlPGK1V08Profile, load_verl_pg_k1_v08_source, pg_compatibility_rule, pg_field_rules_digest, @@ -288,7 +289,7 @@ def summary(self) -> ProfileSummary: ) def config_schema(self) -> dict[str, Any]: - return VerlOPDV08Profile.model_json_schema() + return VerlPGK1V08Profile.model_json_schema() def explain(self, field: str) -> CompatibilityExplanation: rule = pg_compatibility_rule(field) From 6119ed5336b884453677ddcd65b48c204a8d4851 Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Fri, 14 Aug 2026 00:16:31 -0700 Subject: [PATCH 6/7] Mark PG loss tests as torch-optional --- tests/unit/test_verl_pg_k1.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_verl_pg_k1.py b/tests/unit/test_verl_pg_k1.py index fe214c0..89e0e12 100644 --- a/tests/unit/test_verl_pg_k1.py +++ b/tests/unit/test_verl_pg_k1.py @@ -1,9 +1,12 @@ from __future__ import annotations import pytest -import torch -from miniverl.losses.verl_pg import ( +torch = pytest.importorskip("torch") + +pytestmark = pytest.mark.torch + +from miniverl.losses.verl_pg import ( # noqa: E402 VERL_PG_K1_IMPLEMENTATION_VERSION, verl_k1_estimator, verl_pg_k1_loss, From bd03850a63cb1365dc4e5605453da017ce39af0d Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Fri, 14 Aug 2026 00:19:18 -0700 Subject: [PATCH 7/7] Keep PG config validation torch-free --- src/miniverl/bridge/opd_pg_contract.py | 7 +++++++ src/miniverl/config/models.py | 2 +- src/miniverl/losses/verl_pg.py | 4 ++-- src/miniverl/teachers/local.py | 2 +- src/miniverl/training/trainer.py | 2 +- 5 files changed, 12 insertions(+), 5 deletions(-) create mode 100644 src/miniverl/bridge/opd_pg_contract.py diff --git a/src/miniverl/bridge/opd_pg_contract.py b/src/miniverl/bridge/opd_pg_contract.py new file mode 100644 index 0000000..280688c --- /dev/null +++ b/src/miniverl/bridge/opd_pg_contract.py @@ -0,0 +1,7 @@ +"""Torch-free constants for the pinned sampled-k1 OPD contract.""" + +from __future__ import annotations + +VERL_PG_K1_IMPLEMENTATION_VERSION = "verl-v0.8-pg-k1-v1" + +__all__ = ["VERL_PG_K1_IMPLEMENTATION_VERSION"] diff --git a/src/miniverl/config/models.py b/src/miniverl/config/models.py index 867bb99..60b02cb 100644 --- a/src/miniverl/config/models.py +++ b/src/miniverl/config/models.py @@ -918,7 +918,7 @@ def _validate_combination(self) -> RunConfig: ) if self.loss.mode is LossMode.VERL_PG_K1: - from miniverl.losses.verl_pg import VERL_PG_K1_IMPLEMENTATION_VERSION + from miniverl.bridge.opd_pg_contract import VERL_PG_K1_IMPLEMENTATION_VERSION if mode is not TrainingMode.OPD: raise ValueError("loss.mode=verl_pg_k1 requires run.mode=opd") diff --git a/src/miniverl/losses/verl_pg.py b/src/miniverl/losses/verl_pg.py index 3f4b8b6..9d74c8b 100644 --- a/src/miniverl/losses/verl_pg.py +++ b/src/miniverl/losses/verl_pg.py @@ -13,6 +13,8 @@ import torch +from miniverl.bridge.opd_pg_contract import VERL_PG_K1_IMPLEMENTATION_VERSION + __all__ = [ "VERL_PG_K1_IMPLEMENTATION_VERSION", "VerlPGK1Output", @@ -20,8 +22,6 @@ "verl_pg_k1_loss", ] -VERL_PG_K1_IMPLEMENTATION_VERSION = "verl-v0.8-pg-k1-v1" - @dataclass(frozen=True) class VerlPGK1Output: diff --git a/src/miniverl/teachers/local.py b/src/miniverl/teachers/local.py index 01e93dc..178a0d2 100644 --- a/src/miniverl/teachers/local.py +++ b/src/miniverl/teachers/local.py @@ -197,7 +197,7 @@ def score( clip_ratio_c=self.loss.clip_ratio_c, loss_max_clamp=self.loss.loss_max_clamp, ) - from miniverl.losses.verl_pg import VERL_PG_K1_IMPLEMENTATION_VERSION + from miniverl.bridge.opd_pg_contract import VERL_PG_K1_IMPLEMENTATION_VERSION batch = TeacherTargetBatch( trajectory_id=student.trajectory_id, diff --git a/src/miniverl/training/trainer.py b/src/miniverl/training/trainer.py index 7036eda..f6dc555 100644 --- a/src/miniverl/training/trainer.py +++ b/src/miniverl/training/trainer.py @@ -1288,7 +1288,7 @@ def _collect( def _open_cache(self) -> TeacherCache: if self._cache is None: - from miniverl.losses.verl_pg import VERL_PG_K1_IMPLEMENTATION_VERSION + from miniverl.bridge.opd_pg_contract import VERL_PG_K1_IMPLEMENTATION_VERSION from miniverl.losses.verl_topk import VERL_TOPK_SCORE_IMPLEMENTATION assert self.teacher is not None