第五轮:锁/续跑/超时/自动分支/R4 探针与 encode-once/控制台加固;阶段内断点续跑(gen 按条、prep 按 shard、训练与 APG 定时原子存档+体面停止);分歧以实现方意见为准

This commit is contained in:
Hark
2026-10-18 06:40:25 +08:00
committed by cbb
parent 3675fbbd9b
commit f40a40c56b
36 changed files with 1199 additions and 285 deletions
+18 -7
View File
@@ -3,11 +3,12 @@
吞吐摸底: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, time
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
from ml.data import Samples, load_tasks
from ml.model.policy import build, wta_loss
from ml.model.meanflow import meanflow_loss, one_step
@@ -45,11 +46,12 @@ def main():
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 = 0, -1.0
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 = ck.get('step', 0), ck.get('best', -1.0)
step0, best, bad_evals = ck.get('step', 0), ck.get('best', -1.0), ck.get('bad_evals', 0)
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'
@@ -58,11 +60,14 @@ def main():
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
bad_evals, t0, seen = 0, time.time(), 0
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, 'fp': kin.fp}
d = {'model': model.state_dict(), 'opt': opt.state_dict(), 'cfg': cfg, 'step': step, 'best': best, 'bad_evals': bad_evals, 'fp': kin.fp}
if extra: d.update(extra)
torch.save(d, os.path.join(a.out, name))
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'])
@@ -93,8 +98,12 @@ def main():
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); break
print('非有限损失 → 训练发散(M5 门禁:触发 WTA 回退条件)', flush=True); save('diverged.pt', step); diverged = True; break
opt.step(); seen += X.shape[0]
if saver.stop():
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(),
@@ -112,6 +121,8 @@ def main():
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)