ml: 显存预算按 19–20 GB 落地 + ort 本地副本入库 + R4 复跑（feat_dyn）通过
- 显存：APG batch 128→48（实测 d=192 下 128 无论 plans 60/150 都 OOM，崩前已分配 17.85 GiB；48/plans150=13.08、48/60=7.67）；新增 ml/common/vram.py 每进程硬帽 VRAM_FRAC=0.78 并把 vram 峰值写进 train/APG 日志、监控页加「显存峰值」卡与按 20 GB 预算的显存条；smoke.json 补 apg 段并让自检带 --config。 - 自检修复：find_bash 兼容便携版 git 的 mingw64 布局（本机原返回 None → 网页「开始训练」起不动）；selftest_proc.alive 读中文 tasklist 的 GBK 崩溃；rec 值非法 JSON 时静默丢项；gen_resume 判据剔除 plan_ms/cert.ms 两个墙钟字段（旧判据恒 false）。整轮重跑 7/8 通过、3.5 分钟，仅 train_term 为 Windows 真实限制（TERM 送不到原生 python）。 - vendor: onnxruntime-web 1.20.1 本地副本入库（约 33 MB），查看器与 R4/R9 探针不再依赖 npm/CDN；.gitignore 同步移除该目录。 - docs: 第八节 R4 复跑（feat_dyn 新模型）：部署口径 P99 12.67 通过；WebGPU 首次实测 27.09（有界面窗口，adapter nvidia ampere）；viewer_ref 随机模型 planP99 131 / starve 32.7% / lookaheadP99 170 / certFailWhy time 948 / tickGapMax 736，三条判读留给实现方。
This commit is contained in:
@@ -17,6 +17,7 @@ from ml.model.policy import build, wta_loss
|
||||
from ml.model.meanflow import meanflow_loss
|
||||
from ml.sim.closed_loop import evaluate, tasks_to_tensors, goal_err
|
||||
from ml.common.ckpt import Saver, atomic_save, rng_state, set_rng_state
|
||||
from ml.common import vram
|
||||
|
||||
|
||||
def load_pool(paths):
|
||||
@@ -38,6 +39,7 @@ def main():
|
||||
cfg = json.load(open(a.config)); C = cfg['apg']
|
||||
if a.steps: C['steps'] = a.steps
|
||||
dev = a.device; torch.manual_seed(a.seed); random.seed(a.seed)
|
||||
vram.limit()
|
||||
os.makedirs(a.out, exist_ok=True)
|
||||
ck = torch.load(a.init, map_location=dev)
|
||||
mcfg = ck['cfg']
|
||||
@@ -147,7 +149,7 @@ def main():
|
||||
with torch.no_grad():
|
||||
e, oe = goal_err(kin, q, T, R)
|
||||
rec = {'it': it, 'total': C['steps'], 't': time.time(), 'plans': n_plans, 'err_mm_mean': e.mean().item(), 'reach_frac': ((e < 1) & (oe < 0.5)).float().mean().item(), 'bc': l_bc.item(),
|
||||
'gnorm': gn.item(), **stats, 'min': (time.time() - t0) / 60}
|
||||
'gnorm': gn.item(), **stats, 'min': (time.time() - t0) / 60, 'vram': vram.peak_gib()}
|
||||
if it % 10 == 0: print(json.dumps(rec), flush=True)
|
||||
log.write(json.dumps(rec) + '\n'); log.flush()
|
||||
if saver.stop() or (a.stop_at and it + 1 >= a.stop_at):
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
"""显存预算（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
|
||||
+2
-1
@@ -1,7 +1,8 @@
|
||||
{
|
||||
"apg": {
|
||||
"batch": 128, "lr": 1e-4, "steps": 3000, "grad_clip": 1.0, "amp": "bf16",
|
||||
"batch": 48, "lr": 1e-4, "steps": 3000, "grad_clip": 1.0, "amp": "bf16",
|
||||
"plans": 60, "plans_max": 150, "seg_plans": 10,
|
||||
"_vram": "batch 由 128 降到 48：2026-10-18 本机实测（d=192 同结构、每进程 19.2 GB 硬帽）batch=128 无论 plans 取 60 还是 150 都 OOM（崩前已分配 17.85 GiB）；batch=48/plans=150 峰值 13.08 GiB、plans=60 为 7.67 GiB，加 Windows+4K 桌面常驻 4.4 GB 仍在 19–20 GB 预算内。batch=64/plans=150 要 17.66 GiB（不取）。展开长度 plans 没动——降 plans 会把长任务的绕出段截掉。",
|
||||
"margin_train_mm": 3.0,
|
||||
"w_tgt": 1.0, "w_rot": 0.3, "w_col": 2.0, "w_smooth": 0.02, "w_len": 0.001, "w_bc": 0.5,
|
||||
"cat_weights": {"A": 0.4, "B": 1.0, "C": 1.6, "D": 0.6, "Dr": 0.6},
|
||||
|
||||
@@ -31,5 +31,13 @@
|
||||
"log_every": 10,
|
||||
"patience_evals": 5
|
||||
},
|
||||
"_note": "冒烟专用（bash ml/smoke.sh）：只验证流程能跑通，数值无意义"
|
||||
"apg": {
|
||||
"batch": 8, "lr": 0.0001, "steps": 10, "grad_clip": 1.0, "amp": "bf16",
|
||||
"plans": 20, "plans_max": 30, "seg_plans": 5,
|
||||
"margin_train_mm": 3.0,
|
||||
"w_tgt": 1.0, "w_rot": 0.3, "w_col": 2.0, "w_smooth": 0.02, "w_len": 0.001, "w_bc": 0.5,
|
||||
"cat_weights": {"A": 0.4, "B": 1.0, "C": 1.6, "D": 0.6, "Dr": 0.6},
|
||||
"ckpt_min": 10, "eval_every": 1000, "eval_tasks": 4
|
||||
},
|
||||
"_note": "冒烟与自检专用（bash ml/pipeline.sh selftest / smoke）：只验证流程能跑通，数值无意义。apg 段必须有：selftest.sh 第 6 项原先不带 --config，走的是正式 apg.json 的 batch=128，本机实测第一次反向就 OOM；改成正式值 48 也要约 13 GiB，自检不必吃这么多。"
|
||||
}
|
||||
+6
-3
@@ -436,6 +436,7 @@ function renderKpis() {
|
||||
tile('展开内到达比例', pct(r.reach_frac), '误差 < 1 mm 且姿态 < 0.5°', '', last('reach_frac'), COLORS[2]);
|
||||
tile('闭环到达率', pct(e.success) + dlt('success'), `最好 ${pct(best)}`, '', evs('success'), COLORS[2]);
|
||||
tile('碰撞率', pct(e.collision_rate, 2) + dlt('collision_rate', -1), '关安全层', e.collision_rate > 0.03 ? 'warn' : '', evs('collision_rate'), COLORS[3]);
|
||||
tile('显存峰值', isNum(r.vram) ? `${r.vram.toFixed(2)}<small>GiB</small>` : '—', '预算 19–20 GB（整卡 24，桌面占 4–5）', r.vram > 17 ? 'bad' : r.vram > 14 ? 'warn' : '', last('vram'), COLORS[1]);
|
||||
} else {
|
||||
const sps = r.samples_per_s;
|
||||
tile('总损失', fmt(r.loss, 4), `MF ${fmt(r.mf)} · FK ${fmt(r.fk_cm)} cm`, '', ema(last('loss'), 0.9), COLORS[0]);
|
||||
@@ -443,6 +444,7 @@ function renderKpis() {
|
||||
tile('闭环到达率', pct(e.success) + dlt('success'), evals.length ? `最好 ${pct(best)} · @${e.x.toLocaleString()} 步` : '尚未评估', '', evs('success'), COLORS[2]);
|
||||
tile('碰撞率', pct(e.collision_rate, 2) + dlt('collision_rate', -1), '关安全层，部署另有护栏', e.collision_rate > 0.03 ? 'warn' : '', evs('collision_rate'), COLORS[3]);
|
||||
tile('梯度范数', fmt(r.gnorm), '裁剪阈值 1.0', r.gnorm > 5 ? 'warn' : '', ema(last('gnorm'), 0.9), COLORS[4]);
|
||||
tile('显存峰值', isNum(r.vram) ? `${r.vram.toFixed(2)}<small>GiB</small>` : '—', '预算 19–20 GB（整卡 24，桌面占 4–5）', r.vram > 17 ? 'bad' : r.vram > 14 ? 'warn' : '', last('vram'), COLORS[1]);
|
||||
}
|
||||
$('kpis').innerHTML = tiles.map(t => t.html).join('');
|
||||
[...$('kpis').querySelectorAll('canvas')].forEach((cv, i) => spark(cv, tiles[i].sp, tiles[i].color));
|
||||
@@ -477,7 +479,7 @@ async function poll() {
|
||||
render();
|
||||
}
|
||||
const P = lastStatus.pipeline, busy = (job && job.running) || (P && P.current);
|
||||
if (DEMO) { setPill('demo', '示例数据'); if (!$('demobar')) $('pipe').insertAdjacentHTML('beforebegin', `<div class="demobar" id="demobar">这是示例数据，不是你的训练。${location.protocol === 'file:' ? '直接双击 monitor.html 只能看示例；' : ''}看真实训练请双击 ml/console.bat（或 python -m ml.monitor --open）。</div>`); }
|
||||
if (DEMO) { setPill('demo', '示例数据'); if (!$('demobar')) $('pipe').closest('.hero').insertAdjacentHTML('beforebegin', `<div class="demobar" id="demobar">这是示例数据，不是你的训练。${location.protocol === 'file:' ? '直接双击 monitor.html 只能看示例；' : ''}看真实训练请双击 ml/console.bat（或 python -m ml.monitor --open）。</div>`); }
|
||||
else if (busy) setPill('live', `运行中 · ${P && P.current ? SM[P.current]?.n || P.current : '启动中'}`);
|
||||
else if (cur && Date.now() / 1000 - lastMtime < 120) setPill('live', `运行中 · ${ago(lastMtime)}更新`);
|
||||
else setPill('', '空闲');
|
||||
@@ -531,15 +533,16 @@ function renderStatus(s) {
|
||||
|
||||
// GPU
|
||||
const G = s.gpu || [];
|
||||
const BUD = (s.vram_budget_gb || 24) * 1024; // 显存预算（MB）：桌面常驻占 4–5 GB，不能按整卡 24 GB 看
|
||||
$('gputime').textContent = G.length ? '实时' : '';
|
||||
$('gpu').innerHTML = G.length ? G.map(g => {
|
||||
const u = g.util ?? 0, C = 2 * Math.PI * 50, col = u > 85 ? css('--ok') : u > 50 ? css('--accent') : css('--warn');
|
||||
const memP = 100 * g.mem_used / g.mem_total, tc = g.temp > 83 ? css('--bad') : g.temp > 75 ? css('--warn') : css('--c6');
|
||||
const memP = 100 * g.mem_used / BUD, tc = g.temp > 83 ? css('--bad') : g.temp > 75 ? css('--warn') : css('--c6');
|
||||
const row = (k, v, p, c) => `<div class="gr"><div class="l"><span>${k}</span><b>${v}</b></div><div class="mini"><div style="width:${Math.min(100, p)}%;background:${c}"></div></div></div>`;
|
||||
return `<div class="gpuwrap"><div class="ring"><svg width="116" height="116"><circle cx="58" cy="58" r="50" stroke="#1b2433" stroke-width="10" fill="none"/>
|
||||
<circle cx="58" cy="58" r="50" stroke="${col}" stroke-width="10" fill="none" stroke-linecap="round" stroke-dasharray="${C}" stroke-dashoffset="${C * (1 - u / 100)}" style="transition:stroke-dashoffset .6s"/></svg>
|
||||
<div class="c"><div><b>${u}%</b><span>利用率</span></div></div></div>
|
||||
<div class="gpurows">${row('显存', `${(g.mem_used / 1024).toFixed(1)} / ${(g.mem_total / 1024).toFixed(0)} GB`, memP, css('--accent2'))}
|
||||
<div class="gpurows">${row('显存', `${(g.mem_used / 1024).toFixed(1)} / ${(BUD / 1024).toFixed(0)} GB 预算（整卡 ${(g.mem_total / 1024).toFixed(0)}）`, memP, memP > 100 ? css('--bad') : memP > 85 ? css('--warn') : css('--accent2'))}
|
||||
${row('温度', `${g.temp} °C`, g.temp, tc)}${row('功耗', g.power != null ? `${Math.round(g.power)} W` : '—', g.power ? g.power / 3.5 : 0, css('--c7'))}</div></div>
|
||||
<div class="gname">${g.name}</div>`;
|
||||
}).join('<hr style="border:none;border-top:1px solid var(--line);margin:14px 0">') : empty('未检测到 nvidia-smi');
|
||||
|
||||
+9
-3
@@ -192,7 +192,8 @@ def status():
|
||||
pb = jload(P(RUNS, 'policy_bench.json'))
|
||||
if pb: pb = {k: {kk: v.get(kk) for kk in ('n', 'ok', 'reached', 'collisionFrames', 'brakeRate', 'recoveryRate', 'livelock', 'inferMsP99', 'safetyMsP99')}
|
||||
for k, v in pb.items() if isinstance(v, dict) and 'n' in v}
|
||||
return {'now': time.time(), 'gpu': gpu(), 'pipeline': pl, 'check': CHECK.get(), 'lock': lock_info(),
|
||||
return {'now': time.time(), 'gpu': gpu(),
|
||||
'vram_budget_gb': float(os.environ.get('VRAM_BUDGET_GB', 20)), # 本机桌面常驻占 4–5 GB，给训练的预算 19–20 GB 'pipeline': pl, 'check': CHECK.get(), 'lock': lock_info(),
|
||||
'stage_meta': [{'id': k, 'name': n, 'est_h': h, 'timeout_h': pstate.TIMEOUT_H[k]} for k, n, h in pstate.STAGES],
|
||||
'r4': jload(P(RUNS, 'r4.json')),
|
||||
'data': {'consistency': (jload(P(DATA, 'consistency.json')) or {}).get('consistentByGateRule'),
|
||||
@@ -218,8 +219,13 @@ def find_bash():
|
||||
if p and os.path.exists(p): return p
|
||||
g = shutil.which('git')
|
||||
if g:
|
||||
p = os.path.join(os.path.dirname(os.path.dirname(g)), 'bin', 'bash.exe')
|
||||
if os.path.exists(p): return p
|
||||
root = os.path.dirname(os.path.dirname(g)) # 便携版 git 常放在 <根>\mingw64\bin\git.exe，要再上一层才找到 bash
|
||||
for base in dict.fromkeys([root, os.path.dirname(root)]):
|
||||
for sub in ('bin', os.path.join('usr', 'bin')):
|
||||
p = os.path.join(base, sub, 'bash.exe')
|
||||
if os.path.exists(p): return p
|
||||
b = shutil.which('bash')
|
||||
if b and 'windows\\system32' not in b.lower(): return b
|
||||
return None
|
||||
return shutil.which('bash')
|
||||
|
||||
|
||||
+15
-5
@@ -17,14 +17,21 @@ git ls-files -z | tar --null -T - -cf - | (cd "$D" && tar xf -)
|
||||
[ -d ml/node_modules ] && cp -r ml/node_modules "$D/ml/node_modules"
|
||||
cd "$D"; mkdir -p ml/runs ml/data
|
||||
R=ml/runs/selftest_items.jsonl; : > $R
|
||||
rec() { $PY -c "import json,sys;print(json.dumps({'item':sys.argv[1],'pass':json.loads(sys.argv[2]),'detail':sys.argv[3]},ensure_ascii=False))" "$1" "$2" "$3" >> $R; echo "[selftest] $1 → $2 $3"; }
|
||||
rec() { $PY -c "import json,sys
|
||||
try: v=json.loads(sys.argv[2])
|
||||
except Exception: v=False # 值不是合法 JSON 时记 false，不能让整项静默消失（第七轮实测：'none' 抛异常 → 两项没进结果）
|
||||
print(json.dumps({'item':sys.argv[1],'pass':v,'detail':sys.argv[3]},ensure_ascii=False))" "$1" "$2" "$3" >> $R; echo "[selftest] $1 → $2 $3"; }
|
||||
CFGS=ml/configs/smoke.json
|
||||
# 1 生成续写
|
||||
node ml/gen/gen_dataset.cjs --n 6 --seed 11 --out ml/data/a.jsonl 2>/dev/null
|
||||
node ml/gen/gen_dataset.cjs --n 3 --seed 11 --out ml/data/shard_0011.jsonl.part 2>/dev/null
|
||||
sz=$(wc -c < ml/data/shard_0011.jsonl.part); head -c $((sz - 40)) ml/data/shard_0011.jsonl.part > t && mv t ml/data/shard_0011.jsonl.part
|
||||
node ml/gen/gen_dataset.cjs --n 6 --seed 11 --out ml/data/shard_0011.jsonl.part 2>ml/runs/st_gen.log; mv ml/data/shard_0011.jsonl.part ml/data/shard_0011.jsonl
|
||||
cmp -s ml/data/a.jsonl ml/data/shard_0011.jsonl && rec gen_resume true "$(grep -o '续写[^,]*' ml/runs/st_gen.log | head -1)" || rec gen_resume false "续写结果与一口气生成不同"
|
||||
gd=$($PY -c "import json,sys
|
||||
f=lambda p:[json.dumps({**r,'plan_ms':0,'cert':{**r.get('cert',{}),'ms':0}},sort_keys=True) for r in (json.loads(l) for l in open(p,encoding='utf-8') if l.strip())]
|
||||
a,b=f(sys.argv[1]),f(sys.argv[2]);print('OK' if a==b else '不同 %d/%d 条'%(sum(x!=y for x,y in zip(a,b))+abs(len(a)-len(b)),max(len(a),len(b))))" ml/data/a.jsonl ml/data/shard_0011.jsonl)
|
||||
if [ "$gd" = OK ]; then rec gen_resume true "$(grep -o '续写[^,]*' ml/runs/st_gen.log | head -1)（逐条一致；不含 plan_ms/cert.ms 两个墙钟字段，第七轮实测 cmp -s 恒不等）"
|
||||
else rec gen_resume false "续写与一口气生成不同：$gd"; fi
|
||||
rm -f ml/data/a.jsonl
|
||||
node ml/gen/gen_dataset.cjs --n 40 --seed 12 --out ml/data/shard_0012.jsonl 2>/dev/null # 训练自检用的小数据
|
||||
# 2 打包二次运行
|
||||
@@ -49,13 +56,16 @@ if command -v timeout >/dev/null; then
|
||||
grep -q "收到停止请求" ml/runs/st_c.log && rec train_term true "rc=$rc，TERM 送达并存档" || rec train_term false "rc=$rc，TERM 没送到 python（Windows 原生进程常见）→ 阶段超时只能硬杀，进度靠每 10 分钟存档"
|
||||
else rec train_term false "没有 timeout 命令"; fi
|
||||
# 6 APG 续跑
|
||||
$PY -m ml.apg.apg_finetune --init ml/runs/st_a/last.pt --samples ml/data/samples --out ml/runs/st_apg --steps 10 --stop-at 5 > ml/runs/st_apg.log 2>&1; rc=$?
|
||||
$PY -m ml.apg.apg_finetune --init ml/runs/st_a/last.pt --samples ml/data/samples --out ml/runs/st_apg --steps 10 --resume ml/runs/st_apg/last.pt >> ml/runs/st_apg.log 2>&1; rc2=$?
|
||||
$PY -m ml.apg.apg_finetune --config $CFGS --init ml/runs/st_a/last.pt --samples ml/data/samples --out ml/runs/st_apg --steps 10 --stop-at 5 > ml/runs/st_apg.log 2>&1; rc=$?
|
||||
$PY -m ml.apg.apg_finetune --config $CFGS --init ml/runs/st_a/last.pt --samples ml/data/samples --out ml/runs/st_apg --steps 10 --resume ml/runs/st_apg/last.pt >> ml/runs/st_apg.log 2>&1; rc2=$?
|
||||
[ $rc = 143 ] && [ $rc2 = 0 ] && rec apg_resume true "" || rec apg_resume false "rc=$rc 续跑rc=$rc2（日志 $D/ml/runs/st_apg.log）"
|
||||
# 7 进程树 / 关窗口
|
||||
pj=$($PY ml/selftest_proc.py 2>ml/runs/st_proc.err | tail -1)
|
||||
echo "$pj" > ml/runs/st_proc.json
|
||||
k=$($PY -c "import json,sys;d=json.loads(sys.argv[1]);print(str(d.get('tree_killed')).lower(), str(d.get('survive_parent_exit')).lower())" "$pj" 2>/dev/null || echo "null null")
|
||||
k=$($PY -c "import json,sys
|
||||
try: d=json.loads(sys.argv[1])
|
||||
except Exception: d={}
|
||||
print(json.dumps(d.get('tree_killed')), json.dumps(d.get('survive_parent_exit')))" "$pj" 2>/dev/null || echo "null null")
|
||||
rec tree_killed ${k% *} "$pj" ; rec survive_parent_exit ${k#* } ""
|
||||
mkdir -p "$SRC/ml/runs"
|
||||
$PY - "$R" "$SRC/ml/runs/selftest.json" "$D" <<'PY'
|
||||
|
||||
+8
-4
@@ -14,8 +14,9 @@ FLAGS = {'creationflags': subprocess.CREATE_NEW_PROCESS_GROUP | 0x08000000} if o
|
||||
|
||||
def alive(pid):
|
||||
if os.name == 'nt':
|
||||
r = subprocess.run(['tasklist', '/FI', f'PID eq {pid}', '/NH'], capture_output=True, text=True)
|
||||
return str(pid) in r.stdout
|
||||
# 中文 Windows 的 tasklist 表头是 GBK，text=True 会让读线程解码崩掉（stdout 变 None）→ 按字节读再容错解码
|
||||
r = subprocess.run(['tasklist', '/FI', f'PID eq {pid}', '/NH'], capture_output=True, text=True, errors='replace')
|
||||
return str(pid) in (r.stdout or '')
|
||||
try: os.kill(pid, 0); return True
|
||||
except OSError: return False
|
||||
|
||||
@@ -26,11 +27,14 @@ def tree_cmd(d):
|
||||
f'node -e "require(\'fs\').writeFileSync(\'{dd}/n.pid\',String(process.pid));setInterval(()=>{{}},1000)" & wait')
|
||||
|
||||
|
||||
WAITERR = ['']
|
||||
|
||||
|
||||
def wait_pids(d, t=20):
|
||||
for _ in range(t * 4):
|
||||
try: return int(open(os.path.join(d, 'p.pid')).read()), int(open(os.path.join(d, 'n.pid')).read())
|
||||
except Exception: time.sleep(0.25)
|
||||
return None, None
|
||||
except Exception as e: WAITERR[0] = type(e).__name__ + ': ' + str(e)[:90]; time.sleep(0.25)
|
||||
return None, None # 起不来时把最后一次异常带进结果，别让 tree_started=false 变成没有原因的失败
|
||||
|
||||
|
||||
def main():
|
||||
|
||||
+3
-1
@@ -9,6 +9,7 @@ import torch
|
||||
from ml.common.kin_torch import Kin
|
||||
from ml.common.features import STEP, D2R
|
||||
from ml.common.ckpt import Saver, atomic_save, rng_state, set_rng_state
|
||||
from ml.common import vram
|
||||
from ml.data import Samples, load_tasks
|
||||
from ml.model.policy import build, wta_loss
|
||||
from ml.model.meanflow import meanflow_loss, one_step
|
||||
@@ -40,6 +41,7 @@ def main():
|
||||
os.makedirs(a.out, exist_ok=True)
|
||||
torch.manual_seed(a.seed)
|
||||
dev = a.device
|
||||
vram.limit()
|
||||
kin = Kin(device=dev)
|
||||
data = Samples(a.samples, device=dev, split='train')
|
||||
val_tasks = [t for t in load_tasks(a.samples) if t['split'] == 'val'][:tc['eval_tasks']]
|
||||
@@ -111,7 +113,7 @@ def main():
|
||||
el = time.time() - t0
|
||||
rec = {'step': step, 'total': total, 't': time.time(), 'loss': loss.item(), 'mf': mf_raw.item(), 'fk_cm': l_fk.item(), 'smooth': l_sm.item(),
|
||||
'value': l_val.item(), 'sub': l_sub.item(), 'gnorm': gn.item(), 'lr': lr_at(step), 'samples_per_s': seen / max(el, 1e-6),
|
||||
'eta_h': (total - step - 1) * tc['batch'] / max(seen / max(el, 1e-6), 1e-6) / 3600}
|
||||
'eta_h': (total - step - 1) * tc['batch'] / max(seen / max(el, 1e-6), 1e-6) / 3600, 'vram': vram.peak_gib()}
|
||||
print(json.dumps(rec), flush=True); log.write(json.dumps(rec) + '\n'); log.flush()
|
||||
if not a.no_eval and (step + 1) % tc['eval_every'] == 0:
|
||||
model.eval()
|
||||
|
||||
Reference in New Issue
Block a user