align / code /tests /test_cispo_config.py
viai957's picture
Add files using upload-large-folder tool
201dc77 verified
Raw History Blame Contribute Delete
1.55 kB
"""configs/p4_cispo_lean.toml: the draft loads under the stage loader's struct
mode, selects CISPO + FP32 head, passes the P4 gate, and its denominator is the
DERIVED constant (a hand-typed one is refused by GrpoStageConfig.validate)."""
from __future__ import annotations
import pathlib
import pytest
from align.rl.cispo import assert_p4_cispo_conformant
from align.rl.grpo_loss import GRPOLossType
from align.rl.stage_config import load_stage_config
TOML = pathlib.Path(__file__).resolve().parents[1] / "configs" / "p4_cispo_lean.toml"
def _load(tmp_path, edit=lambda s: s):
band = tmp_path / "band.jsonl"
band.write_text("{}\n")
text = TOML.read_text().replace("/apps/development/rlhf/align/data/p4/lean_band0.jsonl", str(band))
p = tmp_path / "p4.toml"
p.write_text(edit(text))
return load_stage_config(str(p))
def test_draft_loads_and_selects_cispo_fp32(tmp_path):
cfg = _load(tmp_path)
g = cfg.grpo
assert g.loss_fn == GRPOLossType.cispo and g.fp32_head is True
assert g.cispo_clip_lower is None # paper: no lower bound
assert g.loss_denominator == 16 * 16 * 1536.0
assert_p4_cispo_conformant(g)
def test_hand_typed_denominator_still_refused(tmp_path):
with pytest.raises(ValueError, match="DERIVED"):
_load(tmp_path, lambda s: s.replace("loss_denominator = 393216.0", "loss_denominator = 100000.0"))
def test_unknown_grpo_key_refused(tmp_path):
with pytest.raises(TypeError):
_load(tmp_path, lambda s: s.replace("fp32_head = true", "fp32_hed = true"))