Files
Robot_Arm/ml/common/vram.py
T
cbb f9fa3e67b6 ml: 显存预算按 19–20 GB 落地(APG batch 128→48 + 每进程硬帽 + vram 上日志/网页);自检四处判据与环境坑修好,整轮 7/8 通过 3.5 分钟
- ml/configs/apg.json:batch 128→48。本机实测(d=192 同结构、每进程 19.2 GB 硬帽、1 步即到位)batch=128 无论 plans 取 60 还是 150 都 OOM(崩前已分配 17.85 GiB;d=32 小模型也 17.54 GiB,说明是展开图吃显存而非模型大小);48/plans=150 为 13.08、48/60 为 7.67。展开长度 plans 没动,降 plans 会把长任务绕出段截掉。
- ml/common/vram.py(新):limit() 设每进程帽 VRAM_FRAC 默认 0.78≈18.7 GiB 并打印上限;peak_gib() 供日志。train.py 与 apg_finetune.py 调用,每条记录新增 vram 字段(实测 train batch=1024 → 2.71,APG 自检档 → 0.33)。没帽时 APG 会把整卡顶到 23.6 GB、连桌面一起挤掉。
- ml/monitor.py / monitor.html:状态带 vram_budget_gb(默认 20,VRAM_BUDGET_GB 可改),GPU 显存条按预算算并超线变色,指标卡加「显存峰值」(>14 黄 >17 红)。headless 加载渲染 DOM 复核过。
- ml/configs/smoke.json 补 apg 段(batch=8)+ selftest.sh 的 APG 两项带 --config:原先不带 → 自检第 6 项按正式 batch=128 跑,本机必 OOM。
- ml/selftest.sh:gen_resume 判据由 cmp -s 改为剔除 plan_ms/cert.ms 两个墙钟字段后逐条比(旧判据恒 false);rec() 值非合法 JSON 时记 false 不再静默丢项。
- ml/selftest_proc.py:alive() 读 tasklist 改 text=True+errors='replace' 并防 None(中文 Windows 表头 GBK 会让读线程抛 UnicodeDecodeError→stdout=None→TypeError);wait_pids 带上最后一次异常。
- ml/monitor.py find_bash():便携版 git 的 mingw64 布局要再上一层找 bash,并兜 which('bash')(排除 System32 的 WSL)。本机原先返回 None → 网页「开始训练」起不动、自检进程两项没跑。
- 整轮自检重跑:gen_resume/prep_rerun/train_stop/train_stopfile/apg_resume/tree_killed/survive_parent_exit 全 true,仅 train_term false(timeout 的 TERM 送不到原生 python,rc=124 硬杀);用时 3.5 分钟。
- 文档:新增第七节(显存实测表、自检 8 项结果、四条判据/环境坑、本机改动清单);A-3 改为「实现方已改、待复跑验证」,C-5/C-6 标为已实测,D 组复核后关闭。
2026-10-18 05:34:07 +08:00

22 lines
1.2 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""显存预算(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。"""
import os
import torch
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)
return True
def peak_gib():
return round(torch.cuda.max_memory_allocated() / 2**30, 2) if torch.cuda.is_available() else None