File size: 3,557 Bytes
5de8603 | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | #!/usr/bin/env python3
"""Pin every complete production checkpoint through the requested stage."""
import argparse
import hashlib
import json
from pathlib import Path
HERE = Path(__file__).resolve().parent
SC = Path('/e/scratch/reformo/schuhmann1_moss')
CHECKPOINTS = SC / 'out/m2_600m_ladder/caption_curriculum_v1/training/checkpoints'
PLAN = HERE / 'plan_caption_production_32nodes.json'
PROMPTS = SC / 'out/m2_20k_score/eval/timed_prompts_v3.json'
def sha(path: Path) -> str:
digest = hashlib.sha256()
with path.open('rb') as stream:
for block in iter(lambda: stream.read(8 << 20), b''):
digest.update(block)
return digest.hexdigest()
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--through-stage', default='S5')
ap.add_argument('--output', default=str(HERE / 'validation_sweep_S1_S5_all_checkpoints.json'))
args = ap.parse_args()
plan = json.loads(PLAN.read_text())
boundaries, end = [], 0
for stage in plan['stages']:
end += int(stage['updates'])
boundaries.append((stage['name'], end))
assert args.through_stage in dict(boundaries)
final_step = dict(boundaries)[args.through_stage]
candidates = []
for directory in sorted(CHECKPOINTS.glob('step-*')):
complete = directory / 'COMPLETE.json'
if not complete.exists():
continue
proof = json.loads(complete.read_text())
step = int(proof['step'])
if step > final_step:
continue
assert proof['world'] == plan['world_size']
assert proof['contract']['plan_sha256'] == sha(PLAN)
assert proof['contract']['objective'] == 'global_token_mean_continuous_ladder_v1'
stage, _ = next((name, bound) for name, bound in boundaries if step <= bound)
state = directory / 'training_state.pt'
assert state.stat().st_size == int(proof['state_bytes'])
candidates.append({'stage': stage, 'step': step, 'checkpoint': str(directory),
'state_sha256': proof['state_sha256'],
'export_available': (directory / 'model_bf16.pt').exists()})
assert candidates and max(row['step'] for row in candidates) == final_step
stage_order = {name: index for index, (name, _) in enumerate(boundaries)}
candidates.sort(key=lambda row: (-stage_order[row['stage']], row['step']))
payload = {'version': 1, 'purpose': 'all complete checkpoints through selected stage',
'metric': 'independent_heldout_validation_loss',
'selection_policy': 'minimum pooled token-weighted heldout loss per stage',
'acting_metrics_used_for_selection': False, 'training_loss_used_for_selection': False,
'world': 4, 'limit': None, 'production_plan': str(PLAN),
'production_plan_sha256': sha(PLAN), 'heldout_prompts': str(PROMPTS),
'heldout_prompts_sha256': sha(PROMPTS), 'heldout_rows': 50,
'tasks_per_checkpoint': 100, 'through_stage': args.through_stage,
'candidate_count': len(candidates), 'candidates': candidates}
output = Path(args.output).resolve()
assert not output.exists(), f'refusing to overwrite existing manifest: {output}'
output.write_text(json.dumps(payload, indent=2) + '\n')
print(json.dumps({'manifest': str(output), 'candidate_count': len(candidates),
'stages': {name: sum(c['stage'] == name for c in candidates)
for name, bound in boundaries if bound <= final_step}}, indent=2))
if __name__ == '__main__':
main()
|