mirror of https://github.com/razor-ai/soup.git
351 lines
12 KiB
Python
351 lines
12 KiB
Python
"""Tests for SimPO training — config, data format, template, routing, sweep."""
|
|
|
|
from unittest.mock import MagicMock
|
|
from unittest.mock import patch as mock_patch
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
from soup_cli.config.schema import TEMPLATES, SoupConfig
|
|
|
|
# ─── Config Tests ───────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestSimPOConfig:
|
|
"""Test SimPO task config validation."""
|
|
|
|
def test_simpo_task_accepted(self):
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="simpo",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
assert cfg.task == "simpo"
|
|
|
|
def test_simpo_gamma_default(self):
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="simpo",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
assert cfg.training.simpo_gamma == 0.5
|
|
|
|
def test_cpo_alpha_default(self):
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="simpo",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
assert cfg.training.cpo_alpha == 1.0
|
|
|
|
def test_simpo_gamma_custom(self):
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="simpo",
|
|
data={"train": "./data.jsonl"},
|
|
training={"simpo_gamma": 1.0},
|
|
)
|
|
assert cfg.training.simpo_gamma == pytest.approx(1.0)
|
|
|
|
def test_cpo_alpha_must_be_positive(self):
|
|
with pytest.raises(ValidationError):
|
|
SoupConfig(
|
|
base="some-model",
|
|
task="simpo",
|
|
data={"train": "./data.jsonl"},
|
|
training={"cpo_alpha": 0},
|
|
)
|
|
|
|
def test_simpo_gamma_zero_accepted(self):
|
|
"""simpo_gamma can be 0 (no margin)."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="simpo",
|
|
data={"train": "./data.jsonl"},
|
|
training={"simpo_gamma": 0},
|
|
)
|
|
assert cfg.training.simpo_gamma == 0
|
|
|
|
def test_simpo_full_config(self):
|
|
cfg = SoupConfig(
|
|
base="meta-llama/Llama-3.1-8B-Instruct",
|
|
task="simpo",
|
|
data={"train": "./data.jsonl", "format": "dpo", "max_length": 2048},
|
|
training={
|
|
"epochs": 3,
|
|
"lr": 1e-5,
|
|
"simpo_gamma": 1.5,
|
|
"cpo_alpha": 0.5,
|
|
"lora": {"r": 64, "alpha": 16},
|
|
"quantization": "4bit",
|
|
},
|
|
)
|
|
assert cfg.task == "simpo"
|
|
assert cfg.training.simpo_gamma == pytest.approx(1.5)
|
|
assert cfg.training.cpo_alpha == pytest.approx(0.5)
|
|
|
|
|
|
# ─── Template Tests ──────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestSimPOTemplate:
|
|
"""Test the SimPO template."""
|
|
|
|
def test_simpo_template_exists(self):
|
|
assert "simpo" in TEMPLATES
|
|
|
|
def test_simpo_template_valid_yaml(self):
|
|
import yaml
|
|
|
|
config = yaml.safe_load(TEMPLATES["simpo"])
|
|
assert config["task"] == "simpo"
|
|
assert config["training"]["simpo_gamma"] == 0.5
|
|
assert config["training"]["cpo_alpha"] == 1.0
|
|
assert config["data"]["format"] == "dpo"
|
|
|
|
def test_simpo_template_valid_config(self):
|
|
import yaml
|
|
|
|
raw = yaml.safe_load(TEMPLATES["simpo"])
|
|
cfg = SoupConfig(**raw)
|
|
assert cfg.task == "simpo"
|
|
assert cfg.training.simpo_gamma == 0.5
|
|
|
|
|
|
# ─── Train Command Routing Tests ─────────────────────────────────────────────
|
|
|
|
|
|
class TestSimPOTrainRouting:
|
|
"""Test that train command routes to SimPO trainer."""
|
|
|
|
def test_simpo_import_exists(self):
|
|
from soup_cli.trainer.simpo import SimPOTrainerWrapper
|
|
|
|
assert SimPOTrainerWrapper is not None
|
|
|
|
def test_simpo_wrapper_init(self):
|
|
from soup_cli.trainer.simpo import SimPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="simpo",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
wrapper = SimPOTrainerWrapper(cfg, device="cpu")
|
|
assert wrapper.config.task == "simpo"
|
|
assert wrapper.device == "cpu"
|
|
assert wrapper.model is None
|
|
assert wrapper.trainer is None
|
|
|
|
def test_simpo_wrapper_init_with_options(self):
|
|
from soup_cli.trainer.simpo import SimPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="simpo",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
wrapper = SimPOTrainerWrapper(
|
|
cfg, device="cuda", report_to="wandb", deepspeed_config="ds.json",
|
|
)
|
|
assert wrapper.report_to == "wandb"
|
|
assert wrapper.deepspeed_config == "ds.json"
|
|
|
|
|
|
# ─── Sweep Shortcut Tests ────────────────────────────────────────────────────
|
|
|
|
|
|
class TestSimPOSweepParams:
|
|
"""Test SimPO parameter shortcuts in sweep."""
|
|
|
|
def test_simpo_gamma_shortcut(self):
|
|
from soup_cli.commands.sweep import _set_nested_param
|
|
|
|
config = {"training": {"simpo_gamma": 0.5}}
|
|
_set_nested_param(config, "simpo_gamma", 1.0)
|
|
assert config["training"]["simpo_gamma"] == 1.0
|
|
|
|
def test_cpo_alpha_shortcut(self):
|
|
from soup_cli.commands.sweep import _set_nested_param
|
|
|
|
config = {}
|
|
_set_nested_param(config, "cpo_alpha", 0.5)
|
|
assert config["training"]["cpo_alpha"] == pytest.approx(0.5)
|
|
|
|
def test_sweep_run_single_routes_to_simpo_trainer(self):
|
|
from soup_cli.commands.sweep import _run_single
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="simpo",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
|
|
fake_dataset = {
|
|
"train": [
|
|
{"prompt": "Q?", "chosen": "A", "rejected": "B"},
|
|
]
|
|
}
|
|
fake_result = {
|
|
"initial_loss": 1.0,
|
|
"final_loss": 0.5,
|
|
"total_steps": 10,
|
|
"duration_secs": 60.0,
|
|
"output_dir": "./output",
|
|
"duration": "1m",
|
|
}
|
|
|
|
fake_gpu_info = {"memory_total": "0 MB", "memory_total_bytes": 0}
|
|
with mock_patch("soup_cli.data.loader.load_dataset", return_value=fake_dataset), \
|
|
mock_patch("soup_cli.utils.gpu.detect_device", return_value=("cpu", "CPU")), \
|
|
mock_patch("soup_cli.utils.gpu.get_gpu_info", return_value=fake_gpu_info), \
|
|
mock_patch("soup_cli.experiment.tracker.ExperimentTracker") as mock_tracker_cls, \
|
|
mock_patch("soup_cli.monitoring.display.TrainingDisplay"), \
|
|
mock_patch("soup_cli.trainer.simpo.SimPOTrainerWrapper.setup"), \
|
|
mock_patch(
|
|
"soup_cli.trainer.simpo.SimPOTrainerWrapper.train", return_value=fake_result
|
|
) as mock_train:
|
|
mock_tracker = MagicMock()
|
|
mock_tracker.start_run.return_value = "run-simpo-1"
|
|
mock_tracker_cls.return_value = mock_tracker
|
|
|
|
result = _run_single(cfg, {}, "simpo_run_1", None)
|
|
|
|
mock_train.assert_called_once()
|
|
assert result["run_id"] == "run-simpo-1"
|
|
|
|
|
|
# ─── Train Method Guard Test ──────────────────────────────────────────────────
|
|
|
|
|
|
class TestSimPOTrainGuard:
|
|
def test_train_before_setup_raises_runtime_error(self):
|
|
from soup_cli.trainer.simpo import SimPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="simpo",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
wrapper = SimPOTrainerWrapper(cfg)
|
|
with pytest.raises(RuntimeError, match="setup\\(dataset\\) first"):
|
|
wrapper.train()
|
|
|
|
|
|
# ─── Train Method Result Structure ───────────────────────────────────────────
|
|
|
|
|
|
class TestSimPOTrainResults:
|
|
def _make_wrapper_with_mock_trainer(self, log_history=None, global_step=20):
|
|
from soup_cli.trainer.simpo import SimPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="simpo",
|
|
data={"train": "./data.jsonl"},
|
|
output="./output",
|
|
)
|
|
wrapper = SimPOTrainerWrapper(cfg, device="cpu")
|
|
mock_trainer = MagicMock()
|
|
mock_trainer.train = MagicMock()
|
|
mock_trainer.state.log_history = log_history if log_history is not None else []
|
|
mock_trainer.state.global_step = global_step
|
|
mock_trainer.save_model = MagicMock()
|
|
wrapper.trainer = mock_trainer
|
|
wrapper.tokenizer = MagicMock()
|
|
wrapper._output_dir = "./output"
|
|
return wrapper, mock_trainer
|
|
|
|
def test_train_returns_expected_keys(self):
|
|
wrapper, _ = self._make_wrapper_with_mock_trainer(
|
|
log_history=[{"loss": 1.5}, {"loss": 0.8}], global_step=10
|
|
)
|
|
result = wrapper.train()
|
|
assert "initial_loss" in result
|
|
assert "final_loss" in result
|
|
assert "duration" in result
|
|
assert "output_dir" in result
|
|
assert "total_steps" in result
|
|
|
|
def test_train_result_losses(self):
|
|
wrapper, _ = self._make_wrapper_with_mock_trainer(
|
|
log_history=[{"loss": 2.0}, {"loss": 0.5}], global_step=20
|
|
)
|
|
result = wrapper.train()
|
|
assert result["initial_loss"] == pytest.approx(2.0)
|
|
assert result["final_loss"] == pytest.approx(0.5)
|
|
|
|
def test_train_calls_save_model(self):
|
|
wrapper, mock_trainer = self._make_wrapper_with_mock_trainer(
|
|
log_history=[{"loss": 1.0}], global_step=5
|
|
)
|
|
wrapper.train()
|
|
mock_trainer.save_model.assert_called_once_with("./output")
|
|
|
|
|
|
# ─── CLI Init SimPO Template Tests ─────────────────────────────────────────────
|
|
|
|
|
|
class TestSimPOInitTemplate:
|
|
def test_init_simpo_template_creates_file(self, tmp_path):
|
|
from typer.testing import CliRunner
|
|
|
|
from soup_cli.cli import app
|
|
|
|
runner = CliRunner()
|
|
output = tmp_path / "soup.yaml"
|
|
result = runner.invoke(app, ["init", "--template", "simpo", "--output", str(output)])
|
|
assert result.exit_code == 0
|
|
assert output.exists()
|
|
content = output.read_text()
|
|
assert "task: simpo" in content
|
|
assert "simpo_gamma" in content
|
|
|
|
def test_init_simpo_template_produces_valid_config(self, tmp_path):
|
|
from pathlib import Path
|
|
|
|
from typer.testing import CliRunner
|
|
|
|
from soup_cli.cli import app
|
|
from soup_cli.config.loader import load_config
|
|
|
|
runner = CliRunner()
|
|
output = tmp_path / "soup.yaml"
|
|
runner.invoke(app, ["init", "--template", "simpo", "--output", str(output)])
|
|
cfg = load_config(Path(output))
|
|
assert cfg.task == "simpo"
|
|
|
|
|
|
# ─── Wizard SimPO Path Tests ────────────────────────────────────────────────────
|
|
|
|
|
|
class TestSimPOWizardPath:
|
|
def test_wizard_simpo_task_sets_dpo_format(self):
|
|
from soup_cli.commands.init import _interactive_wizard
|
|
|
|
with mock_patch("soup_cli.commands.init.Prompt.ask", side_effect=[
|
|
"some-model",
|
|
"simpo",
|
|
"./data.jsonl",
|
|
"3",
|
|
"yes",
|
|
]):
|
|
config_text = _interactive_wizard()
|
|
|
|
assert "task: simpo" in config_text
|
|
assert "format: dpo" in config_text
|
|
assert "simpo_gamma: 0.5" in config_text
|
|
assert "cpo_alpha: 1.0" in config_text
|
|
|
|
|
|
# ─── Config Loader Round-trip Tests ──────────────────────────────────────────
|
|
|
|
|
|
class TestSimPOConfigLoaderRoundTrip:
|
|
def test_simpo_template_round_trip(self):
|
|
from soup_cli.config.loader import load_config_from_string
|
|
|
|
cfg = load_config_from_string(TEMPLATES["simpo"])
|
|
assert cfg.task == "simpo"
|
|
assert cfg.training.simpo_gamma == pytest.approx(0.5)
|