49 lines
2.8 KiB
Python
49 lines
2.8 KiB
Python
"""显存预算（2026-10-18 本机实测后加；同日两次修正）。
|
||
口径：整卡留 1 GB 给系统（VRAM_RESERVE_GIB），本进程上限 = CUDA 可用 − 余量。
|
||
为什么不能按"整卡 − nvidia-smi 已用"扣：WDDM 下 nvidia-smi 的"已用"含**驱动可以换出**的量，
|
||
把它当硬占用会把自己卡死——实测预训练第一次 evaluate 峰值要 ~15 GiB，那样算出的上限只有
|
||
14.45 GiB，torch 就报 OOM，而同一时刻卡上其实还有 9.98 GiB 空闲（train.py step 2000 的评估里挂掉）。
|
||
CUDA 的 mem_get_info 才是"我实际能拿到多少"：它已经把可换出的部分算进去了。
|
||
limit() 上限 = min(CUDA 可用, 整卡) − 余量；显式给 VRAM_FRAC 时再取较小值（跟旧口径 0.78 对拍用）。
|
||
别的应用占多少照样打印，但只作参考、不从上限里扣。
|
||
peak_gib() 记进日志的 vram 字段：网页指标卡与 ml/selftest.sh 都按它核对，不用再去猜。
|
||
实测参考（d=192 同结构）：预训练 batch=1024 的训练步峰值只有 2.76 GiB，但第一次 evaluate 就要 ~15 GiB
|
||
（kin.gaps 一次开 2.35 GiB）；APG batch=48/plans=150 → 13.08 GiB。"""
|
||
import os
|
||
import subprocess
|
||
import torch
|
||
|
||
|
||
def _smi():
|
||
"""(整卡已用 GiB, 整卡总量 GiB)；拿不到就 (None, None)。"""
|
||
try:
|
||
q = 'memory.used,memory.total'
|
||
r = subprocess.run(['nvidia-smi', f'--query-gpu={q}', '--format=csv,noheader,nounits'],
|
||
capture_output=True, text=True, timeout=3)
|
||
u, t = [float(x) for x in r.stdout.strip().splitlines()[0].split(',')[:2]]
|
||
return u / 1024, t / 1024
|
||
except Exception:
|
||
return None, None
|
||
|
||
|
||
def limit():
|
||
if not torch.cuda.is_available(): return False
|
||
tot = torch.cuda.get_device_properties(0).total_memory / 2**30
|
||
torch.empty(1, device='cuda') # 让本进程的 CUDA 上下文先落地，再读可用（否则会把自家上下文算进可用）
|
||
free = torch.cuda.mem_get_info()[0] / 2**30
|
||
res = float(os.environ.get('VRAM_RESERVE_GIB', '1'))
|
||
cap = min(free, tot) - res
|
||
if os.environ.get('VRAM_FRAC'): # 显式硬帽（对拍旧口径用）
|
||
cap = min(cap, tot * float(os.environ['VRAM_FRAC']))
|
||
cap = max(1.0, min(cap, tot - res))
|
||
torch.cuda.set_per_process_memory_fraction(cap / tot)
|
||
used, _ = _smi()
|
||
note = '' if used is None else '；别的应用当前占 %.1f GiB，需要时驱动自行换出，不从上限里扣' % used
|
||
print('显存上限：本进程 %.1f GiB（CUDA 可用 %.1f − 系统余量 %.1f，= 整卡 %.1f 的 %.0f%%）%s'
|
||
% (cap, free, res, tot, 100 * cap / tot, note), flush=True)
|
||
return True
|
||
|
||
|
||
def peak_gib():
|
||
return round(torch.cuda.max_memory_allocated() / 2**30, 2) if torch.cuda.is_available() else None
|