Buckets:
| """Compare fixed-rank and dynamic GRPO rollout scheduling on identical draws.""" | |
| import argparse | |
| import hashlib | |
| import json | |
| from pathlib import Path | |
| import os | |
| import time | |
| import torch | |
| from rl.collection import collect_logical_batch | |
| from rl.config import load_config | |
| from rl.distributed import DistributedContext, launch_configured_workers | |
| from rl.evaluation import make_backend | |
| from rl.pipeline import protocol_limits | |
| from rl.rollout import GRPORolloutController | |
| from shared.artifacts import acquire_run_writer, claim_run_directory, seed_everything, write_json | |
| from shared.data import PackedTrajectoryCache, query_indices_for_update | |
| from shared.modeling import load_model, load_processor | |
| def worker(config): | |
| parallel = DistributedContext.initialize(config.trainer.world_size) | |
| device = torch.device("cuda", parallel.local_rank) | |
| torch.cuda.set_device(device) | |
| seed_everything(config.run.seed) | |
| run = Path(config.run.output_dir) | |
| writer = None | |
| try: | |
| def claim(): | |
| nonlocal writer | |
| claim_run_directory(run, resume=(run / ".run-claim.json").exists()) | |
| writer = acquire_run_writer(run) | |
| parallel.main_call(claim) | |
| cache = PackedTrajectoryCache(config.data.cache_root) | |
| indices = query_indices_for_update(cache.indices("train"), batch_size=2, | |
| seed=config.run.seed, update=0, shuffle=config.data.shuffle) | |
| processor = load_processor(config.model.path) | |
| actor = load_model(config, device, for_training=True) | |
| actor.eval() | |
| records = [] | |
| for scheduling, cached in (("static", False), ("static", True), ("dynamic", True)): | |
| path = run / f"{scheduling}-autocast-{cached}.json" | |
| completed = parallel.main_call(lambda: json.loads(path.read_text()) if path.exists() else None) | |
| if completed is not None: | |
| records.append(completed) | |
| continue | |
| config.rollout.scheduling = scheduling | |
| config.rollout.cache_autocast = cached | |
| torch.cuda.reset_peak_memory_stats(device) | |
| controller = GRPORolloutController(make_backend(actor, processor, config), protocol_limits(config)) | |
| groups, timing = collect_logical_batch(controller=controller, cache=cache, | |
| prompt_indices=indices, config=config, scorer=None, tokenizer=processor.tokenizer, | |
| update=0, parallel=parallel, run_dir=run) | |
| signatures = [] | |
| for group in groups: | |
| for scored in group.trajectories: | |
| trajectory = scored.trajectory | |
| digest = hashlib.sha256() | |
| for value in (trajectory.input_ids, trajectory.sampled_visual_codes, | |
| trajectory.text_old_log_probs, trajectory.visual_old_factor_log_probs): | |
| digest.update(value.numpy().tobytes()) | |
| signatures.append(digest.hexdigest()) | |
| peak = max(parallel.all_gather_object(torch.cuda.max_memory_allocated(device) / 1024**3)) | |
| record = {"scheduling": scheduling, "cache_autocast": cached, "timing": timing, | |
| "peak_gpu_memory_gib": peak, "signatures": signatures} | |
| parallel.main_call(lambda: write_json(path, record)) | |
| records.append(record) | |
| print("RL_COLLECTION_BENCHMARK " + json.dumps(record), flush=True) | |
| if any(record["signatures"] != records[0]["signatures"] for record in records): | |
| raise RuntimeError("execution optimization changed sampled tokens or old log-probabilities") | |
| parallel.main_call(lambda: write_json(run / "benchmark.json", {"passed": True, "records": records})) | |
| finally: | |
| if writer: | |
| writer.close() | |
| parallel.close() | |
| def main(): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--config", required=True) | |
| parser.add_argument("--output-dir", required=True) | |
| args = parser.parse_args() | |
| config = load_config(args.config, [f"run.output_dir={args.output_dir}"]) | |
| launch_configured_workers(config, worker) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 4.17 kB
- Xet hash:
- ee2412ebe857a3ece5274e7cec03e73a43e8650f6d53e1b72e972459b673cce4
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.