File size: 3,363 Bytes
026f63e | 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 | """Hardware detection and automatic runtime profile (CUDA / Apple MPS / CPU).
The *effective* batch size is part of the experiment (it changes results); the
micro-batch that fits in memory and the gradient accumulation to reach the
effective batch are hardware details chosen here, so the same config gives
comparable results on a laptop and on a rented GPU.
"""
from __future__ import annotations
import os
import platform
from dataclasses import asdict, dataclass
import torch
@dataclass
class RuntimeProfile:
device: str
device_name: str
vram_gb: float | None
amp_dtype: str | None # "bf16" | "fp16" | None
micro_batch: int
grad_accum: int
num_workers: int
compile: bool
threads: int
def to_dict(self) -> dict:
return asdict(self)
def detect_device(prefer: str = "auto") -> str:
if prefer != "auto":
return prefer
if torch.cuda.is_available():
return "cuda"
if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
return "mps"
return "cpu"
def resolve_runtime(train_cfg: dict, model_params: int | None = None) -> RuntimeProfile:
hw = train_cfg.get("hardware", {}) or {}
device = detect_device(hw.get("device", "auto"))
eff = int(train_cfg.get("batch_size", 256))
vram = None
name = platform.processor() or platform.machine()
if device == "cuda":
props = torch.cuda.get_device_properties(0)
vram = props.total_memory / 1024**3
name = props.name
# activations dominate for these tiny models; ~1 GB per 4k samples of a ~5M model is a
# conservative rule of thumb. `max_micro_batch` in the config always wins.
per_gb = hw.get("samples_per_gb", 2048 if (model_params or 0) > 3e6 else 8192)
cap = int(max(64, (vram - 1.5) * per_gb))
elif device == "mps":
cap = 1024
else:
cap = 512
cap = min(cap, int(hw.get("max_micro_batch", cap)))
micro = eff
while micro > cap and micro % 2 == 0:
micro //= 2
accum = max(1, eff // micro)
amp = hw.get("amp", "auto")
if amp == "auto":
if device == "cuda":
amp = "bf16" if torch.cuda.is_bf16_supported() else "fp16"
else:
amp = None # CPU autocast rarely helps small models; MPS fp16 is flaky for training
elif amp in ("none", False):
amp = None
cpus = os.cpu_count() or 2
workers = hw.get("num_workers", "auto")
if workers == "auto":
workers = min(8, max(0, cpus // 2 - 1)) if device != "cpu" else min(4, max(0, cpus // 4))
threads = hw.get("threads", "auto")
if threads == "auto":
threads = max(1, cpus - int(workers)) if device == "cpu" else max(1, min(4, cpus))
compile_ = hw.get("compile", "auto")
if compile_ == "auto":
compile_ = device == "cuda" # launch-overhead bound tiny models gain most on GPU
return RuntimeProfile(device, name, round(vram, 1) if vram else None, amp, micro, accum,
int(workers), bool(compile_), int(threads))
def autocast(profile: RuntimeProfile):
if profile.amp_dtype is None:
return torch.autocast(device_type="cpu", enabled=False)
dtype = torch.bfloat16 if profile.amp_dtype == "bf16" else torch.float16
return torch.autocast(device_type=profile.device, dtype=dtype)
|