mirror of https://github.com/razor-ai/soup.git
597 lines
20 KiB
Python
597 lines
20 KiB
Python
"""Tests for soup profile command — memory/speed estimator."""
|
|
|
|
import json
|
|
|
|
import pytest
|
|
from typer.testing import CliRunner
|
|
|
|
from soup_cli.cli import app
|
|
|
|
runner = CliRunner()
|
|
|
|
|
|
# --- Memory estimation tests ---
|
|
|
|
|
|
class TestEstimateModelMemory:
|
|
"""Test model memory estimation for various sizes and quantizations."""
|
|
|
|
def test_7b_4bit(self):
|
|
from soup_cli.utils.profiler import estimate_model_memory
|
|
|
|
mem_gb = estimate_model_memory(7.0, "4bit")
|
|
assert 3.0 <= mem_gb <= 5.0 # ~3.5 GB for 7B 4bit
|
|
|
|
def test_7b_8bit(self):
|
|
from soup_cli.utils.profiler import estimate_model_memory
|
|
|
|
mem_gb = estimate_model_memory(7.0, "8bit")
|
|
assert 6.0 <= mem_gb <= 9.0 # ~7 GB for 7B 8bit
|
|
|
|
def test_7b_none(self):
|
|
from soup_cli.utils.profiler import estimate_model_memory
|
|
|
|
mem_gb = estimate_model_memory(7.0, "none")
|
|
assert 13.0 <= mem_gb <= 16.0 # ~14 GB for 7B FP16
|
|
|
|
def test_1b_4bit(self):
|
|
from soup_cli.utils.profiler import estimate_model_memory
|
|
|
|
mem_gb = estimate_model_memory(1.0, "4bit")
|
|
assert 0.3 <= mem_gb <= 1.5
|
|
|
|
def test_13b_4bit(self):
|
|
from soup_cli.utils.profiler import estimate_model_memory
|
|
|
|
mem_gb = estimate_model_memory(13.0, "4bit")
|
|
assert 5.0 <= mem_gb <= 9.0
|
|
|
|
def test_70b_4bit(self):
|
|
from soup_cli.utils.profiler import estimate_model_memory
|
|
|
|
mem_gb = estimate_model_memory(70.0, "4bit")
|
|
assert 30.0 <= mem_gb <= 45.0
|
|
|
|
|
|
class TestEstimateLoraMemory:
|
|
"""Test LoRA parameter memory estimation."""
|
|
|
|
def test_default_lora(self):
|
|
from soup_cli.utils.profiler import estimate_lora_memory
|
|
|
|
mem_gb = estimate_lora_memory(7.0, lora_r=64, lora_alpha=16)
|
|
assert 0.01 <= mem_gb <= 1.0
|
|
|
|
def test_higher_rank(self):
|
|
from soup_cli.utils.profiler import estimate_lora_memory
|
|
|
|
low = estimate_lora_memory(7.0, lora_r=16, lora_alpha=16)
|
|
high = estimate_lora_memory(7.0, lora_r=128, lora_alpha=32)
|
|
assert high > low
|
|
|
|
def test_small_model(self):
|
|
from soup_cli.utils.profiler import estimate_lora_memory
|
|
|
|
mem_gb = estimate_lora_memory(1.0, lora_r=8, lora_alpha=16)
|
|
assert mem_gb >= 0
|
|
|
|
|
|
class TestEstimateOptimizerMemory:
|
|
"""Test optimizer memory estimation."""
|
|
|
|
def test_adam_optimizer(self):
|
|
from soup_cli.utils.profiler import estimate_optimizer_memory
|
|
|
|
# ~13M trainable params for 7B with LoRA r=64
|
|
mem_gb = estimate_optimizer_memory(13_000_000, "adamw_torch")
|
|
assert 0.01 <= mem_gb <= 0.5
|
|
|
|
def test_adam_8bit(self):
|
|
from soup_cli.utils.profiler import estimate_optimizer_memory
|
|
|
|
adam = estimate_optimizer_memory(13_000_000, "adamw_torch")
|
|
adam8 = estimate_optimizer_memory(13_000_000, "adamw_bnb_8bit")
|
|
assert adam8 < adam
|
|
|
|
def test_sgd_optimizer(self):
|
|
from soup_cli.utils.profiler import estimate_optimizer_memory
|
|
|
|
adam = estimate_optimizer_memory(13_000_000, "adamw_torch")
|
|
sgd = estimate_optimizer_memory(13_000_000, "sgd")
|
|
assert sgd < adam
|
|
|
|
|
|
class TestEstimateActivationMemory:
|
|
"""Test activation memory estimation."""
|
|
|
|
def test_basic(self):
|
|
from soup_cli.utils.profiler import estimate_activation_memory
|
|
|
|
mem_gb = estimate_activation_memory(
|
|
batch_size=4, seq_len=2048, hidden_size=4096, num_layers=32
|
|
)
|
|
assert mem_gb > 0
|
|
|
|
def test_larger_batch_more_memory(self):
|
|
from soup_cli.utils.profiler import estimate_activation_memory
|
|
|
|
small = estimate_activation_memory(
|
|
batch_size=1, seq_len=2048, hidden_size=4096, num_layers=32
|
|
)
|
|
large = estimate_activation_memory(
|
|
batch_size=8, seq_len=2048, hidden_size=4096, num_layers=32
|
|
)
|
|
assert large > small
|
|
|
|
def test_longer_seq_more_memory(self):
|
|
from soup_cli.utils.profiler import estimate_activation_memory
|
|
|
|
short = estimate_activation_memory(
|
|
batch_size=4, seq_len=512, hidden_size=4096, num_layers=32
|
|
)
|
|
long = estimate_activation_memory(
|
|
batch_size=4, seq_len=4096, hidden_size=4096, num_layers=32
|
|
)
|
|
assert long > short
|
|
|
|
def test_gradient_checkpointing_reduces_memory(self):
|
|
from soup_cli.utils.profiler import estimate_activation_memory
|
|
|
|
normal = estimate_activation_memory(
|
|
batch_size=4, seq_len=2048, hidden_size=4096, num_layers=32,
|
|
gradient_checkpointing=False,
|
|
)
|
|
checkpointed = estimate_activation_memory(
|
|
batch_size=4, seq_len=2048, hidden_size=4096, num_layers=32,
|
|
gradient_checkpointing=True,
|
|
)
|
|
assert checkpointed < normal
|
|
|
|
|
|
class TestEstimateTotal:
|
|
"""Test total profile estimation from SoupConfig."""
|
|
|
|
def test_basic_profile(self):
|
|
from soup_cli.utils.profiler import estimate_total
|
|
|
|
result = estimate_total(
|
|
model_name="meta-llama/Llama-3.1-8B-Instruct",
|
|
model_params_b=8.0,
|
|
quantization="4bit",
|
|
lora_r=64,
|
|
lora_alpha=16,
|
|
batch_size=4,
|
|
seq_len=2048,
|
|
optimizer="adamw_torch",
|
|
gradient_checkpointing=False,
|
|
)
|
|
assert "model_memory_gb" in result
|
|
assert "lora_memory_gb" in result
|
|
assert "optimizer_memory_gb" in result
|
|
assert "activation_memory_gb" in result
|
|
assert "total_memory_gb" in result
|
|
assert result["total_memory_gb"] > 0
|
|
|
|
def test_total_is_sum_of_parts(self):
|
|
from soup_cli.utils.profiler import estimate_total
|
|
|
|
result = estimate_total(
|
|
model_name="meta-llama/Llama-3.1-8B-Instruct",
|
|
model_params_b=8.0,
|
|
quantization="4bit",
|
|
lora_r=64,
|
|
lora_alpha=16,
|
|
batch_size=4,
|
|
seq_len=2048,
|
|
optimizer="adamw_torch",
|
|
gradient_checkpointing=False,
|
|
)
|
|
parts_sum = (
|
|
result["model_memory_gb"]
|
|
+ result["lora_memory_gb"]
|
|
+ result["optimizer_memory_gb"]
|
|
+ result["activation_memory_gb"]
|
|
)
|
|
# Total includes overhead, so should be >= parts_sum
|
|
assert result["total_memory_gb"] >= parts_sum
|
|
|
|
|
|
class TestEstimateSpeed:
|
|
"""Test speed estimation."""
|
|
|
|
def test_basic_speed(self):
|
|
from soup_cli.utils.profiler import estimate_speed
|
|
|
|
tokens_per_sec = estimate_speed(model_params_b=7.0, quantization="4bit", batch_size=4)
|
|
assert tokens_per_sec > 0
|
|
|
|
def test_larger_model_slower(self):
|
|
from soup_cli.utils.profiler import estimate_speed
|
|
|
|
small = estimate_speed(model_params_b=7.0, quantization="4bit", batch_size=4)
|
|
large = estimate_speed(model_params_b=70.0, quantization="4bit", batch_size=4)
|
|
assert small > large
|
|
|
|
def test_no_quantization_slower(self):
|
|
from soup_cli.utils.profiler import estimate_speed
|
|
|
|
quantized = estimate_speed(model_params_b=7.0, quantization="4bit", batch_size=4)
|
|
full = estimate_speed(model_params_b=7.0, quantization="none", batch_size=4)
|
|
assert quantized > full
|
|
|
|
|
|
class TestEstimateTrainingTime:
|
|
"""Test training time estimation."""
|
|
|
|
def test_basic_time(self):
|
|
from soup_cli.utils.profiler import estimate_training_time
|
|
|
|
minutes = estimate_training_time(
|
|
dataset_size=10000, epochs=3, samples_per_sec=0.5
|
|
)
|
|
assert minutes > 0
|
|
|
|
def test_more_epochs_more_time(self):
|
|
from soup_cli.utils.profiler import estimate_training_time
|
|
|
|
time_1 = estimate_training_time(dataset_size=10000, epochs=1, samples_per_sec=0.5)
|
|
time_3 = estimate_training_time(dataset_size=10000, epochs=3, samples_per_sec=0.5)
|
|
assert time_3 == pytest.approx(time_1 * 3, rel=0.01)
|
|
|
|
def test_zero_speed_returns_inf(self):
|
|
from soup_cli.utils.profiler import estimate_training_time
|
|
|
|
minutes = estimate_training_time(
|
|
dataset_size=10000, epochs=3, samples_per_sec=0.0
|
|
)
|
|
assert minutes == float("inf")
|
|
|
|
|
|
class TestRecommendBatchSize:
|
|
"""Test batch size recommendation."""
|
|
|
|
def test_basic_recommendation(self):
|
|
from soup_cli.utils.profiler import recommend_batch_size
|
|
|
|
bs = recommend_batch_size(total_memory_gb=6.0, gpu_memory_gb=24.0)
|
|
assert bs >= 1
|
|
|
|
def test_tight_memory_returns_1(self):
|
|
from soup_cli.utils.profiler import recommend_batch_size
|
|
|
|
bs = recommend_batch_size(total_memory_gb=23.0, gpu_memory_gb=24.0)
|
|
assert bs >= 1
|
|
|
|
def test_more_gpu_memory_higher_batch(self):
|
|
from soup_cli.utils.profiler import recommend_batch_size
|
|
|
|
small_gpu = recommend_batch_size(total_memory_gb=6.0, gpu_memory_gb=8.0)
|
|
big_gpu = recommend_batch_size(total_memory_gb=6.0, gpu_memory_gb=80.0)
|
|
assert big_gpu >= small_gpu
|
|
|
|
|
|
class TestRecommendGpu:
|
|
"""Test GPU recommendation logic."""
|
|
|
|
def test_small_model(self):
|
|
from soup_cli.utils.profiler import recommend_gpu
|
|
|
|
gpus = recommend_gpu(total_memory_gb=5.0)
|
|
assert len(gpus) > 0
|
|
# Should include most GPUs
|
|
assert any("RTX" in g or "A100" in g or "H100" in g for g in gpus)
|
|
|
|
def test_large_model(self):
|
|
from soup_cli.utils.profiler import recommend_gpu
|
|
|
|
gpus = recommend_gpu(total_memory_gb=50.0)
|
|
# Only large GPUs should be recommended
|
|
assert all("RTX 3060" not in g for g in gpus)
|
|
|
|
def test_very_large_model(self):
|
|
from soup_cli.utils.profiler import recommend_gpu
|
|
|
|
gpus = recommend_gpu(total_memory_gb=200.0)
|
|
# Should suggest multi-GPU or mention no single GPU fits
|
|
assert isinstance(gpus, list)
|
|
|
|
|
|
class TestModelArchLookup:
|
|
"""Test model architecture lookup for hidden_size/num_layers."""
|
|
|
|
def test_known_model_8b(self):
|
|
from soup_cli.utils.profiler import get_model_arch
|
|
|
|
arch = get_model_arch("meta-llama/Llama-3.1-8B-Instruct", 8.0)
|
|
assert arch["hidden_size"] > 0
|
|
assert arch["num_layers"] > 0
|
|
|
|
def test_unknown_model_uses_params(self):
|
|
from soup_cli.utils.profiler import get_model_arch
|
|
|
|
arch = get_model_arch("unknown/model-7b", 7.0)
|
|
assert arch["hidden_size"] > 0
|
|
assert arch["num_layers"] > 0
|
|
|
|
def test_70b_model(self):
|
|
from soup_cli.utils.profiler import get_model_arch
|
|
|
|
arch = get_model_arch("meta-llama/Llama-3.1-70B-Instruct", 70.0)
|
|
assert arch["hidden_size"] >= 8192
|
|
assert arch["num_layers"] >= 80
|
|
|
|
|
|
# --- CLI tests ---
|
|
|
|
|
|
class TestProfileCLI:
|
|
"""Test the profile CLI command."""
|
|
|
|
def test_profile_with_config_file(self, tmp_path):
|
|
"""soup profile --config soup.yaml shows profile panel."""
|
|
config_file = tmp_path / "soup.yaml"
|
|
config_file.write_text(
|
|
"base: meta-llama/Llama-3.1-8B-Instruct\n"
|
|
"task: sft\n"
|
|
"data:\n"
|
|
" train: ./data/train.jsonl\n"
|
|
" max_length: 2048\n"
|
|
"training:\n"
|
|
" epochs: 3\n"
|
|
" batch_size: 4\n"
|
|
" quantization: 4bit\n"
|
|
" lora:\n"
|
|
" r: 64\n"
|
|
" alpha: 16\n"
|
|
"output: ./output\n"
|
|
)
|
|
result = runner.invoke(app, ["profile", "--config", str(config_file)])
|
|
assert result.exit_code == 0
|
|
assert "Training Profile" in result.output
|
|
|
|
def test_profile_shows_memory(self, tmp_path):
|
|
"""Profile output includes memory breakdown."""
|
|
config_file = tmp_path / "soup.yaml"
|
|
config_file.write_text(
|
|
"base: meta-llama/Llama-3.1-8B-Instruct\n"
|
|
"task: sft\n"
|
|
"data:\n"
|
|
" train: ./data/train.jsonl\n"
|
|
" max_length: 2048\n"
|
|
"training:\n"
|
|
" batch_size: 4\n"
|
|
" quantization: 4bit\n"
|
|
" lora:\n"
|
|
" r: 64\n"
|
|
"output: ./output\n"
|
|
)
|
|
result = runner.invoke(app, ["profile", "--config", str(config_file)])
|
|
assert result.exit_code == 0
|
|
assert "Model" in result.output
|
|
assert "LoRA" in result.output
|
|
assert "Total" in result.output
|
|
|
|
def test_profile_shows_recommendations(self, tmp_path):
|
|
"""Profile output includes recommendations."""
|
|
config_file = tmp_path / "soup.yaml"
|
|
config_file.write_text(
|
|
"base: meta-llama/Llama-3.1-8B-Instruct\n"
|
|
"task: sft\n"
|
|
"data:\n"
|
|
" train: ./data/train.jsonl\n"
|
|
" max_length: 2048\n"
|
|
"training:\n"
|
|
" batch_size: 4\n"
|
|
" quantization: 4bit\n"
|
|
" lora:\n"
|
|
" r: 64\n"
|
|
"output: ./output\n"
|
|
)
|
|
result = runner.invoke(app, ["profile", "--config", str(config_file)])
|
|
assert result.exit_code == 0
|
|
assert "Recommend" in result.output or "GPU" in result.output
|
|
|
|
def test_profile_json_output(self, tmp_path):
|
|
"""soup profile --json outputs valid JSON."""
|
|
config_file = tmp_path / "soup.yaml"
|
|
config_file.write_text(
|
|
"base: meta-llama/Llama-3.1-8B-Instruct\n"
|
|
"task: sft\n"
|
|
"data:\n"
|
|
" train: ./data/train.jsonl\n"
|
|
" max_length: 2048\n"
|
|
"training:\n"
|
|
" batch_size: 4\n"
|
|
" quantization: 4bit\n"
|
|
" lora:\n"
|
|
" r: 64\n"
|
|
"output: ./output\n"
|
|
)
|
|
result = runner.invoke(app, ["profile", "--config", str(config_file), "--json"])
|
|
assert result.exit_code == 0
|
|
data = json.loads(result.output)
|
|
assert "total_memory_gb" in data
|
|
assert "model_memory_gb" in data
|
|
|
|
def test_profile_with_gpu_flag(self, tmp_path):
|
|
"""soup profile --gpu a100 uses specified GPU memory."""
|
|
config_file = tmp_path / "soup.yaml"
|
|
config_file.write_text(
|
|
"base: meta-llama/Llama-3.1-8B-Instruct\n"
|
|
"task: sft\n"
|
|
"data:\n"
|
|
" train: ./data/train.jsonl\n"
|
|
" max_length: 2048\n"
|
|
"training:\n"
|
|
" batch_size: 4\n"
|
|
" quantization: 4bit\n"
|
|
" lora:\n"
|
|
" r: 64\n"
|
|
"output: ./output\n"
|
|
)
|
|
result = runner.invoke(
|
|
app, ["profile", "--config", str(config_file), "--gpu", "a100"]
|
|
)
|
|
assert result.exit_code == 0
|
|
assert "A100" in result.output or "80" in result.output
|
|
|
|
def test_profile_auto_batch_size(self, tmp_path):
|
|
"""Profile with batch_size: auto estimates batch size."""
|
|
config_file = tmp_path / "soup.yaml"
|
|
config_file.write_text(
|
|
"base: meta-llama/Llama-3.1-8B-Instruct\n"
|
|
"task: sft\n"
|
|
"data:\n"
|
|
" train: ./data/train.jsonl\n"
|
|
" max_length: 2048\n"
|
|
"training:\n"
|
|
" batch_size: auto\n"
|
|
" quantization: 4bit\n"
|
|
" lora:\n"
|
|
" r: 64\n"
|
|
"output: ./output\n"
|
|
)
|
|
result = runner.invoke(app, ["profile", "--config", str(config_file)])
|
|
assert result.exit_code == 0
|
|
|
|
def test_profile_missing_config(self):
|
|
"""soup profile --config nonexistent.yaml fails gracefully."""
|
|
result = runner.invoke(app, ["profile", "--config", "nonexistent.yaml"])
|
|
assert result.exit_code != 0
|
|
|
|
def test_profile_no_lora(self, tmp_path):
|
|
"""Profile works for full-parameter training (no LoRA quantization=none)."""
|
|
config_file = tmp_path / "soup.yaml"
|
|
config_file.write_text(
|
|
"base: meta-llama/Llama-3.1-8B-Instruct\n"
|
|
"task: sft\n"
|
|
"data:\n"
|
|
" train: ./data/train.jsonl\n"
|
|
" max_length: 2048\n"
|
|
"training:\n"
|
|
" batch_size: 4\n"
|
|
" quantization: none\n"
|
|
" lora:\n"
|
|
" r: 64\n"
|
|
"output: ./output\n"
|
|
)
|
|
result = runner.invoke(app, ["profile", "--config", str(config_file)])
|
|
assert result.exit_code == 0
|
|
|
|
def test_profile_gradient_checkpointing(self, tmp_path):
|
|
"""Profile respects gradient_checkpointing flag."""
|
|
config_file = tmp_path / "soup.yaml"
|
|
config_file.write_text(
|
|
"base: meta-llama/Llama-3.1-8B-Instruct\n"
|
|
"task: sft\n"
|
|
"data:\n"
|
|
" train: ./data/train.jsonl\n"
|
|
" max_length: 2048\n"
|
|
"training:\n"
|
|
" batch_size: 4\n"
|
|
" quantization: 4bit\n"
|
|
" gradient_checkpointing: true\n"
|
|
" lora:\n"
|
|
" r: 64\n"
|
|
"output: ./output\n"
|
|
)
|
|
result = runner.invoke(app, ["profile", "--config", str(config_file)])
|
|
assert result.exit_code == 0
|
|
assert "gradient_checkpointing" in result.output or "checkpointing" in result.output.lower()
|
|
|
|
|
|
# --- GPU lookup table tests ---
|
|
|
|
|
|
class TestGpuLookup:
|
|
"""Test GPU memory lookup table."""
|
|
|
|
def test_known_gpus(self):
|
|
from soup_cli.utils.profiler import GPU_MEMORY
|
|
|
|
assert "a100" in GPU_MEMORY
|
|
assert "h100" in GPU_MEMORY
|
|
assert "rtx4090" in GPU_MEMORY
|
|
assert GPU_MEMORY["a100"] == 80
|
|
assert GPU_MEMORY["h100"] == 80
|
|
|
|
def test_rtx_gpus(self):
|
|
from soup_cli.utils.profiler import GPU_MEMORY
|
|
|
|
assert "rtx3090" in GPU_MEMORY
|
|
assert GPU_MEMORY["rtx3090"] == 24
|
|
|
|
def test_invalid_gpu_flag(self, tmp_path):
|
|
"""Unknown GPU name shows error."""
|
|
config_file = tmp_path / "soup.yaml"
|
|
config_file.write_text(
|
|
"base: meta-llama/Llama-3.1-8B-Instruct\n"
|
|
"task: sft\n"
|
|
"data:\n"
|
|
" train: ./data/train.jsonl\n"
|
|
" max_length: 2048\n"
|
|
"training:\n"
|
|
" batch_size: 4\n"
|
|
" quantization: 4bit\n"
|
|
" lora:\n"
|
|
" r: 64\n"
|
|
"output: ./output\n"
|
|
)
|
|
result = runner.invoke(
|
|
app, ["profile", "--config", str(config_file), "--gpu", "nonexistent_gpu"]
|
|
)
|
|
assert result.exit_code != 0
|
|
|
|
|
|
class TestEstimateModelMemoryEdge:
|
|
"""Edge case tests for model memory estimation."""
|
|
|
|
def test_unknown_quantization_falls_back_to_fp16(self):
|
|
from soup_cli.utils.profiler import estimate_model_memory
|
|
|
|
mem = estimate_model_memory(7.0, "bfloat16")
|
|
fp16_mem = estimate_model_memory(7.0, "none")
|
|
assert mem == fp16_mem # unknown falls back to 2.0 bytes/param
|
|
|
|
def test_negative_samples_per_sec(self):
|
|
from soup_cli.utils.profiler import estimate_training_time
|
|
|
|
result = estimate_training_time(1000, 3, -1.0)
|
|
assert result == float("inf")
|
|
|
|
|
|
class TestResolveGpuMemory:
|
|
"""Test GPU memory resolution."""
|
|
|
|
def test_fallback_to_default_24gb(self, monkeypatch):
|
|
from soup_cli.commands.profile import _resolve_gpu_memory
|
|
|
|
# Force GPU detection to fail by patching inside the utils module
|
|
monkeypatch.setattr(
|
|
"soup_cli.utils.gpu.get_gpu_info",
|
|
lambda: (_ for _ in ()).throw(RuntimeError("no GPU")),
|
|
)
|
|
mem = _resolve_gpu_memory(None)
|
|
assert mem == 24.0
|
|
|
|
|
|
class TestTrainableParamsEstimate:
|
|
"""Test trainable parameter count estimation."""
|
|
|
|
def test_lora_params(self):
|
|
from soup_cli.utils.profiler import estimate_trainable_params
|
|
|
|
params = estimate_trainable_params(
|
|
model_params_b=7.0, lora_r=64, hidden_size=4096
|
|
)
|
|
assert params > 0
|
|
# LoRA should be a small fraction of total
|
|
assert params < 7_000_000_000 * 0.05
|
|
|
|
def test_higher_rank_more_params(self):
|
|
from soup_cli.utils.profiler import estimate_trainable_params
|
|
|
|
low = estimate_trainable_params(model_params_b=7.0, lora_r=16, hidden_size=4096)
|
|
high = estimate_trainable_params(model_params_b=7.0, lora_r=128, hidden_size=4096)
|
|
assert high > low
|