File size: 3,653 Bytes
df31653 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | from __future__ import annotations
import io
from collections.abc import Callable
from typing import Any
import pytest
from fitcheck import repl
from fitcheck.config_parser import ModelConfig, fetch_model_config
from fitcheck.estimator import TrainingConfig
from fitcheck.gpu_db import GpuSpec, get_gpu
from rich.console import Console
@pytest.fixture
def llama_model(
fake_config_download: Callable[[dict[str, Any]], None],
llama_31_8b_config: dict[str, Any],
) -> ModelConfig:
fake_config_download(llama_31_8b_config)
return fetch_model_config("meta-llama/Llama-3.1-8B")
@pytest.fixture
def qlora_training() -> TrainingConfig:
return TrainingConfig(
precision="bf16",
quantization="nf4",
optimizer="adamw",
optimizer_dtype="fp32",
batch_size=1,
seq_len=2048,
lora_rank=64,
lora_targets=["q_proj", "k_proj", "v_proj", "o_proj"],
grad_checkpoint=True,
flash_attn=True,
)
class _Repl:
def __init__(self, model: ModelConfig, training: TrainingConfig, gpu: GpuSpec) -> None:
self.buffer = io.StringIO()
console = Console(file=self.buffer, width=200, no_color=True, legacy_windows=False)
self.session = repl._Session(
console=console,
glyphs=repl._ASCII_GLYPHS,
ascii_only=True,
model=model,
model_id="meta-llama/Llama-3.1-8B",
gpu=gpu,
training=training,
)
def run(self, line: str) -> str:
before = self.buffer.tell()
repl._dispatch(self.session, line)
self.buffer.seek(before)
return " ".join(self.buffer.read().split())
def test_gpu_flag_uses_the_database_key(llama_model: ModelConfig) -> None:
assert repl._gpu_flag(get_gpu("4090")) == "--gpu 4090"
def test_gpu_flag_reproduces_a_custom_card(llama_model: ModelConfig) -> None:
assert repl._gpu_flag(get_gpu("mycard", 24_000)) == "--gpu mycard --vram-mib 24000"
def test_optimize_command_names_the_session_gpu(
llama_model: ModelConfig, qlora_training: TrainingConfig
) -> None:
session = _Repl(llama_model, qlora_training, get_gpu("4090"))
assert "memory --gpu 4090 --batch-size" in session.run("optimize")
def test_optimize_command_keeps_a_one_shot_gpu_override(
llama_model: ModelConfig, qlora_training: TrainingConfig
) -> None:
session = _Repl(llama_model, qlora_training, get_gpu("a100-80"))
session.run("memory --gpu 4090")
assert "memory --gpu 4090 --batch-size" in session.run("optimize")
assert session.session.gpu == get_gpu("a100-80")
def test_rescue_command_keeps_the_gpu_override(
llama_model: ModelConfig, qlora_training: TrainingConfig
) -> None:
session = _Repl(llama_model, qlora_training, get_gpu("a100-80"))
session.run("memory --gpu t4")
output = session.run("optimize")
assert "Nothing fits at batch_size 1" in output
assert "memory --gpu t4 --batch-size 1" in output
def test_compare_prices_every_card_separately(
llama_model: ModelConfig, qlora_training: TrainingConfig
) -> None:
session = _Repl(llama_model, qlora_training, get_gpu("t4"))
output = session.run("compare 4090")
assert "Peak (MiB)" in output
assert "Peak is identical on every card" not in output
assert "Peak moves with the card" in output
def test_compare_still_states_equality_when_the_peaks_match(
llama_model: ModelConfig, qlora_training: TrainingConfig
) -> None:
session = _Repl(llama_model, qlora_training, get_gpu("4090"))
output = session.run("compare a100-80")
assert "Peak is identical on every card" in output
|