| """Atomic full-state DDP checkpoints with per-rank RNG and pinned contracts.""" |
| from datetime import datetime |
| import json |
| import os |
| from pathlib import Path |
| import random |
| import numpy as np |
| import torch |
| import torch.distributed as dist |
| from continuation_control import sha |
|
|
|
|
| def ranks(): |
| return (dist.get_rank(), dist.get_world_size()) if dist.is_initialized() else (0, 1) |
|
|
|
|
| def barrier(): |
| if dist.is_initialized(): |
| dist.barrier() |
|
|
|
|
| def unwrap(model): |
| return model.module if hasattr(model, 'module') else model |
|
|
|
|
| def save(root, model, optimizer, scheduler, step, contract, export=False): |
| rank, world = ranks() |
| root = Path(root) |
| rng = {'torch': torch.get_rng_state(), 'cuda': torch.cuda.get_rng_state(), |
| 'numpy': np.random.get_state(), 'python': random.getstate()} |
| all_rng = [None] * world if rank == 0 else None |
| if world > 1: |
| dist.gather_object(rng, all_rng, dst=0) |
| else: |
| all_rng = [rng] |
| directory = root / f'step-{step:08d}' |
| if rank == 0: |
| temporary = root / f'.step-{step:08d}-{os.environ.get("SLURM_JOB_ID", "local")}.partial' |
| temporary.mkdir(parents=True, exist_ok=True) |
| torch.save({'model': unwrap(model).state_dict(), 'optimizer': optimizer.state_dict(), |
| 'scheduler': scheduler.state_dict(), 'rng': all_rng, 'step': step, 'world': world, |
| 'contract': contract}, temporary / 'training_state.pt') |
| metadata = {'step': step, 'world': world, 'contract': contract, |
| 'created_at': datetime.now().astimezone().isoformat(), |
| 'state_bytes': (temporary / 'training_state.pt').stat().st_size, |
| 'state_sha256': sha(temporary / 'training_state.pt')} |
| if export: |
| base = unwrap(model).base |
| cache, state = {}, {} |
| for name, value in base.state_dict().items(): |
| key = (value.data_ptr(), tuple(value.shape), value.dtype) |
| if key not in cache: |
| cache[key] = value.detach().to(dtype=torch.bfloat16, device='cpu') |
| state[name] = cache[key] |
| torch.save(state, temporary / 'model_bf16.pt') |
| metadata['export_bytes'] = (temporary / 'model_bf16.pt').stat().st_size |
| metadata['export_sha256'] = sha(temporary / 'model_bf16.pt') |
| (temporary / 'COMPLETE.json').write_text(json.dumps(metadata, indent=2) + '\n') |
| if directory.exists(): |
| raise RuntimeError(f'Refusing to overwrite existing checkpoint {directory}') |
| temporary.replace(directory) |
| pointer = root / 'current.json.tmp' |
| pointer.write_text(json.dumps({'directory': directory.name, 'step': step}) + '\n') |
| pointer.replace(root / 'current.json') |
| barrier() |
| return directory |
|
|
|
|
| def load(directory, model, optimizer, scheduler, contract): |
| directory = Path(directory) |
| rank, world = ranks() |
| metadata = json.loads((directory / 'COMPLETE.json').read_text()) |
| assert metadata['contract'] == contract and metadata['world'] == world |
| assert (directory / 'training_state.pt').stat().st_size == metadata['state_bytes'] |
| state = torch.load(directory / 'training_state.pt', map_location='cpu', weights_only=False) |
| assert state['contract'] == contract and state['world'] == world and len(state['rng']) == world |
| unwrap(model).load_state_dict(state['model'], strict=True) |
| optimizer.load_state_dict(state['optimizer']); scheduler.load_state_dict(state['scheduler']) |
| rng = state['rng'][rank] |
| torch.set_rng_state(rng['torch']); torch.cuda.set_rng_state(rng['cuda']) |
| np.random.set_state(rng['numpy']); random.setstate(rng['python']) |
| step = int(state['step']) |
| assert step==metadata['step'] |
| del state |
| barrier() |
| return step |
|
|
