File size: 7,027 Bytes
9b42c40 a413ecf 9b42c40 a413ecf 9b42c40 | 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | #!/usr/bin/env python3
"""make_peer_compare.py - A/B charts: v20 (16 private PEER pools) vs v21 (3 shared).
Reads the two training logs and emits the comparison figures + a metrics json.
baseline : peer_compare/baseline_peer_pipeline.log (v20, instance 45557221)
v21 : peer_compare/v21_peer21.log (v21, instance 45785944)
"""
import json, re, sys
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
D = Path(__file__).parent / 'peer_compare'
STEP_RE = re.compile(r'step=(\d+) tokens=(\d+) ce=([\d.]+).*?tok_s=([\d,]+)')
CKPT_RE = re.compile(r'CKPT step=(\d+) valid_ce=([\d.]+)')
def parse(path):
steps, toks, ce, tps, vsteps, vtok, vce = [], [], [], [], [], [], []
tok_at = {}
for ln in Path(path).read_text(errors='ignore').splitlines():
m = STEP_RE.search(ln)
if m:
s, t, c, r = int(m[1]), int(m[2]), float(m[3]), int(m[4].replace(',', ''))
steps.append(s); toks.append(t); ce.append(c); tps.append(r); tok_at[s] = t
continue
m = CKPT_RE.search(ln)
if m:
s, v = int(m[1]), float(m[2])
if s in tok_at: vsteps.append(s); vtok.append(tok_at[s]); vce.append(v)
return dict(steps=steps, tokens=toks, ce=ce, tok_s=tps,
vsteps=vsteps, vtokens=vtok, vce=vce)
def med(x):
x = sorted(x); return x[len(x)//2] if x else 0
def ce_at(d, lo, hi):
"""Mean valid CE over a token window - the only fair quality comparison when the
two runs stopped at different token counts."""
v = [c for t, c in zip(d['vtokens'], d['vce']) if lo <= t <= hi]
return sum(v)/len(v) if v else None
def main():
a = parse(D / 'baseline_peer_pipeline.log') # v20
b = parse(D / 'v21_peer21.log') # v21
# steady-state throughput ignores the first few logged points (compile warmup)
a_tp = a['tok_s'][5:] or a['tok_s']
b_tp = b['tok_s'][5:] or b['tok_s']
A = sum(a_tp)/len(a_tp); B = sum(b_tp)/len(b_tp)
M = dict(
v20=dict(avg_tok_s=A, median_tok_s=med(a_tp), max_tok_s=max(a_tp),
samples=len(a_tp), final_valid_ce=(a['vce'][-1] if a['vce'] else None),
tokens=(a['tokens'][-1] if a['tokens'] else 0),
experts_per_layer=30976, peer_layers=16, pools='16 private',
ctrl_params=569357312, h_times_k=8),
v21=dict(avg_tok_s=B, median_tok_s=med(b_tp), max_tok_s=max(b_tp),
samples=len(b_tp), final_valid_ce=(b['vce'][-1] if b['vce'] else None),
tokens=(b['tokens'][-1] if b['tokens'] else 0),
experts_per_layer=495616, peer_layers=3, pools='1 shared',
ctrl_params=598145536, h_times_k=8),
speedup=B/A if A else 0)
# same-token quality: compare both runs over the window they BOTH covered
hi = min(max(a['vtokens'] or [0]), max(b['vtokens'] or [0]))
lo = 0.5 * hi
if hi > 0:
M['same_token_window'] = {'lo_tokens': lo, 'hi_tokens': hi,
'v20_mean_valid_ce': ce_at(a, lo, hi),
'v21_mean_valid_ce': ce_at(b, lo, hi)}
(D / 'compare_metrics.json').write_text(json.dumps(M, indent=2))
print(json.dumps(M, indent=2))
C20, C21 = '#c44e52', '#4c72b0'
# ---- 1. throughput over training ----
fig, ax = plt.subplots(figsize=(11, 5.5))
ax.plot([t/1e6 for t in a['tokens']], a['tok_s'], color=C20, lw=.8, alpha=.55)
ax.plot([t/1e6 for t in b['tokens']], b['tok_s'], color=C21, lw=.8, alpha=.55)
ax.axhline(A, color=C20, ls='--', lw=2, label=f'v20 (16 private pools) avg {A:,.0f} tok/s')
ax.axhline(B, color=C21, ls='--', lw=2, label=f'v21 (3 shared pools) avg {B:,.0f} tok/s')
ax.axhline(45700, color='#55a868', ls=':', lw=2, label='dense DNA-2B reference 45,700 tok/s')
ax.set_xlabel('tokens seen (millions)'); ax.set_ylabel('training throughput (tok/s)')
ax.set_title(f'PEER training throughput: v21 is {B/A:.2f}x faster than v20\n'
f'(identical GPU: RTX 5060 Ti 16GB)', fontweight='bold')
ax.legend(loc='lower right'); ax.grid(alpha=.3); ax.set_ylim(0, 50000)
fig.tight_layout(); fig.savefig(D / 'throughput_v20_vs_v21.png', dpi=130); plt.close(fig)
# ---- 2. bar summary ----
fig, axes = plt.subplots(1, 3, figsize=(13, 4.6))
ax = axes[0]
ax.bar(['v20\n16 private', 'v21\n3 shared'], [A, B], color=[C20, C21])
ax.axhline(45700, color='#55a868', ls=':', lw=2)
for i, v in enumerate([A, B]): ax.text(i, v*1.02, f'{v:,.0f}', ha='center', fontweight='bold')
ax.set_ylabel('avg tok/s'); ax.set_title(f'Throughput ({B/A:.2f}x)'); ax.grid(alpha=.3, axis='y')
ax = axes[1]
ax.bar(['v20', 'v21'], [30976, 495616], color=[C20, C21])
ax.set_yscale('log'); ax.set_ylabel('experts per PEER layer (log)')
for i, v in enumerate([30976, 495616]): ax.text(i, v*1.15, f'{v:,}', ha='center', fontweight='bold')
ax.set_title('Expert granularity (16x)'); ax.grid(alpha=.3, axis='y')
ax = axes[2]
ax.bar(['v20', 'v21'], [16, 3], color=[C20, C21])
ax.set_ylabel('PEER layers (gathers / token)')
for i, v in enumerate([16, 3]): ax.text(i, v+.2, str(v), ha='center', fontweight='bold')
ax.set_title('Gathers per token (5.3x fewer)'); ax.grid(alpha=.3, axis='y')
fig.suptitle('v20 vs v21 at identical ~507M expert capacity', fontweight='bold')
fig.tight_layout(); fig.savefig(D / 'summary_v20_vs_v21.png', dpi=130); plt.close(fig)
# ---- 3. loss vs tokens and vs wall-clock ----
fig, axes = plt.subplots(1, 2, figsize=(13, 5))
ax = axes[0]
ax.plot([t/1e6 for t in a['tokens']], a['ce'], color=C20, lw=.7, alpha=.4)
ax.plot([t/1e6 for t in b['tokens']], b['ce'], color=C21, lw=.7, alpha=.4)
if a['vce']: ax.plot([t/1e6 for t in a['vtokens']], a['vce'], color=C20, lw=2.2, marker='o', ms=3, label='v20 valid CE')
if b['vce']: ax.plot([t/1e6 for t in b['vtokens']], b['vce'], color=C21, lw=2.2, marker='o', ms=3, label='v21 valid CE')
ax.set_xlabel('tokens seen (millions)'); ax.set_ylabel('cross-entropy')
ax.set_title('Quality vs tokens (sample efficiency)'); ax.legend(); ax.grid(alpha=.3)
ax = axes[1]
ah = [t/A/3600 for t in a['tokens']]; bh = [t/B/3600 for t in b['tokens']]
ax.plot(ah, a['ce'], color=C20, lw=.7, alpha=.4)
ax.plot(bh, b['ce'], color=C21, lw=.7, alpha=.4)
if a['vce']: ax.plot([t/A/3600 for t in a['vtokens']], a['vce'], color=C20, lw=2.2, marker='o', ms=3, label='v20')
if b['vce']: ax.plot([t/B/3600 for t in b['vtokens']], b['vce'], color=C21, lw=2.2, marker='o', ms=3, label='v21')
ax.set_xlabel('GPU-hours at measured throughput'); ax.set_ylabel('cross-entropy')
ax.set_title('Quality vs wall-clock (what the speedup buys)'); ax.legend(); ax.grid(alpha=.3)
fig.tight_layout(); fig.savefig(D / 'loss_v20_vs_v21.png', dpi=130); plt.close(fig)
print('\nwrote 3 charts + compare_metrics.json to', D)
if __name__ == '__main__':
main()
|