feat(vram): 显存上限改为浮动——本进程上限 = min(整卡 − nvidia-smi 已用 − 1 GB 余量, CUDA 空闲 − 余量),别人占多少自动扣掉,不再写死 0.78(实测当前口径 14.1 GiB;VRAM_FRAC 仍可显式对拍旧口径);网页显存行改成"8.7 / 23 GB · 空 14.3"并把解释放悬停(长文本不再把行挤成两行,.gr 行加了省略号保护);README/审查文档同步口径

This commit is contained in:
cbb
2026-10-18 17:23:57 +08:00
parent e431862793
commit b1cf4126fd
5 changed files with 48 additions and 16 deletions
+38 -9
View File
@@ -1,19 +1,48 @@
"""显存预算(2026-10-18 本机实测后加)。
用户定的预算:这台 3090 只有 19–20 GB 能给训练用——Windows + 4K 桌面常驻就占掉 4–5 GB(空载 nvidia-smi 4.4 GB)。
limit() 给本进程设硬帽(VRAM_FRAC,默认 0.78 ≈ 18.7 GiB)。超限时在自己的反向里早失败,
而不是把桌面挤到没有显存(整机卡顿、别的应用崩),这也是 APG batch=128 原先 OOM 时把显存顶到 23.6 GB 的原因。
peak_gib() 记进日志的 vram 字段:网页指标卡与 ml/selftest.sh 都按它核对预算,不用再去猜。
实测参考(d=192 同结构、1 步即到位):预训练 batch=1024 → 2.76;APG batch=48/plans=150 → 13.08、batch=128 → OOM。"""
"""显存预算(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 / 浏览器这类占显存的收一收——上限会跟着启动时的占用变小。"""
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
f = float(os.environ.get('VRAM_FRAC', '0.78'))
tot = torch.cuda.get_device_properties(0).total_memory / 2**30
torch.cuda.set_per_process_memory_fraction(f)
print('显存上限:本进程 %.1f GiB(整卡 %.1f GiB × %.2f,其余留给桌面)' % (f * tot, tot, f), flush=True)
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)
return True