gpu.py

"""GPU detection, memory calculation, and auto batch size."""

import math
import re


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()

    # Whisper ASR checkpoints carry the size in the name suffix, not an "Nb"
    # token — check these first so a 39M whisper-tiny isn't mistaken for the
    # 7B default (v0.71.32: the default guess blocked ASR training on the
    # hardware-fit gate).
    whisper_markers = [
        ("whisper-large", 1.55), ("whisper-medium", 0.769),
        ("whisper-small", 0.244), ("whisper-base", 0.074),
        ("whisper-tiny", 0.039),
    ]
    for marker, size in whisper_markers:
        if marker in name_lower:
            return size

    # Longer markers first: "1.7b" contains "7b", so a naive scan would call a
    # 1.7B model a 7B one (and over-predict its VRAM by 4x).
    size_markers = [
        ("70b", 70), ("65b", 65), ("34b", 34), ("33b", 33),
        ("13b", 13), ("8b", 8), ("3b", 3),
        ("1.5b", 1.5), ("1.7b", 1.7), ("0.5b", 0.5), ("0.6b", 0.6),
        ("7b", 7), ("1b", 1),
    ]

    for marker, size in size_markers:
        if marker in name_lower:
            return size

    # Sub-billion checkpoints carry their size in MILLIONS (SmolLM2-135M,
    # SmolVLM-256M, ...). Without this they fell through to the 7B default and
    # the hardware-fit gate refused to train them — which blocked `soup draft`
    # for exactly the tiny models drafts are made of (v0.71.33 live smoke;
    # same class as the v0.71.32 whisper fix).
    #
    # Checked AFTER the "b" markers on purpose: `Qwen2.5-7B-Instruct-1M` is a
    # 7B model with a 1M *context*, not a 1M-parameter model.
    million = re.search(r"(?<![a-z0-9.])(\d+(?:\.\d+)?)m(?![a-z0-9])", name_lower)
    if million:
        return float(million.group(1)) / 1000.0

    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