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