vramcalc 0.1.0__py3-none-any.whl

vramcalc/__init__.py

@@ -0,0 +1,9 @@
+"""vramcalc: predict GPU VRAM usage from config.json."""
+
+from .arch import Arch, UnsupportedModelError
+from .bench import BenchResult, bench
+from .estimate import Estimate, EstimateParams, FitResult, Suggestion, estimate
+
+__version__ = "0.1.0"
+__all__ = ["Arch", "BenchResult", "Estimate", "EstimateParams", "FitResult", "Suggestion",
+           "UnsupportedModelError", "bench", "estimate"]

vramcalc/arch.py

@@ -0,0 +1,111 @@
+"""Normalize a Hugging Face config.json into the fields the memory formulas need."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+SUPPORTED_MODEL_TYPES = ("llama", "mistral", "qwen2", "gemma", "gemma2")
+
+_REQUIRED = ("hidden_size", "num_hidden_layers", "num_attention_heads",
+             "intermediate_size", "vocab_size")
+
+
+class UnsupportedModelError(ValueError):
+    """Raised when model_type is outside the supported allowlist."""
+
+
+def _positive_int(cfg: dict, key: str) -> int:
+    value = cfg[key]
+    if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
+        raise ValueError(f"{key} must be a positive integer, got {value!r}")
+    return value
+
+
+@dataclass(frozen=True)
+class Arch:
+    hidden_size: int
+    num_layers: int
+    num_heads: int
+    num_kv_heads: int
+    head_dim: int
+    intermediate_size: int
+    vocab_size: int
+    tie_word_embeddings: bool
+    attention_bias: bool
+    norms_per_layer: int
+    name: str = ""
+
+    @property
+    def q(self) -> int:
+        """Width of the query projection output."""
+        return self.num_heads * self.head_dim
+
+    @property
+    def kv(self) -> int:
+        """Width of each of the key and value projection outputs."""
+        return self.num_kv_heads * self.head_dim
+
+    @classmethod
+    def from_config(cls, cfg: dict) -> Arch:
+        model_type = cfg.get("model_type")
+        if model_type not in SUPPORTED_MODEL_TYPES:
+            raise UnsupportedModelError(
+                f"model_type {model_type!r} is not supported; "
+                f"supported: {', '.join(SUPPORTED_MODEL_TYPES)}"
+            )
+        missing = [k for k in _REQUIRED if k not in cfg]
+        if missing:
+            raise ValueError(f"config is missing required keys: {', '.join(missing)}")
+
+        hidden = _positive_int(cfg, "hidden_size")
+        layers = _positive_int(cfg, "num_hidden_layers")
+        heads = _positive_int(cfg, "num_attention_heads")
+        kv_heads = _positive_int(cfg, "num_key_value_heads") if "num_key_value_heads" in cfg else heads
+        if heads % kv_heads != 0:
+            raise ValueError(
+                f"num_key_value_heads={kv_heads} must divide num_attention_heads={heads}"
+            )
+        if "head_dim" in cfg:
+            head_dim = _positive_int(cfg, "head_dim")
+        else:
+            if hidden % heads != 0:
+                raise ValueError(
+                    f"hidden_size={hidden} is not divisible by num_attention_heads={heads}"
+                )
+            head_dim = hidden // heads
+
+        is_gemma = model_type in ("gemma", "gemma2")
+        name = str(cfg.get("_name_or_path", "")).rstrip("/").split("/")[-1]
+        return cls(
+            hidden_size=hidden,
+            num_layers=layers,
+            num_heads=heads,
+            num_kv_heads=kv_heads,
+            head_dim=head_dim,
+            intermediate_size=_positive_int(cfg, "intermediate_size"),
+            vocab_size=_positive_int(cfg, "vocab_size"),
+            tie_word_embeddings=bool(cfg.get("tie_word_embeddings", is_gemma)),
+            attention_bias=(model_type == "qwen2") or bool(cfg.get("attention_bias", False)),
+            norms_per_layer=4 if model_type == "gemma2" else 2,
+            name=name,
+        )
+
+    # Parameter-count pieces are exposed so memory.py can reuse them for P_cast.
+    def attn_params(self) -> int:
+        h = self.hidden_size
+        return h * self.q + 2 * h * self.kv + self.q * h
+
+    def attn_bias_params(self) -> int:
+        return self.q + 2 * self.kv if self.attention_bias else 0
+
+    def mlp_params(self) -> int:
+        return 3 * self.hidden_size * self.intermediate_size
+
+    def embed_params(self) -> int:
+        return self.vocab_size * self.hidden_size
+
+    def param_count(self) -> int:
+        per_layer = (self.attn_params() + self.attn_bias_params() + self.mlp_params()
+                     + self.norms_per_layer * self.hidden_size)
+        lm_head = 0 if self.tie_word_embeddings else self.embed_params()
+        return self.embed_params() + self.num_layers * per_layer + self.hidden_size + lm_head

vramcalc/bench.py

@@ -0,0 +1,220 @@
+"""Measure real GPU memory with torch and compare against the prediction. Spec section 7."""
+
+from __future__ import annotations
+
+import contextlib
+import dataclasses
+import datetime as dt
+import gc
+import json
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+from .estimate import GB, Estimate, EstimateParams, estimate
+
+BUCKETS = ("weights", "states", "peak", "total", "overhead")
+LORA_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj"]
+_ATTN_IMPL = {"eager": "eager", "sdpa": "sdpa", "flash": "flash_attention_2"}
+
+
+def buckets(est: Estimate) -> dict[str, int]:
+    inference = est.params.mode == "inference"
+    return {
+        "weights": est.weights,
+        "states": est.gradients + est.optimizer,
+        "peak": est.peak_extra + (est.kv_cache if inference else 0),
+        "total": est.total,
+        "overhead": est.overhead,
+    }
+
+
+@dataclass(frozen=True)
+class BenchResult:
+    status: str                   # "ok" | "oom" | "failed"
+    error: str | None
+    params: EstimateParams
+    predicted: dict[str, int]
+    measured: dict[str, int]
+    gpu_name: str
+    versions: dict[str, str]
+    timestamp: str
+
+    def to_dict(self) -> dict:
+        d = dataclasses.asdict(self)
+        d["params"]["arch"] = dataclasses.asdict(self.params.arch)
+        return d
+
+    def table(self) -> str:
+        lines = [f"{'bucket':<10}{'predicted':>12}{'measured':>12}{'error':>9}"]
+        for k in BUCKETS:
+            p = self.predicted.get(k, 0)
+            m = self.measured.get(k)
+            err = "N/A" if not m else f"{(p - m) / m * 100:+.1f}%"
+            m_txt = f"{m / GB:.2f}" if m is not None else "-"
+            lines.append(f"{k:<10}{p / GB:>12.2f}{m_txt:>12}{err:>9}")
+        if self.status != "ok":
+            lines.append(f"status: {self.status}: {self.error}")
+        return "\n".join(lines)
+
+
+def write_record(result: BenchResult, path: str | Path) -> None:
+    with open(path, "a", encoding="utf-8") as f:
+        f.write(json.dumps(result.to_dict(), ensure_ascii=False) + "\n")
+
+
+_ABS_PATH = re.compile(r"(?<![\w])/[^\s'\"]+")
+
+
+def _sanitize(message: str, model_path) -> str:
+    """Keep local paths out of recorded errors (spec 7.4)."""
+    message = message.replace(str(model_path), "<model>")
+    return _ABS_PATH.sub("<path>", message)[:500]
+
+
+def _require_bench_deps():
+    try:
+        import peft
+        import torch
+        import transformers
+    except ImportError as exc:
+        raise ImportError(
+            "bench needs torch, transformers and peft; install vramcalc[bench]") from exc
+    if not torch.cuda.is_available():
+        raise RuntimeError("bench needs a CUDA device")
+    return torch, transformers, peft
+
+
+def bench(model_path, *, mode="inference", dtype="bf16", batch=1, seq_len=2048, lora_rank=None,
+          grad_checkpoint=False, optimizer="adamw", precision="amp", attn_impl="sdpa",
+          overhead_gb=0.75, steps=3, record=None) -> BenchResult:
+    # Pre-validation: everything here raises; nothing after model load does.
+    compute_fp32 = dtype == "fp32" and (mode == "inference" or precision == "pure")
+    if attn_impl == "sdpa" and compute_fp32:
+        # The pinned efficient/flash backends do not serve fp32 GQA; math would store scores.
+        raise ValueError("attn_impl='sdpa' cannot be pinned for fp32 compute; use attn_impl='eager'")
+    torch, transformers, peft = _require_bench_deps()
+    if attn_impl == "flash":
+        try:
+            import flash_attn  # noqa: F401
+        except ImportError as exc:
+            raise ImportError("attn_impl='flash' needs the flash-attn package") from exc
+    if not isinstance(steps, int) or steps <= 0:
+        raise ValueError(f"steps must be a positive integer, got {steps!r}")
+    est = estimate(model_path, mode=mode, dtype=dtype, batch=batch, seq_len=seq_len,
+                   lora_rank=lora_rank, grad_checkpoint=grad_checkpoint, optimizer=optimizer,
+                   precision=precision, attn_impl=attn_impl, overhead_gb=overhead_gb)
+    p = est.params
+    versions = {"torch": torch.__version__, "transformers": transformers.__version__,
+                "peft": peft.__version__}
+    gpu_name = torch.cuda.get_device_name(torch.cuda.current_device())
+    stamp = dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds")
+
+    status, error, measured = "ok", None, {}
+    try:
+        measured = _measure(torch, transformers, peft, model_path, p, steps)
+    except torch.cuda.OutOfMemoryError as exc:
+        status, error = "oom", _sanitize((str(exc).splitlines() or [""])[0], model_path)
+    except Exception as exc:  # noqa: BLE001 - runtime failure after model load, spec 7.3
+        status, error = "failed", _sanitize(f"{type(exc).__name__}: {exc}", model_path)
+    finally:
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    result = BenchResult(status, error, p, buckets(est), measured, gpu_name, versions, stamp)
+    if record is not None:
+        write_record(result, record)
+    return result
+
+
+def _reset(torch):
+    torch.cuda.synchronize()
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+
+
+def _measure(torch, transformers, peft, model_path, p: EstimateParams, steps: int) -> dict[str, int]:
+    torch_dtype = {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16}[p.dtype]
+    load_dtype = torch.float32 if (p.mode == "train" and p.precision == "amp") else torch_dtype
+    measured: dict[str, int] = {}
+
+    _reset(torch)
+    preexisting = torch.cuda.memory_allocated()  # tensors the caller already holds
+    model = transformers.AutoModelForCausalLM.from_pretrained(
+        model_path, torch_dtype=load_dtype, attn_implementation=_ATTN_IMPL[p.attn_impl],
+    ).cuda()
+    if p.lora_rank is not None:
+        cfg = peft.LoraConfig(r=p.lora_rank, lora_alpha=p.lora_rank, lora_dropout=0.0,
+                              target_modules=LORA_TARGETS, task_type="CAUSAL_LM")
+        model = peft.get_peft_model(model, cfg)
+    torch.cuda.synchronize()
+    measured["weights"] = torch.cuda.memory_allocated() - preexisting
+
+    input_ids = torch.randint(0, p.arch.vocab_size, (p.batch, p.seq_len), device="cuda")
+
+    if p.attn_impl == "sdpa":
+        from torch.nn.attention import SDPBackend, sdpa_kernel
+
+        def attn_ctx():
+            return sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION, SDPBackend.FLASH_ATTENTION])
+    else:
+        attn_ctx = contextlib.nullcontext
+
+    if p.mode == "inference":
+        model.eval()
+        _reset(torch)
+        with torch.inference_mode(), attn_ctx():
+            out = model(input_ids=input_ids, use_cache=True)
+        torch.cuda.synchronize()
+        measured["peak"] = torch.cuda.max_memory_allocated() - preexisting - measured["weights"]
+        measured["states"] = 0
+        del out
+    else:
+        model.train()
+        model.config.use_cache = False
+        if p.grad_checkpoint:
+            model.gradient_checkpointing_enable(
+                gradient_checkpointing_kwargs={"use_reentrant": False})
+        trainable = [q for q in model.parameters() if q.requires_grad]
+        if p.optimizer == "adamw":
+            opt = torch.optim.AdamW(trainable, lr=1e-5, foreach=True)
+        else:
+            opt = torch.optim.SGD(trainable, lr=1e-5, momentum=0.9, foreach=True)
+        amp = p.precision == "amp"
+        if amp:
+            def autocast():
+                return torch.autocast("cuda", dtype=torch_dtype)
+        else:
+            autocast = contextlib.nullcontext
+
+        def step():
+            # attn_ctx stays active through backward so checkpoint recomputation
+            # uses the same SDPA backend as the forward pass
+            with attn_ctx():
+                with autocast():
+                    out = model(input_ids=input_ids, labels=input_ids, use_cache=False)
+                    loss = out.loss
+                loss.backward()
+            del out, loss  # release logits before step() so workspace does not overlap (spec 5.6)
+            opt.step()
+            opt.zero_grad(set_to_none=False)
+
+        step()  # warm-up
+        torch.cuda.synchronize()
+        baseline = torch.cuda.memory_allocated()
+        measured["states"] = baseline - preexisting - measured["weights"]
+        peak = 0
+        for _ in range(steps):
+            _reset(torch)
+            step()
+            torch.cuda.synchronize()
+            peak = max(peak, torch.cuda.max_memory_allocated() - baseline)
+        measured["peak"] = peak
+        del opt
+
+    measured["total"] = measured["weights"] + measured["states"] + measured["peak"]
+    torch.cuda.synchronize()
+    free, total = torch.cuda.mem_get_info()
+    measured["overhead"] = (total - free) - torch.cuda.memory_reserved()
+    return measured  # model and inputs are released when this frame returns

vramcalc/cli.py

@@ -0,0 +1,126 @@
+"""Command-line interface: `vramcalc MODEL ...`, `vramcalc bench ...`, `vramcalc gpus`."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+
+from . import gpus
+from .estimate import GB, estimate
+
+_SUBCOMMANDS = ("bench", "gpus")
+
+
+def _add_estimate_args(p: argparse.ArgumentParser) -> None:
+    p.add_argument("--train", action="store_true")
+    p.add_argument("--lora", type=int, metavar="R")
+    p.add_argument("--dtype", default="bf16", choices=("fp32", "bf16", "fp16"))
+    p.add_argument("--batch", type=int, default=1)
+    p.add_argument("--seq", type=int, default=2048)
+    p.add_argument("--grad-checkpoint", action="store_true")
+    p.add_argument("--optimizer", default="adamw", choices=("adamw", "sgd"))
+    p.add_argument("--precision", default="amp", choices=("amp", "pure"))
+    p.add_argument("--attn", default="sdpa", choices=("eager", "sdpa", "flash"))
+    p.add_argument("--overhead-gb", type=float, default=0.75)
+
+
+def _estimate_kwargs(a: argparse.Namespace) -> dict:
+    return {
+        "mode": "train" if a.train else "inference", "dtype": a.dtype, "batch": a.batch,
+        "seq_len": a.seq, "lora_rank": a.lora, "grad_checkpoint": a.grad_checkpoint,
+        "optimizer": a.optimizer, "precision": a.precision, "attn_impl": a.attn,
+        "overhead_gb": a.overhead_gb,
+    }
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="vramcalc", description=__doc__)
+    sub = parser.add_subparsers(dest="command")
+
+    est = sub.add_parser("estimate", help="predict memory (default when MODEL is given)")
+    est.add_argument("model")
+    _add_estimate_args(est)
+    est.add_argument("--gpu", help="GPU name or GB number; exit 1 if the model does not fit")
+    est.add_argument("--json", action="store_true")
+
+    sub.add_parser("gpus", help="list built-in GPU table")
+
+    bench = sub.add_parser(
+        "bench", help="measure on the local GPU and compare (needs vramcalc[bench])")
+    bench.add_argument("model_path")
+    _add_estimate_args(bench)
+    bench.add_argument("--steps", type=int, default=3)
+    bench.add_argument("--record", metavar="PATH")
+    bench.add_argument("--json", action="store_true")
+    return parser
+
+
+def _parse_gpu(value: str | None):
+    """Table names win over numbers so that `--gpu 3090` means the card, not 3090 GB."""
+    if value is None:
+        return None
+    if gpus._key(value) in gpus._INDEX:
+        return value
+    try:
+        return float(value)
+    except ValueError:
+        return value
+
+
+def cmd_estimate(a: argparse.Namespace) -> int:
+    e = estimate(a.model, **_estimate_kwargs(a))
+    fit = e.fits(_parse_gpu(a.gpu)) if a.gpu is not None else None
+    if a.json:
+        out = {"estimate": e.to_dict()}
+        if fit is not None:
+            out["fit"] = fit.to_dict()
+        print(json.dumps(out, indent=2))
+    else:
+        print(e.breakdown())
+        if fit is not None:
+            print()
+            print(fit.message)
+            for s in fit.suggestions:
+                tag = " (changes training objective)" if s.changes_objective else ""
+                change = ", ".join(f"{k}={v}" for k, v in s.change.items())
+                need = (s.estimate.total + s.estimate.overhead) / GB
+                print(f"  -> {change}: {need:.1f} GB{tag}")
+    return 0 if fit is None or fit.ok else 1
+
+
+def cmd_gpus(_: argparse.Namespace) -> int:
+    for name, size in gpus.GPUS.items():
+        print(f"{name:<20} {size / gpus.GIB:>5.0f} GiB")
+    return 0
+
+
+def cmd_bench(a: argparse.Namespace) -> int:
+    from .bench import bench
+
+    kwargs = _estimate_kwargs(a)
+    r = bench(a.model_path, steps=a.steps, record=a.record, **kwargs)
+    print(json.dumps(r.to_dict(), indent=2) if a.json else r.table())
+    return 0 if r.status == "ok" else 1
+
+
+def main(argv: list[str] | None = None) -> int:
+    argv = list(sys.argv[1:] if argv is None else argv)
+    if argv and argv[0] not in _SUBCOMMANDS and argv[0] != "estimate" and not argv[0].startswith("-"):
+        argv.insert(0, "estimate")
+    parser = build_parser()
+    a = parser.parse_args(argv)
+    if a.command is None:
+        parser.print_help()
+        return 2
+    handler = {"estimate": cmd_estimate, "gpus": cmd_gpus, "bench": cmd_bench}[a.command]
+    try:
+        return handler(a)
+    except (ValueError, KeyError, FileNotFoundError, ImportError, RuntimeError) as exc:
+        msg = exc.args[0] if exc.args else str(exc)
+        print(f"vramcalc: {msg}", file=sys.stderr)
+        return 2
+
+
+if __name__ == "__main__":
+    sys.exit(main())

