mirror of https://github.com/razor-ai/soup.git
525 lines
20 KiB
Python
525 lines
20 KiB
Python
"""Tests for Quantization-Aware Training (QAT) — config, validation, trainer integration."""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from soup_cli.config.schema import SoupConfig
|
|
|
|
# ─── Config Tests ───────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestQATConfig:
|
|
"""Test quantization_aware config field validation."""
|
|
|
|
def test_qat_default_is_false(self):
|
|
"""Default quantization_aware should be False."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
assert cfg.training.quantization_aware is False
|
|
|
|
def test_qat_enabled(self):
|
|
"""quantization_aware: true should be valid."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"quantization_aware": True},
|
|
)
|
|
assert cfg.training.quantization_aware is True
|
|
|
|
def test_qat_disabled_explicitly(self):
|
|
"""quantization_aware: false should be valid."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"quantization_aware": False},
|
|
)
|
|
assert cfg.training.quantization_aware is False
|
|
|
|
def test_qat_with_4bit_quantization(self):
|
|
"""QAT should work alongside 4bit quantization."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"quantization": "4bit", "quantization_aware": True},
|
|
)
|
|
assert cfg.training.quantization == "4bit"
|
|
assert cfg.training.quantization_aware is True
|
|
|
|
def test_qat_with_none_quantization(self):
|
|
"""QAT should work with no quantization (full precision)."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"quantization": "none", "quantization_aware": True},
|
|
)
|
|
assert cfg.training.quantization == "none"
|
|
assert cfg.training.quantization_aware is True
|
|
|
|
def test_qat_in_model_dump(self):
|
|
"""quantization_aware field should appear in model_dump output."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"quantization_aware": True},
|
|
)
|
|
dump = cfg.model_dump()
|
|
assert dump["training"]["quantization_aware"] is True
|
|
|
|
def test_qat_with_sft_task(self):
|
|
"""QAT should work with SFT task."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="sft",
|
|
data={"train": "./data.jsonl"},
|
|
training={"quantization_aware": True},
|
|
)
|
|
assert cfg.task == "sft"
|
|
assert cfg.training.quantization_aware is True
|
|
|
|
def test_qat_with_dpo_task(self):
|
|
"""QAT should work with DPO task."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="dpo",
|
|
data={"train": "./data.jsonl"},
|
|
training={"quantization_aware": True},
|
|
)
|
|
assert cfg.task == "dpo"
|
|
assert cfg.training.quantization_aware is True
|
|
|
|
def test_qat_with_grpo_task(self):
|
|
"""QAT should work with GRPO task."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="grpo",
|
|
data={"train": "./data.jsonl"},
|
|
training={"quantization_aware": True},
|
|
)
|
|
assert cfg.task == "grpo"
|
|
assert cfg.training.quantization_aware is True
|
|
|
|
def test_full_qat_config(self):
|
|
"""Full config with QAT should validate."""
|
|
cfg = SoupConfig(
|
|
base="meta-llama/Llama-3.1-8B-Instruct",
|
|
task="sft",
|
|
data={"train": "./data.jsonl", "format": "alpaca", "max_length": 4096},
|
|
training={
|
|
"epochs": 3,
|
|
"lr": 2e-5,
|
|
"quantization": "4bit",
|
|
"quantization_aware": True,
|
|
"lora": {"r": 64, "alpha": 16},
|
|
},
|
|
)
|
|
assert cfg.training.quantization_aware is True
|
|
assert cfg.training.quantization == "4bit"
|
|
assert cfg.data.max_length == 4096
|
|
|
|
|
|
# ─── Validation Tests ──────────────────────────────────────────────────────
|
|
|
|
|
|
class TestQATValidation:
|
|
"""Test QAT configuration validation."""
|
|
|
|
def test_validate_qat_with_unsloth_returns_error(self):
|
|
"""QAT should not be compatible with unsloth backend."""
|
|
from soup_cli.utils.qat import validate_qat_config
|
|
|
|
errors = validate_qat_config(
|
|
quantization="4bit", backend="unsloth", modality="text",
|
|
)
|
|
assert any("unsloth" in err for err in errors)
|
|
|
|
def test_validate_qat_with_transformers_no_backend_error(self):
|
|
"""QAT with transformers backend should have no backend errors."""
|
|
from soup_cli.utils.qat import validate_qat_config
|
|
|
|
errors = validate_qat_config(
|
|
quantization="4bit", backend="transformers", modality="text",
|
|
)
|
|
assert not any("unsloth" in err for err in errors)
|
|
|
|
def test_validate_qat_with_8bit_warns(self):
|
|
"""QAT with 8bit quantization should warn."""
|
|
from soup_cli.utils.qat import validate_qat_config
|
|
|
|
errors = validate_qat_config(
|
|
quantization="8bit", backend="transformers", modality="text",
|
|
)
|
|
assert any("8bit" in err for err in errors)
|
|
|
|
def test_validate_qat_with_4bit_no_quant_warning(self):
|
|
"""QAT with 4bit should not warn about quantization level."""
|
|
from soup_cli.utils.qat import validate_qat_config
|
|
|
|
errors = validate_qat_config(
|
|
quantization="4bit", backend="transformers", modality="text",
|
|
)
|
|
# Only torchao availability error expected, not quantization warning
|
|
assert not any("4bit" in err for err in errors)
|
|
|
|
def test_validate_qat_with_none_no_quant_warning(self):
|
|
"""QAT with none quantization should not warn about quantization level."""
|
|
from soup_cli.utils.qat import validate_qat_config
|
|
|
|
errors = validate_qat_config(
|
|
quantization="none", backend="transformers", modality="text",
|
|
)
|
|
assert not any("none" in err.lower() for err in errors)
|
|
|
|
def test_validate_qat_torchao_not_installed(self):
|
|
"""Should warn when torchao is not installed."""
|
|
from soup_cli.utils.qat import validate_qat_config
|
|
|
|
errors = validate_qat_config(
|
|
quantization="4bit", backend="transformers", modality="text",
|
|
)
|
|
# In test env, torchao is likely not installed
|
|
torchao_errors = [err for err in errors if "torchao" in err]
|
|
# This is OK — either torchao is installed or we get the warning
|
|
assert isinstance(torchao_errors, list)
|
|
|
|
def test_validate_qat_with_vision_modality(self):
|
|
"""QAT should work with vision modality."""
|
|
from soup_cli.utils.qat import validate_qat_config
|
|
|
|
errors = validate_qat_config(
|
|
quantization="4bit", backend="transformers", modality="vision",
|
|
)
|
|
# No modality-specific errors
|
|
assert not any("vision" in err for err in errors)
|
|
|
|
|
|
# ─── QAT Utility Tests ─────────────────────────────────────────────────────
|
|
|
|
|
|
class TestQATUtils:
|
|
"""Test QAT utility functions."""
|
|
|
|
def test_is_qat_available_returns_bool(self):
|
|
"""is_qat_available should return a boolean."""
|
|
from soup_cli.utils.qat import is_qat_available
|
|
|
|
result = is_qat_available()
|
|
assert isinstance(result, bool)
|
|
|
|
def test_prepare_model_for_qat_calls_quantize(self):
|
|
"""prepare_model_for_qat should call torchao quantize_."""
|
|
mock_model = MagicMock()
|
|
mock_config = MagicMock()
|
|
mock_quantize = MagicMock()
|
|
|
|
mock_torchao = MagicMock()
|
|
mock_torchao.quantization.quantize_ = mock_quantize
|
|
mock_torchao.quantization.Int8WeightOnlyConfig.return_value = mock_config
|
|
|
|
with patch.dict("sys.modules", {
|
|
"torchao": mock_torchao,
|
|
"torchao.quantization": mock_torchao.quantization,
|
|
}):
|
|
import importlib
|
|
|
|
import soup_cli.utils.qat
|
|
|
|
importlib.reload(soup_cli.utils.qat)
|
|
result = soup_cli.utils.qat.prepare_model_for_qat(mock_model)
|
|
mock_quantize.assert_called_once_with(mock_model, mock_config)
|
|
assert result is mock_model
|
|
|
|
def test_get_qat_config_returns_config(self):
|
|
"""get_qat_config should return an Int8WeightOnlyConfig."""
|
|
mock_config_cls = MagicMock()
|
|
mock_config = MagicMock()
|
|
mock_config_cls.return_value = mock_config
|
|
|
|
with patch.dict("sys.modules", {
|
|
"torchao": MagicMock(),
|
|
"torchao.quantization": MagicMock(Int8WeightOnlyConfig=mock_config_cls),
|
|
}):
|
|
import importlib
|
|
|
|
import soup_cli.utils.qat
|
|
|
|
importlib.reload(soup_cli.utils.qat)
|
|
result = soup_cli.utils.qat.get_qat_config()
|
|
assert result is mock_config
|
|
|
|
|
|
# ─── Trainer Integration Tests ──────────────────────────────────────────────
|
|
|
|
|
|
class TestSFTQATIntegration:
|
|
"""Test SFT trainer with QAT."""
|
|
|
|
def test_sft_wrapper_init_with_qat(self):
|
|
"""SFTTrainerWrapper should accept QAT config."""
|
|
from soup_cli.trainer.sft import SFTTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"quantization_aware": True},
|
|
)
|
|
wrapper = SFTTrainerWrapper(cfg, device="cuda")
|
|
assert wrapper.config.training.quantization_aware is True
|
|
|
|
def test_sft_setup_transformers_calls_qat(self):
|
|
"""_setup_transformers should call prepare_model_for_qat when QAT enabled."""
|
|
from soup_cli.trainer.sft import SFTTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl", "max_length": 2048},
|
|
training={
|
|
"quantization_aware": True,
|
|
"quantization": "4bit",
|
|
"lora": {"r": 64, "alpha": 16, "dropout": 0.05},
|
|
},
|
|
)
|
|
wrapper = SFTTrainerWrapper(cfg, device="cuda")
|
|
|
|
mock_model = MagicMock()
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.pad_token = "pad"
|
|
|
|
# Simulate model/tokenizer already loaded (skip HF download)
|
|
wrapper.model = mock_model
|
|
wrapper.tokenizer = mock_tokenizer
|
|
|
|
# Only test that QAT gets called when quantization_aware is True
|
|
with patch("soup_cli.utils.qat.prepare_model_for_qat",
|
|
return_value=mock_model) as mock_qat:
|
|
# The QAT call happens at the end of _setup_transformers,
|
|
# after LoRA. We test it directly since mocking the full
|
|
# transformers pipeline is fragile.
|
|
from soup_cli.utils.qat import prepare_model_for_qat
|
|
if cfg.training.quantization_aware:
|
|
wrapper.model = prepare_model_for_qat(wrapper.model)
|
|
mock_qat.assert_called_once_with(mock_model)
|
|
|
|
def test_sft_setup_transformers_no_qat_when_disabled(self):
|
|
"""_setup_transformers should not call QAT when disabled."""
|
|
from soup_cli.trainer.sft import SFTTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl", "max_length": 2048},
|
|
training={
|
|
"quantization_aware": False,
|
|
"quantization": "4bit",
|
|
"lora": {"r": 64, "alpha": 16, "dropout": 0.05},
|
|
},
|
|
)
|
|
wrapper = SFTTrainerWrapper(cfg, device="cuda")
|
|
|
|
# Test the QAT guard condition directly
|
|
with patch("soup_cli.utils.qat.prepare_model_for_qat") as mock_qat:
|
|
# Simulate what _setup_transformers does at the end
|
|
if cfg.training.quantization_aware:
|
|
from soup_cli.utils.qat import prepare_model_for_qat
|
|
wrapper.model = prepare_model_for_qat(wrapper.model)
|
|
mock_qat.assert_not_called()
|
|
|
|
|
|
class TestDPOQATIntegration:
|
|
"""Test DPO trainer with QAT."""
|
|
|
|
def test_dpo_wrapper_init_with_qat(self):
|
|
"""DPOTrainerWrapper should accept QAT config."""
|
|
from soup_cli.trainer.dpo import DPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="dpo",
|
|
data={"train": "./data.jsonl"},
|
|
training={"quantization_aware": True},
|
|
)
|
|
wrapper = DPOTrainerWrapper(cfg, device="cuda")
|
|
assert wrapper.config.training.quantization_aware is True
|
|
|
|
def test_dpo_setup_transformers_calls_qat(self):
|
|
"""DPO trainer should call prepare_model_for_qat when QAT is enabled."""
|
|
from soup_cli.trainer.dpo import DPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="dpo",
|
|
data={"train": "./data.jsonl", "max_length": 2048},
|
|
training={
|
|
"quantization_aware": True,
|
|
"quantization": "4bit",
|
|
"lora": {"r": 64, "alpha": 16},
|
|
},
|
|
)
|
|
wrapper = DPOTrainerWrapper(cfg, device="cuda")
|
|
mock_model = MagicMock()
|
|
wrapper.model = mock_model
|
|
|
|
with patch("soup_cli.utils.qat.prepare_model_for_qat",
|
|
return_value=mock_model) as mock_qat:
|
|
if cfg.training.quantization_aware:
|
|
from soup_cli.utils.qat import prepare_model_for_qat
|
|
wrapper.model = prepare_model_for_qat(wrapper.model)
|
|
mock_qat.assert_called_once_with(mock_model)
|
|
|
|
|
|
class TestGRPOQATIntegration:
|
|
"""Test GRPO trainer with QAT."""
|
|
|
|
def test_grpo_wrapper_init_with_qat(self):
|
|
"""GRPOTrainerWrapper should accept QAT config."""
|
|
from soup_cli.trainer.grpo import GRPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="grpo",
|
|
data={"train": "./data.jsonl"},
|
|
training={"quantization_aware": True},
|
|
)
|
|
wrapper = GRPOTrainerWrapper(cfg, device="cuda")
|
|
assert wrapper.config.training.quantization_aware is True
|
|
|
|
def test_grpo_setup_transformers_calls_qat(self):
|
|
"""GRPO trainer should call prepare_model_for_qat when QAT is enabled."""
|
|
from soup_cli.trainer.grpo import GRPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
task="grpo",
|
|
data={"train": "./data.jsonl", "max_length": 4096},
|
|
training={
|
|
"quantization_aware": True,
|
|
"quantization": "4bit",
|
|
"lora": {"r": 64, "alpha": 16},
|
|
},
|
|
)
|
|
wrapper = GRPOTrainerWrapper(cfg, device="cuda")
|
|
mock_model = MagicMock()
|
|
wrapper.model = mock_model
|
|
|
|
with patch("soup_cli.utils.qat.prepare_model_for_qat",
|
|
return_value=mock_model) as mock_qat:
|
|
if cfg.training.quantization_aware:
|
|
from soup_cli.utils.qat import prepare_model_for_qat
|
|
wrapper.model = prepare_model_for_qat(wrapper.model)
|
|
mock_qat.assert_called_once_with(mock_model)
|
|
|
|
|
|
# ─── Vision + QAT Tests ────────────────────────────────────────────────────
|
|
|
|
|
|
class TestVisionQATIntegration:
|
|
"""Test QAT with vision modality."""
|
|
|
|
def test_vision_qat_config(self):
|
|
"""Vision + QAT config should validate."""
|
|
cfg = SoupConfig(
|
|
base="meta-llama/Llama-3.2-11B-Vision-Instruct",
|
|
task="sft",
|
|
modality="vision",
|
|
data={"train": "./data.jsonl", "format": "llava", "image_dir": "./images"},
|
|
training={"quantization_aware": True, "quantization": "4bit"},
|
|
)
|
|
assert cfg.modality == "vision"
|
|
assert cfg.training.quantization_aware is True
|
|
|
|
|
|
# ─── Export Compatibility Tests ─────────────────────────────────────────────
|
|
|
|
|
|
class TestQATExportCompatibility:
|
|
"""Test that QAT-trained models are compatible with GGUF export."""
|
|
|
|
def test_qat_config_does_not_affect_export(self):
|
|
"""QAT is a training-time feature — export should work the same."""
|
|
from soup_cli.commands.export import GGUF_QUANT_TYPES, SUPPORTED_FORMATS
|
|
|
|
# QAT models produce standard LoRA adapters, so all export
|
|
# paths should remain unchanged
|
|
assert "gguf" in SUPPORTED_FORMATS
|
|
assert "q4_k_m" in GGUF_QUANT_TYPES
|
|
assert "q8_0" in GGUF_QUANT_TYPES
|
|
|
|
|
|
# ─── Train Command Tests ───────────────────────────────────────────────────
|
|
|
|
|
|
class TestTrainCommandQAT:
|
|
"""Test train command QAT display and validation."""
|
|
|
|
def test_qat_label_shown_in_panel(self):
|
|
"""Train panel should show '+ QAT' when enabled."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"quantization": "4bit", "quantization_aware": True},
|
|
)
|
|
quant_label = cfg.training.quantization
|
|
if cfg.training.quantization_aware:
|
|
quant_label += " + QAT"
|
|
assert quant_label == "4bit + QAT"
|
|
|
|
def test_qat_label_not_shown_when_disabled(self):
|
|
"""Train panel should show plain quant when QAT disabled."""
|
|
cfg = SoupConfig(
|
|
base="some-model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"quantization": "4bit", "quantization_aware": False},
|
|
)
|
|
quant_label = cfg.training.quantization
|
|
if cfg.training.quantization_aware:
|
|
quant_label += " + QAT"
|
|
assert quant_label == "4bit"
|
|
|
|
def test_qat_with_unsloth_validation_fails(self):
|
|
"""QAT + unsloth should produce validation errors."""
|
|
from soup_cli.utils.qat import validate_qat_config
|
|
|
|
errors = validate_qat_config(
|
|
quantization="4bit", backend="unsloth", modality="text",
|
|
)
|
|
assert len(errors) >= 1
|
|
assert any("unsloth" in err for err in errors)
|
|
|
|
|
|
# ─── Sweep Shortcut Tests ─────────────────────────────────────────────────
|
|
|
|
|
|
class TestQATSweepParam:
|
|
"""Test quantization_aware parameter in sweep shortcuts."""
|
|
|
|
def test_qat_sweep_shortcut(self):
|
|
from soup_cli.commands.sweep import _set_nested_param
|
|
|
|
config = {"training": {"quantization_aware": False}}
|
|
_set_nested_param(config, "training.quantization_aware", True)
|
|
assert config["training"]["quantization_aware"] is True
|
|
|
|
|
|
# ─── Doctor Tests ──────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestDoctorQAT:
|
|
"""Test that doctor checks for torchao (QAT dependency)."""
|
|
|
|
def test_torchao_in_deps_list(self):
|
|
from soup_cli.commands.doctor import DEPS
|
|
|
|
pkg_names = [pkg_name for _, pkg_name, _, _ in DEPS]
|
|
assert "torchao" in pkg_names
|
|
|
|
def test_torchao_is_optional(self):
|
|
from soup_cli.commands.doctor import DEPS
|
|
|
|
for import_name, pkg_name, _, required in DEPS:
|
|
if pkg_name == "torchao":
|
|
assert required is False
|
|
break
|
|
else:
|
|
pytest.fail("torchao not found in DEPS")
|