mirror of https://github.com/razor-ai/soup.git
379 lines
13 KiB
Python
379 lines
13 KiB
Python
"""v0.64.0 Part D — Hardware-fit calculator tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
import math
|
|
|
|
import pytest
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Module imports
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_module_imports():
|
|
from soup_cli.utils import hardware_fit
|
|
|
|
assert hasattr(hardware_fit, "HardwareFitInput")
|
|
assert hasattr(hardware_fit, "VRAMBreakdown")
|
|
assert hasattr(hardware_fit, "HardwareFitReport")
|
|
assert hasattr(hardware_fit, "estimate_peak_vram_gb")
|
|
assert hasattr(hardware_fit, "decide_hardware_fit")
|
|
assert hasattr(hardware_fit, "validate_seq_len")
|
|
assert hasattr(hardware_fit, "validate_batch_size")
|
|
assert hasattr(hardware_fit, "VRAM_SAFETY_MARGIN")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# VRAM_SAFETY_MARGIN
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_safety_margin_is_10pct():
|
|
from soup_cli.utils.hardware_fit import VRAM_SAFETY_MARGIN
|
|
|
|
assert VRAM_SAFETY_MARGIN == pytest.approx(0.10, abs=1e-6)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# validate_seq_len
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize("v", [64, 1024, 8192, 1_048_576])
|
|
def test_validate_seq_len_happy(v):
|
|
from soup_cli.utils.hardware_fit import validate_seq_len
|
|
|
|
assert validate_seq_len(v) == v
|
|
|
|
|
|
@pytest.mark.parametrize("bad", [True, False, "1024", -1, 0, 63, 1_048_577, 1.5])
|
|
def test_validate_seq_len_rejects(bad):
|
|
from soup_cli.utils.hardware_fit import validate_seq_len
|
|
|
|
with pytest.raises((TypeError, ValueError)):
|
|
validate_seq_len(bad)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# validate_batch_size
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize("v", [1, 4, 64, 1024])
|
|
def test_validate_batch_size_happy(v):
|
|
from soup_cli.utils.hardware_fit import validate_batch_size
|
|
|
|
assert validate_batch_size(v) == v
|
|
|
|
|
|
@pytest.mark.parametrize("bad", [True, False, "4", -1, 0, 1025, 1.5])
|
|
def test_validate_batch_size_rejects(bad):
|
|
from soup_cli.utils.hardware_fit import validate_batch_size
|
|
|
|
with pytest.raises((TypeError, ValueError)):
|
|
validate_batch_size(bad)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# HardwareFitInput
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_input_frozen():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput
|
|
|
|
inp = HardwareFitInput(
|
|
params_b=7.0, seq_len=2048, batch_size=4,
|
|
optimizer="adamw_torch", quant="4bit", peft="lora",
|
|
gradient_checkpointing=True,
|
|
)
|
|
with pytest.raises(dataclasses.FrozenInstanceError):
|
|
inp.batch_size = 1 # type: ignore[misc]
|
|
|
|
|
|
def test_input_rejects_negative_params():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput
|
|
|
|
with pytest.raises(ValueError, match="params"):
|
|
HardwareFitInput(
|
|
params_b=-1.0, seq_len=2048, batch_size=4,
|
|
optimizer="adamw_torch", quant="4bit", peft="lora",
|
|
gradient_checkpointing=True,
|
|
)
|
|
|
|
|
|
def test_input_rejects_invalid_quant():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput
|
|
|
|
with pytest.raises(ValueError, match="quant"):
|
|
HardwareFitInput(
|
|
params_b=7.0, seq_len=2048, batch_size=4,
|
|
optimizer="adamw_torch", quant="bogus", peft="lora",
|
|
gradient_checkpointing=True,
|
|
)
|
|
|
|
|
|
def test_input_rejects_invalid_peft():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput
|
|
|
|
with pytest.raises(ValueError, match="peft"):
|
|
HardwareFitInput(
|
|
params_b=7.0, seq_len=2048, batch_size=4,
|
|
optimizer="adamw_torch", quant="4bit", peft="bogus",
|
|
gradient_checkpointing=True,
|
|
)
|
|
|
|
|
|
def test_input_rejects_bool_params():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput
|
|
|
|
with pytest.raises(TypeError, match="bool"):
|
|
HardwareFitInput(
|
|
params_b=True, seq_len=2048, batch_size=4, # type: ignore[arg-type]
|
|
optimizer="adamw_torch", quant="4bit", peft="lora",
|
|
gradient_checkpointing=True,
|
|
)
|
|
|
|
|
|
def test_input_rejects_non_bool_gradient_ckpt():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput
|
|
|
|
with pytest.raises(TypeError, match="gradient_checkpointing"):
|
|
HardwareFitInput(
|
|
params_b=7.0, seq_len=2048, batch_size=4,
|
|
optimizer="adamw_torch", quant="4bit", peft="lora",
|
|
gradient_checkpointing=1, # type: ignore[arg-type]
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# VRAMBreakdown
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_vram_breakdown_frozen():
|
|
from soup_cli.utils.hardware_fit import VRAMBreakdown
|
|
|
|
v = VRAMBreakdown(
|
|
weights_gb=2.0,
|
|
optimizer_gb=0.5,
|
|
gradients_gb=0.5,
|
|
activations_gb=1.0,
|
|
overhead_gb=0.5,
|
|
)
|
|
with pytest.raises(dataclasses.FrozenInstanceError):
|
|
v.weights_gb = 0.0 # type: ignore[misc]
|
|
|
|
|
|
def test_vram_breakdown_rejects_negative():
|
|
from soup_cli.utils.hardware_fit import VRAMBreakdown
|
|
|
|
with pytest.raises(ValueError, match="negative"):
|
|
VRAMBreakdown(
|
|
weights_gb=-1.0,
|
|
optimizer_gb=0.5,
|
|
gradients_gb=0.5,
|
|
activations_gb=1.0,
|
|
overhead_gb=0.5,
|
|
)
|
|
|
|
|
|
def test_vram_breakdown_total():
|
|
from soup_cli.utils.hardware_fit import VRAMBreakdown
|
|
|
|
v = VRAMBreakdown(
|
|
weights_gb=2.0,
|
|
optimizer_gb=0.5,
|
|
gradients_gb=0.5,
|
|
activations_gb=1.0,
|
|
overhead_gb=0.5,
|
|
)
|
|
assert v.total_gb == pytest.approx(4.5)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# estimate_peak_vram_gb
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_estimate_peak_vram_returns_breakdown():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput, VRAMBreakdown, estimate_peak_vram_gb
|
|
|
|
inp = HardwareFitInput(
|
|
params_b=7.0, seq_len=2048, batch_size=4,
|
|
optimizer="adamw_torch", quant="4bit", peft="lora",
|
|
gradient_checkpointing=True,
|
|
)
|
|
breakdown = estimate_peak_vram_gb(inp)
|
|
assert isinstance(breakdown, VRAMBreakdown)
|
|
assert breakdown.weights_gb > 0
|
|
assert breakdown.total_gb > 0
|
|
|
|
|
|
def test_estimate_peak_vram_4bit_smaller_than_fp16():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput, estimate_peak_vram_gb
|
|
|
|
base = dict(params_b=7.0, seq_len=1024, batch_size=1,
|
|
optimizer="adamw_torch", peft="lora",
|
|
gradient_checkpointing=True)
|
|
fp16 = estimate_peak_vram_gb(HardwareFitInput(quant="none", **base))
|
|
q4 = estimate_peak_vram_gb(HardwareFitInput(quant="4bit", **base))
|
|
assert q4.weights_gb < fp16.weights_gb
|
|
|
|
|
|
def test_estimate_peak_vram_lora_smaller_than_full():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput, estimate_peak_vram_gb
|
|
|
|
base = dict(params_b=7.0, seq_len=1024, batch_size=1,
|
|
optimizer="adamw_torch", quant="4bit",
|
|
gradient_checkpointing=True)
|
|
full = estimate_peak_vram_gb(HardwareFitInput(peft="full", **base))
|
|
lora = estimate_peak_vram_gb(HardwareFitInput(peft="lora", **base))
|
|
assert lora.optimizer_gb < full.optimizer_gb
|
|
|
|
|
|
def test_estimate_peak_vram_seq_len_scales_activations():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput, estimate_peak_vram_gb
|
|
|
|
base = dict(params_b=1.0, batch_size=1,
|
|
optimizer="adamw_torch", quant="4bit", peft="lora",
|
|
gradient_checkpointing=False)
|
|
short = estimate_peak_vram_gb(HardwareFitInput(seq_len=512, **base))
|
|
long_ = estimate_peak_vram_gb(HardwareFitInput(seq_len=8192, **base))
|
|
assert long_.activations_gb > short.activations_gb
|
|
|
|
|
|
def test_estimate_peak_vram_grad_ckpt_reduces_activations():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput, estimate_peak_vram_gb
|
|
|
|
base = dict(params_b=7.0, seq_len=4096, batch_size=1,
|
|
optimizer="adamw_torch", quant="4bit", peft="lora")
|
|
off = estimate_peak_vram_gb(HardwareFitInput(gradient_checkpointing=False, **base))
|
|
on = estimate_peak_vram_gb(HardwareFitInput(gradient_checkpointing=True, **base))
|
|
assert on.activations_gb < off.activations_gb
|
|
|
|
|
|
def test_estimate_peak_vram_finite():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput, estimate_peak_vram_gb
|
|
|
|
inp = HardwareFitInput(
|
|
params_b=70.0, seq_len=8192, batch_size=8,
|
|
optimizer="adamw_torch", quant="none", peft="full",
|
|
gradient_checkpointing=False,
|
|
)
|
|
bd = estimate_peak_vram_gb(inp)
|
|
assert math.isfinite(bd.total_gb)
|
|
|
|
|
|
def test_estimate_peak_vram_rejects_non_input():
|
|
from soup_cli.utils.hardware_fit import estimate_peak_vram_gb
|
|
|
|
with pytest.raises(TypeError):
|
|
estimate_peak_vram_gb("not an input") # type: ignore[arg-type]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# decide_hardware_fit
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_decide_hardware_fit_ok(tmp_path):
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput, decide_hardware_fit
|
|
|
|
inp = HardwareFitInput(
|
|
params_b=1.0, seq_len=1024, batch_size=1,
|
|
optimizer="adamw_torch", quant="4bit", peft="lora",
|
|
gradient_checkpointing=True,
|
|
)
|
|
# Plenty of headroom
|
|
report = decide_hardware_fit(inp, available_vram_gb=24.0)
|
|
assert report.ok is True
|
|
assert report.peak_vram_gb > 0
|
|
assert report.required_with_margin_gb >= report.peak_vram_gb
|
|
|
|
|
|
def test_decide_hardware_fit_oom():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput, decide_hardware_fit
|
|
|
|
inp = HardwareFitInput(
|
|
params_b=70.0, seq_len=8192, batch_size=8,
|
|
optimizer="adamw_torch", quant="none", peft="full",
|
|
gradient_checkpointing=False,
|
|
)
|
|
report = decide_hardware_fit(inp, available_vram_gb=8.0)
|
|
assert report.ok is False
|
|
assert "exceed" in report.reason.lower() or "oom" in report.reason.lower() or \
|
|
"available" in report.reason.lower()
|
|
|
|
|
|
def test_decide_hardware_fit_rejects_negative_vram():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput, decide_hardware_fit
|
|
|
|
inp = HardwareFitInput(
|
|
params_b=1.0, seq_len=1024, batch_size=1,
|
|
optimizer="adamw_torch", quant="4bit", peft="lora",
|
|
gradient_checkpointing=True,
|
|
)
|
|
with pytest.raises(ValueError, match="vram"):
|
|
decide_hardware_fit(inp, available_vram_gb=-1.0)
|
|
|
|
|
|
def test_decide_hardware_fit_rejects_bool_vram():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput, decide_hardware_fit
|
|
|
|
inp = HardwareFitInput(
|
|
params_b=1.0, seq_len=1024, batch_size=1,
|
|
optimizer="adamw_torch", quant="4bit", peft="lora",
|
|
gradient_checkpointing=True,
|
|
)
|
|
with pytest.raises(TypeError, match="bool"):
|
|
decide_hardware_fit(inp, available_vram_gb=True) # type: ignore[arg-type]
|
|
|
|
|
|
def test_decide_hardware_fit_rejects_non_finite_vram():
|
|
from soup_cli.utils.hardware_fit import HardwareFitInput, decide_hardware_fit
|
|
|
|
inp = HardwareFitInput(
|
|
params_b=1.0, seq_len=1024, batch_size=1,
|
|
optimizer="adamw_torch", quant="4bit", peft="lora",
|
|
gradient_checkpointing=True,
|
|
)
|
|
with pytest.raises(ValueError, match="finite"):
|
|
decide_hardware_fit(inp, available_vram_gb=float("nan"))
|
|
|
|
|
|
def test_hardware_fit_report_frozen():
|
|
from soup_cli.utils.hardware_fit import HardwareFitReport, VRAMBreakdown
|
|
|
|
bd = VRAMBreakdown(
|
|
weights_gb=1.0, optimizer_gb=0.5, gradients_gb=0.5,
|
|
activations_gb=1.0, overhead_gb=0.5,
|
|
)
|
|
rep = HardwareFitReport(
|
|
ok=True,
|
|
peak_vram_gb=3.5,
|
|
required_with_margin_gb=3.85,
|
|
available_vram_gb=24.0,
|
|
breakdown=bd,
|
|
reason="ok",
|
|
)
|
|
with pytest.raises(dataclasses.FrozenInstanceError):
|
|
rep.ok = False # type: ignore[misc]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Source-wiring regression
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_no_heavy_top_level_imports():
|
|
from pathlib import Path
|
|
|
|
src = Path(__file__).resolve().parent.parent / "src" / "soup_cli" / "utils" / "hardware_fit.py"
|
|
text = src.read_text(encoding="utf-8")
|
|
import re
|
|
for bad in ["^import torch", "^from torch", "^import transformers", "^from transformers"]:
|
|
assert not re.search(bad, text, re.MULTILINE)
|