scripts/make_peer_compare.py

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