cbb
f9fa3e67b6
- 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 组复核后关闭。
138 lines
8.2 KiB
Python
138 lines
8.2 KiB
Python
"""预训练（方案 §6）：MeanFlow 一步头 + FK 损失 + 平滑 + 价值头 + 子目标头。
|
||
python -m ml.train --samples ml/data/samples --out ml/runs/base [--config ml/configs/base.json]
|
||
吞吐摸底：python -m ml.train --samples ml/data/samples --out ml/runs/bench_base --steps 2000 --no-eval
|
||
通过线（方案 R10）：全程 60000 步 ≤ 3.5 h，即 samples_per_s ≥ 5000（3090 实测 6.2 k/s ≈ 2.8 h）；日志 eta_h 为剩余小时数。
|
||
每 eval_every 步做一次闭环验证（ml/sim/closed_loop.py，部署节奏），按 success 存 best.pt；连续 patience_evals 次不升即停。"""
|
||
import argparse, json, math, os, sys, time
|
||
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
|
||
from ml.sim.closed_loop import evaluate
|
||
|
||
|
||
def fk_points(kin, q, idx):
|
||
"""选定胶囊端点 + TCP 的世界坐标 [B, 2*len(idx)+1, 3]"""
|
||
xf = kin.cum_xf(q)
|
||
a, b = kin.caps_world(q, xf)
|
||
return torch.cat([a[:, idx], b[:, idx], kin.fk(q, xf)[:, None]], 1)
|
||
|
||
|
||
def q_from_feat(X):
|
||
s, c = X[:, 0:12:2], X[:, 1:12:2]
|
||
return torch.atan2(s, c) / D2R
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument('--samples', nargs='+', default=['ml/data/samples']); ap.add_argument('--out', default='ml/runs/base')
|
||
ap.add_argument('--config', default='ml/configs/base.json'); ap.add_argument('--steps', type=int, default=None)
|
||
ap.add_argument('--resume', default=None); ap.add_argument('--no-eval', action='store_true')
|
||
ap.add_argument('--stop-at', type=int, default=0, help='自检用：跑到第 N 步按停止请求处理（存 last.pt、rc 143）')
|
||
ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu'); ap.add_argument('--seed', type=int, default=0)
|
||
a = ap.parse_args()
|
||
cfg = json.load(open(a.config)); tc = cfg['train']
|
||
if a.steps: tc['steps'] = a.steps
|
||
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']]
|
||
print(f'train samples {data.n}, val tasks {len(val_tasks)}, device {dev}', flush=True)
|
||
model = build(cfg).to(dev)
|
||
print(f'params {sum(p.numel() for p in model.parameters()) / 1e6:.2f} M', flush=True)
|
||
opt = torch.optim.AdamW(model.parameters(), lr=tc['lr'], betas=tuple(tc['betas']), weight_decay=tc['wd'])
|
||
step0, best, bad_evals = 0, -1.0, 0
|
||
if a.resume:
|
||
ck = torch.load(a.resume, map_location=dev); model.load_state_dict(ck['model'])
|
||
if 'opt' in ck: opt.load_state_dict(ck['opt'])
|
||
step0, best, bad_evals = ck.get('step', 0), ck.get('best', -1.0), ck.get('bad_evals', 0)
|
||
if set_rng_state(ck.get('rng')): print('续训：随机数状态已恢复（batch 顺序与不断点一致）', flush=True)
|
||
print(f'续训：从第 {step0} 步继续（共 {tc["steps"]} 步），当前最好 success {best:.4f}', flush=True)
|
||
total = tc['steps']
|
||
lr_at = lambda s: tc['lr'] * (s / tc['warmup'] if s < tc['warmup'] else 0.5 * (1 + math.cos(math.pi * (s - tc['warmup']) / max(1, total - tc['warmup']))))
|
||
amp = dev.startswith('cuda') and tc.get('amp') == 'bf16'
|
||
stepv = torch.tensor(STEP, device=dev)
|
||
fk_idx = torch.linspace(0, kin.cap_a.shape[0] - 1, tc['fk_caps'], device=dev).long()
|
||
log = open(os.path.join(a.out, 'log.jsonl'), 'a' if a.resume else 'w') # 新训练覆盖旧日志，续训才追加
|
||
if not val_tasks and not a.no_eval:
|
||
print('没有 val 任务（样本太少），本次跳过闭环验证', flush=True); a.no_eval = True
|
||
t0, seen, diverged = time.time(), 0, False
|
||
saver = Saver(tc.get('ckpt_min', 10)) # 2026-10-18：每 10 分钟 + 收到停止请求时存 last.pt（原子写）
|
||
def save(name, step, extra=None):
|
||
d = {'model': model.state_dict(), 'opt': opt.state_dict(), 'cfg': cfg, 'step': step, 'best': best, 'bad_evals': bad_evals, 'fp': kin.fp}
|
||
if name == 'last.pt': d['rng'] = rng_state()
|
||
if extra: d.update(extra)
|
||
atomic_save(d, os.path.join(a.out, name))
|
||
if name == 'last.pt': saver.mark()
|
||
step = step0 - 1 # 续训时已训满：循环一次不跑，下面照常收尾
|
||
for step in range(step0, total):
|
||
for gpar in opt.param_groups: gpar['lr'] = lr_at(step)
|
||
X, Y, V, SG = data.batch(tc['batch'])
|
||
model.train()
|
||
with torch.autocast('cuda', dtype=torch.bfloat16, enabled=amp):
|
||
mem = model.encode(X)
|
||
if model.head == 'wta':
|
||
l_mf, xh = wta_loss(model, mem, Y); mf_raw = l_mf.detach(); xh = xh.float()
|
||
else:
|
||
l_mf, mf_raw = meanflow_loss(model, mem, Y, p_neq=tc['p_neq'])
|
||
e = torch.randn_like(Y)
|
||
xh = one_step(model, mem, e).float()
|
||
v, s = model.heads(mem)
|
||
l_val = (v.float() - V).abs().mean(); l_sub = (s.float() - SG).abs().mean()
|
||
d2 = xh[:, 2:] - 2 * xh[:, 1:-1] + xh[:, :-2]
|
||
l_sm = d2.pow(2).mean()
|
||
l_fk = torch.zeros((), device=dev)
|
||
if tc['w_fk'] > 0:
|
||
q = q_from_feat(X)
|
||
ks = torch.tensor([3, 7, 11, 15], device=dev)
|
||
qh = q[:, None] + torch.cumsum(xh.clamp(-1, 1) * stepv, 1)[:, ks]
|
||
qt = q[:, None] + torch.cumsum(Y * stepv, 1)[:, ks]
|
||
B4 = qh.shape[0] * 4
|
||
ph = fk_points(kin, qh.reshape(B4, 6), fk_idx); pt = fk_points(kin, qt.reshape(B4, 6), fk_idx)
|
||
l_fk = ((ph - pt).norm(dim=-1) / 10).mean()
|
||
loss = tc['w_mf'] * l_mf + tc['w_fk'] * l_fk + tc['w_smooth'] * l_sm + tc['w_value'] * l_val + tc['w_sub'] * l_sub
|
||
opt.zero_grad(set_to_none=True)
|
||
loss.backward()
|
||
gn = torch.nn.utils.clip_grad_norm_(model.parameters(), tc['grad_clip'])
|
||
if not torch.isfinite(loss):
|
||
print('非有限损失 → 训练发散（M5 门禁：触发 WTA 回退条件）', flush=True); save('diverged.pt', step); diverged = True; break
|
||
opt.step(); seen += X.shape[0]
|
||
if saver.stop() or (a.stop_at and step + 1 >= a.stop_at):
|
||
save('last.pt', step + 1); log.close()
|
||
print(f'收到停止请求：已存档 last.pt（第 {step + 1} 步），下次运行自动续训', flush=True); sys.exit(143)
|
||
if saver.due(): save('last.pt', step + 1)
|
||
if step % tc['log_every'] == 0:
|
||
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, '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()
|
||
res, _ = evaluate(model, kin, val_tasks, dev)
|
||
res['step'] = step + 1; res['t'] = time.time()
|
||
print('EVAL', json.dumps(res), flush=True); log.write(json.dumps({'eval': res}) + '\n'); log.flush()
|
||
save('last.pt', step + 1)
|
||
if res['success'] > best + 1e-4:
|
||
best = res['success']; bad_evals = 0; save('best.pt', step + 1, {'eval': res})
|
||
else:
|
||
bad_evals += 1
|
||
if bad_evals >= tc['patience_evals']: print('闭环到达率不再上升，停止', flush=True); break
|
||
if diverged: # 发散的权重不能覆盖 last.pt / best.pt：last.pt 保留发散前最后一次存档，M5 门禁会因此判不过 → 自动 WTA
|
||
print('训练发散，last.pt / best.pt 保留发散前的存档', flush=True); return
|
||
save('last.pt', step + 1)
|
||
if a.no_eval: save('best.pt', step + 1)
|
||
print(f'done, best success {best:.4f}, {(time.time() - t0) / 3600:.2f} h', flush=True)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|