#!/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 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) (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()