vramcalc/estimate.py

@@ -0,0 +1,219 @@
+"""Compose per-term formulas into an Estimate for a given mode."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import asdict, dataclass, field
+
+from . import memory
+from .arch import Arch
+from .loaders import load_arch
+
+GB = 10**9
+
+MODES = ("inference", "train")
+ATTN_IMPLS = ("eager", "sdpa", "flash")
+OPTIMIZERS = ("adamw", "sgd")
+PRECISIONS = ("amp", "pure")
+
+
+@dataclass(frozen=True)
+class EstimateParams:
+    arch: Arch
+    mode: str
+    dtype: str
+    batch: int
+    seq_len: int
+    lora_rank: int | None
+    grad_checkpoint: bool
+    optimizer: str
+    precision: str
+    attn_impl: str
+    overhead_gb: float
+
+    def replace(self, **changes) -> EstimateParams:
+        return EstimateParams(**{**asdict_shallow(self), **changes})
+
+
+def asdict_shallow(params: EstimateParams) -> dict:
+    """asdict() without recursing into Arch, so Arch stays an Arch."""
+    return {f: getattr(params, f) for f in EstimateParams.__dataclass_fields__}
+
+
+@dataclass(frozen=True)
+class Suggestion:
+    change: dict
+    changes_objective: bool
+    estimate: Estimate
+    margin_bytes: int
+
+
+@dataclass(frozen=True)
+class FitResult:
+    ok: bool
+    gpu_name: str
+    gpu_bytes: int
+    usable_bytes: int
+    required_bytes: int
+    margin_bytes: int
+    suggestions: list[Suggestion]
+    message: str
+
+    def to_dict(self) -> dict:
+        d = {k: getattr(self, k) for k in ("ok", "gpu_name", "gpu_bytes", "usable_bytes",
+                                            "required_bytes", "margin_bytes", "message")}
+        d["suggestions"] = [
+            {"change": s.change, "changes_objective": s.changes_objective,
+             "margin_bytes": s.margin_bytes, "total": s.estimate.total}
+            for s in self.suggestions
+        ]
+        return d
+
+
+@dataclass(frozen=True)
+class Estimate:
+    weights: int
+    kv_cache: int
+    gradients: int
+    optimizer: int
+    activations: int
+    autocast_cache: int
+    step_workspace: int
+    overhead: int
+    params: EstimateParams = field(repr=False)
+
+    @property
+    def persistent(self) -> int:
+        return self.weights + self.kv_cache + self.gradients + self.optimizer
+
+    @property
+    def peak_extra(self) -> int:
+        return max(self.activations + self.autocast_cache, self.step_workspace)
+
+    @property
+    def total(self) -> int:
+        return self.persistent + self.peak_extra
+
+    def to_dict(self) -> dict:
+        d = {k: getattr(self, k) for k in _TERMS}
+        d.update(persistent=self.persistent, peak_extra=self.peak_extra, total=self.total)
+        params = asdict_shallow(self.params)
+        params["arch"] = asdict(self.params.arch)
+        d["params"] = params
+        return d
+
+    def fits(self, gpu=None) -> FitResult:
+        from . import gpus  # local import keeps gpus optional for pure estimation
+        from .suggest import suggest
+
+        name, size = gpus.resolve(gpu)
+        usable = int(size * gpus.USABLE_RATIO)
+        required = self.total + self.overhead
+        margin = usable - required
+        ok = required <= usable
+        suggestions = [] if ok else suggest(self, usable)
+        if ok:
+            message = f"fits {name}: {required / GB:.1f} GB of {usable / GB:.1f} GB usable"
+        elif suggestions:
+            n = len(suggestions)
+            message = (f"short by {-margin / GB:.1f} GB on {name}; "
+                       f"{n} single adjustment{'s fit' if n != 1 else ' fits'}")
+        else:
+            message = (f"short by {-margin / GB:.1f} GB on {name}; "
+                       "no tested single adjustment fits. "
+                       "Combine adjustments or consider quantization or distributed training")
+        return FitResult(ok, name, size, usable, required, margin, suggestions, message)
+
+    def breakdown(self) -> str:
+        rows = [(k, getattr(self, k)) for k in _TERMS if k != "overhead" and getattr(self, k)]
+        rows += [("total", self.total), ("overhead", self.overhead)]
+        width = max(len(k) for k, _ in rows)
+        return "\n".join(f"{k:<{width}}  {v / GB:8.2f} GB" for k, v in rows)
+
+
+_TERMS = ("weights", "kv_cache", "gradients", "optimizer",
+          "activations", "autocast_cache", "step_workspace", "overhead")
+
+
+def _validate(p: EstimateParams) -> None:
+    if p.mode not in MODES:
+        raise ValueError(f"mode must be one of {MODES}, got {p.mode!r}")
+    if p.dtype not in memory.DTYPE_BYTES:
+        raise ValueError(f"dtype must be one of {tuple(memory.DTYPE_BYTES)}, got {p.dtype!r}")
+    if p.attn_impl not in ATTN_IMPLS:
+        raise ValueError(f"attn_impl must be one of {ATTN_IMPLS}, got {p.attn_impl!r}")
+    if p.optimizer not in OPTIMIZERS:
+        raise ValueError(f"optimizer must be one of {OPTIMIZERS}, got {p.optimizer!r}")
+    if p.precision not in PRECISIONS:
+        raise ValueError(f"precision must be one of {PRECISIONS}, got {p.precision!r}")
+    for name in ("batch", "seq_len"):
+        v = getattr(p, name)
+        if not isinstance(v, int) or isinstance(v, bool) or v <= 0:
+            raise ValueError(f"{name} must be a positive integer, got {v!r}")
+    if p.lora_rank is not None:
+        if p.mode != "train":
+            raise ValueError("lora_rank requires mode='train'")
+        if not isinstance(p.lora_rank, int) or p.lora_rank <= 0:
+            raise ValueError(f"lora_rank must be a positive integer, got {p.lora_rank!r}")
+    if (not isinstance(p.overhead_gb, (int, float)) or not math.isfinite(p.overhead_gb)
+            or p.overhead_gb < 0):
+        raise ValueError(f"overhead_gb must be a finite number >= 0, got {p.overhead_gb!r}")
+
+
+def estimate(
+    model,
+    *,
+    mode: str = "inference",
+    dtype: str = "bf16",
+    batch: int = 1,
+    seq_len: int = 2048,
+    lora_rank: int | None = None,
+    grad_checkpoint: bool = False,
+    optimizer: str = "adamw",
+    precision: str = "amp",
+    attn_impl: str = "sdpa",
+    overhead_gb: float = 0.75,
+) -> Estimate:
+    params = EstimateParams(
+        arch=load_arch(model), mode=mode, dtype=dtype, batch=batch, seq_len=seq_len,
+        lora_rank=lora_rank, grad_checkpoint=grad_checkpoint, optimizer=optimizer,
+        precision=precision, attn_impl=attn_impl, overhead_gb=overhead_gb,
+    )
+    return estimate_from_params(params)
+
+
+def estimate_from_params(p: EstimateParams) -> Estimate:
+    _validate(p)
+    B = memory.DTYPE_BYTES[p.dtype]
+    overhead = int(p.overhead_gb * GB)
+    if p.mode == "inference":
+        return Estimate(
+            weights=memory.weights(p.arch, B),
+            kv_cache=memory.kv_cache(p.arch, p.batch, p.seq_len, B),
+            gradients=0, optimizer=0,
+            activations=memory.inference_activations(p.arch, p.batch, p.seq_len, p.attn_impl, B),
+            autocast_cache=0, step_workspace=0,
+            overhead=overhead, params=p,
+        )
+    if p.dtype == "fp32" and p.precision == "amp":
+        p = p.replace(precision="pure")  # nothing to autocast
+    S = 4 if p.precision == "amp" else B
+    lora = p.lora_rank is not None
+    trainable = memory.lora_params(p.arch, p.lora_rank) if lora else p.arch.param_count()
+    adapter_bytes = 4 if lora else S
+    weights = memory.weights(p.arch, S) + (trainable * 4 if lora else 0)
+    return Estimate(
+        weights=weights,
+        kv_cache=0,
+        gradients=memory.gradients(trainable, adapter_bytes),
+        optimizer=memory.optimizer_states(trainable, p.optimizer, adapter_bytes),
+        activations=memory.train_activations(
+            p.arch, p.batch, p.seq_len, p.attn_impl, B,
+            grad_checkpoint=p.grad_checkpoint, lora_rank=p.lora_rank,
+        ),
+        autocast_cache=memory.autocast_cache(
+            p.arch, p.precision, B, lora=lora, grad_checkpoint=p.grad_checkpoint,
+        ),
+        step_workspace=memory.step_workspace(trainable, adapter_bytes),
+        overhead=overhead, params=p,
+    )

vramcalc/gpus.py

@@ -0,0 +1,76 @@
+"""Built-in GPU capacity table and resolution of the `gpu` argument of fits()."""
+
+from __future__ import annotations
+
+import math
+import re
+
+GIB = 2**30
+
+# Nominal capacity in GiB. Keys are canonical display names.
+GPUS: dict[str, int] = {name: gib * GIB for name, gib in {
+    "RTX 3060": 12, "RTX 3070": 8, "RTX 3080": 10, "RTX 3080 Ti": 12,
+    "RTX 3090": 24, "RTX 3090 Ti": 24,
+    "RTX 4060": 8, "RTX 4060 Ti": 16, "RTX 4070": 12, "RTX 4070 Ti": 12,
+    "RTX 4070 Ti Super": 16, "RTX 4080": 16, "RTX 4080 Super": 16, "RTX 4090": 24,
+    "RTX 5070": 12, "RTX 5070 Ti": 16, "RTX 5080": 16, "RTX 5090": 32,
+    "T4": 16, "V100 16GB": 16, "V100 32GB": 32,
+    "A10": 24, "A40": 48, "A100 40GB": 40, "A100 80GB": 80,
+    "H100 80GB": 80, "H200": 141, "L4": 24, "L40S": 48, "B200": 192,
+}.items()}
+
+# Provisional safety reserve for an otherwise idle GPU. Confirm with bench.
+USABLE_RATIO = 0.97
+
+_STRIP = re.compile(r"[\s\-_]+")
+_PREFIXES = ("nvidia", "geforce", "tesla", "rtx")
+_DETECT_HINT = "cannot detect GPU; pass a GPU name or GB number"
+
+
+def _key(name: str) -> str:
+    k = _STRIP.sub("", name.lower())
+    for p in _PREFIXES:
+        k = k.removeprefix(p)
+    return k
+
+
+_INDEX = {_key(n): n for n in GPUS}
+
+
+def lookup(name: str) -> tuple[str, int]:
+    k = _key(name)
+    if k in _INDEX:
+        canonical = _INDEX[k]
+        return canonical, GPUS[canonical]
+    candidates = [n for key, n in _INDEX.items() if k and (k in key or key in k)]
+    hint = f"; did you mean {', '.join(candidates)}?" if candidates else ""
+    raise KeyError(f"unknown GPU {name!r}{hint}")
+
+
+def detect() -> tuple[str, int]:
+    try:
+        import pynvml
+    except ImportError as exc:
+        raise RuntimeError(_DETECT_HINT) from exc
+    if pynvml is None:  # sys.modules poisoned in tests
+        raise RuntimeError(_DETECT_HINT)
+    try:
+        pynvml.nvmlInit()
+        handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+        name = pynvml.nvmlDeviceGetName(handle)
+        total = pynvml.nvmlDeviceGetMemoryInfo(handle).total
+    except Exception as exc:  # pynvml raises NVMLError subclasses
+        raise RuntimeError(_DETECT_HINT) from exc
+    return (name.decode() if isinstance(name, bytes) else name), int(total)
+
+
+def resolve(gpu) -> tuple[str, int]:
+    if gpu is None:
+        return detect()
+    if isinstance(gpu, bool) or not isinstance(gpu, (int, float, str)):
+        raise ValueError(f"gpu must be None, a GB number or a name, got {gpu!r}")  # noqa: TRY004
+    if isinstance(gpu, (int, float)):
+        if not math.isfinite(gpu) or gpu <= 0:
+            raise ValueError(f"gpu GB must be a positive finite number, got {gpu!r}")
+        return f"{gpu:g} GB", int(gpu * 10**9)
+    return lookup(gpu)

vramcalc/loaders.py

@@ -0,0 +1,57 @@
+"""Resolve the `model` argument of estimate() into an Arch."""
+
+from __future__ import annotations
+
+import dataclasses
+import json
+from pathlib import Path
+
+from .arch import Arch
+
+_HF_HINT = "is not a local path; install vramcalc[hf] to load from the Hub"
+
+
+def load_arch(model) -> Arch:
+    if isinstance(model, Arch):
+        return model
+    if isinstance(model, dict):
+        return Arch.from_config(model)
+    path = Path(model)
+    if path.is_dir():
+        config = path / "config.json"
+        if not config.is_file():
+            raise FileNotFoundError(f"no config.json in directory {path}")
+        return _named(Arch.from_config(_read_json(config)), path.name)
+    if path.is_file():
+        return _named(Arch.from_config(_read_json(path)), path.parent.name)
+    repo_id = str(model)
+    return _named(Arch.from_config(_download_config(repo_id)), repo_id.split("/")[-1])
+
+
+def _named(arch: Arch, fallback: str) -> Arch:
+    """Configs without _name_or_path get the input's own name so records stay identifiable."""
+    return arch if arch.name else dataclasses.replace(arch, name=fallback)
+
+
+def _read_json(path: Path) -> dict:
+    try:
+        return json.loads(path.read_text())
+    except json.JSONDecodeError as exc:
+        raise ValueError(f"{path} is not valid JSON: {exc}") from exc
+
+
+def _download_config(repo_id: str) -> dict:
+    try:
+        import huggingface_hub
+    except ImportError as exc:
+        raise ImportError(f"{repo_id!r} {_HF_HINT}") from exc
+    if huggingface_hub is None:  # sys.modules poisoned in tests
+        raise ImportError(f"{repo_id!r} {_HF_HINT}")
+    try:
+        local = huggingface_hub.hf_hub_download(repo_id, "config.json")
+    except Exception as exc:
+        raise FileNotFoundError(
+            f"cannot fetch config.json for {repo_id!r}: {exc}"
+            " (gated repo? run `huggingface-cli login`; offline? pass a local path)"
+        ) from exc
+    return _read_json(Path(local))

