mirror of https://github.com/razor-ai/soup.git
354 lines
13 KiB
Python
354 lines
13 KiB
Python
"""Part C — v0.28.1 speed/memory live (#43, #44, #47) for v0.33.0.
|
|
|
|
Covers:
|
|
- #43 Multi-trainer wiring: apply_v028_speed_memory helper +
|
|
supports_v028_features extension to dpo + pretrain. Schema gate
|
|
softens to allow sft / dpo / pretrain.
|
|
- #44 install_selective_hooks: per-layer checkpoint hooks across the
|
|
three granularities, with a fake transformer-shaped model.
|
|
- #47 CrossDocCollator: block-diagonal mask injection via the
|
|
underlying build_cross_doc_mask helper.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# #43 — apply_v028_speed_memory helper
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestApplyV028SpeedMemory:
|
|
def test_no_features_returns_all_false(self):
|
|
from soup_cli.utils.v028_features import apply_v028_speed_memory
|
|
|
|
tcfg = SimpleNamespace(
|
|
use_cut_ce=False, quantization_aware=False,
|
|
kernel_auto_compose=False,
|
|
)
|
|
result = apply_v028_speed_memory(
|
|
model=MagicMock(), tcfg=tcfg, base_model="x/y",
|
|
)
|
|
assert result == {
|
|
"cut_ce": False, "fp8": False, "kernel_auto_compose": False,
|
|
}
|
|
|
|
def test_cut_ce_failure_logs_yellow(self, monkeypatch, capsys):
|
|
from soup_cli.utils import v028_features as vf
|
|
|
|
# Force apply_cut_ce import to raise
|
|
def _boom(_name):
|
|
raise RuntimeError("no cut_cross_entropy")
|
|
|
|
monkeypatch.setattr(
|
|
"soup_cli.utils.cut_ce.apply_cut_ce", _boom, raising=False,
|
|
)
|
|
from rich.console import Console
|
|
console = Console()
|
|
tcfg = SimpleNamespace(
|
|
use_cut_ce=True, quantization_aware=False,
|
|
kernel_auto_compose=False,
|
|
)
|
|
result = vf.apply_v028_speed_memory(
|
|
model=MagicMock(), tcfg=tcfg,
|
|
base_model="x/y", console=console,
|
|
)
|
|
assert result["cut_ce"] is False
|
|
|
|
def test_supports_v028_features_extends_to_all_transformer_trainers(self):
|
|
"""v0.35.0 #60 expanded support from {sft, dpo, pretrain} to every
|
|
transformer-backend trainer."""
|
|
from soup_cli.utils.v028_features import supports_v028_features
|
|
|
|
for task in (
|
|
"sft", "dpo", "pretrain", "grpo", "kto", "orpo",
|
|
"simpo", "ipo", "ppo", "reward_model", "embedding",
|
|
):
|
|
assert supports_v028_features(task) is True
|
|
# Unknown / future tasks default to False
|
|
assert supports_v028_features("nonexistent") is False
|
|
|
|
def test_warn_unsupported_returns_none_for_supported(self):
|
|
from soup_cli.utils.v028_features import warn_unsupported_features
|
|
|
|
tcfg = SimpleNamespace(
|
|
use_cut_ce=True, quantization_aware="fp8",
|
|
kernel_auto_compose=True, activation_offloading="cpu",
|
|
)
|
|
for task in (
|
|
"sft", "dpo", "pretrain", "grpo", "kto", "orpo",
|
|
"simpo", "ipo", "ppo", "reward_model", "embedding",
|
|
):
|
|
assert warn_unsupported_features(tcfg, task) is None
|
|
|
|
def test_warn_unsupported_lists_offenders_for_unknown_task(self):
|
|
from soup_cli.utils.v028_features import warn_unsupported_features
|
|
|
|
tcfg = SimpleNamespace(
|
|
use_cut_ce=True, quantization_aware="fp8",
|
|
kernel_auto_compose=False, activation_offloading=None,
|
|
)
|
|
msg = warn_unsupported_features(tcfg, "future_task")
|
|
assert msg is not None
|
|
assert "use_cut_ce" in msg
|
|
assert "fp8" in msg
|
|
|
|
|
|
class TestSchemaGateExpanded:
|
|
def _config(self, task: str, **training_extra):
|
|
import yaml
|
|
|
|
from soup_cli.config.loader import load_config_from_string
|
|
|
|
body = {
|
|
"base": "test/model",
|
|
"task": task,
|
|
"data": {"train": "data.jsonl", "format": "alpaca"}
|
|
if task != "pretrain"
|
|
else {"train": "data.jsonl", "format": "plaintext"},
|
|
"training": {"epochs": 1, "lr": 1e-4, "batch_size": 1, **training_extra},
|
|
}
|
|
if task in ("dpo", "kto", "orpo", "simpo", "ipo", "grpo"):
|
|
body["data"]["format"] = "dpo"
|
|
return load_config_from_string(yaml.safe_dump(body))
|
|
|
|
def test_dpo_now_accepts_v028_features(self):
|
|
cfg = self._config("dpo", use_cut_ce=True)
|
|
assert cfg.task == "dpo"
|
|
assert cfg.training.use_cut_ce is True
|
|
|
|
def test_pretrain_now_accepts_v028_features(self):
|
|
cfg = self._config("pretrain", use_cut_ce=True)
|
|
assert cfg.task == "pretrain"
|
|
|
|
def test_kto_now_accepts_v028_features(self):
|
|
"""v0.35.0 #60 lifted the SFT-only schema gate; KTO now accepts."""
|
|
cfg = self._config("kto", use_cut_ce=True)
|
|
assert cfg.task == "kto"
|
|
assert cfg.training.use_cut_ce is True
|
|
|
|
def test_grpo_now_accepts_v028_features(self):
|
|
cfg = self._config("grpo", quantization_aware="fp8")
|
|
assert cfg.task == "grpo"
|
|
assert cfg.training.quantization_aware == "fp8"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# #44 — install_selective_hooks
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestInstallSelectiveHooks:
|
|
def test_unknown_granularity_rejected(self):
|
|
from soup_cli.utils.gradient_ckpt import install_selective_hooks
|
|
|
|
with pytest.raises(ValueError, match="must be one of"):
|
|
install_selective_hooks(MagicMock(), "weird")
|
|
|
|
def test_full_wraps_every_block(self):
|
|
"""A 4-block fake model with 'medium' wraps 2 blocks (every other)."""
|
|
from soup_cli.utils.gradient_ckpt import install_selective_hooks
|
|
|
|
class FakeBlock:
|
|
def __init__(self, idx):
|
|
self.idx = idx
|
|
|
|
def forward(self, *args, **kwargs):
|
|
return ("orig", self.idx)
|
|
|
|
# Simulate having attention sub-modules
|
|
def named_modules(self, prefix=""):
|
|
yield (f"{prefix}", self)
|
|
yield (f"{prefix}.self_attn", _AttnStub(self.idx))
|
|
|
|
class _AttnStub:
|
|
def __init__(self, idx):
|
|
self.idx = idx
|
|
|
|
def forward(self, *args, **kwargs):
|
|
return ("attn", self.idx)
|
|
|
|
class FakeModel:
|
|
def __init__(self):
|
|
self.blocks = [FakeBlock(i) for i in range(4)]
|
|
|
|
def named_modules(self):
|
|
# HF-style names: ``model.layers.<i>``
|
|
for idx, blk in enumerate(self.blocks):
|
|
yield (f"model.layers.{idx}", blk)
|
|
for child_name, child in blk.named_modules(
|
|
prefix=f"model.layers.{idx}",
|
|
):
|
|
if child_name != f"model.layers.{idx}":
|
|
yield (child_name, child)
|
|
|
|
# Full granularity wraps every numbered child
|
|
model = FakeModel()
|
|
hooked_full = install_selective_hooks(model, "full")
|
|
assert hooked_full == 4
|
|
|
|
def test_medium_wraps_every_other_block(self):
|
|
from soup_cli.utils.gradient_ckpt import install_selective_hooks
|
|
|
|
class FakeBlock:
|
|
def forward(self, *args, **kwargs):
|
|
return None
|
|
|
|
def named_modules(self, prefix=""):
|
|
yield (prefix, self)
|
|
|
|
class FakeModel:
|
|
def __init__(self, n):
|
|
self.blocks = [FakeBlock() for _ in range(n)]
|
|
|
|
def named_modules(self):
|
|
for i, blk in enumerate(self.blocks):
|
|
yield (f"model.layers.{i}", blk)
|
|
|
|
hooked = install_selective_hooks(FakeModel(6), "medium")
|
|
# 6 blocks, every-other → 3
|
|
assert hooked == 3
|
|
|
|
def test_selective_wraps_attention_only(self):
|
|
from soup_cli.utils.gradient_ckpt import install_selective_hooks
|
|
|
|
class _Attn:
|
|
def forward(self, *args, **kwargs):
|
|
return None
|
|
|
|
class FakeBlock:
|
|
def __init__(self):
|
|
self._attn = _Attn()
|
|
self._mlp = MagicMock()
|
|
self._mlp.forward = lambda *_a, **_k: None
|
|
|
|
def forward(self, *args, **kwargs):
|
|
return None
|
|
|
|
def named_modules(self, prefix=""):
|
|
yield (prefix, self)
|
|
yield (f"{prefix}.self_attn", self._attn)
|
|
yield (f"{prefix}.mlp", self._mlp)
|
|
|
|
class FakeModel:
|
|
def __init__(self):
|
|
self.blocks = [FakeBlock() for _ in range(3)]
|
|
|
|
def named_modules(self):
|
|
for i, blk in enumerate(self.blocks):
|
|
yield (f"model.layers.{i}", blk)
|
|
for child_name, child in blk.named_modules(
|
|
prefix=f"model.layers.{i}",
|
|
):
|
|
if child_name != f"model.layers.{i}":
|
|
yield (child_name, child)
|
|
|
|
hooked = install_selective_hooks(FakeModel(), "selective")
|
|
# 3 blocks, each with one attention child = 3 hooks
|
|
assert hooked == 3
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# #47 — CrossDocCollator
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCrossDocCollator:
|
|
def test_requires_base_collator(self):
|
|
from soup_cli.data.collators import CrossDocCollator
|
|
|
|
with pytest.raises(ValueError, match="base_collator"):
|
|
CrossDocCollator(base_collator=None)
|
|
|
|
def test_passes_through_when_no_doc_lengths(self):
|
|
from soup_cli.data.collators import CrossDocCollator
|
|
|
|
base = MagicMock(return_value={"input_ids": MagicMock()})
|
|
collator = CrossDocCollator(base_collator=base)
|
|
result = collator([{"input_ids": [1, 2, 3]}])
|
|
assert "input_ids" in result
|
|
# No cross_doc_attn_mask since no doc_lengths supplied
|
|
assert "cross_doc_attn_mask" not in result
|
|
|
|
def test_strips_doc_lengths_before_base_call(self):
|
|
from soup_cli.data.collators import CrossDocCollator
|
|
|
|
captured: list[dict] = []
|
|
|
|
def _base(features):
|
|
captured.extend(features)
|
|
return {"input_ids": MagicMock()}
|
|
|
|
collator = CrossDocCollator(base_collator=_base)
|
|
collator([{"input_ids": [1, 2], "doc_lengths": [1, 1]}])
|
|
# base collator must NOT see the doc_lengths key
|
|
assert "doc_lengths" not in captured[0]
|
|
|
|
def test_mismatched_doc_lengths_falls_back_to_causal(self):
|
|
"""If sum(doc_lengths) > seq_length the collator must NOT crash —
|
|
it falls back to a plain lower-triangular mask."""
|
|
from soup_cli.data.collators import CrossDocCollator
|
|
|
|
seq_len = 4
|
|
|
|
class _Tensor:
|
|
shape = (1, seq_len)
|
|
|
|
def _base(features):
|
|
return {"input_ids": _Tensor()}
|
|
|
|
collator = CrossDocCollator(base_collator=_base)
|
|
# 3 + 3 = 6 > seq_len 4 → fallback path
|
|
result = collator([{"doc_lengths": [3, 3]}])
|
|
# Either no mask emitted (graceful skip) or fallback mask present.
|
|
# Either is acceptable; what matters is no exception.
|
|
assert "input_ids" in result
|
|
|
|
def test_does_not_mutate_input_dict(self):
|
|
"""Regression: collator pops doc_lengths in v0.33.0 wave; HF Dataset
|
|
rows are cached, so mutation breaks subsequent batches."""
|
|
from soup_cli.data.collators import CrossDocCollator
|
|
|
|
def _base(features):
|
|
return {"input_ids": object()}
|
|
|
|
collator = CrossDocCollator(base_collator=_base)
|
|
original = {"doc_lengths": [2, 2], "input_ids": [1, 2, 3, 4]}
|
|
collator([original])
|
|
# The original dict must still have doc_lengths after the call.
|
|
assert original.get("doc_lengths") == [2, 2]
|
|
|
|
def test_injects_block_diag_mask_with_doc_lengths(self):
|
|
"""End-to-end: when doc_lengths are present, cross_doc_attn_mask
|
|
appears on the batch with the right shape."""
|
|
import numpy as np
|
|
|
|
from soup_cli.data.collators import CrossDocCollator
|
|
|
|
# 2 documents, each 2 tokens long, packed into seq_length=4
|
|
seq_len = 4
|
|
|
|
class _Tensor:
|
|
shape = (1, seq_len)
|
|
|
|
def _base(features):
|
|
return {"input_ids": _Tensor()}
|
|
|
|
collator = CrossDocCollator(base_collator=_base)
|
|
result = collator([{"doc_lengths": [2, 2]}])
|
|
|
|
assert "cross_doc_attn_mask" in result
|
|
mask = result["cross_doc_attn_mask"]
|
|
# Shape: (batch=1, seq, seq)
|
|
assert mask.shape == (1, seq_len, seq_len)
|
|
# Token in doc 0 should NOT attend to token in doc 1 (positions 0,1 vs 2,3)
|
|
np_mask = mask[0].numpy() if hasattr(mask[0], "numpy") else np.array(mask[0])
|
|
# Position 0 attending to position 2 → should be 0
|
|
assert np_mask[0, 2] == 0
|
|
# Position 2 attending to position 2 → 1 (causal within doc 1)
|
|
assert np_mask[2, 2] == 1
|