mirror of https://github.com/razor-ai/soup.git
131 lines
3.9 KiB
Python
131 lines
3.9 KiB
Python
"""GPU detection, memory calculation, and auto batch size."""
|
|
|
|
import math
|
|
|
|
|
|
def detect_device() -> tuple[str, str]:
|
|
"""Detect available device. Returns (device_string, human_name)."""
|
|
try:
|
|
import torch
|
|
|
|
if torch.cuda.is_available():
|
|
name = torch.cuda.get_device_name(0)
|
|
return "cuda", name
|
|
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
|
return "mps", "Apple Silicon (MPS)"
|
|
except ImportError:
|
|
pass
|
|
|
|
return "cpu", "CPU (no GPU detected)"
|
|
|
|
|
|
def get_gpu_info() -> dict:
|
|
"""Get GPU memory info."""
|
|
try:
|
|
import torch
|
|
|
|
if torch.cuda.is_available():
|
|
total = torch.cuda.get_device_properties(0).total_memory
|
|
total_gb = total / (1024**3)
|
|
return {
|
|
"memory_total": f"{total_gb:.1f} GB",
|
|
"memory_total_bytes": total,
|
|
"gpu_count": torch.cuda.device_count(),
|
|
}
|
|
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
|
# MPS doesn't expose memory easily, estimate from system
|
|
return {
|
|
"memory_total": "shared (Apple Silicon)",
|
|
"memory_total_bytes": 0,
|
|
"gpu_count": 1,
|
|
}
|
|
except ImportError:
|
|
pass
|
|
|
|
return {
|
|
"memory_total": "N/A (CPU mode)",
|
|
"memory_total_bytes": 0,
|
|
"gpu_count": 0,
|
|
}
|
|
|
|
|
|
def estimate_batch_size(
|
|
model_params_b: float,
|
|
seq_length: int,
|
|
gpu_memory_bytes: int,
|
|
quantization: str = "4bit",
|
|
lora_r: int = 64,
|
|
) -> int:
|
|
"""Estimate max batch size that fits in GPU memory.
|
|
|
|
Conservative estimate — better to start smaller and gradient accumulate.
|
|
"""
|
|
if gpu_memory_bytes == 0:
|
|
return 1 # CPU fallback
|
|
|
|
gpu_gb = gpu_memory_bytes / (1024**3)
|
|
|
|
# Rough memory per param based on quantization
|
|
bytes_per_param = {"4bit": 0.5, "8bit": 1.0, "none": 2.0} # FP16
|
|
bpp = bytes_per_param.get(quantization, 2.0)
|
|
|
|
# Model memory (static)
|
|
model_mem_gb = model_params_b * bpp
|
|
|
|
# LoRA trainable params (usually ~1-3% of total)
|
|
lora_ratio = min(lora_r * 2 / 4096, 0.05) # rough estimate
|
|
trainable_mem_gb = model_params_b * 2 * lora_ratio # FP16 for trainable
|
|
|
|
# Optimizer states (Adam: 2x params)
|
|
optimizer_mem_gb = trainable_mem_gb * 2
|
|
|
|
# Available for activations
|
|
overhead_gb = 1.5 # CUDA overhead, fragmentation
|
|
available_gb = gpu_gb - model_mem_gb - trainable_mem_gb - optimizer_mem_gb - overhead_gb
|
|
|
|
if available_gb <= 0:
|
|
return 1
|
|
|
|
# Rough activation memory per sample per token
|
|
# ~2 bytes per hidden dim per layer per token for a transformer
|
|
activation_per_sample_gb = (seq_length * model_params_b * 0.001) # very rough
|
|
activation_per_sample_gb = max(activation_per_sample_gb, 0.5) # minimum 0.5 GB
|
|
|
|
batch_size = max(1, int(available_gb / activation_per_sample_gb))
|
|
# Clamp to power of 2 (common practice)
|
|
batch_size = 2 ** int(math.log2(batch_size)) if batch_size > 1 else 1
|
|
|
|
return min(batch_size, 32) # cap at 32
|
|
|
|
|
|
def model_size_from_name(model_name: str) -> float:
|
|
"""Guess model size in billions from model name."""
|
|
name_lower = model_name.lower()
|
|
|
|
size_markers = [
|
|
("70b", 70), ("65b", 65), ("34b", 34), ("33b", 33),
|
|
("13b", 13), ("8b", 8), ("7b", 7), ("3b", 3),
|
|
("1.5b", 1.5), ("1b", 1), ("0.5b", 0.5),
|
|
]
|
|
|
|
for marker, size in size_markers:
|
|
if marker in name_lower:
|
|
return size
|
|
|
|
return 7.0 # default guess
|
|
|
|
|
|
def get_compute_dtype():
|
|
"""Return the best compute dtype for the current device.
|
|
|
|
Uses bfloat16 on CUDA GPUs that support it, float16 otherwise.
|
|
On CPU, uses float32 to avoid dtype mismatch errors.
|
|
"""
|
|
import torch
|
|
|
|
if torch.cuda.is_available():
|
|
if torch.cuda.is_bf16_supported():
|
|
return torch.bfloat16
|
|
return torch.float16
|
|
return torch.float32
|