vramcalc/memory.py

@@ -0,0 +1,114 @@
+"""Per-term memory formulas. Every function returns bytes unless its name says elements.
+
+See spec section 5. Symbols: L layers, h hidden, a heads, q = a*head_dim,
+kv = kv_heads*head_dim, i intermediate, V vocab, b batch, s seq_len.
+"""
+
+from __future__ import annotations
+
+from .arch import Arch
+
+DTYPE_BYTES = {"fp32": 4, "bf16": 2, "fp16": 2}
+
+# RMSNorm keeps fp32 intermediates for backward; count per layer. Confirm with bench.
+NORM_FP32_TENSORS = 2
+
+
+def weights(arch: Arch, dtype_bytes: int) -> int:
+    return arch.param_count() * dtype_bytes
+
+
+def kv_cache(arch: Arch, batch: int, seq_len: int, dtype_bytes: int) -> int:
+    return 2 * arch.num_layers * batch * seq_len * arch.kv * dtype_bytes
+
+
+def per_layer_act(arch: Arch, batch: int, seq_len: int, attn_impl: str, dtype_bytes: int) -> int:
+    """Elements saved per layer for backward (also the transient per-layer footprint in inference)."""
+    h, q, kv, i, a = arch.hidden_size, arch.q, arch.kv, arch.intermediate_size, arch.num_heads
+    tokens = seq_len * batch
+    elements = tokens * (2 * h + q + 2 * kv + q + h + 4 * i)
+    # The +2h term is empirical. A6000 runs on Llama-3.2-1B/3B and Qwen2.5-1.5B leave a
+    # residual of about two hidden-sized tensors per layer after the logits term is fixed.
+    # Split it into named tensors when a model family shows >10% peak error.
+    elements += tokens * 2 * h
+    if attn_impl in ("sdpa", "eager") and kv < q:
+        elements += 2 * tokens * q  # repeat_kv materializes K and V at query width
+    if attn_impl == "eager":
+        # score matrix in dtype plus fp32 softmax, expressed in dtype-sized elements
+        elements += a * seq_len * seq_len * batch * (dtype_bytes + 4) // dtype_bytes
+    return elements
+
+
+def inference_activations(
+    arch: Arch, batch: int, seq_len: int, attn_impl: str, dtype_bytes: int
+) -> int:
+    """Forward-only peak: prefill logits plus a small working set.
+
+    Per-layer buffers are freed as the forward proceeds, so they do not survive to the peak
+    (measured: kv_cache + logits explains the peak within 30 MB on 1B to 8B models).
+    """
+    tokens = seq_len * batch
+    logits = tokens * arch.vocab_size * dtype_bytes
+    working = 2 * tokens * arch.hidden_size * dtype_bytes
+    if attn_impl == "eager":
+        # score matrix in dtype plus fp32 softmax for one layer, and repeat_kv K/V expansion
+        working += arch.num_heads * seq_len * seq_len * batch * (dtype_bytes + 4)
+        working += 2 * tokens * arch.q * dtype_bytes
+    return logits + working
+
+
+LORA_TARGETS = 4  # q_proj, k_proj, v_proj, o_proj
+
+
+def lora_params(arch: Arch, rank: int) -> int:
+    h, q, kv = arch.hidden_size, arch.q, arch.kv
+    # in+out per target: q (h,q), k (h,kv), v (h,kv), o (q,h)
+    return arch.num_layers * rank * (4 * h + 2 * q + 2 * kv)
+
+
+def cast_params(arch: Arch) -> int:
+    """Linear weights autocast copies; the output linear is cast once whether tied or not."""
+    return arch.num_layers * (arch.attn_params() + arch.mlp_params()) + arch.embed_params()
+
+
+def gradients(trainable: int, state_bytes: int) -> int:
+    return trainable * state_bytes
+
+
+def optimizer_states(trainable: int, optimizer: str, state_bytes: int) -> int:
+    moments = 2 if optimizer == "adamw" else 1
+    return trainable * moments * state_bytes
+
+
+def step_workspace(trainable: int, state_bytes: int) -> int:
+    """Temporary buffer of the foreach optimizer implementation, alive only during step()."""
+    return trainable * state_bytes
+
+
+def autocast_cache(
+    arch: Arch, precision: str, dtype_bytes: int, *, lora: bool, grad_checkpoint: bool
+) -> int:
+    if precision != "amp":
+        return 0
+    full = cast_params(arch) * dtype_bytes
+    if lora and grad_checkpoint:
+        # frozen casts are not cached; only the layer being recomputed holds them
+        return full // arch.num_layers
+    return full
+
+
+def train_activations(
+    arch: Arch, batch: int, seq_len: int, attn_impl: str, dtype_bytes: int,
+    *, grad_checkpoint: bool, lora_rank: int | None,
+) -> int:
+    L, h, V = arch.num_layers, arch.hidden_size, arch.vocab_size
+    tokens = seq_len * batch
+    layer = per_layer_act(arch, batch, seq_len, attn_impl, dtype_bytes) * dtype_bytes
+    norm_fp32 = NORM_FP32_TENSORS * 4 * tokens * h * L
+    lora_branch = LORA_TARGETS * L * tokens * lora_rank * 4 if lora_rank else 0
+    # logits peak sits in log_softmax backward: dtype logits + saved fp32 log_softmax
+    # + fp32 grad_out + fp32 grad_in (the fp32 upcast itself is freed before backward)
+    logits = tokens * V * (dtype_bytes + 4 + 4 + 4)
+    if grad_checkpoint:
+        return L * tokens * h * dtype_bytes + layer + norm_fp32 // L + lora_branch + logits
+    return L * layer + norm_fp32 + lora_branch + logits

