Buckets:
| """Launch the unchanged OPD objective and summarize every GPU's phase peaks.""" | |
| import argparse | |
| from collections import Counter | |
| import json | |
| from pathlib import Path | |
| import subprocess | |
| import sys | |
| def summarize(run, config, returncode): | |
| import torch | |
| from bagel_mmcot.shared.trajectory import ImageBlock, TextBlock | |
| from bagel_mmcot.opd.runtime import atomic_json | |
| records = [json.loads(path.read_text()) for path in (run/'memory').glob('*/*.json')] | |
| memory = {} | |
| for rank in range(config.trainer.processes): | |
| rows = [row for row in records if row['rank'] == rank] | |
| memory[str(rank)] = { | |
| 'peak_allocated_gib': max((row['peak_allocated_gib'] for row in rows), default=0), | |
| 'peak_reserved_gib': max((row['peak_reserved_gib'] for row in rows), default=0), | |
| 'phases_completed': sorted({row['phase'] for row in rows if row['phase'].endswith('_done')}), | |
| 'errors': [row['error'] for row in rows if 'error' in row], | |
| } | |
| coverage = dict(samples=0, valid=0, awaiting_verdict=0, privileged_wrong=0, max_images=0, | |
| max_text_block_tokens=0, max_generated_image_pixels=0) | |
| reasons = Counter() | |
| for path in (run/'rollouts/train').glob('update-*/*.pt'): | |
| row = torch.load(path, map_location='cpu', weights_only=False) | |
| trajectory, verdict = row['trajectory'], row['verdict'] | |
| images = [block for block in trajectory.blocks if isinstance(block, ImageBlock)] | |
| texts = [block for block in trajectory.blocks if isinstance(block, TextBlock)] | |
| coverage['samples'] += 1 | |
| coverage['valid'] += int(trajectory.valid) | |
| coverage['awaiting_verdict'] += int(verdict is None) | |
| coverage['privileged_wrong'] += int(trajectory.valid and verdict is not None and not verdict.correct | |
| and config.routing.mode == 'privileged_wrong') | |
| coverage['max_images'] = max(coverage['max_images'], len(images)) | |
| coverage['max_text_block_tokens'] = max(coverage['max_text_block_tokens'], | |
| max((len(block.tokens)-1 for block in texts), default=0)) | |
| coverage['max_generated_image_pixels'] = max(coverage['max_generated_image_pixels'], | |
| max((block.image.width * block.image.height for block in images), default=0)) | |
| reasons[trajectory.stop_reason] += 1 | |
| coverage['stop_reasons'] = dict(reasons) | |
| metrics = [json.loads(path.read_text()) for path in (run/'metrics').glob('*.json')] | |
| needed = ['opd/' + '/'.join(key) for key,term in config.components().items() if term.weight > 0] | |
| exercised = [key for key in needed if any(row.get(key,0) > 0 for row in metrics)] | |
| complete = (run/'COMPLETE').is_file() | |
| passed = (returncode == 0 and complete and set(exercised) == set(needed) | |
| and all('optimizer_done' in row['phases_completed'] for row in memory.values())) | |
| result = dict(passed=passed, returncode=returncode, complete=complete, memory_per_rank=memory, | |
| components_exercised=exercised, trajectory_coverage=coverage, | |
| configured_max_text_tokens=config.rollout.max_text_tokens, | |
| configured_max_images=config.rollout.max_images, | |
| limitation='Natural rollout coverage only; configured length/image caps are not forced.') | |
| atomic_json(run/'smoke_report.json', result) | |
| return result | |
| def main(): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument('--config', type=Path, required=True) | |
| parser.add_argument('--output-dir', type=Path, required=True) | |
| args = parser.parse_args() | |
| from bagel_mmcot.opd.config import load_config | |
| config = load_config(args.config) | |
| if not config.trainer.memory_profile: | |
| raise ValueError('memory smoke requires trainer.memory_profile=true') | |
| result = subprocess.run([sys.executable, '-m', 'bagel_mmcot', 'opd', 'train', | |
| '--config', str(args.config), '--output-dir', str(args.output_dir)]) | |
| report = summarize(args.output_dir, config, result.returncode) | |
| print(json.dumps(report, indent=2), flush=True) | |
| return 0 if report['passed'] else 1 | |
| if __name__ == '__main__': | |
| raise SystemExit(main()) | |
Xet Storage Details
- Size:
- 4.27 kB
- Xet hash:
- e9597f9b25eb0d614fc2e319fc317dc63ee57b7234088e0e954a60c5f981e7e3
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.