Buckets:
| """Fixed-query GH200 rollout timing and hard-trajectory equivalence checks.""" | |
| import argparse | |
| import json | |
| from pathlib import Path | |
| import torch | |
| from opd.config import load_config | |
| from opd.distributed import OPDDistributedContext | |
| from opd.rollout_artifacts import load_rollout | |
| from opd.workers import RolloutWorkers | |
| from opd.verifier import build_verifier | |
| from shared.artifacts import claim_run_directory, seed_everything, write_json | |
| from shared.data import PackedTrajectoryCache, rollout_indices_for_update | |
| from shared.modeling import load_model, load_processor | |
| def main(): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--config", required=True) | |
| parser.add_argument("--output-dir", required=True) | |
| parser.add_argument("--adapter", required=True) | |
| parser.add_argument("--update", type=int, default=18) | |
| parser.add_argument("--workers", type=int, nargs="+", default=[1, 2]) | |
| parser.add_argument("--reference-dir") | |
| args = parser.parse_args() | |
| config = load_config(args.config) | |
| config.model.initial_adapter = args.adapter | |
| config.rollout.workers_per_gpu = max(args.workers) | |
| parallel = OPDDistributedContext.initialize(config.trainer) | |
| directory = Path(args.output_dir) | |
| if parallel.is_main: | |
| claim_run_directory(directory, resume=(directory / ".run-claim.json").exists()) | |
| parallel.barrier() | |
| seed_everything(config.run.seed) | |
| torch.backends.cuda.matmul.allow_tf32 = config.trainer.tf32 | |
| model = load_model(config, torch.device("cuda", parallel.local_rank), for_training=True) | |
| parallel.broadcast_parameters(parameter for parameter in model.parameters() if parameter.requires_grad) | |
| processor = load_processor(config.model.path) | |
| cache = PackedTrajectoryCache(config.data.cache_root) | |
| indices = rollout_indices_for_update(cache.indices(config.data.train_split), batch_size=16, seed=config.run.seed, update=args.update, shuffle=True) | |
| verifier = build_verifier(config.verifier, processor.tokenizer) | |
| pool = RolloutWorkers(model, processor, cache, verifier, config, parallel) | |
| try: | |
| variants = ([] if args.reference_dir else [("old_autocast_w1", False, 1)]) | |
| variants += [(f"cached_autocast_w{workers}", True, workers) for workers in args.workers] | |
| for name, autocast, workers in variants: | |
| if parallel.broadcast_object((directory / name / "summary.json").exists() if parallel.is_main else None): | |
| continue | |
| config.rollout.cache_autocast = autocast | |
| # A partially completed benchmark phase resumes its same sample set. | |
| from opd.evaluation import committed_sample | |
| stage = directory / name | |
| if parallel.is_main: | |
| pending = tuple(index for index in indices if not committed_sample(stage, index)) | |
| else: | |
| pending = None | |
| pending = parallel.broadcast_object(pending) | |
| result = pool.run(pending, stage, phase="train", update=args.update, workers_per_gpu=workers) | |
| result["partial_resume"] = len(pending) != len(indices) | |
| result["sample_indices"] = list(indices) | |
| if parallel.is_main: | |
| if name != "old_autocast_w1": | |
| equal_tokens = equal_codes = 0 | |
| for index in indices: | |
| reference = Path(args.reference_dir) if args.reference_dir else directory / "old_autocast_w1" | |
| baseline = load_rollout(reference / "samples" / f"{index:08d}" / "trajectory.pt") | |
| current = load_rollout(stage / "samples" / f"{index:08d}" / "trajectory.pt") | |
| equal_tokens += int(torch.equal(baseline.input_ids, current.input_ids)) | |
| equal_codes += int(torch.equal(baseline.sampled_visual_codes, current.sampled_visual_codes)) | |
| result.update(exact_token_trajectories=equal_tokens, exact_visual_trajectories=equal_codes) | |
| write_json(stage / "summary.json", result) | |
| print("OPD_BENCHMARK " + json.dumps({"variant": name, **result}), flush=True) | |
| parallel.barrier() | |
| finally: | |
| pool.close() | |
| parallel.close() | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 4.27 kB
- Xet hash:
- ac772c370670925d3770f644a06355976db0d827cb019c89c064eb212a4256a5
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.