mirror of https://github.com/razor-ai/soup.git
241 lines
8.7 KiB
Python
241 lines
8.7 KiB
Python
"""Tests for NEFTune (neftune_alpha) and rsLoRA (use_rslora) support."""
|
|
|
|
import pytest
|
|
|
|
from soup_cli.config.schema import LoraConfig, SoupConfig, TrainingConfig
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# NEFTune config tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestNEFTuneConfig:
|
|
"""NEFTune config validation."""
|
|
|
|
def test_neftune_alpha_default_none(self):
|
|
"""neftune_alpha defaults to None."""
|
|
tcfg = TrainingConfig()
|
|
assert tcfg.neftune_alpha is None
|
|
|
|
def test_neftune_alpha_valid(self):
|
|
"""neftune_alpha accepts valid float."""
|
|
tcfg = TrainingConfig(neftune_alpha=5.0)
|
|
assert tcfg.neftune_alpha == 5.0
|
|
|
|
def test_neftune_alpha_zero(self):
|
|
"""neftune_alpha accepts 0 (disabled)."""
|
|
tcfg = TrainingConfig(neftune_alpha=0.0)
|
|
assert tcfg.neftune_alpha == 0.0
|
|
|
|
def test_neftune_alpha_max(self):
|
|
"""neftune_alpha accepts 50.0 (max)."""
|
|
tcfg = TrainingConfig(neftune_alpha=50.0)
|
|
assert tcfg.neftune_alpha == 50.0
|
|
|
|
def test_neftune_alpha_negative_rejected(self):
|
|
"""neftune_alpha rejects negative values."""
|
|
with pytest.raises(Exception):
|
|
TrainingConfig(neftune_alpha=-1.0)
|
|
|
|
def test_neftune_alpha_too_high_rejected(self):
|
|
"""neftune_alpha rejects values > 50."""
|
|
with pytest.raises(Exception):
|
|
TrainingConfig(neftune_alpha=51.0)
|
|
|
|
def test_neftune_in_full_config(self):
|
|
"""neftune_alpha works in full SoupConfig."""
|
|
cfg = SoupConfig(
|
|
base="test-model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"neftune_alpha": 5.0},
|
|
)
|
|
assert cfg.training.neftune_alpha == 5.0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# NEFTune trainer integration tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestNEFTuneTrainer:
|
|
"""NEFTune trainer argument passing."""
|
|
|
|
def test_sft_neftune_in_training_kwargs(self):
|
|
"""SFT trainer includes neftune_noise_alpha in training kwargs."""
|
|
tcfg = TrainingConfig(neftune_alpha=5.0)
|
|
assert tcfg.neftune_alpha == 5.0
|
|
# Verify the pattern: if neftune_alpha is not None, it should be passed
|
|
training_kwargs = {}
|
|
if tcfg.neftune_alpha is not None:
|
|
training_kwargs["neftune_noise_alpha"] = tcfg.neftune_alpha
|
|
assert training_kwargs["neftune_noise_alpha"] == 5.0
|
|
|
|
def test_neftune_none_not_passed(self):
|
|
"""When neftune_alpha is None, neftune_noise_alpha is NOT added."""
|
|
tcfg = TrainingConfig()
|
|
training_kwargs = {}
|
|
if tcfg.neftune_alpha is not None:
|
|
training_kwargs["neftune_noise_alpha"] = tcfg.neftune_alpha
|
|
assert "neftune_noise_alpha" not in training_kwargs
|
|
|
|
def test_dpo_neftune_forwarding(self):
|
|
"""DPO trainer config includes neftune_noise_alpha when set."""
|
|
tcfg = TrainingConfig(neftune_alpha=10.0)
|
|
extra = (
|
|
{"neftune_noise_alpha": tcfg.neftune_alpha}
|
|
if tcfg.neftune_alpha is not None else {}
|
|
)
|
|
assert extra == {"neftune_noise_alpha": 10.0}
|
|
|
|
def test_kto_neftune_forwarding(self):
|
|
"""KTO trainer config includes neftune_noise_alpha when set."""
|
|
tcfg = TrainingConfig(neftune_alpha=7.5)
|
|
extra = (
|
|
{"neftune_noise_alpha": tcfg.neftune_alpha}
|
|
if tcfg.neftune_alpha is not None else {}
|
|
)
|
|
assert extra == {"neftune_noise_alpha": 7.5}
|
|
|
|
def test_orpo_neftune_forwarding(self):
|
|
"""ORPO trainer config includes neftune_noise_alpha when set."""
|
|
tcfg = TrainingConfig(neftune_alpha=3.0)
|
|
extra = (
|
|
{"neftune_noise_alpha": tcfg.neftune_alpha}
|
|
if tcfg.neftune_alpha is not None else {}
|
|
)
|
|
assert extra == {"neftune_noise_alpha": 3.0}
|
|
|
|
def test_simpo_neftune_forwarding(self):
|
|
"""SimPO trainer config includes neftune_noise_alpha when set."""
|
|
tcfg = TrainingConfig(neftune_alpha=15.0)
|
|
extra = (
|
|
{"neftune_noise_alpha": tcfg.neftune_alpha}
|
|
if tcfg.neftune_alpha is not None else {}
|
|
)
|
|
assert extra == {"neftune_noise_alpha": 15.0}
|
|
|
|
def test_ipo_neftune_forwarding(self):
|
|
"""IPO trainer config includes neftune_noise_alpha when set."""
|
|
tcfg = TrainingConfig(neftune_alpha=2.0)
|
|
extra = (
|
|
{"neftune_noise_alpha": tcfg.neftune_alpha}
|
|
if tcfg.neftune_alpha is not None else {}
|
|
)
|
|
assert extra == {"neftune_noise_alpha": 2.0}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# NEFTune sweep integration tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestNEFTuneSweep:
|
|
"""NEFTune sweep parameter support."""
|
|
|
|
def test_neftune_in_sweep_shortcuts(self):
|
|
"""neftune_alpha is a valid sweep parameter."""
|
|
from soup_cli.commands.sweep import _set_nested_param
|
|
|
|
config_dict = {
|
|
"base": "test",
|
|
"task": "sft",
|
|
"data": {"train": "./data.jsonl"},
|
|
"training": {},
|
|
}
|
|
_set_nested_param(config_dict, "neftune_alpha", 5.0)
|
|
assert config_dict["training"]["neftune_alpha"] == 5.0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# rsLoRA config tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestRsLoRAConfig:
|
|
"""rsLoRA config validation."""
|
|
|
|
def test_rslora_default_false(self):
|
|
"""use_rslora defaults to False."""
|
|
lora = LoraConfig()
|
|
assert lora.use_rslora is False
|
|
|
|
def test_rslora_enable(self):
|
|
"""use_rslora can be enabled."""
|
|
lora = LoraConfig(use_rslora=True)
|
|
assert lora.use_rslora is True
|
|
|
|
def test_rslora_in_full_config(self):
|
|
"""use_rslora works in full SoupConfig."""
|
|
cfg = SoupConfig(
|
|
base="test-model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"lora": {"use_rslora": True}},
|
|
)
|
|
assert cfg.training.lora.use_rslora is True
|
|
|
|
def test_rslora_with_dora(self):
|
|
"""use_rslora can be used alongside use_dora."""
|
|
lora = LoraConfig(use_rslora=True, use_dora=True)
|
|
assert lora.use_rslora is True
|
|
assert lora.use_dora is True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# rsLoRA sweep integration tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestRsLoRATrainer:
|
|
"""rsLoRA trainer integration — use_rslora flows to LoraConfig."""
|
|
|
|
def test_rslora_in_lora_config_kwargs(self):
|
|
"""use_rslora is passed to peft.LoraConfig constructor."""
|
|
lcfg = LoraConfig(use_rslora=True, r=16, alpha=32)
|
|
# Verify the value that would be passed to peft.LoraConfig
|
|
assert lcfg.use_rslora is True
|
|
kwargs = {
|
|
"r": lcfg.r,
|
|
"lora_alpha": lcfg.alpha,
|
|
"lora_dropout": lcfg.dropout,
|
|
"use_dora": lcfg.use_dora,
|
|
"use_rslora": lcfg.use_rslora,
|
|
}
|
|
assert kwargs["use_rslora"] is True
|
|
assert kwargs["use_dora"] is False
|
|
|
|
def test_rslora_false_in_lora_config_kwargs(self):
|
|
"""use_rslora=False is passed correctly."""
|
|
lcfg = LoraConfig(use_rslora=False)
|
|
kwargs = {"use_rslora": lcfg.use_rslora}
|
|
assert kwargs["use_rslora"] is False
|
|
|
|
def test_rslora_with_all_tasks(self):
|
|
"""use_rslora works in config for every task type."""
|
|
for task in [
|
|
"sft", "dpo", "grpo", "ppo", "reward_model",
|
|
"kto", "orpo", "simpo", "ipo", "pretrain", "embedding",
|
|
]:
|
|
cfg = SoupConfig(
|
|
base="test-model",
|
|
task=task,
|
|
data={"train": "./data.jsonl"},
|
|
training={"lora": {"use_rslora": True}},
|
|
)
|
|
assert cfg.training.lora.use_rslora is True, f"Failed for task={task}"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# rsLoRA sweep integration tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestRsLoRASweep:
|
|
"""rsLoRA sweep parameter support."""
|
|
|
|
def test_rslora_in_sweep_shortcuts(self):
|
|
"""use_rslora is a valid sweep parameter."""
|
|
from soup_cli.commands.sweep import _set_nested_param
|
|
|
|
config_dict = {
|
|
"base": "test",
|
|
"task": "sft",
|
|
"data": {"train": "./data.jsonl"},
|
|
"training": {"lora": {}},
|
|
}
|
|
_set_nested_param(config_dict, "use_rslora", True)
|
|
assert config_dict["training"]["lora"]["use_rslora"] is True
|