2026-10-18 17:23:57 +08:00
|
|
|
|
"""显存预算（2026-10-18 本机实测后加；同日按用户要求改成浮动）。
|
|
|
|
|
|
口径：整卡留 1 GB 给系统，其余"能给应用用多少"按启动那一刻的实际情况算——桌面 / 浏览器 / OBS /
|
|
|
|
|
|
直播软件占了多少都自动扣掉，不再写死一个比例（写死有两个毛病：别人占得多时上限仍很高 → 顶满整卡、
|
|
|
|
|
|
桌面卡顿；占得少时又白白少用）。
|
|
|
|
|
|
两个口径取较小（2026-10-18 实测差的不是一点）：
|
|
|
|
|
|
nvidia-smi：整卡 24.0 − 已用 8.6 − 余量 1.0 → 14.4 GiB（别人手里的显存先扣掉，别把它们挤去系统内存）
|
|
|
|
|
|
CUDA free：mem_get_info 报空闲 22.8 GiB —— WDDM 下这部分含"驱动能换出"的量，单独信它会顶满整卡
|
|
|
|
|
|
（APG batch=128 那次顶到 23.6 GB 就是整机卡顿的根源）。
|
|
|
|
|
|
limit() 本进程上限 = min(上面两个) ；显式给 VRAM_FRAC 时再取较小值，用来跟旧口径（0.78）对拍。
|
|
|
|
|
|
peak_gib() 记进日志的 vram 字段：网页指标卡与 ml/selftest.sh 都按它核对，不用再去猜。
|
|
|
|
|
|
实测参考（d=192 同结构、1 步即到位）：预训练 batch=1024 → 2.76 GiB；APG batch=48/plans=150 → 13.08 GiB。
|
|
|
|
|
|
要跑大 batch 的步骤（APG）之前先把 OBS / 浏览器这类占显存的收一收——上限会跟着启动时的占用变小。"""
|
2026-10-18 06:34:21 +08:00
|
|
|
|
import os
|
2026-10-18 17:23:57 +08:00
|
|
|
|
import subprocess
|
2026-10-18 06:34:21 +08:00
|
|
|
|
import torch
|
|
|
|
|
|
|
|
|
|
|
|
|
2026-10-18 17:23:57 +08:00
|
|
|
|
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
|
|
|
|
|
|
|
|
|
|
|
|
|
2026-10-18 06:34:21 +08:00
|
|
|
|
def limit():
|
|
|
|
|
|
if not torch.cuda.is_available(): return False
|
|
|
|
|
|
tot = torch.cuda.get_device_properties(0).total_memory / 2**30
|
2026-10-18 17:23:57 +08:00
|
|
|
|
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, why = free - res, 'CUDA 空闲 %.1f − 余量 %.1f' % (free, res)
|
|
|
|
|
|
used, _ = _smi()
|
|
|
|
|
|
if used is not None:
|
|
|
|
|
|
head = tot - used - res
|
|
|
|
|
|
if head < cap: cap, why = head, '整卡 %.1f − 已用 %.1f − 余量 %.1f' % (tot, used, 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)
|
|
|
|
|
|
print('显存上限：本进程 %.1f GiB（%s；= 整卡的 %.0f%%）' % (cap, why, 100 * cap / tot), flush=True)
|
2026-10-18 06:34:21 +08:00
|
|
|
|
return True
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def peak_gib():
|
|
|
|
|
|
return round(torch.cuda.max_memory_allocated() / 2**30, 2) if torch.cuda.is_available() else None
|
