diff --git a/ml/common/vram.py b/ml/common/vram.py index e1ca3b0..4765d01 100644 --- a/ml/common/vram.py +++ b/ml/common/vram.py @@ -1,15 +1,14 @@ -"""显存预算(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)对拍。 +"""显存预算(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 同结构、1 步即到位):预训练 batch=1024 → 2.76 GiB;APG batch=48/plans=150 → 13.08 GiB。 -要跑大 batch 的步骤(APG)之前先把 OBS / 浏览器这类占显存的收一收——上限会跟着启动时的占用变小。""" +实测参考(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 @@ -30,19 +29,18 @@ def _smi(): 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 上下文先落地,再读空闲(否则会把自己的上下文算进空闲) + 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) + 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) - print('显存上限:本进程 %.1f GiB(%s;= 整卡的 %.0f%%)' % (cap, why, 100 * cap / tot), flush=True) + 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 diff --git a/ml/monitor.html b/ml/monitor.html index 25a2c6c..9a4e757 100644 --- a/ml/monitor.html +++ b/ml/monitor.html @@ -420,7 +420,7 @@ function renderProgress() {