| """Pure control rules shared by the real trainer and its recovery verifier.""" |
| import hashlib |
| import json |
| import math |
| from pathlib import Path |
|
|
|
|
| def sha(path): |
| result=hashlib.sha256() |
| with Path(path).open('rb') as stream: |
| for block in iter(lambda:stream.read(8<<20),b''):result.update(block) |
| return result.hexdigest() |
|
|
|
|
| def complete_checkpoint(directory,contract=None,world=None,content=False): |
| directory=Path(directory);meta=json.loads((directory/'COMPLETE.json').read_text()) |
| assert meta['step']>=1 |
| if contract is not None:assert meta['contract']==contract |
| if world is not None:assert meta['world']==world |
| assert (directory/'training_state.pt').stat().st_size==meta['state_bytes']>0 |
| if content:assert sha(directory/'training_state.pt')==meta['state_sha256'] |
| if 'export_bytes' in meta: |
| assert (directory/'model_bf16.pt').stat().st_size==meta['export_bytes']>0 |
| if content:assert sha(directory/'model_bf16.pt')==meta['export_sha256'] |
| return meta |
|
|
|
|
| def latest_checkpoint(root,contract,world): |
| """A crash after directory rename but before pointer update is recoverable.""" |
| found=[] |
| for directory in Path(root).glob('step-*'): |
| |
| |
| meta=complete_checkpoint(directory,contract,world) |
| assert directory.name==f'step-{meta["step"]:08d}' |
| found.append((meta['step'],directory)) |
| return max(found)[1] if found else None |
|
|
|
|
| def epoch_plan(samples,world,per_gpu): |
| assert samples>0 and world>0 and per_gpu>0 |
| batch=world*per_gpu;steps=math.ceil(samples/batch);tail=samples-(steps-1)*batch |
| assert tail>=world,'Every final rank must receive a real target; no silent tail drop/padding' |
| quarters=sorted({math.ceil(steps*q/4) for q in (1,2,3,4)}) |
| return {'samples':samples,'world':world,'global_batch':batch,'steps':steps,'tail_samples':tail, |
| 'quarter_steps':quarters,'epochs':1,'presentations_per_target':1,'dropped_targets':0,'duplicated_targets':0} |
|
|
|
|
| def contract(plan_path,plan,world): |
| return {'plan_sha256':sha(plan_path),'schema_sha256':plan['schema_sha256'],'manifest_sha256':plan['manifest_sha256'], |
| 'world':world,'objective':'global_token_mean_v1','seed':plan['seed']} |
|
|
