mirror of https://github.com/razor-ai/soup.git
236 lines
7.4 KiB
Python
236 lines
7.4 KiB
Python
"""Tests for v0.37.0 Part B — multipack config wiring + sampler builder.
|
|
|
|
Covers:
|
|
- ``TrainingConfig.multipack`` Pydantic field default + type
|
|
- Cross-validator: ``multipack`` and ``packing`` are mutually exclusive
|
|
- Cross-validator: SoupConfig restricts ``multipack`` to sft / pretrain
|
|
- Cross-validator: MLX backend rejects multipack (sampler injection is HF Trainer-specific)
|
|
- ``build_multipack_sampler_for_lengths`` helper — returns a configured
|
|
:class:`MultipackBatchSampler` from a list of sample lengths
|
|
- ``supports_multipack`` task allowlist
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
from soup_cli.config.schema import SoupConfig, TrainingConfig
|
|
from soup_cli.utils.multipack import (
|
|
build_multipack_sampler_for_lengths,
|
|
supports_multipack,
|
|
)
|
|
from soup_cli.utils.multipack_sampler import MultipackBatchSampler
|
|
|
|
# ---- TrainingConfig.multipack field --------------------------------------
|
|
|
|
|
|
def test_multipack_default_false():
|
|
tcfg = TrainingConfig()
|
|
assert tcfg.multipack is False
|
|
|
|
|
|
def test_multipack_accepts_true():
|
|
tcfg = TrainingConfig(multipack=True)
|
|
assert tcfg.multipack is True
|
|
|
|
|
|
def test_multipack_rejects_non_bool():
|
|
# Pydantic v2 coerces "true"/"false" strings to bool, but rejects
|
|
# arbitrary objects. A list cannot be coerced to bool.
|
|
with pytest.raises(ValidationError):
|
|
TrainingConfig(multipack=[1, 2]) # type: ignore[arg-type]
|
|
|
|
|
|
# ---- mutually exclusive with packing -------------------------------------
|
|
|
|
|
|
def test_multipack_packing_mutually_exclusive():
|
|
with pytest.raises(ValidationError, match="mutually exclusive"):
|
|
TrainingConfig(multipack=True, packing=True)
|
|
|
|
|
|
def test_multipack_alone_ok():
|
|
tcfg = TrainingConfig(multipack=True, packing=False)
|
|
assert tcfg.multipack is True
|
|
assert tcfg.packing is False
|
|
|
|
|
|
def test_packing_alone_ok():
|
|
tcfg = TrainingConfig(packing=True, multipack=False)
|
|
assert tcfg.packing is True
|
|
assert tcfg.multipack is False
|
|
|
|
|
|
# ---- SoupConfig task gate ------------------------------------------------
|
|
|
|
|
|
def _base_soup_kwargs(task: str = "sft", **overrides):
|
|
cfg = {
|
|
"base": "fake-org/fake-model",
|
|
"task": task,
|
|
"data": {"train": "data.jsonl", "format": "alpaca"},
|
|
"training": {"epochs": 1, "lr": 1e-4, "multipack": True},
|
|
"output": "./out",
|
|
}
|
|
cfg.update(overrides)
|
|
return cfg
|
|
|
|
|
|
def test_multipack_allowed_for_sft():
|
|
cfg = SoupConfig(**_base_soup_kwargs(task="sft"))
|
|
assert cfg.training.multipack is True
|
|
|
|
|
|
def test_multipack_allowed_for_pretrain():
|
|
kwargs = _base_soup_kwargs(task="pretrain")
|
|
kwargs["data"]["format"] = "plaintext"
|
|
cfg = SoupConfig(**kwargs)
|
|
assert cfg.training.multipack is True
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"task",
|
|
["dpo", "grpo", "kto", "orpo", "simpo", "ipo", "ppo",
|
|
"reward_model", "embedding"],
|
|
)
|
|
def test_multipack_rejected_for_non_sft_pretrain(task):
|
|
kwargs = _base_soup_kwargs(task=task)
|
|
# Adjust data format to satisfy each task's data validator before our
|
|
# multipack guard fires.
|
|
if task in {"dpo", "kto", "orpo", "simpo", "ipo"}:
|
|
kwargs["data"]["format"] = "dpo"
|
|
elif task == "embedding":
|
|
kwargs["data"]["format"] = "embedding"
|
|
elif task == "reward_model":
|
|
kwargs["data"]["format"] = "dpo"
|
|
with pytest.raises(ValidationError, match="multipack"):
|
|
SoupConfig(**kwargs)
|
|
|
|
|
|
def test_multipack_off_does_not_trip_task_gate():
|
|
# multipack=False on a non-sft task should NOT raise.
|
|
kwargs = _base_soup_kwargs(task="dpo")
|
|
kwargs["training"]["multipack"] = False
|
|
kwargs["data"]["format"] = "dpo"
|
|
cfg = SoupConfig(**kwargs)
|
|
assert cfg.training.multipack is False
|
|
|
|
|
|
def test_multipack_rejected_on_mlx_backend():
|
|
kwargs = _base_soup_kwargs(task="sft")
|
|
kwargs["backend"] = "mlx"
|
|
with pytest.raises(ValidationError, match="mlx"):
|
|
SoupConfig(**kwargs)
|
|
|
|
|
|
# ---- supports_multipack helper -------------------------------------------
|
|
|
|
|
|
def test_supports_multipack_allowed_tasks():
|
|
assert supports_multipack("sft") is True
|
|
assert supports_multipack("pretrain") is True
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"task",
|
|
["dpo", "grpo", "kto", "orpo", "simpo", "ipo", "ppo",
|
|
"reward_model", "embedding"],
|
|
)
|
|
def test_supports_multipack_rejected_tasks(task):
|
|
assert supports_multipack(task) is False
|
|
|
|
|
|
def test_supports_multipack_unknown_task():
|
|
assert supports_multipack("nonexistent_task") is False
|
|
|
|
|
|
# ---- build_multipack_sampler_for_lengths ---------------------------------
|
|
|
|
|
|
def test_build_sampler_returns_multipack_sampler():
|
|
tcfg = TrainingConfig(
|
|
multipack=True, batch_size=2, packing=False,
|
|
)
|
|
lengths = [3, 5, 2, 4, 1, 6]
|
|
sampler = build_multipack_sampler_for_lengths(
|
|
lengths=lengths, tcfg=tcfg, max_seq_length=10, seed=0,
|
|
)
|
|
assert isinstance(sampler, MultipackBatchSampler)
|
|
|
|
|
|
def test_build_sampler_real_batches_uses_batch_size():
|
|
tcfg = TrainingConfig(
|
|
multipack=True, batch_size=4, packing=False,
|
|
)
|
|
lengths = [3] * 16
|
|
sampler = build_multipack_sampler_for_lengths(
|
|
lengths=lengths, tcfg=tcfg, max_seq_length=12,
|
|
real_batches=True, seed=0,
|
|
)
|
|
for batch in sampler:
|
|
assert len(batch) <= 4
|
|
|
|
|
|
def test_build_sampler_flat_mode():
|
|
# real_batches=False yields flat index lists, max_len = batch_size * max_seq_length
|
|
tcfg = TrainingConfig(
|
|
multipack=True, batch_size=4, packing=False,
|
|
)
|
|
lengths = [10, 8, 6, 4, 2]
|
|
sampler = build_multipack_sampler_for_lengths(
|
|
lengths=lengths, tcfg=tcfg, max_seq_length=8,
|
|
real_batches=False, seed=0,
|
|
)
|
|
# max bin len in flat mode = 4 * 8 = 32, so total of all lengths (=30)
|
|
# should fit in one bin given FFD.
|
|
batches = list(sampler)
|
|
assert len(batches) == 1, "expected single bin given budget"
|
|
flat = sorted(idx for batch in batches for idx in batch)
|
|
assert flat == list(range(5))
|
|
|
|
|
|
def test_build_sampler_requires_multipack_enabled():
|
|
tcfg = TrainingConfig(multipack=False)
|
|
with pytest.raises(ValueError, match="multipack"):
|
|
build_multipack_sampler_for_lengths(
|
|
lengths=[3, 4], tcfg=tcfg, max_seq_length=10, seed=0,
|
|
)
|
|
|
|
|
|
def test_build_sampler_rejects_non_positive_max_seq_length():
|
|
tcfg = TrainingConfig(multipack=True)
|
|
with pytest.raises(ValueError, match="max_seq_length"):
|
|
build_multipack_sampler_for_lengths(
|
|
lengths=[3, 4], tcfg=tcfg, max_seq_length=0, seed=0,
|
|
)
|
|
|
|
|
|
def test_build_sampler_rejects_bool_max_seq_length():
|
|
# bool is subclass of int — reject explicitly per v0.30.0+ policy.
|
|
tcfg = TrainingConfig(multipack=True)
|
|
with pytest.raises(TypeError, match="bool"):
|
|
build_multipack_sampler_for_lengths(
|
|
lengths=[3, 4], tcfg=tcfg, max_seq_length=True, seed=0,
|
|
)
|
|
|
|
|
|
def test_build_sampler_rejects_auto_batch_size():
|
|
tcfg = TrainingConfig(multipack=True, batch_size="auto")
|
|
with pytest.raises(ValueError, match="auto"):
|
|
build_multipack_sampler_for_lengths(
|
|
lengths=[3, 4], tcfg=tcfg, max_seq_length=10, seed=0,
|
|
)
|
|
|
|
|
|
def test_build_sampler_seed_determinism():
|
|
tcfg = TrainingConfig(multipack=True, batch_size=2)
|
|
lengths = [3, 5, 2, 4, 1, 6, 7, 2]
|
|
s1 = build_multipack_sampler_for_lengths(
|
|
lengths=lengths, tcfg=tcfg, max_seq_length=10, seed=42,
|
|
)
|
|
s2 = build_multipack_sampler_for_lengths(
|
|
lengths=lengths, tcfg=tcfg, max_seq_length=10, seed=42,
|
|
)
|
|
assert list(s1) == list(s2)
|