mirror of https://github.com/razor-ai/soup.git
915 lines
33 KiB
Python
915 lines
33 KiB
Python
"""v0.71.11 — GRPO / RL callbacks live wiring.
|
|
|
|
Closes #235 (reward-hack), #236 (ULD), #237 (MiniLLM), #238 (RL checkpoint),
|
|
#239 (iterative-DPO), #240 (echo-trap), #159 (variant fallback warning),
|
|
#160 (in-place GRPO EMA).
|
|
|
|
These tests lift the v0.70.0 deferred-stub family to live behaviour and
|
|
exercise the math + wiring on CPU / tiny fakes (no GPU). The Step-6 smoke
|
|
runs a real SmolLM2-135M GRPO loop separately.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
import soup_cli
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Shared helpers
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
class _FakeState:
|
|
def __init__(self, global_step: int = 1):
|
|
self.global_step = global_step
|
|
self.log_history: list[dict] = []
|
|
|
|
|
|
class _FakeControl:
|
|
def __init__(self):
|
|
self.should_training_stop = False
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Shared RL signal buffer
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
class TestRLSignalBuffer:
|
|
def test_record_and_snapshot(self):
|
|
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer
|
|
|
|
buf = RLSignalBuffer()
|
|
buf.record(
|
|
func_name="reward", completions=["a", "b", "c", "d"], rewards=[1, 2, 3, 4]
|
|
)
|
|
snap = buf.snapshot()
|
|
assert snap["completions"] == ["a", "b", "c", "d"]
|
|
assert snap["rewards"] == [1.0, 2.0, 3.0, 4.0]
|
|
assert "reward" in snap["per_func"]
|
|
|
|
def test_aggregate_sums_across_funcs(self):
|
|
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer
|
|
|
|
buf = RLSignalBuffer()
|
|
buf.record(func_name="a", completions=["x", "y"], rewards=[1, 2])
|
|
buf.record(func_name="b", completions=["x", "y"], rewards=[10, 20])
|
|
snap = buf.snapshot()
|
|
assert snap["rewards"] == [11.0, 22.0]
|
|
assert set(snap["per_func"]) == {"a", "b"}
|
|
|
|
def test_non_finite_reward_dropped(self):
|
|
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer
|
|
|
|
buf = RLSignalBuffer()
|
|
buf.record(func_name="r", completions=["a"], rewards=[float("nan")])
|
|
snap = buf.snapshot()
|
|
# NaN coerced to None → aggregate position is 0.0 (no finite value).
|
|
assert snap["per_func"]["r"] == [None]
|
|
|
|
def test_conversational_completion_extracted(self):
|
|
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer
|
|
|
|
buf = RLSignalBuffer()
|
|
buf.record(
|
|
func_name="r",
|
|
completions=[[{"role": "assistant", "content": "hello world"}]],
|
|
rewards=[1.0],
|
|
)
|
|
snap = buf.snapshot()
|
|
assert snap["completions"] == ["hello world"]
|
|
|
|
def test_wrap_preserves_name_and_captures(self):
|
|
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer, wrap_reward_funcs
|
|
|
|
buf = RLSignalBuffer()
|
|
|
|
def my_reward(prompts=None, completions=None, **kwargs):
|
|
return [float(len(c)) for c in completions]
|
|
|
|
wrapped = wrap_reward_funcs(my_reward, buf)
|
|
assert wrapped.__name__ == "my_reward"
|
|
result = wrapped(prompts=["p"], completions=["aa", "bbb"])
|
|
assert result == [2.0, 3.0]
|
|
snap = buf.snapshot()
|
|
assert snap["completions"] == ["aa", "bbb"]
|
|
assert snap["rewards"] == [2.0, 3.0]
|
|
|
|
def test_wrap_list_shape_preserved(self):
|
|
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer, wrap_reward_funcs
|
|
|
|
buf = RLSignalBuffer()
|
|
fns = [lambda completions=None, **k: [1.0]]
|
|
wrapped = wrap_reward_funcs(fns, buf)
|
|
assert isinstance(wrapped, list)
|
|
assert len(wrapped) == 1
|
|
|
|
def test_capture_never_breaks_reward(self):
|
|
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer, wrap_reward_funcs
|
|
|
|
buf = RLSignalBuffer()
|
|
|
|
def reward(prompts=None, completions=None, **kwargs):
|
|
return [1.0, 2.0]
|
|
|
|
wrapped = wrap_reward_funcs(reward, buf)
|
|
# Bizarre completions that the normaliser can't handle must not raise.
|
|
out = wrapped(prompts=None, completions=object())
|
|
assert out == [1.0, 2.0]
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# #235 — reward-hack callback
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
class TestRewardHackCallback:
|
|
def test_build_returns_callback_not_notimplemented(self):
|
|
from soup_cli.utils.reward_hacking import (
|
|
RewardHackCallback,
|
|
build_reward_hack_callback,
|
|
)
|
|
|
|
cb = build_reward_hack_callback(detector="info_rm")
|
|
assert isinstance(cb, RewardHackCallback)
|
|
|
|
def test_build_rejects_unknown_detector(self):
|
|
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
|
|
|
with pytest.raises(ValueError, match="not supported"):
|
|
build_reward_hack_callback(detector="evil")
|
|
|
|
def test_build_rejects_non_bool_halt(self):
|
|
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
|
|
|
with pytest.raises(TypeError, match="halt_on_hack"):
|
|
build_reward_hack_callback(detector="info_rm", halt_on_hack="yes")
|
|
|
|
def test_info_rm_compute_signal(self):
|
|
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
|
|
|
cb = build_reward_hack_callback(detector="info_rm")
|
|
snap = {"rewards": [0.0, 0.0, 5.0, 5.0], "per_func": {}}
|
|
sig = cb.compute_signal(snap)
|
|
assert sig is not None and sig > 0.0
|
|
|
|
def test_info_rm_insufficient_data_returns_none(self):
|
|
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
|
|
|
cb = build_reward_hack_callback(detector="info_rm")
|
|
assert cb.compute_signal({"rewards": [1.0, 2.0], "per_func": {}}) is None
|
|
|
|
def test_observe_baseline_then_drop_to_hack(self):
|
|
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
|
|
|
cb = build_reward_hack_callback(detector="info_rm")
|
|
r0 = cb.observe_signal(10.0, step=1) # baseline separation 10
|
|
assert r0.verdict == "OK"
|
|
r1 = cb.observe_signal(2.0, step=2) # dropped 80% → HACK
|
|
assert r1.verdict == "HACK"
|
|
assert r1.signal == 2.0
|
|
|
|
def test_on_step_end_halts_on_hack(self):
|
|
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
|
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer
|
|
|
|
buf = RLSignalBuffer()
|
|
cb = build_reward_hack_callback(
|
|
detector="info_rm", halt_on_hack=True, buffer=buf
|
|
)
|
|
state, control = _FakeState(1), _FakeControl()
|
|
# Step 1 — high separation = baseline.
|
|
buf.record(func_name="r", completions=["a"] * 4, rewards=[0, 0, 9, 9])
|
|
cb.on_step_end(None, state, control)
|
|
# Step 2 — bunched rewards = HACK.
|
|
buf.record(func_name="r", completions=["a"] * 4, rewards=[5, 5, 5, 5])
|
|
state.global_step = 2
|
|
cb.on_step_end(None, state, control)
|
|
assert control.should_training_stop is True
|
|
assert any("reward_hack_verdict" in e for e in state.log_history)
|
|
|
|
def test_rm_ensemble_needs_two_funcs(self):
|
|
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
|
|
|
cb = build_reward_hack_callback(detector="rm_ensemble")
|
|
# One func → None.
|
|
assert cb.compute_signal({"rewards": [], "per_func": {"a": [1.0, 2.0]}}) is None
|
|
# Two funcs → divergence.
|
|
sig = cb.compute_signal(
|
|
{"rewards": [], "per_func": {"a": [1.0, 2.0], "b": [3.0, 0.0]}}
|
|
)
|
|
assert sig is not None and sig >= 0.0
|
|
|
|
def test_on_log_fallback_without_buffer(self):
|
|
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
|
|
|
cb = build_reward_hack_callback(detector="info_rm", buffer=None)
|
|
state, control = _FakeState(1), _FakeControl()
|
|
cb.on_log(None, state, control, logs={"reward": 5.0, "reward_std": 0.1})
|
|
assert cb.last_report() is not None
|
|
|
|
def test_compute_separation_from_stats(self):
|
|
from soup_cli.utils.reward_hacking import compute_separation_from_stats
|
|
|
|
high = compute_separation_from_stats(5.0, 0.1)
|
|
low = compute_separation_from_stats(5.0, 10.0)
|
|
assert high > low
|
|
|
|
def test_separation_stats_rejects_bool(self):
|
|
from soup_cli.utils.reward_hacking import compute_separation_from_stats
|
|
|
|
with pytest.raises(ValueError, match="bool"):
|
|
compute_separation_from_stats(True, 1.0)
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# #240 — echo-trap callback
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
class _FakeTokenizer:
|
|
def encode(self, text, add_special_tokens=False):
|
|
# deterministic id per whitespace token
|
|
return [abs(hash(t)) % 1000 for t in text.split()]
|
|
|
|
|
|
class TestEchoTrapCallback:
|
|
def test_build_returns_callback_not_notimplemented(self):
|
|
from soup_cli.utils.echo_trap import EchoTrapCallback, build_echo_trap_callback
|
|
|
|
cb = build_echo_trap_callback(threshold=0.5)
|
|
assert isinstance(cb, EchoTrapCallback)
|
|
|
|
def test_build_rejects_bad_threshold(self):
|
|
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
|
|
|
with pytest.raises(ValueError):
|
|
build_echo_trap_callback(threshold=2.0)
|
|
|
|
def test_build_rejects_bool_threshold(self):
|
|
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
|
|
|
with pytest.raises(ValueError):
|
|
build_echo_trap_callback(threshold=True)
|
|
|
|
def test_build_rejects_non_bool_halt(self):
|
|
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
|
|
|
with pytest.raises(TypeError):
|
|
build_echo_trap_callback(threshold=0.5, halt_on_trap="yes")
|
|
|
|
def test_compute_signal_repetitive_high(self):
|
|
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
|
|
|
cb = build_echo_trap_callback(threshold=0.5)
|
|
# "a a a a a" — every 2-gram repeats.
|
|
snap = {"completions": ["a a a a a", "b b b b b"]}
|
|
sig = cb.compute_signal(snap)
|
|
assert sig is not None and sig > 0.5
|
|
|
|
def test_compute_signal_no_completions(self):
|
|
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
|
|
|
cb = build_echo_trap_callback(threshold=0.5)
|
|
assert cb.compute_signal({"completions": []}) is None
|
|
|
|
def test_tokenizer_aware_path(self):
|
|
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
|
|
|
cb = build_echo_trap_callback(
|
|
threshold=0.5, tokenizer_aware=True, tokenizer=_FakeTokenizer()
|
|
)
|
|
sig = cb.compute_signal({"completions": ["x x x x", "y y y y"]})
|
|
assert sig is not None and sig >= 0.0
|
|
|
|
def test_on_step_end_halts_on_trap(self):
|
|
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
|
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer
|
|
|
|
buf = RLSignalBuffer()
|
|
cb = build_echo_trap_callback(threshold=0.3, halt_on_trap=True, buffer=buf)
|
|
buf.record(
|
|
func_name="r",
|
|
completions=["a a a a a a", "b b b b b b"],
|
|
rewards=[1.0, 1.0],
|
|
)
|
|
state, control = _FakeState(1), _FakeControl()
|
|
cb.on_step_end(None, state, control)
|
|
assert control.should_training_stop is True
|
|
assert any("echo_trap_verdict" in e for e in state.log_history)
|
|
|
|
def test_observe_classifies_ok(self):
|
|
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
|
|
|
cb = build_echo_trap_callback(threshold=0.5)
|
|
report = cb.observe_signal(0.0, step=1, n_trajectories=3)
|
|
assert report.verdict == "OK"
|
|
assert report.trajectories_seen == 3
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# #238 — RL checkpoint callback
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
class _FakeSavableModel:
|
|
def __init__(self):
|
|
self.saved_to = None
|
|
|
|
def save_pretrained(self, path):
|
|
import os
|
|
|
|
os.makedirs(path, exist_ok=True)
|
|
with open(os.path.join(path, "adapter_model.safetensors"), "wb") as fh:
|
|
fh.write(b"\x00")
|
|
self.saved_to = path
|
|
|
|
|
|
class TestRLCheckpointCallback:
|
|
def test_build_requires_output_dir(self):
|
|
from soup_cli.utils.rl_checkpoint import (
|
|
RLCheckpointConfig,
|
|
build_rl_checkpoint_callback,
|
|
)
|
|
|
|
cfg = RLCheckpointConfig(save_every_steps=1)
|
|
with pytest.raises(ValueError, match="output_dir"):
|
|
build_rl_checkpoint_callback(cfg)
|
|
|
|
def test_build_rejects_non_config(self):
|
|
from soup_cli.utils.rl_checkpoint import build_rl_checkpoint_callback
|
|
|
|
with pytest.raises(TypeError):
|
|
build_rl_checkpoint_callback({"save_every_steps": 1}, output_dir="x")
|
|
|
|
def test_save_checkpoint_writes_manifest(self, tmp_path, monkeypatch):
|
|
import torch
|
|
|
|
from soup_cli.utils.rl_checkpoint import (
|
|
RLCheckpointConfig,
|
|
build_rl_checkpoint_callback,
|
|
)
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
cfg = RLCheckpointConfig(save_every_steps=1)
|
|
cb = build_rl_checkpoint_callback(cfg, output_dir="run", task="grpo")
|
|
model = _FakeSavableModel()
|
|
opt = torch.optim.SGD([torch.nn.Parameter(torch.zeros(2))], lr=0.1)
|
|
ckpt = cb.save_checkpoint(step=2, model=model, optimizer=opt)
|
|
manifest = Path(ckpt) / "manifest.json"
|
|
assert manifest.is_file()
|
|
data = json.loads(manifest.read_text())
|
|
assert data["step"] == 2 and data["task"] == "grpo"
|
|
assert data["has_optimizer"] is True
|
|
assert (Path(ckpt) / "optimizer.pt").is_file()
|
|
|
|
def test_prune_keeps_keep_last(self, tmp_path, monkeypatch):
|
|
from soup_cli.utils.rl_checkpoint import (
|
|
RLCheckpointConfig,
|
|
build_rl_checkpoint_callback,
|
|
)
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
cfg = RLCheckpointConfig(save_every_steps=1, keep_last=2)
|
|
cb = build_rl_checkpoint_callback(cfg, output_dir="run", task="grpo")
|
|
for step in (1, 2, 3):
|
|
cb.save_checkpoint(step=step, model=_FakeSavableModel(), optimizer=None)
|
|
root = tmp_path / "run" / "rl-checkpoints"
|
|
dirs = sorted(p.name for p in root.iterdir())
|
|
assert dirs == ["step-2", "step-3"] # step-1 pruned
|
|
|
|
def test_on_step_end_cadence(self, tmp_path, monkeypatch):
|
|
from soup_cli.utils.rl_checkpoint import (
|
|
RLCheckpointConfig,
|
|
build_rl_checkpoint_callback,
|
|
)
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
cfg = RLCheckpointConfig(save_every_steps=2)
|
|
cb = build_rl_checkpoint_callback(cfg, output_dir="run", task="grpo")
|
|
# step 1 → no save; step 2 → save.
|
|
cb.on_step_end(None, _FakeState(1), _FakeControl(), model=_FakeSavableModel())
|
|
assert not (tmp_path / "run" / "rl-checkpoints").exists()
|
|
cb.on_step_end(None, _FakeState(2), _FakeControl(), model=_FakeSavableModel())
|
|
assert (tmp_path / "run" / "rl-checkpoints" / "step-2").is_dir()
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# #236 — ULD
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
class TestULD:
|
|
def test_build_returns_projection(self):
|
|
from soup_cli.utils.uld import ULDConfig, ULDProjection, build_uld_projection
|
|
|
|
proj = build_uld_projection(
|
|
ULDConfig(strategy="wasserstein", student_vocab_size=10, teacher_vocab_size=12)
|
|
)
|
|
assert isinstance(proj, ULDProjection)
|
|
|
|
def test_build_rejects_non_config(self):
|
|
from soup_cli.utils.uld import build_uld_projection
|
|
|
|
with pytest.raises(TypeError):
|
|
build_uld_projection({"strategy": "wasserstein"})
|
|
|
|
def test_wasserstein_loss_different_vocab(self):
|
|
import torch
|
|
|
|
from soup_cli.utils.uld import ULDConfig, uld_distill_loss
|
|
|
|
cfg = ULDConfig(strategy="wasserstein", student_vocab_size=8, teacher_vocab_size=12)
|
|
s = torch.randn(2, 3, 8, requires_grad=True)
|
|
t = torch.randn(2, 3, 12)
|
|
loss = uld_distill_loss(s, t, config=cfg)
|
|
assert torch.isfinite(loss)
|
|
loss.backward()
|
|
assert s.grad is not None
|
|
|
|
def test_topk_align_loss(self):
|
|
import torch
|
|
|
|
from soup_cli.utils.uld import ULDConfig, uld_distill_loss
|
|
|
|
cfg = ULDConfig(
|
|
strategy="topk_align", student_vocab_size=8, teacher_vocab_size=12, top_k=4
|
|
)
|
|
s = torch.randn(2, 3, 8, requires_grad=True)
|
|
t = torch.randn(2, 3, 12)
|
|
loss = uld_distill_loss(s, t, config=cfg)
|
|
assert torch.isfinite(loss)
|
|
loss.backward()
|
|
|
|
def test_identical_distributions_low_loss(self):
|
|
import torch
|
|
|
|
from soup_cli.utils.uld import ULDConfig, uld_distill_loss
|
|
|
|
cfg = ULDConfig(strategy="wasserstein", student_vocab_size=8, teacher_vocab_size=8)
|
|
logits = torch.randn(2, 3, 8)
|
|
loss = uld_distill_loss(logits, logits.clone(), config=cfg)
|
|
assert float(loss) < 1e-5
|
|
|
|
def test_attention_mask_applied(self):
|
|
import torch
|
|
|
|
from soup_cli.utils.uld import ULDConfig, uld_distill_loss
|
|
|
|
cfg = ULDConfig(strategy="wasserstein", student_vocab_size=8, teacher_vocab_size=8)
|
|
s = torch.randn(2, 3, 8)
|
|
t = torch.randn(2, 3, 8)
|
|
mask = torch.tensor([[1, 1, 0], [1, 0, 0]])
|
|
loss = uld_distill_loss(s, t, config=cfg, attention_mask=mask)
|
|
assert torch.isfinite(loss)
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# #237 — MiniLLM
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
class TestMiniLLM:
|
|
def test_build_returns_callback(self):
|
|
from soup_cli.utils.minillm import (
|
|
MiniLLMCallback,
|
|
MiniLLMConfig,
|
|
build_minillm_callback,
|
|
)
|
|
|
|
cb = build_minillm_callback(MiniLLMConfig(teacher_mix_ratio=0.5))
|
|
assert isinstance(cb, MiniLLMCallback)
|
|
|
|
def test_build_rejects_non_config(self):
|
|
from soup_cli.utils.minillm import build_minillm_callback
|
|
|
|
with pytest.raises(TypeError):
|
|
build_minillm_callback({})
|
|
|
|
def test_distill_term_finite_and_differentiable(self):
|
|
import torch
|
|
|
|
from soup_cli.utils.minillm import MiniLLMConfig, minillm_distill_term
|
|
|
|
cfg = MiniLLMConfig(teacher_mix_ratio=0.5, length_normalize=True)
|
|
s = torch.randn(2, 4, 16, requires_grad=True)
|
|
t = torch.randn(2, 4, 16)
|
|
labels = torch.randint(0, 16, (2, 4))
|
|
labels[0, 0] = -100 # masked
|
|
loss = minillm_distill_term(s, t, labels, config=cfg)
|
|
assert torch.isfinite(loss)
|
|
loss.backward()
|
|
assert s.grad is not None
|
|
|
|
def test_teacher_mix_ratio_zero_gives_near_zero(self):
|
|
import torch
|
|
|
|
from soup_cli.utils.minillm import MiniLLMConfig, minillm_distill_term
|
|
|
|
cfg = MiniLLMConfig(teacher_mix_ratio=0.0)
|
|
s = torch.randn(2, 4, 16)
|
|
t = torch.randn(2, 4, 16)
|
|
labels = torch.randint(0, 16, (2, 4))
|
|
loss = minillm_distill_term(s, t, labels, config=cfg)
|
|
# ratio=0 → target = student_detached → reverse-KL ≈ 0.
|
|
assert abs(float(loss)) < 1e-4
|
|
|
|
def test_anchor_term_with_file(self, tmp_path, monkeypatch):
|
|
import torch
|
|
|
|
from soup_cli.utils.minillm import MiniLLMConfig, build_minillm_callback
|
|
|
|
# Loading a tokenizer/model from the Hub flakes on CI runners that get
|
|
# HF-rate-limited (the cache-warm step is best-effort). Skip on network
|
|
# failure — the anchor math is covered network-free by
|
|
# ``test_anchor_term_with_fake_model`` below.
|
|
try:
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
tok = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
"hf-internal-testing/tiny-random-gpt2"
|
|
)
|
|
except OSError as exc: # pragma: no cover — network-dependent
|
|
pytest.skip(f"HF model unavailable (offline / rate-limited): {exc}")
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
anchor = tmp_path / "anchor.jsonl"
|
|
anchor.write_text(
|
|
"\n".join(json.dumps({"text": f"sentence number {i}"}) for i in range(4))
|
|
)
|
|
if tok.pad_token is None:
|
|
tok.pad_token = tok.eos_token
|
|
cb = build_minillm_callback(
|
|
MiniLLMConfig(pretrain_anchor_weight=0.1, pretrain_anchor_path="anchor.jsonl"),
|
|
tokenizer=tok,
|
|
)
|
|
term = cb.anchor_term(model)
|
|
assert term is not None
|
|
assert torch.isfinite(term)
|
|
|
|
def test_anchor_term_with_fake_model(self, tmp_path, monkeypatch):
|
|
"""Network-free coverage of ``_load_anchor`` + ``anchor_term`` — a fake
|
|
tokenizer + tiny ``nn.Module`` exercise the same lines as the Hub-backed
|
|
test above, so the coverage gate does not depend on HF availability."""
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from soup_cli.utils.minillm import MiniLLMConfig, build_minillm_callback
|
|
|
|
class _FakeOut:
|
|
def __init__(self, logits):
|
|
self.logits = logits
|
|
|
|
class _TinyLM(nn.Module):
|
|
def __init__(self, vocab=16):
|
|
super().__init__()
|
|
self.emb = nn.Embedding(vocab, 4)
|
|
self.head = nn.Linear(4, vocab)
|
|
|
|
def forward(self, input_ids, attention_mask=None): # noqa: ARG002
|
|
return _FakeOut(self.head(self.emb(input_ids)))
|
|
|
|
class _FakeTok:
|
|
def __call__(
|
|
self, texts, return_tensors=None, padding=None,
|
|
truncation=None, max_length=None,
|
|
): # noqa: ARG002
|
|
ids = torch.tensor([[1, 2, 3, 4] for _ in texts])
|
|
return {"input_ids": ids, "attention_mask": torch.ones_like(ids)}
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
anchor = tmp_path / "anchor.jsonl"
|
|
anchor.write_text(
|
|
"\n".join(json.dumps({"text": f"sentence {i}"}) for i in range(4))
|
|
)
|
|
cb = build_minillm_callback(
|
|
MiniLLMConfig(pretrain_anchor_weight=0.25, pretrain_anchor_path="anchor.jsonl"),
|
|
tokenizer=_FakeTok(),
|
|
)
|
|
term = cb.anchor_term(_TinyLM())
|
|
assert term is not None
|
|
assert torch.isfinite(term)
|
|
|
|
def test_anchor_term_disabled_returns_none(self):
|
|
import torch.nn as nn
|
|
|
|
from soup_cli.utils.minillm import MiniLLMConfig, build_minillm_callback
|
|
|
|
cb = build_minillm_callback(MiniLLMConfig(teacher_mix_ratio=0.3))
|
|
assert cb.anchor_term(nn.Linear(2, 2)) is None
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# #239 — iterative DPO
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
class TestIterativeDPO:
|
|
def test_build_pairs_from_scored(self):
|
|
from soup_cli.utils.iterative_dpo import build_pairs_from_scored
|
|
|
|
assert build_pairs_from_scored([("a", 1.0), ("b", 3.0)]) == ("b", "a")
|
|
assert build_pairs_from_scored([("a", 1.0)]) is None
|
|
assert build_pairs_from_scored([("a", 2.0), ("b", 2.0)]) is None
|
|
|
|
def test_run_iterative_dpo_with_fakes(self, tmp_path, monkeypatch):
|
|
from soup_cli.utils.iterative_dpo import (
|
|
IterativeDPOResult,
|
|
build_iterative_dpo_plan,
|
|
run_iterative_dpo,
|
|
)
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
prompts = tmp_path / "prompts.jsonl"
|
|
prompts.write_text(
|
|
"\n".join(json.dumps({"prompt": f"q{i}"}) for i in range(3))
|
|
)
|
|
|
|
plan = build_iterative_dpo_plan(
|
|
base_model="tiny",
|
|
reward_model="rm",
|
|
prompts_path="prompts.jsonl",
|
|
output_dir="out",
|
|
rounds=2,
|
|
pairs_per_round=10,
|
|
)
|
|
|
|
calls = {"sample_adapters": [], "score": 0, "train": []}
|
|
|
|
def fake_sample(*, base_model, adapter_path, prompts, num_samples,
|
|
max_new_tokens, device):
|
|
calls["sample_adapters"].append(adapter_path)
|
|
return [[f"{p}-a", f"{p}-b"] for p in prompts]
|
|
|
|
def fake_score(*, reward_model, prompt, completions, device):
|
|
calls["score"] += 1
|
|
return [float(len(c)) for c in completions]
|
|
|
|
def fake_train(*, base_model, pairs_path, adapter_path):
|
|
calls["train"].append((base_model, adapter_path))
|
|
import os
|
|
|
|
os.makedirs(adapter_path, exist_ok=True)
|
|
|
|
result = run_iterative_dpo(
|
|
plan, sample_fn=fake_sample, score_fn=fake_score, train_fn=fake_train
|
|
)
|
|
assert isinstance(result, IterativeDPOResult)
|
|
assert result.rounds_completed == 2
|
|
# Training ALWAYS starts from the plan's base (never an adapter dir).
|
|
assert calls["train"][0][0] == "tiny"
|
|
assert calls["train"][1][0] == "tiny"
|
|
# Round 0 samples from base (None adapter); round 1 from round-0 adapter.
|
|
assert calls["sample_adapters"][0] is None
|
|
assert calls["sample_adapters"][1].endswith("round-00/adapter")
|
|
# pairs written
|
|
assert (tmp_path / "out" / "round-00" / "pairs.jsonl").is_file()
|
|
|
|
def test_run_rejects_non_plan(self):
|
|
from soup_cli.utils.iterative_dpo import run_iterative_dpo
|
|
|
|
with pytest.raises(TypeError):
|
|
run_iterative_dpo({"rounds": 1})
|
|
|
|
def test_cli_plan_only_still_exits_zero(self, tmp_path, monkeypatch):
|
|
from typer.testing import CliRunner
|
|
|
|
from soup_cli.commands.iterative_dpo import app
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
(tmp_path / "p.jsonl").write_text(json.dumps({"prompt": "q"}))
|
|
runner = CliRunner()
|
|
res = runner.invoke(
|
|
app,
|
|
[
|
|
"--base-model", "b", "--reward-model", "rm",
|
|
"--prompts", "p.jsonl", "--output-dir", "o",
|
|
"--rounds", "1", "--pairs-per-round", "10", "--plan-only",
|
|
],
|
|
)
|
|
assert res.exit_code == 0, res.output
|
|
|
|
def test_cli_runs_with_monkeypatched_runner(self, tmp_path, monkeypatch):
|
|
from typer.testing import CliRunner
|
|
|
|
import soup_cli.utils.iterative_dpo as idpo
|
|
from soup_cli.commands.iterative_dpo import app
|
|
from soup_cli.utils.iterative_dpo import IterativeDPOResult
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
(tmp_path / "p.jsonl").write_text(json.dumps({"prompt": "q"}))
|
|
|
|
def fake_run(plan, **kwargs):
|
|
return IterativeDPOResult(
|
|
rounds_completed=1, final_adapter="o/round-00/adapter",
|
|
per_round_pairs=(1,),
|
|
)
|
|
|
|
monkeypatch.setattr(idpo, "run_iterative_dpo", fake_run)
|
|
runner = CliRunner()
|
|
res = runner.invoke(
|
|
app,
|
|
[
|
|
"--base-model", "b", "--reward-model", "rm",
|
|
"--prompts", "p.jsonl", "--output-dir", "o",
|
|
"--rounds", "1", "--pairs-per-round", "10",
|
|
],
|
|
)
|
|
assert res.exit_code == 0, res.output
|
|
assert "Done" in res.output
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# #159 — GRPO variant fallback warning
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
class _FakeGRPOBase:
|
|
"""Minimal stand-in for trl.GRPOTrainer for the variant subclass."""
|
|
|
|
class _Args:
|
|
beta = 0.1
|
|
|
|
def __init__(self):
|
|
self.args = self._Args()
|
|
self.super_called = 0
|
|
|
|
def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
|
|
self.super_called += 1
|
|
return 0.0
|
|
|
|
|
|
class TestGrpoVariantFallbackWarning:
|
|
def test_fallback_warns_once(self, caplog):
|
|
from soup_cli.trainer.grpo import make_grpo_trainer_variant
|
|
|
|
cls = make_grpo_trainer_variant(_FakeGRPOBase, "gspo")
|
|
inst = cls()
|
|
# inputs missing per-token logps → fallback path.
|
|
with caplog.at_level(logging.WARNING, logger="soup_cli.trainer.grpo"):
|
|
inst.compute_loss(None, {})
|
|
inst.compute_loss(None, {})
|
|
warnings = [r for r in caplog.records if "fell back" in r.message]
|
|
assert len(warnings) == 1 # one-shot
|
|
assert inst.super_called == 2 # but fallback happened both times
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# #160 — in-place GRPO EMA
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
class TestGrpoEmaInPlace:
|
|
def test_in_place_blends_toward_policy(self):
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from soup_cli.monitoring.grpo_stability_callback import update_ema_in_place
|
|
|
|
ref = nn.Linear(3, 3)
|
|
pol = nn.Linear(3, 3)
|
|
with torch.no_grad():
|
|
ref.weight.fill_(0.0)
|
|
pol.weight.fill_(1.0)
|
|
update_ema_in_place(ref, pol, 0.25)
|
|
# ref = 0.75*0 + 0.25*1 = 0.25
|
|
assert torch.allclose(ref.weight, torch.full_like(ref.weight, 0.25))
|
|
|
|
def test_returns_updated_count(self):
|
|
import torch.nn as nn
|
|
|
|
from soup_cli.monitoring.grpo_stability_callback import update_ema_in_place
|
|
|
|
# nn.Linear(2, 2) has weight + bias → 2 shared params updated.
|
|
n = update_ema_in_place(nn.Linear(2, 2), nn.Linear(2, 2), 0.5)
|
|
assert n == 2
|
|
|
|
def test_zero_overlap_returns_zero(self):
|
|
import torch.nn as nn
|
|
|
|
from soup_cli.monitoring.grpo_stability_callback import update_ema_in_place
|
|
|
|
# Disjoint parameter names → no overlap → count 0.
|
|
class _A(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.alpha = nn.Linear(2, 2)
|
|
|
|
class _B(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.beta = nn.Linear(2, 2)
|
|
|
|
assert update_ema_in_place(_A(), _B(), 0.5) == 0
|
|
|
|
def test_rejects_bool_alpha(self):
|
|
import torch.nn as nn
|
|
|
|
from soup_cli.monitoring.grpo_stability_callback import update_ema_in_place
|
|
|
|
with pytest.raises(TypeError):
|
|
update_ema_in_place(nn.Linear(2, 2), nn.Linear(2, 2), True)
|
|
|
|
def test_rejects_out_of_range_alpha(self):
|
|
import torch.nn as nn
|
|
|
|
from soup_cli.monitoring.grpo_stability_callback import update_ema_in_place
|
|
|
|
with pytest.raises(ValueError):
|
|
update_ema_in_place(nn.Linear(2, 2), nn.Linear(2, 2), 1.5)
|
|
|
|
def test_shape_mismatch_skipped(self):
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from soup_cli.monitoring.grpo_stability_callback import update_ema_in_place
|
|
|
|
ref = nn.Linear(3, 3)
|
|
pol = nn.Linear(2, 2) # different shapes for the same param name
|
|
with torch.no_grad():
|
|
ref.weight.fill_(7.0)
|
|
update_ema_in_place(ref, pol, 0.5)
|
|
# shape mismatch → ref untouched.
|
|
assert torch.allclose(ref.weight, torch.full_like(ref.weight, 7.0))
|
|
|
|
def test_callback_warns_once_on_zero_overlap(self, caplog):
|
|
import logging
|
|
|
|
import torch.nn as nn
|
|
|
|
from soup_cli.monitoring.grpo_stability_callback import (
|
|
GRPOStabilityCallback,
|
|
)
|
|
|
|
class _A(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.alpha = nn.Linear(2, 2)
|
|
|
|
class _B(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.beta = nn.Linear(2, 2)
|
|
|
|
cb = GRPOStabilityCallback(ref_model_ema_alpha=0.5)
|
|
cb._policy_model = _A()
|
|
cb._ref_model = _B()
|
|
with caplog.at_level(logging.WARNING):
|
|
cb.on_step_end(args=None, state=None, control=None, model=cb._policy_model)
|
|
cb.on_step_end(args=None, state=None, control=None, model=cb._policy_model)
|
|
warnings = [r for r in caplog.records if "0 shared parameters" in r.message]
|
|
assert len(warnings) == 1 # one-shot
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Source wiring + patch invariants
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
class TestSourceWiring:
|
|
def _read(self, rel: str) -> str:
|
|
root = Path(__file__).resolve().parent.parent
|
|
return (root / "src" / "soup_cli" / rel).read_text(encoding="utf-8")
|
|
|
|
def test_grpo_wires_rl_callbacks(self):
|
|
src = self._read("trainer/grpo.py")
|
|
assert "attach_rl_callbacks" in src
|
|
assert "wrap_reward_funcs" in src
|
|
|
|
def test_distill_wires_uld_and_minillm(self):
|
|
src = self._read("trainer/distill.py")
|
|
assert "build_uld_projection" in src
|
|
assert "build_minillm_callback" in src
|
|
|
|
def test_stability_callback_uses_in_place_ema(self):
|
|
src = self._read("monitoring/grpo_stability_callback.py")
|
|
assert "update_ema_in_place" in src
|
|
# the old full-state_dict round-trip should be gone from on_step_end.
|
|
assert "self._ref_model.load_state_dict(ref_sd" not in src
|
|
|
|
def test_no_top_level_torch_in_new_utils(self):
|
|
for rel in (
|
|
"utils/rl_signal_buffer.py",
|
|
"utils/reward_hacking.py",
|
|
"utils/echo_trap.py",
|
|
"utils/uld.py",
|
|
"utils/minillm.py",
|
|
"utils/iterative_dpo.py",
|
|
"utils/rl_checkpoint.py",
|
|
):
|
|
src = self._read(rel)
|
|
assert "\nimport torch" not in src, rel
|
|
assert "\nfrom torch" not in src, rel
|
|
|
|
|
|
class TestPatchInvariants:
|
|
def test_version_bumped(self):
|
|
parts = soup_cli.__version__.split(".")
|
|
assert (int(parts[0]), int(parts[1]), int(parts[2])) >= (0, 71, 11)
|