1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
|
"""Retro + live wandb sync from trainer logs (runs launched before wandb default-on).
History logs replayed once; fw72m.log followed until DONE. Multi-run files (nsize_ladder)
switch tag on '[tag] cascade-EP' headers. Project: ept-cascade."""
import re, time
from pathlib import Path
import wandb
ENTITY = 'eqprop-llm-training'
R = Path('/home/yurenh2/ept/ep_run/runs')
EP_RE = re.compile(
r'^step\s+(\d+)/\d+ \| train ([\d.]+) val ([\d.]+) \(best ([\d.]+)\) '
r'\| beta=([\d.e+-]+) K=(\d+) skips=(\d+)\(d(\d+)/g(\d+)/r(\d+)\)(?: cos=(-?[\d.]+))? '
r'drift=([\d.]+) gn=([\d.e+-]+) sig=([\d.]+) \| ([\d.]+) it/s')
BP_RE = re.compile(r'^step\s+(\d+)/\d+ \| train ([\d.]+) val ([\d.]+) \(best ([\d.]+)\) \| ([\d.]+) it/s')
HDR = re.compile(r'^\[(\w+)\] cascade-')
def parse(line):
m = EP_RE.match(line)
if m:
g = m.groups()
d = {'train_ce': float(g[1]), 'val_ce': float(g[2]), 'best': float(g[3]),
'beta_t': float(g[4]), 'K': int(g[5]), 'skips': int(g[6]),
'drift': float(g[11]), 'gn': float(g[12]), 'sig': float(g[13]), 'it_s': float(g[14])}
if g[10]: d['gate_cos'] = float(g[10])
return int(g[0]), d
m = BP_RE.match(line)
if m:
g = m.groups()
return int(g[0]), {'train_ce': float(g[1]), 'val_ce': float(g[2]), 'best': float(g[3]), 'it_s': float(g[4])}
return None
def sync_file(path, default_tag, project, follow=False):
tag, run, last = default_tag, None, -1
def ensure(t):
nonlocal run, last
if run is not None: run.finish()
import os as _os
run = wandb.init(entity=ENTITY, project=project, name=t, id=t + _os.environ.get('SYNC_ID_SUFFIX', ''), resume='allow', reinit=True)
last = -1
f = open(path, errors='replace')
ensure(tag)
while True:
line = f.readline()
if not line:
if not follow: break
if 'DONE best' in open(path, errors='replace').read()[-2000:]: break
time.sleep(60); continue
h = HDR.match(line)
if h and h.group(1) != tag:
tag = h.group(1); ensure(tag)
p = parse(line)
if p and p[0] > last:
run.log(p[1], step=p[0]); last = p[0]
if run is not None: run.finish()
print(f'[sync] {path} done', flush=True)
TS, FW = 'ept-tinystories-42m', 'ept-fineweb-72m'
HISTORY = [
('stage1b_ep_muon.log', 'stage1b_ep_muon', TS), ('stage1b_amp.log', 'stage1b_amp', TS),
('stage1b_bp_muon.log', 'stage1b_bp_muon', TS),
('amp_s1.log', 'amp_s1', TS), ('amp_s2.log', 'amp_s2', TS), ('amp_s3.log', 'amp_s3', TS),
('nsize_ladder.log', 'nsize_c256', TS),
('bsign_s1.log', 'bsign_s1', TS), ('bsign_s2.log', 'bsign_s2', TS), ('bsign_s3.log', 'bsign_s3', TS),
('stage1b_f3e3cont.log', 'f3e3cont', TS),
('arm_ctl.log', 'arm_ctl', TS), ('arm_mom99_f3e4.log', 'arm_mom99_f3e4', TS),
('arm_mom995_f1e4.log', 'arm_mom995_f1e4', TS), ('arm_cent_f3e3.log', 'arm_cent_f3e3', TS),
('arm_rich_f1e3.log', 'arm_rich_f1e3', TS), ('arm_mom99_f1e3.log', 'arm_mom99_f1e3', TS),
('arm_adamw_f1e3.log', 'arm_adamw_f1e3', TS),
('fw_smoke.log', 'fw_smoke', FW),
('fw72m.log', 'fw72m', FW), ('fw72m_bp.log', 'fw72m_bp', FW), ('fw72m_f3e3.log', 'fw72m_f3e3', FW),
('fw72m_c.log', 'fw72m_c', FW), ('fw72m_c2.log', 'fw72m_c2', FW),
('w2_ctl.log', 'w2_ctl', FW), ('w2_blow.log', 'w2_blow', FW), ('w2_geta07.log', 'w2_geta07', FW),
('w2_geta05_k5.log', 'w2_geta05_k5', FW), ('w2_bcap.log', 'w2_bcap', FW), ('w2_adapt.log', 'w2_adapt', FW),
]
import sys
if len(sys.argv) > 1:
HISTORY = [h for h in HISTORY if h[1] in sys.argv[1:]]
for fn, tag, prj in HISTORY:
p = R / fn
if p.exists():
try: sync_file(p, tag, prj)
except Exception as e: print(f'[sync] {fn} FAILED: {e}', flush=True)
print('[sync] ALL DONE', flush=True)
|
