File size: 4,203 Bytes
dc9acb9 | 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 116 117 118 119 120 121 122 123 | """
Mide el uso de VRAM (pico) de un modelo en distintas presentaciones:
- transformers bf16 (19.3 GB → NO cabe en 4 GB VRAM)
- transformers 4-bit (bitsandbytes NF4, ~6 GB)
- GGUF Q4_K_M vía llama.cpp (llama-cli / llama-server)
Modo transformers:
python measure_vram.py --hf-dir <dir> [--nf4]
Modo GGUF:
python measure_vram.py --gguf <file.gguf> [--n-gpu-layers 999]
Reporta pico de VRAM (torch.cuda.max_memory_allocated + nvidia-smi) y footprint.
"""
import argparse
import subprocess
import time
from pathlib import Path
PROMPTS = [
"What is the capital of France?",
"Explain the water cycle in three sentences.",
"Write a haiku about winter.",
]
def nvidia_mem_used_mb():
out = subprocess.run(
["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
capture_output=True, text=True, check=False,
).stdout.strip()
try:
return int(out.splitlines()[0].strip())
except (ValueError, IndexError):
return -1
def measure_transformers(path, nf4=False, max_new=24):
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
if torch.cuda.is_available():
torch.cuda.reset_peak_memory_stats()
kwargs = {"torch_dtype": torch.bfloat16 if not nf4 else None}
if nf4:
from transformers import BitsAndBytesConfig
kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
print(f"Cargando {'NF4 4-bit' if nf4 else 'bf16'} desde {path} ...")
model = AutoModelForCausalLM.from_pretrained(str(path), **kwargs, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(str(path))
print(f" footprint: {model.get_memory_footprint() / 1e9:.2f} GB")
before = nvidia_mem_used_mb()
device = next(model.parameters()).device
with torch.inference_mode():
for p in PROMPTS:
inputs = tokenizer(p, return_tensors="pt").to(device)
out = model.generate(**inputs, max_new_tokens=max_new, do_sample=False)
print(f" → {tokenizer.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)[:60]}")
if torch.cuda.is_available():
peak = torch.cuda.max_memory_allocated() / 1e6
else:
peak = -1
after = nvidia_mem_used_mb()
print(f" Pico VRAM (torch): {peak:.0f} MB")
print(f" VRAM nvidia-smi: antes={before} MB → tras generar={after} MB (delta {after - before} MB)")
def measure_gguf(path, n_gpu_layers=999, max_new=24, ckpt="/tmp/vram.ckpt"):
binary = "/home/methodwhite/.local/bin/llama-cli"
if not Path(binary).exists():
print(f"llama-cli no encontrado en {binary}")
return
cmd = [
binary, "-m", str(path), "-n", str(max_new), "-p", "What is the capital of France?",
"-ngl", str(n_gpu_layers), "--no-display-prompt", "--no-warmup", "-c", "512",
]
print(f"Ejecutando: {' '.join(cmd)}")
before = nvidia_mem_used_mb()
proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
peak = before
try:
while True:
line = proc.stdout.readline()
if not line:
break
line = line.strip()
if line and "log" not in line.lower() and not line.startswith("llama_"):
pass
used = nvidia_mem_used_mb()
if used > peak:
peak = used
if proc.poll() is not None:
break
time.sleep(0.2)
finally:
proc.kill()
print(f" VRAM pico (nvidia-smi): {peak} MB (antes {before} MB)")
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--hf-dir", type=str)
parser.add_argument("--nf4", action="store_true", help="Cargar en 4-bit NF4 en vez de bf16")
parser.add_argument("--gguf", type=str)
parser.add_argument("--n-gpu-layers", type=int, default=999)
args = parser.parse_args()
if args.hf_dir:
measure_transformers(args.hf_dir, nf4=args.nf4)
elif args.gguf:
measure_gguf(args.gguf, args.n_gpu_layers)
else:
parser.print_help()
if __name__ == "__main__":
main()
|