import torch
ckpt = torch.load('ckpts/step_259567/archon.pt', map_location='cpu', weights_only=False)
print(f'step: {ckpt["step"]}')
print(f'wat_buffer_size: {ckpt["wat_buffer_size"]}')
print(f'optimizer keys: {list(ckpt["optimizer"].keys())}')
m = ckpt['model']
print(f'\nmodel state_dict first 15 keys + dtypes:')
for k in list(m.keys())[:15]:
    v = m[k]
    shape = tuple(v.shape) if hasattr(v, 'shape') else type(v).__name__
    dt = v.dtype if hasattr(v, 'dtype') else ''
    print(f'  {k:50s} {str(shape):25s} {dt}')
print(f'\nmodel state_dict last 10 keys:')
for k in list(m.keys())[-10:]:
    v = m[k]
    shape = tuple(v.shape) if hasattr(v, 'shape') else type(v).__name__
    dt = v.dtype if hasattr(v, 'dtype') else ''
    print(f'  {k:50s} {str(shape):25s} {dt}')
total = sum(v.numel() for v in m.values() if hasattr(v, 'numel'))
print(f'\nTotal params: {total/1e6:.1f}M')
print(f'\n--- config.py ---')
with open('source/config.py') as f:
    print(f.read())