File size: 2,877 Bytes
7bcd331 | 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 | #!/usr/bin/env python3
"""Plot recorded logs; requires matplotlib 3.10.8 (not needed for inference)."""
import csv
import os
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
ROOT = Path(__file__).resolve().parents[1]
def read(name):
rows = list(csv.DictReader((ROOT/'results'/name).open()))
return [dict(row, **{k:float(v) for k,v in row.items() if k!='phase'}) for row in rows]
def main():
train = read('pretraining_steps.csv'); validation = read('pretraining_validation.csv')
train = sorted({r['step']:r for r in train}.values(), key=lambda r:r['step'])
validation = sorted({r['step']:r for r in validation}.values(), key=lambda r:r['step'])
plt.rcParams.update({'font.family':'DejaVu Sans','font.size':10,'axes.spines.top':False,'axes.spines.right':False})
output = ROOT/'docs/figures'; output.mkdir(parents=True, exist_ok=True)
for name, lower, upper in [('pretraining-overview',0,42.1),('cooldown-detail',35,42.1)]:
fig, (ax, lr) = plt.subplots(2,1,figsize=(10,6.8),sharex=True,gridspec_kw={'height_ratios':[2,1]},layout='constrained')
ts=[r for r in train if r['tokens']/1e9>=lower]
vs=[r for r in validation if r['tokens']/1e9>=lower]
ax.plot([r['tokens']/1e9 for r in ts],[r['ema'] for r in ts],color='#8898a3',lw=.9,alpha=.8,label='Training loss EMA (changing mixture)')
ax.plot([r['tokens']/1e9 for r in vs],[r['reported_val_loss'] for r in vs],color='#1557a0',lw=1.5,label='Held-out prose validation (logged BF16 path)')
if lower==0: ax.set_ylim(2.2,5.6)
ax.set_ylabel('Token cross-entropy loss'); ax.legend(loc='upper right',fontsize=9)
lr.plot([r['tokens']/1e9 for r in ts],[r['lr'] for r in ts],color='#b25819',lw=1.5)
lr.set_ylabel('Learning rate'); lr.ticklabel_format(axis='y',style='sci',scilimits=(0,0))
lr.set_xlabel('Processed pretraining tokens (billions)')
for panel in (ax,lr):
panel.set_xlim(lower,upper); panel.grid(alpha=.16)
panel.axvline(36.801216512,color='#5f6984',lw=.8,ls='--')
panel.axvspan(40.001339392,42.001367040,color='#ead9ac',alpha=.4)
fig.suptitle('Solstice-1 路 175M | '+('Pretraining history' if lower==0 else 'Late training and 2B-token cooldown'),fontsize=16,fontweight='bold')
fig.supxlabel('Dashed line: Python/SQL introduced 路 Shaded region: final cooldown\nTraining and validation use different data; early training loss above 5.6 is outside the overview scale.' if lower==0 else 'Dashed line: Python/SQL introduced 路 Shaded region: final cooldown\nMetrics are transcribed from saved logs, with their original reported precision.',fontsize=9,color='#44505c')
for ext in ('png','svg'): fig.savefig(output/f'{name}.{ext}',dpi=180)
plt.close(fig)
if __name__ == '__main__':
main()
|