| """Smoke test: load ArchonBrain + state_dict + forward pass V100.""" |
| import sys, time, torch |
| sys.path.insert(0, 'source') |
| from config import ArchonBrainConfig |
| import model as model_module |
|
|
| print('=== model.py classes ===') |
| print([c for c in dir(model_module) if not c.startswith('_') and c[0].isupper()]) |
|
|
| |
| ArchonBrain = None |
| for name in ['ArchonBrain', 'Brain', 'ArchonBrainModel', 'Model']: |
| if hasattr(model_module, name): |
| ArchonBrain = getattr(model_module, name) |
| print(f'Found: {name}') |
| break |
|
|
| cfg = ArchonBrainConfig() |
| print(f'\nConfig param_count estimate: {cfg.param_count_human}') |
|
|
| print('\n[1] Instantiating model...') |
| m = ArchonBrain(cfg) |
| n_params = sum(p.numel() for p in m.parameters()) |
| print(f' Instantiated: {n_params/1e6:.1f}M params') |
|
|
| print('\n[2] Loading ckpt state_dict...') |
| ckpt = torch.load('ckpts/step_259567/archon.pt', map_location='cpu', weights_only=False) |
| missing, unexpected = m.load_state_dict(ckpt['model'], strict=False) |
| print(f' missing: {len(missing)} (first 5: {missing[:5]})') |
| print(f' unexpected: {len(unexpected)} (first 5: {unexpected[:5]})') |
|
|
| print('\n[3] Moving to GPU bf16...') |
| m = m.to(device='cuda', dtype=torch.bfloat16) |
| torch.cuda.synchronize() |
| free, total = torch.cuda.mem_get_info() |
| print(f' VRAM used: {(total-free)/1e9:.2f}GB / {total/1e9:.2f}GB') |
|
|
| print('\n[4] Forward pass batch=4 seq=4096...') |
| x = torch.randint(0, 32000, (4, 4096), device='cuda') |
| with torch.no_grad(): |
| torch.cuda.synchronize() |
| t0 = time.time() |
| out = m(x) if not hasattr(m, 'forward_impl') else m.forward(x) |
| torch.cuda.synchronize() |
| dt = time.time() - t0 |
| print(f' output type: {type(out).__name__}') |
| if isinstance(out, torch.Tensor): |
| print(f' output shape: {tuple(out.shape)} dtype: {out.dtype}') |
| elif isinstance(out, (tuple, list)): |
| for i, o in enumerate(out): |
| if isinstance(o, torch.Tensor): |
| print(f' out[{i}]: {tuple(o.shape)} {o.dtype}') |
| elif isinstance(out, dict): |
| for k, v in out.items(): |
| if isinstance(v, torch.Tensor): |
| print(f' out[{k}]: {tuple(v.shape)} {v.dtype}') |
| print(f' forward time: {dt*1000:.1f}ms') |
|
|
| print('\n[5] Throughput estimate (forward only, no grad):') |
| tokens = 4 * 4096 |
| print(f' {tokens/dt:.0f} tok/s forward bf16') |
| print(f' (training est ~3x slower = ~{tokens/(dt*3):.0f} tok/s with backward+optim)') |
|
|
| print('\nSMOKE_TEST_OK') |
|
|