vramcalc/suggest.py

@@ -0,0 +1,45 @@
+"""Adjustment suggestions when an estimate does not fit. Spec section 6.3."""
+
+from __future__ import annotations
+
+from .estimate import Estimate, EstimateParams, Suggestion, estimate_from_params
+
+MIN_SEQ = 512
+MIN_BATCH = 1
+
+
+def _halving(p: EstimateParams, key: str, floor: int, usable: int):
+    """Return the first halved value of `key` whose estimate fits, or None."""
+    value = getattr(p, key)
+    while value > floor:
+        value = max(value // 2, floor)
+        e = estimate_from_params(p.replace(**{key: value}))
+        if e.total + e.overhead <= usable:
+            return {key: value}
+    return None
+
+
+def _candidates(p: EstimateParams, usable: int):
+    train = p.mode == "train"
+    if train and not p.grad_checkpoint:
+        yield {"grad_checkpoint": True}, False
+    for key, floor in (("batch", MIN_BATCH), ("seq_len", MIN_SEQ)):
+        hit = _halving(p, key, floor, usable)
+        if hit:
+            yield hit, train
+    if train and p.lora_rank is None:
+        yield {"lora_rank": 16}, True
+    if train and p.optimizer == "adamw":
+        yield {"optimizer": "sgd"}, True
+    if train and p.precision == "amp":
+        yield {"precision": "pure"}, True
+
+
+def suggest(est: Estimate, usable_bytes: int) -> list[Suggestion]:
+    out: list[Suggestion] = []
+    for change, changes_objective in _candidates(est.params, usable_bytes):
+        e = estimate_from_params(est.params.replace(**change))
+        margin = usable_bytes - (e.total + e.overhead)
+        if margin >= 0:
+            out.append(Suggestion(change, changes_objective, e, margin))
+    return out

vramcalc-0.1.0.dist-info/METADATA

@@ -0,0 +1,237 @@
+Metadata-Version: 2.5
+Name: vramcalc
+Version: 0.1.0
+Summary: Predict GPU VRAM usage of Llama-family models from config.json and check whether they fit a GPU
+License-Expression: MIT
+License-File: LICENSE
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Provides-Extra: bench
+Requires-Dist: nvidia-ml-py; extra == 'bench'
+Requires-Dist: peft>=0.12; extra == 'bench'
+Requires-Dist: torch>=2.3; extra == 'bench'
+Requires-Dist: transformers<5,>=4.46; extra == 'bench'
+Provides-Extra: dev
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Provides-Extra: hf
+Requires-Dist: huggingface-hub; extra == 'hf'
+Description-Content-Type: text/markdown
+
+# vramcalc
+
+[English](#english) | [한국어](#한국어)
+
+---
+
+## English
+
+Predict how much GPU memory an LLM needs before you download the weights or rent the GPU. vramcalc reads only the Hugging Face `config.json`, tells you whether the model fits a given GPU, and when it does not, which single change would make it fit. A `bench` subcommand measures real memory on your GPU and compares it with the prediction, bucket by bucket.
+
+```python
+from vramcalc import estimate
+
+e = estimate("meta-llama/Llama-3.1-8B", mode="train", lora_rank=16,
+             precision="pure", seq_len=4096)
+print(e.breakdown())
+r = e.fits("RTX 3090")
+print(r.message)
+for s in r.suggestions:
+    print(s.change, s.margin_bytes)
+```
+
+```
+$ vramcalc meta-llama/Llama-3.1-8B --train --lora 16 --precision pure --seq 4096 --gpu "RTX 3090"
+weights            16.12 GB
+gradients           0.05 GB
+optimizer           0.11 GB
+activations        36.92 GB
+step_workspace      0.05 GB
+total              53.19 GB
+overhead            0.75 GB
+
+short by 28.9 GB on RTX 3090; 1 single adjustment fits
+  -> seq_len=512: 21.6 GB (changes training objective)
+```
+
+### Install
+
+| Command | What you get |
+|---|---|
+| `pip install vramcalc` | Estimator and CLI. No dependencies. Takes a local `config.json` or a model directory |
+| `pip install "vramcalc[hf]"` | Hugging Face Hub repo ids as input (`huggingface_hub`) |
+| `pip install "vramcalc[bench]"` | `vramcalc bench` for real measurements (`torch`, `transformers`, `peft`, `nvidia-ml-py`) |
+
+Offline or on-premises: pass a model directory, `estimate("/path/to/model")`. With `HF_HUB_OFFLINE=1`, Hub ids resolve from the local cache only.
+
+### Why
+
+"Will this model run on this GPU with these settings?" is a question you want answered before downloading 16 GB of weights or paying for an H100 hour. Existing calculators mostly size the weights and treat activations, optimizer state and autocast copies loosely, which is where training estimates drift by several gigabytes. vramcalc keeps one formula per memory term and ships a `bench` command that measures the same configuration on a real GPU, so every term can be checked against the allocator. The formulas are fixed; there is no learned correction factor. Measurements are evidence for fixing a formula, not a fudge applied on top.
+
+### Scope
+
+- Models: dense decoder-only models with `model_type` of `llama`, `mistral`, `qwen2`, `gemma` or `gemma2`. GQA, QKV bias (Qwen2), tied embeddings and Gemma 2's four norms per layer are read from the config.
+- Modes: inference (with KV cache), full fine-tuning, LoRA on `q_proj`, `k_proj`, `v_proj`, `o_proj`.
+- Knobs: `dtype` (fp32/bf16/fp16), `batch`, `seq_len`, `grad_checkpoint`, `optimizer` (adamw/sgd), `precision` (amp/pure), `attn_impl` (eager/sdpa/flash).
+
+`precision="amp"` models HF Trainer with `bf16=True` (fp32 weights, bf16 autocast). `precision="pure"` models loading with `torch_dtype=bf16` and training as is. The two differ by almost 2x in memory, so pick the one you actually run.
+
+Out of scope for 0.1: MoE, MLA, quantization (int8/int4/QLoRA), distributed training (ZeRO/FSDP), multimodal and encoder models, and sliding-window KV caches (currently counted as global attention, which is conservative).
+
+### What is counted
+
+The prediction targets PyTorch's caching allocator (`allocated` memory). CUDA context and other memory outside the allocator is kept separate as `overhead` (default 0.75 GB) and only added in `fits()`.
+
+| Term | Meaning |
+|---|---|
+| `weights` | parameter count x storage dtype; for LoRA, frozen base plus fp32 adapters |
+| `kv_cache` | inference only, `2 x L x b x s x kv x bytes` |
+| `gradients`, `optimizer` | over trainable parameters; AdamW keeps two moments, SGD one |
+| `activations` | per-layer tensors saved for backward (SwiGLU, RMSNorm) plus the fp32 logits chain at the `log_softmax` backward peak |
+| `autocast_cache` | bf16 weight copies alive during an amp forward pass |
+| `step_workspace` | temporary buffer of the foreach optimizer during `step()` |
+
+`total = weights + kv_cache + gradients + optimizer + max(activations + autocast_cache, step_workspace)`. The full formulas are in `docs/design.md`, section 5.
+
+### Validation
+
+`vramcalc bench` loads the model with the same arguments, records `torch.cuda.max_memory_allocated()`, and compares five buckets: weights, states, peak, total, overhead.
+
+```
+vramcalc bench /path/to/Llama-3.2-1B --train --precision pure --record benchmarks/results.jsonl
+```
+
+Environment: NVIDIA RTX A6000 48 GB, torch 2.14.1+cu130, transformers 4.57.6, peft 0.21.2. Llama weights are gated, so the `unsloth/` mirrors of the same checkpoints were used. All 19 runs land within +/-4.3% on `total`; `weights` is exact in every run.
+
+| model | mode | setting | GPU | predicted | measured | error |
+|---|---|---|---|---|---|---|
+| Llama-3.2-1B | inference | bf16 b=1 s=2048 sdpa | RTX A6000 | 3.08 GB | 3.08 GB | -0.1% |
+| Llama-3.2-1B | inference | bf16 b=1 s=4096 sdpa | RTX A6000 | 3.69 GB | 3.68 GB | +0.2% |
+| Llama-3.2-1B | inference | bf16 b=1 s=2048 eager | RTX A6000 | 3.90 GB | 3.95 GB | -1.2% |
+| Llama-3.2-3B | inference | bf16 b=1 s=2048 sdpa | RTX A6000 | 7.21 GB | 7.21 GB | +0.0% |
+| Llama-3.1-8B | inference | bf16 b=1 s=2048 sdpa | RTX A6000 | 16.89 GB | 16.88 GB | +0.0% |
+| Qwen2.5-1.5B | inference | bf16 b=1 s=2048 sdpa | RTX A6000 | 3.78 GB | 3.78 GB | -0.1% |
+| Llama-3.2-1B | train | bf16 b=1 s=2048 sdpa full pure | RTX A6000 | 17.52 GB | 17.51 GB | +0.0% |
+| Llama-3.2-1B | train | bf16 b=1 s=2048 sdpa full pure ckpt | RTX A6000 | 13.95 GB | 13.76 GB | +1.4% |
+| Llama-3.2-1B | train | bf16 b=1 s=2048 sdpa full amp | RTX A6000 | 29.88 GB | 30.55 GB | -2.2% |
+| Llama-3.2-1B | train | bf16 b=1 s=2048 sdpa full amp ckpt | RTX A6000 | 26.30 GB | 25.22 GB | +4.3% |
+| Llama-3.2-1B | train | bf16 b=1 s=2048 eager full pure | RTX A6000 | 30.41 GB | 30.26 GB | +0.5% |
+| Llama-3.2-3B | train | bf16 b=1 s=2048 sdpa full pure ckpt | RTX A6000 | 32.13 GB | 32.15 GB | -0.1% |
+| Qwen2.5-1.5B | train | bf16 b=1 s=2048 sdpa full pure | RTX A6000 | 23.17 GB | 23.40 GB | -1.0% |
+| Llama-3.2-1B | train | bf16 b=1 s=2048 sdpa lora=16 pure | RTX A6000 | 10.17 GB | 10.13 GB | +0.4% |
+| Llama-3.2-1B | train | bf16 b=1 s=2048 sdpa lora=16 amp | RTX A6000 | 15.11 GB | 14.52 GB | +4.1% |
+| Llama-3.2-1B | train | bf16 b=1 s=2048 sdpa lora=16 amp ckpt | RTX A6000 | 9.22 GB | 9.51 GB | -3.1% |
+| Llama-3.2-3B | train | bf16 b=1 s=2048 sdpa lora=16 pure | RTX A6000 | 18.84 GB | 18.98 GB | -0.8% |
+| Llama-3.1-8B | train | bf16 b=1 s=2048 sdpa lora=16 pure ckpt | RTX A6000 | 20.97 GB | 20.55 GB | +2.0% |
+| Qwen2.5-1.5B | train | bf16 b=1 s=2048 sdpa lora=16 pure | RTX A6000 | 13.99 GB | 13.78 GB | +1.5% |
+
+`pure` is `torch_dtype=bf16` training, `amp` is fp32 weights with bf16 autocast, `ckpt` is gradient checkpointing. Per-bucket numbers are in `benchmarks/results.jsonl`; regenerate the table with `python benchmarks/make_table.py`.
+
+What the measurements changed in the formulas: the inference peak is just `kv_cache + logits` (per-layer buffers are freed as the forward proceeds); the training logits peak is 14 bytes per logit, not 10, because `log_softmax` backward holds its saved output and both gradient tensors at once; and each layer saves about two more hidden-sized tensors than the textbook count.
+
+### Limits
+
+- `bench` uses `torch.autocast` directly, not HF Trainer. Extra conversions done by Accelerate are not modelled.
+- `fits()` assumes nothing else is using the GPU and takes 97% of the nominal capacity as usable.
+- `attn_impl="sdpa"` assumes the memory-efficient or flash backend. If your environment falls back to the math backend, use `eager`.
+- Overhead was measured on an A6000 with CUDA 13; other cards and driver versions may differ by a few hundred megabytes.
+
+### License
+
+MIT
+
+---
+
+## 한국어
+
+모델 가중치를 받기 전, GPU를 빌리기 전에 LLM이 얼마만큼의 GPU 메모리를 필요로 하는지 예측하는 라이브러리입니다. Hugging Face `config.json`만 읽어 지정한 GPU에 들어가는지 판정하고, 들어가지 않으면 무엇을 하나 바꾸면 들어가는지 제안합니다. `bench` 명령은 실제 GPU에서 메모리를 측정해 예측과 항목별로 비교합니다.
+
+```python
+from vramcalc import estimate
+
+e = estimate("meta-llama/Llama-3.1-8B", mode="train", lora_rank=16,
+             precision="pure", seq_len=4096)
+print(e.breakdown())
+r = e.fits("RTX 3090")
+print(r.message)
+for s in r.suggestions:
+    print(s.change, s.margin_bytes)
+```
+
+```
+$ vramcalc meta-llama/Llama-3.1-8B --train --lora 16 --precision pure --seq 4096 --gpu "RTX 3090"
+weights            16.12 GB
+gradients           0.05 GB
+optimizer           0.11 GB
+activations        36.92 GB
+step_workspace      0.05 GB
+total              53.19 GB
+overhead            0.75 GB
+
+short by 28.9 GB on RTX 3090; 1 single adjustment fits
+  -> seq_len=512: 21.6 GB (changes training objective)
+```
+
+### 설치
+
+| 명령 | 포함 내용 |
+|---|---|
+| `pip install vramcalc` | 계산기와 CLI. 의존성 없음. 로컬 `config.json` 또는 모델 디렉터리 입력 |
+| `pip install "vramcalc[hf]"` | Hugging Face Hub repo id 입력 (`huggingface_hub`) |
+| `pip install "vramcalc[bench]"` | 실측 비교 `vramcalc bench` (`torch`, `transformers`, `peft`, `nvidia-ml-py`) |
+
+오프라인이나 온프레미스 환경에서는 `estimate("/path/to/model")`처럼 모델 디렉터리를 넘기면 됩니다. `HF_HUB_OFFLINE=1`이면 Hub id도 로컬 캐시만 조회합니다.
+
+### 만든 이유
+
+"이 모델을 이 설정으로 이 GPU에서 돌릴 수 있는가"는 16GB 가중치를 받거나 H100을 한 시간 빌리기 전에 알아야 하는 질문입니다. 기존 계산기는 대부분 가중치 크기 위주이고 activation, optimizer 상태, autocast 복사본은 대략적으로만 다루기 때문에 학습 모드에서 수 GB씩 어긋납니다. vramcalc는 메모리 항목마다 식을 하나씩 두고, 같은 설정을 실제 GPU에서 측정하는 `bench` 명령을 함께 제공해 항목별로 allocator 값과 대조할 수 있게 했습니다. 식은 고정이며 보정 계수를 학습하지 않습니다. 실측은 식을 고치는 근거이지, 위에 덧씌우는 보정이 아닙니다.
+
+### 지원 범위
+
+- 모델: `model_type`이 `llama`, `mistral`, `qwen2`, `gemma`, `gemma2`인 dense decoder-only 모델. GQA, QKV bias(Qwen2), 임베딩 공유, Gemma 2의 레이어당 norm 4개를 config에서 읽어 반영합니다.
+- 모드: 추론(KV cache 포함), 전체 학습, LoRA(`q_proj`, `k_proj`, `v_proj`, `o_proj` 대상)
+- 옵션: `dtype`(fp32/bf16/fp16), `batch`, `seq_len`, `grad_checkpoint`, `optimizer`(adamw/sgd), `precision`(amp/pure), `attn_impl`(eager/sdpa/flash)
+
+`precision="amp"`는 HF Trainer `bf16=True` 방식(fp32 가중치, bf16 autocast)이고, `precision="pure"`는 `torch_dtype=bf16`으로 로드해 그대로 학습하는 방식입니다. 둘의 메모리는 2배 가까이 차이가 나므로 실제로 사용하는 방식을 지정해야 합니다.
+
+0.1에서 제외된 항목: MoE, MLA, 양자화(int8/int4/QLoRA), 분산 학습(ZeRO/FSDP), 멀티모달과 encoder 모델, sliding window KV cache(현재는 전부 global attention으로 보수적으로 계산).
+
+### 계산 항목
+
+예측 대상은 PyTorch caching allocator 기준 `allocated` 메모리입니다. CUDA context 등 allocator 밖 메모리는 `overhead`(기본 0.75GB)로 분리해 `fits()` 판정에서만 더합니다.
+
+| 항목 | 내용 |
+|---|---|
+| `weights` | 파라미터 수 × 저장 dtype. LoRA는 frozen base + fp32 어댑터 |
+| `kv_cache` | 추론 전용. `2 × L × b × s × kv × bytes` |
+| `gradients`, `optimizer` | 학습 대상 파라미터 기준. AdamW는 상태 2개, SGD는 momentum 1개 |
+| `activations` | 레이어당 역전파 저장분(SwiGLU, RMSNorm) + `log_softmax` 역전파 시점의 fp32 logits 체인 |
+| `autocast_cache` | amp에서 forward 동안 살아 있는 bf16 가중치 복사본 |
+| `step_workspace` | foreach optimizer의 `step()` 중 임시 버퍼 |
+
+`total = weights + kv_cache + gradients + optimizer + max(activations + autocast_cache, step_workspace)`. 상세 식은 `docs/design.md` 5절에 있습니다.
+
+### 실측 검증
+
+`vramcalc bench`는 같은 인자로 모델을 실제로 로드해 `torch.cuda.max_memory_allocated()`를 기록하고 weights, states, peak, total, overhead 다섯 묶음으로 예측과 비교합니다.
+
+```
+vramcalc bench /path/to/Llama-3.2-1B --train --precision pure --record benchmarks/results.jsonl
+```
+
+측정 환경: NVIDIA RTX A6000 48GB, torch 2.14.1+cu130, transformers 4.57.6, peft 0.21.2. Llama 가중치는 gated라 동일 가중치의 `unsloth/` 미러를 사용했습니다. 19개 조합 모두 `total` 오차 ±4.3% 안이며 `weights`는 전부 0.0%입니다. 표는 위 영어 절과 같고, 항목별 수치는 `benchmarks/results.jsonl`에 있습니다. `python benchmarks/make_table.py`로 표를 다시 생성할 수 있습니다.
+
+실측으로 바뀐 식: 추론 peak는 `kv_cache + logits`만으로 설명되고(레이어 버퍼는 forward 진행 중 해제됨), 학습 logits peak는 logit당 10바이트가 아니라 14바이트이며(`log_softmax` 역전파가 저장 출력과 두 gradient를 동시에 보유), 레이어마다 교과서 계산보다 hidden 크기 텐서 약 2개가 더 저장됩니다.
+
+### 한계
+
+- `bench`는 HF Trainer가 아니라 `torch.autocast`를 직접 사용합니다. Accelerate가 추가로 수행하는 변환은 반영하지 않습니다.
+- `fits()`는 GPU에 다른 프로세스가 없다고 가정하고 표기 용량의 97%를 사용 가능 용량으로 봅니다.
+- `attn_impl="sdpa"`는 memory-efficient 또는 flash backend를 전제로 합니다. math backend로 떨어지는 환경이면 `eager`로 지정해야 합니다.
+- overhead는 A6000과 CUDA 13에서 측정한 값이라 다른 카드와 드라이버에서는 수백 MB 차이가 날 수 있습니다.
+
+### 라이선스
+
+MIT

vramcalc-0.1.0.dist-info/RECORD

@@ -0,0 +1,14 @@
+vramcalc/__init__.py,sha256=DUN1u42xUL8MEdOvxGS8DCqCVIIMIaZrcyDh51qOGWs,393
+vramcalc/arch.py,sha256=NMnvXXJR8bAK1qdkeU_8LXCfwHgrz4QfKGhkoKYquKc,4092
+vramcalc/bench.py,sha256=ri4dUM_eibBv4BJb6FF9s6Z_8Nt0TXIL7s1gXvy-yME,8764
+vramcalc/cli.py,sha256=pqXR3kuepRLpCVuHv9Y9spWzr5goaumZCsl6D07zpQE,4548
+vramcalc/estimate.py,sha256=B4_gVJpVLMLeu8LcmscUOr05dRL2azNb7NgasO4xwAI,7811
+vramcalc/gpus.py,sha256=ONH4QdllhkX9tq2CZHZNkYuI-QAT8vAwGlzoO-_R8Qs,2709
+vramcalc/loaders.py,sha256=8LZhso1SeMfIg5KMxJlS7oibMRBLcjIP_mVHVYzSge0,1948
+vramcalc/memory.py,sha256=rg2bJDeuuSEHDzyfJxOZF4M1KKg4KR_gab8SC8MCgz8,4730
+vramcalc/suggest.py,sha256=8quT_oHviKZCVYd3PnEzvhmmLMbfwtQJW1fDAxnQk-c,1589
+vramcalc-0.1.0.dist-info/METADATA,sha256=FVI9Lor-t74oD2yro-5FQw1Uk0g2BPa0UxY3-o_I0co,14832
+vramcalc-0.1.0.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
+vramcalc-0.1.0.dist-info/entry_points.txt,sha256=9gfnZdZWjOVUTQmszl3qMLvQyUs83bskk-IiFNSwi7g,47
+vramcalc-0.1.0.dist-info/licenses/LICENSE,sha256=sZhDRicBLDYEeQNMKQ5o6sRsujSsfddgm_NG8u0a5iY,1068
+vramcalc-0.1.0.dist-info/RECORD,,

vramcalc-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.4
+Root-Is-Purelib: true
+Tag: py3-none-any

vramcalc-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+vramcalc = vramcalc.cli:main

vramcalc-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Seonho Hong
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.