mirror of https://github.com/razor-ai/soup.git
838 lines
31 KiB
Python
838 lines
31 KiB
Python
"""v0.71.20 — Modality II trainers (BETA, hw-gated): #131 TTS, #134 BitNet,
|
||
#136 MoE expert quant + router-only training.
|
||
|
||
Lifts the v0.52.0 schema-only stubs:
|
||
- #131 build_tts_trainer → live TTSTrainerWrapper (SFT-CE + per-family
|
||
emotion templating + codec special-token registration). Pre-encoded chat
|
||
mode is the live/smokeable path; the live-codec (audio→tokens at train time)
|
||
path is hardware/dependency-gated with a friendly per-family RuntimeError.
|
||
- #134 build_bitnet_trainer + export_bitnet_gguf → real plumbing behind
|
||
friendly onebitllms / llama.cpp gates.
|
||
- #136 apply_moe_expert_quant + freeze_experts_train_router + SFT wiring.
|
||
|
||
torch-needing tests use ``pytest.importorskip``; the pure validators/templating
|
||
+ source-wiring guards run everywhere.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
import sys
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
|
||
import pytest
|
||
|
||
from soup_cli import __version__
|
||
|
||
_SRC = Path(__file__).resolve().parent.parent / "src" / "soup_cli"
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# #131 TTS — pure helpers
|
||
# --------------------------------------------------------------------------- #
|
||
class TestTtsCodecPackage:
|
||
def test_per_family_packages(self):
|
||
from soup_cli.utils.tts import TTS_CODEC_PACKAGES, tts_codec_package
|
||
|
||
for fam in ("orpheus", "sesame_csm", "llasa", "spark", "oute"):
|
||
assert tts_codec_package(fam) == TTS_CODEC_PACKAGES[fam]
|
||
assert tts_codec_package("orpheus") == "snac"
|
||
assert tts_codec_package("spark") == "sparktts"
|
||
|
||
def test_case_insensitive(self):
|
||
from soup_cli.utils.tts import tts_codec_package
|
||
|
||
assert tts_codec_package("OrPheus") == "snac"
|
||
|
||
def test_unknown_rejected(self):
|
||
from soup_cli.utils.tts import tts_codec_package
|
||
|
||
with pytest.raises(ValueError, match="not supported"):
|
||
tts_codec_package("bark")
|
||
|
||
def test_packages_immutable(self):
|
||
from soup_cli.utils.tts import TTS_CODEC_PACKAGES
|
||
|
||
with pytest.raises(TypeError):
|
||
TTS_CODEC_PACKAGES["x"] = "y" # type: ignore[index]
|
||
|
||
|
||
class TestFormatTtsMessages:
|
||
def _msgs(self):
|
||
return [
|
||
{"role": "system", "content": "Speak."},
|
||
{"role": "user", "content": "Hello there."},
|
||
{"role": "assistant", "content": "<codec>123 456</codec>"},
|
||
]
|
||
|
||
def test_no_emotion_passthrough(self):
|
||
from soup_cli.utils.tts import format_tts_messages
|
||
|
||
out = format_tts_messages(self._msgs(), "spark", emotion=None)
|
||
assert out[1]["content"] == "Hello there."
|
||
|
||
def test_orpheus_emotion_prefixes_user(self):
|
||
from soup_cli.utils.tts import format_tts_messages
|
||
|
||
out = format_tts_messages(self._msgs(), "orpheus", emotion="happy")
|
||
assert out[1]["content"].startswith("<|emotion|>happy<|/emotion|>")
|
||
assert "Hello there." in out[1]["content"]
|
||
|
||
def test_oute_emotion_template(self):
|
||
from soup_cli.utils.tts import format_tts_messages
|
||
|
||
out = format_tts_messages(self._msgs(), "oute", emotion="calm")
|
||
assert out[1]["content"].startswith("[emotion: calm]")
|
||
|
||
def test_caller_list_not_mutated(self):
|
||
from soup_cli.utils.tts import format_tts_messages
|
||
|
||
msgs = self._msgs()
|
||
original = msgs[1]["content"]
|
||
_ = format_tts_messages(msgs, "orpheus", emotion="happy")
|
||
assert msgs[1]["content"] == original
|
||
|
||
def test_emotion_on_unsupported_family_rejected(self):
|
||
from soup_cli.utils.tts import format_tts_messages
|
||
|
||
with pytest.raises(ValueError, match="does not support emotion"):
|
||
format_tts_messages(self._msgs(), "spark", emotion="happy")
|
||
|
||
def test_non_list_rejected(self):
|
||
from soup_cli.utils.tts import format_tts_messages
|
||
|
||
with pytest.raises(TypeError, match="messages must be a list"):
|
||
format_tts_messages("nope", "spark")
|
||
|
||
def test_non_dict_message_rejected(self):
|
||
from soup_cli.utils.tts import format_tts_messages
|
||
|
||
with pytest.raises(TypeError, match="must be a dict"):
|
||
format_tts_messages(["nope"], "spark")
|
||
|
||
def test_only_first_user_turn_prefixed(self):
|
||
from soup_cli.utils.tts import format_tts_messages
|
||
|
||
msgs = [
|
||
{"role": "user", "content": "one"},
|
||
{"role": "user", "content": "two"},
|
||
]
|
||
out = format_tts_messages(msgs, "orpheus", emotion="sad")
|
||
assert out[0]["content"].startswith("<|emotion|>sad")
|
||
assert out[1]["content"] == "two"
|
||
|
||
def test_non_str_content_left_unchanged(self):
|
||
from soup_cli.utils.tts import format_tts_messages
|
||
|
||
parts = [{"type": "text", "text": "hi"}]
|
||
msgs = [{"role": "user", "content": parts}]
|
||
out = format_tts_messages(msgs, "orpheus", emotion="happy")
|
||
# multimodal content list not string-prefixed; deep-copied (not same obj)
|
||
assert out[0]["content"] == parts
|
||
assert out[0]["content"] is not parts
|
||
|
||
def test_no_user_turn_with_emotion(self):
|
||
from soup_cli.utils.tts import format_tts_messages
|
||
|
||
msgs = [
|
||
{"role": "system", "content": "Speak."},
|
||
{"role": "assistant", "content": "<codec>1</codec>"},
|
||
]
|
||
out = format_tts_messages(msgs, "orpheus", emotion="happy")
|
||
assert out[0]["content"] == "Speak."
|
||
assert out[1]["content"] == "<codec>1</codec>"
|
||
|
||
def test_nested_value_not_shared(self):
|
||
from soup_cli.utils.tts import format_tts_messages
|
||
|
||
msgs = [{"role": "user", "content": "hi", "tool_calls": [{"x": 1}]}]
|
||
out = format_tts_messages(msgs, "orpheus", emotion="happy")
|
||
out[0]["tool_calls"][0]["x"] = 99
|
||
assert msgs[0]["tool_calls"][0]["x"] == 1
|
||
|
||
|
||
class TestBuildTtsTrainer:
|
||
def test_no_arg_typeerror(self):
|
||
from soup_cli.utils.tts import build_tts_trainer
|
||
|
||
with pytest.raises(TypeError):
|
||
build_tts_trainer()
|
||
|
||
def test_returns_wrapper(self):
|
||
from soup_cli.config.loader import load_config_from_string
|
||
from soup_cli.trainer.tts import TTSTrainerWrapper
|
||
from soup_cli.utils.tts import build_tts_trainer
|
||
|
||
yaml_str = (
|
||
"base: hf-internal-testing/tiny-random-gpt2\n"
|
||
"task: tts\n"
|
||
"modality: audio_out\n"
|
||
"data:\n"
|
||
" train: ./x.jsonl\n"
|
||
" format: auto\n"
|
||
"training:\n"
|
||
" tts_family: spark\n"
|
||
" epochs: 1\n"
|
||
" lr: 0.0001\n"
|
||
" batch_size: 1\n"
|
||
)
|
||
cfg = load_config_from_string(yaml_str)
|
||
wrapper = build_tts_trainer(cfg, device="cpu")
|
||
assert isinstance(wrapper, TTSTrainerWrapper)
|
||
|
||
|
||
class TestTtsTrainerSetup:
|
||
def _wrapper(self, *, data_format, family="spark", emotion=None):
|
||
from soup_cli.trainer.tts import TTSTrainerWrapper
|
||
|
||
w = object.__new__(TTSTrainerWrapper)
|
||
w.config = SimpleNamespace(
|
||
training=SimpleNamespace(tts_family=family, tts_emotion=emotion),
|
||
data=SimpleNamespace(format=data_format, new_special_tokens=None),
|
||
)
|
||
w._tts_family = None
|
||
return w
|
||
|
||
def test_live_codec_mode_gates_when_codec_absent(self):
|
||
# data.format='audio' + family whose codec is not installed.
|
||
w = self._wrapper(data_format="audio", family="spark")
|
||
with pytest.raises(RuntimeError, match="sparktts"):
|
||
w.setup({"train": []})
|
||
|
||
def test_live_codec_mode_orpheus_names_snac(self, monkeypatch):
|
||
# v0.71.22 #265 lifted the unconditional not-yet-validated gate —
|
||
# orpheus is live when snac IS installed, so force the missing-dep
|
||
# branch to keep exercising the friendly pip hint.
|
||
import importlib.util as ilu
|
||
|
||
real_find_spec = ilu.find_spec
|
||
monkeypatch.setattr(
|
||
ilu,
|
||
"find_spec",
|
||
lambda name, *a, **k: None
|
||
if name == "snac"
|
||
else real_find_spec(name, *a, **k),
|
||
)
|
||
w = self._wrapper(data_format="audio", family="orpheus")
|
||
with pytest.raises(RuntimeError, match="snac"):
|
||
w.setup({"train": []})
|
||
|
||
def test_chat_mode_applies_templating_and_delegates(self):
|
||
# In pre-encoded chat mode, setup() emotion-templates the dataset then
|
||
# calls super().setup(). Stub super().setup to capture the dataset.
|
||
from soup_cli.trainer.sft import SFTTrainerWrapper
|
||
|
||
w = self._wrapper(data_format="auto", family="orpheus", emotion="happy")
|
||
captured = {}
|
||
|
||
def _fake_super_setup(self, dataset): # noqa: ANN001
|
||
captured["dataset"] = dataset
|
||
|
||
orig = SFTTrainerWrapper.setup
|
||
try:
|
||
SFTTrainerWrapper.setup = _fake_super_setup # type: ignore[assignment]
|
||
w.setup({
|
||
"train": [
|
||
{"messages": [{"role": "user", "content": "hi"}]},
|
||
]
|
||
})
|
||
finally:
|
||
SFTTrainerWrapper.setup = orig # type: ignore[assignment]
|
||
|
||
msg = captured["dataset"]["train"][0]["messages"][0]
|
||
assert msg["content"].startswith("<|emotion|>happy")
|
||
|
||
def test_apply_tts_templating_non_message_rows_passthrough(self):
|
||
w = self._wrapper(data_format="auto", family="spark")
|
||
ds = w._apply_tts_templating(
|
||
{"train": [{"text": "raw"}], "val": []}, "spark", None
|
||
)
|
||
assert ds["train"][0] == {"text": "raw"}
|
||
|
||
def test_apply_tts_templating_val_split_templated(self):
|
||
w = self._wrapper(data_format="auto", family="orpheus", emotion="happy")
|
||
ds = w._apply_tts_templating(
|
||
{
|
||
"train": [],
|
||
"val": [{"messages": [{"role": "user", "content": "v"}]}],
|
||
},
|
||
"orpheus",
|
||
"happy",
|
||
)
|
||
assert ds["val"][0]["messages"][0]["content"].startswith("<|emotion|>happy")
|
||
|
||
def test_apply_tts_templating_no_train_key(self):
|
||
w = self._wrapper(data_format="auto", family="spark")
|
||
ds = w._apply_tts_templating({}, "spark", None)
|
||
assert ds == {}
|
||
|
||
def test_register_special_tokens_adds_and_resizes(self):
|
||
w = self._wrapper(data_format="auto")
|
||
w.config.data.new_special_tokens = ["<aud_0>", "<aud_1>"]
|
||
calls = {"resize": None, "added": None}
|
||
|
||
class _Tok:
|
||
def get_vocab(self):
|
||
return {"<aud_0>": 1}
|
||
|
||
def add_special_tokens(self, d):
|
||
calls["added"] = d["additional_special_tokens"]
|
||
return len(d["additional_special_tokens"])
|
||
|
||
def __len__(self):
|
||
return 100
|
||
|
||
class _Model:
|
||
def resize_token_embeddings(self, n):
|
||
calls["resize"] = n
|
||
|
||
w.tokenizer = _Tok()
|
||
w.model = _Model()
|
||
w._register_tts_special_tokens(w.config)
|
||
# <aud_0> already present → only <aud_1> added.
|
||
assert calls["added"] == ["<aud_1>"]
|
||
assert calls["resize"] == 100
|
||
|
||
def test_register_special_tokens_all_present_no_resize(self):
|
||
w = self._wrapper(data_format="auto")
|
||
w.config.data.new_special_tokens = ["<aud_0>", "<aud_0>"]
|
||
calls = {"resize": False, "added": False}
|
||
|
||
class _Tok:
|
||
def get_vocab(self):
|
||
return {"<aud_0>": 1}
|
||
|
||
def add_special_tokens(self, d):
|
||
calls["added"] = True
|
||
return 0
|
||
|
||
class _Model:
|
||
def resize_token_embeddings(self, n):
|
||
calls["resize"] = True
|
||
|
||
w.tokenizer = _Tok()
|
||
w.model = _Model()
|
||
w._register_tts_special_tokens(w.config)
|
||
# all tokens already present (and deduped) → never add or resize.
|
||
assert calls["added"] is False
|
||
assert calls["resize"] is False
|
||
|
||
def test_register_special_tokens_noop_when_none(self):
|
||
w = self._wrapper(data_format="auto")
|
||
w.config.data.new_special_tokens = None
|
||
w.tokenizer = object()
|
||
w.model = object()
|
||
# Should not raise (no tokenizer access).
|
||
w._register_tts_special_tokens(w.config)
|
||
|
||
|
||
class TestTtsSchemaAndRouting:
|
||
def test_tts_config_loads(self):
|
||
from soup_cli.config.loader import load_config_from_string
|
||
|
||
yaml_str = (
|
||
"base: hf-internal-testing/tiny-random-gpt2\n"
|
||
"task: tts\n"
|
||
"modality: audio_out\n"
|
||
"data:\n"
|
||
" train: ./x.jsonl\n"
|
||
" format: auto\n"
|
||
"training:\n"
|
||
" tts_family: orpheus\n"
|
||
" tts_emotion: happy\n"
|
||
)
|
||
cfg = load_config_from_string(yaml_str)
|
||
assert cfg.task == "tts"
|
||
assert cfg.training.tts_family == "orpheus"
|
||
|
||
def test_train_routes_tts(self):
|
||
src = (_SRC / "commands" / "train.py").read_text(encoding="utf-8")
|
||
assert 'cfg.task == "tts"' in src
|
||
assert "TTSTrainerWrapper" in src
|
||
|
||
def test_utils_tts_no_top_level_torch(self):
|
||
src = (_SRC / "utils" / "tts.py").read_text(encoding="utf-8")
|
||
assert "\nimport torch" not in src
|
||
assert "\nfrom torch" not in src
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# #134 BitNet
|
||
# --------------------------------------------------------------------------- #
|
||
class TestBuildBitnetTrainer:
|
||
def test_no_arg_typeerror(self):
|
||
from soup_cli.utils.bitnet import build_bitnet_trainer
|
||
|
||
with pytest.raises(TypeError):
|
||
build_bitnet_trainer()
|
||
|
||
def test_returns_wrapper(self):
|
||
from soup_cli.config.loader import load_config_from_string
|
||
from soup_cli.trainer.bitnet import BitNetTrainerWrapper
|
||
from soup_cli.utils.bitnet import build_bitnet_trainer
|
||
|
||
yaml_str = (
|
||
"base: hf-internal-testing/tiny-random-gpt2\n"
|
||
"task: sft\n"
|
||
"data:\n"
|
||
" train: ./x.jsonl\n"
|
||
" format: auto\n"
|
||
"training:\n"
|
||
" quantization: bitnet_1.58\n"
|
||
)
|
||
cfg = load_config_from_string(yaml_str)
|
||
wrapper = build_bitnet_trainer(cfg, device="cpu")
|
||
assert isinstance(wrapper, BitNetTrainerWrapper)
|
||
|
||
def test_setup_gates_on_onebitllms(self):
|
||
from soup_cli.trainer.bitnet import BitNetTrainerWrapper
|
||
|
||
# onebitllms is not installed on the maintainer's box → friendly gate.
|
||
if "onebitllms" in sys.modules or _spec("onebitllms"):
|
||
pytest.skip("onebitllms installed; gate path not exercised")
|
||
w = object.__new__(BitNetTrainerWrapper)
|
||
with pytest.raises(RuntimeError, match="onebitllms"):
|
||
w.setup({"train": []})
|
||
|
||
|
||
class TestBitnetGgufQuantArg:
|
||
def test_maps_to_tq1_0(self):
|
||
from soup_cli.utils.bitnet import _BITNET_GGUF_QUANT_ARG
|
||
|
||
assert _BITNET_GGUF_QUANT_ARG["bitnet"] == "TQ1_0"
|
||
assert _BITNET_GGUF_QUANT_ARG["tq1_0"] == "TQ1_0"
|
||
|
||
|
||
class TestExportBitnetGguf:
|
||
def test_no_arg_typeerror(self):
|
||
from soup_cli.utils.bitnet import export_bitnet_gguf
|
||
|
||
with pytest.raises(TypeError):
|
||
export_bitnet_gguf()
|
||
|
||
def test_unknown_format_rejected(self, tmp_path):
|
||
from soup_cli.utils.bitnet import export_bitnet_gguf
|
||
|
||
with pytest.raises(ValueError, match="not supported"):
|
||
export_bitnet_gguf(
|
||
model_dir="model",
|
||
output_path="out.gguf",
|
||
export_format="weird",
|
||
llama_cpp_dir="llama",
|
||
)
|
||
|
||
def test_outside_cwd_model_rejected(self, tmp_path):
|
||
from soup_cli.utils.bitnet import export_bitnet_gguf
|
||
|
||
# tmp_path is outside cwd → containment rejection.
|
||
with pytest.raises(ValueError, match="under cwd"):
|
||
export_bitnet_gguf(
|
||
model_dir=str(tmp_path / "model"),
|
||
output_path="out.gguf",
|
||
export_format="bitnet",
|
||
llama_cpp_dir="llama",
|
||
)
|
||
|
||
def test_missing_dir_in_cwd(self, monkeypatch, tmp_path):
|
||
from soup_cli.utils.bitnet import export_bitnet_gguf
|
||
|
||
monkeypatch.chdir(tmp_path)
|
||
# model dir under cwd but does not exist → FileNotFoundError.
|
||
with pytest.raises(FileNotFoundError, match="model_dir"):
|
||
export_bitnet_gguf(
|
||
model_dir="nope",
|
||
output_path="out.gguf",
|
||
export_format="bitnet",
|
||
llama_cpp_dir="llama",
|
||
)
|
||
|
||
def test_no_llama_cpp_dir(self, monkeypatch, tmp_path):
|
||
from soup_cli.utils.bitnet import export_bitnet_gguf
|
||
|
||
monkeypatch.chdir(tmp_path)
|
||
(tmp_path / "model").mkdir()
|
||
with pytest.raises(FileNotFoundError, match="llama_cpp_dir"):
|
||
export_bitnet_gguf(
|
||
model_dir="model",
|
||
output_path="out.gguf",
|
||
export_format="bitnet",
|
||
llama_cpp_dir="missing_llama",
|
||
)
|
||
|
||
@pytest.mark.skipif(sys.platform == "win32", reason="POSIX symlink")
|
||
def test_symlink_model_rejected(self, monkeypatch, tmp_path):
|
||
import os
|
||
|
||
from soup_cli.utils.bitnet import export_bitnet_gguf
|
||
|
||
monkeypatch.chdir(tmp_path)
|
||
target = tmp_path / "outside"
|
||
target.mkdir()
|
||
link = tmp_path / "model_link"
|
||
os.symlink(str(target), str(link))
|
||
with pytest.raises(ValueError, match="symlink"):
|
||
export_bitnet_gguf(
|
||
model_dir="model_link",
|
||
output_path="out.gguf",
|
||
export_format="bitnet",
|
||
llama_cpp_dir="llama",
|
||
)
|
||
|
||
|
||
class TestBitnetExportCli:
|
||
def _runner(self):
|
||
import typer
|
||
from typer.testing import CliRunner
|
||
|
||
from soup_cli.commands.export import export
|
||
|
||
app = typer.Typer()
|
||
app.command()(export)
|
||
return CliRunner(), app
|
||
|
||
def test_help_lists_bitnet(self):
|
||
from soup_cli.commands.export import SUPPORTED_FORMATS
|
||
|
||
assert "bitnet" in SUPPORTED_FORMATS
|
||
assert "tq1_0" in SUPPORTED_FORMATS
|
||
|
||
def test_bitnet_export_no_llama_exits_nonzero(self, monkeypatch, tmp_path):
|
||
runner, app = self._runner()
|
||
monkeypatch.chdir(tmp_path)
|
||
(tmp_path / "model").mkdir()
|
||
result = runner.invoke(
|
||
app,
|
||
["--model", "model", "--format", "bitnet", "--llama-cpp-path", "nope"],
|
||
)
|
||
assert result.exit_code != 0, result.output
|
||
assert "llama" in result.output.lower() or "failed" in result.output.lower()
|
||
|
||
def test_utils_bitnet_no_top_level_torch(self):
|
||
src = (_SRC / "utils" / "bitnet.py").read_text(encoding="utf-8")
|
||
assert "\nimport torch" not in src
|
||
assert "\nfrom torch" not in src
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# #136 MoE expert quant + router-only
|
||
# --------------------------------------------------------------------------- #
|
||
class TestMoeRouterDetection:
|
||
@pytest.mark.parametrize(
|
||
"name,is_router",
|
||
[
|
||
("model.layers.0.block_sparse_moe.gate.weight", True),
|
||
("model.layers.0.block_sparse_moe.gate.bias", True),
|
||
("model.layers.0.mlp.gate.weight", True),
|
||
("model.layers.0.mlp.router.weight", True),
|
||
("model.layers.0.block_sparse_moe.experts.0.gate_proj.weight", False),
|
||
("model.layers.0.block_sparse_moe.experts.0.w1.weight", False),
|
||
("model.layers.0.self_attn.q_proj.weight", False),
|
||
("model.layers.0.mlp.gate_proj.weight", False),
|
||
# an expert sub-projection named 'router' must NOT win (expert wins)
|
||
("model.layers.0.block_sparse_moe.experts.0.router_proj.weight", False),
|
||
],
|
||
)
|
||
def test_is_router_param(self, name, is_router):
|
||
from soup_cli.utils.moe_quant import _is_router_param
|
||
|
||
assert _is_router_param(name) is is_router
|
||
|
||
|
||
class TestMoeExpertHelpers:
|
||
def _fake_moe(self):
|
||
torch = pytest.importorskip("torch")
|
||
import torch.nn as nn
|
||
|
||
class Block(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.self_attn = nn.ModuleDict({"q_proj": nn.Linear(4, 4)})
|
||
experts = nn.ModuleList([
|
||
nn.ModuleDict({
|
||
"w1": nn.Linear(4, 8),
|
||
"w2": nn.Linear(8, 4),
|
||
})
|
||
for _ in range(2)
|
||
])
|
||
self.block_sparse_moe = nn.ModuleDict({
|
||
"gate": nn.Linear(4, 2),
|
||
"experts": experts,
|
||
})
|
||
|
||
class Model(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.layers = nn.ModuleList([Block()])
|
||
|
||
def get_submodule(self, target):
|
||
return dict(self.named_modules())[target]
|
||
|
||
return torch, Model()
|
||
|
||
def test_find_expert_linears(self):
|
||
from soup_cli.utils.moe_quant import _find_expert_linears
|
||
|
||
_torch, model = self._fake_moe()
|
||
found = _find_expert_linears(model)
|
||
names = {n for n, _ in found}
|
||
# 2 experts × 2 linears = 4 expert Linears.
|
||
assert len(found) == 4
|
||
assert all(".experts." in n for n in names)
|
||
# the gate (router) Linear is NOT in the found set.
|
||
assert not any("gate" in n for n in names)
|
||
|
||
def test_find_expert_linears_none(self):
|
||
pytest.importorskip("torch")
|
||
import torch.nn as nn
|
||
|
||
from soup_cli.utils.moe_quant import _find_expert_linears
|
||
|
||
class Plain(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.lin = nn.Linear(4, 4)
|
||
|
||
assert _find_expert_linears(Plain()) == []
|
||
|
||
def test_freeze_experts_train_router(self):
|
||
from soup_cli.utils.moe_quant import freeze_experts_train_router
|
||
|
||
_torch, model = self._fake_moe()
|
||
frozen, trainable = freeze_experts_train_router(model)
|
||
assert frozen > 0
|
||
assert trainable > 0
|
||
for name, param in model.named_parameters():
|
||
if ".experts." in name:
|
||
assert param.requires_grad is False
|
||
elif ".gate." in name:
|
||
assert param.requires_grad is True
|
||
|
||
|
||
class TestApplyMoeExpertQuant:
|
||
def test_unknown_format_rejected(self):
|
||
from soup_cli.utils.moe_quant import apply_moe_expert_quant
|
||
|
||
with pytest.raises(ValueError, match="not supported"):
|
||
apply_moe_expert_quant(object(), "weird")
|
||
|
||
def test_no_experts_returns_zero(self):
|
||
pytest.importorskip("torch")
|
||
import torch.nn as nn
|
||
|
||
from soup_cli.utils.moe_quant import apply_moe_expert_quant
|
||
|
||
class Plain(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.lin = nn.Linear(4, 4)
|
||
|
||
assert apply_moe_expert_quant(Plain(), "nf4") == 0
|
||
|
||
def _moe_with_experts(self):
|
||
import torch.nn as nn
|
||
|
||
class Model(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
experts = nn.ModuleList([
|
||
nn.ModuleDict({"w1": nn.Linear(8, 16), "w2": nn.Linear(16, 8)})
|
||
for _ in range(2)
|
||
])
|
||
self.block_sparse_moe = nn.ModuleDict({
|
||
"gate": nn.Linear(8, 2),
|
||
"experts": experts,
|
||
})
|
||
|
||
def get_submodule(self, target):
|
||
return dict(self.named_modules())[target]
|
||
|
||
return Model()
|
||
|
||
def test_gates_without_bnb_or_cuda(self):
|
||
torch = pytest.importorskip("torch")
|
||
if torch.cuda.is_available() and _spec("bitsandbytes"):
|
||
pytest.skip("CUDA + bitsandbytes present; real quant path used")
|
||
from soup_cli.utils.moe_quant import apply_moe_expert_quant
|
||
|
||
with pytest.raises(RuntimeError, match="bitsandbytes|CUDA"):
|
||
apply_moe_expert_quant(self._moe_with_experts(), "nf4")
|
||
|
||
def test_real_nf4_quant_swaps_experts(self):
|
||
"""When CUDA + bitsandbytes are present (dev box; skipped on CI),
|
||
the real per-expert quant replaces each expert Linear with a bnb
|
||
Linear4bit. Validated live on an RTX 3050 (Step 6)."""
|
||
torch = pytest.importorskip("torch")
|
||
if not (torch.cuda.is_available() and _spec("bitsandbytes")):
|
||
pytest.skip("requires CUDA + bitsandbytes")
|
||
import bitsandbytes as bnb # noqa: PLC0415
|
||
|
||
from soup_cli.utils.moe_quant import apply_moe_expert_quant
|
||
|
||
model = self._moe_with_experts()
|
||
n = apply_moe_expert_quant(model, "nf4")
|
||
assert n == 4 # 2 experts x 2 linears
|
||
# Every expert Linear is now a bnb Linear4bit (a subclass of nn.Linear).
|
||
expert_mods = [
|
||
m for nm, m in model.named_modules() if ".experts." in nm and ".weight" not in nm
|
||
]
|
||
quantized = [m for m in expert_mods if isinstance(m, bnb.nn.Linear4bit)]
|
||
assert len(quantized) == 4
|
||
# The router (gate) is left in full precision.
|
||
gate = model.get_submodule("block_sparse_moe.gate")
|
||
assert not isinstance(gate, bnb.nn.Linear4bit)
|
||
|
||
|
||
class TestApplyMoeFeaturesIfConfigured:
|
||
def _model(self):
|
||
import torch.nn as nn
|
||
|
||
class Model(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
experts = nn.ModuleList([nn.ModuleDict({"w1": nn.Linear(4, 4)})])
|
||
self.block_sparse_moe = nn.ModuleDict({
|
||
"gate": nn.Linear(4, 1),
|
||
"experts": experts,
|
||
})
|
||
|
||
def get_submodule(self, target):
|
||
return dict(self.named_modules())[target]
|
||
|
||
return Model()
|
||
|
||
def test_quant_noop_when_unset(self):
|
||
from soup_cli.utils.moe_quant import apply_moe_expert_quant_if_configured
|
||
|
||
tcfg = SimpleNamespace(moe_expert_quant=None)
|
||
# Should not touch the model (object() would error on any access).
|
||
apply_moe_expert_quant_if_configured(object(), tcfg)
|
||
|
||
def test_freeze_noop_when_unset(self):
|
||
from soup_cli.utils.moe_quant import apply_router_only_freeze_if_configured
|
||
|
||
tcfg = SimpleNamespace(train_router_only=False)
|
||
apply_router_only_freeze_if_configured(object(), tcfg)
|
||
|
||
def test_router_only_freezes_and_prints(self):
|
||
pytest.importorskip("torch")
|
||
|
||
from soup_cli.utils.moe_quant import apply_router_only_freeze_if_configured
|
||
|
||
model = self._model()
|
||
printed = []
|
||
fake_console = SimpleNamespace(print=lambda m: printed.append(m))
|
||
tcfg = SimpleNamespace(train_router_only=True)
|
||
apply_router_only_freeze_if_configured(model, tcfg, fake_console)
|
||
for name, param in model.named_parameters():
|
||
if ".experts." in name:
|
||
assert param.requires_grad is False
|
||
elif ".gate." in name:
|
||
assert param.requires_grad is True
|
||
assert any("router-only" in m for m in printed)
|
||
|
||
def test_quant_prints_when_set(self, monkeypatch):
|
||
pytest.importorskip("torch")
|
||
|
||
from soup_cli.utils import moe_quant
|
||
|
||
model = self._model()
|
||
printed = []
|
||
fake_console = SimpleNamespace(print=lambda m: printed.append(m))
|
||
# Stub the real quant so this runs without CUDA/bnb.
|
||
monkeypatch.setattr(moe_quant, "apply_moe_expert_quant", lambda m, f: 1)
|
||
tcfg = SimpleNamespace(moe_expert_quant="nf4")
|
||
moe_quant.apply_moe_expert_quant_if_configured(model, tcfg, fake_console)
|
||
assert any("expert quant" in m for m in printed)
|
||
|
||
|
||
class TestMoeSchemaAndWiring:
|
||
def test_moe_features_config_loads(self):
|
||
from soup_cli.config.loader import load_config_from_string
|
||
|
||
yaml_str = (
|
||
"base: hf-internal-testing/tiny-random-gpt2\n"
|
||
"task: sft\n"
|
||
"data:\n"
|
||
" train: ./x.jsonl\n"
|
||
" format: auto\n"
|
||
"training:\n"
|
||
" moe_lora: true\n"
|
||
" moe_expert_quant: nf4\n"
|
||
" train_router_only: true\n"
|
||
)
|
||
cfg = load_config_from_string(yaml_str)
|
||
assert cfg.training.moe_expert_quant == "nf4"
|
||
assert cfg.training.train_router_only is True
|
||
|
||
def test_moe_expert_quant_requires_moe_lora(self):
|
||
from soup_cli.config.loader import load_config_from_string
|
||
|
||
yaml_str = (
|
||
"base: hf-internal-testing/tiny-random-gpt2\n"
|
||
"task: sft\n"
|
||
"data:\n"
|
||
" train: ./x.jsonl\n"
|
||
" format: auto\n"
|
||
"training:\n"
|
||
" moe_expert_quant: nf4\n"
|
||
)
|
||
with pytest.raises(ValueError, match="moe_lora"):
|
||
load_config_from_string(yaml_str)
|
||
|
||
def test_sft_wires_moe_features(self):
|
||
src = (_SRC / "trainer" / "sft.py").read_text(encoding="utf-8")
|
||
# Quant runs pre-LoRA, freeze post-LoRA. Assert both wired + ordering.
|
||
assert "apply_moe_expert_quant_if_configured" in src
|
||
assert "apply_router_only_freeze_if_configured" in src
|
||
quant_at = src.index("apply_moe_expert_quant_if_configured(self.model")
|
||
peft_at = src.index("get_peft_model(self.model")
|
||
freeze_at = src.index("apply_router_only_freeze_if_configured(self.model")
|
||
assert quant_at < peft_at < freeze_at
|
||
|
||
def test_validate_moe_expert_quant_compat_matrix(self):
|
||
from soup_cli.utils.moe_quant import validate_moe_expert_quant_compat
|
||
|
||
# happy
|
||
validate_moe_expert_quant_compat(backend="transformers", moe_lora=True)
|
||
with pytest.raises(ValueError, match="mlx"):
|
||
validate_moe_expert_quant_compat(backend="mlx", moe_lora=True)
|
||
with pytest.raises(ValueError, match="moe_lora"):
|
||
validate_moe_expert_quant_compat(backend="transformers", moe_lora=False)
|
||
with pytest.raises(TypeError, match="must not be bool"):
|
||
validate_moe_expert_quant_compat(backend=True, moe_lora=True)
|
||
with pytest.raises(TypeError, match="must be bool"):
|
||
validate_moe_expert_quant_compat(backend="transformers", moe_lora="yes")
|
||
|
||
def test_validate_train_router_only_compat_matrix(self):
|
||
from soup_cli.utils.moe_quant import validate_train_router_only_compat
|
||
|
||
validate_train_router_only_compat(backend="transformers", moe_lora=True)
|
||
with pytest.raises(ValueError, match="mlx"):
|
||
validate_train_router_only_compat(backend="mlx", moe_lora=True)
|
||
with pytest.raises(ValueError, match="moe_lora"):
|
||
validate_train_router_only_compat(backend="transformers", moe_lora=False)
|
||
|
||
def test_utils_moe_quant_no_top_level_torch(self):
|
||
src = (_SRC / "utils" / "moe_quant.py").read_text(encoding="utf-8")
|
||
assert "\nimport torch" not in src
|
||
assert "\nfrom torch" not in src
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# Patch invariants
|
||
# --------------------------------------------------------------------------- #
|
||
class TestPatchInvariants:
|
||
def test_version_bumped(self):
|
||
parts = tuple(int(p) for p in re.findall(r"\d+", __version__)[:3])
|
||
assert parts >= (0, 71, 20)
|
||
|
||
|
||
def _spec(name: str) -> bool:
|
||
import importlib.util
|
||
|
||
try:
|
||
return importlib.util.find_spec(name) is not None
|
||
except (ImportError, ValueError, ModuleNotFoundError):
|
||
return False
|