| |
| """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() |
|
|
