Buckets:
| """Recreate update-2 draws at the committed actor and compare replay memory.""" | |
| import argparse | |
| import gc | |
| import json | |
| from pathlib import Path | |
| import traceback | |
| import torch | |
| from rl.config import load_config | |
| from rl.distributed import DistributedContext, launch_configured_workers | |
| from rl.evaluation import make_backend | |
| from rl.objective import multimodal_grpo_loss | |
| from rl.pipeline import protocol_limits | |
| from rl.policy import PolicyScorer | |
| from rl.rollout import GRPORolloutController | |
| from rl.runtime import load_reference_state, verify_rl_checkpoint | |
| from shared.artifacts import claim_run_directory, acquire_run_writer, seed_everything, write_json | |
| from shared.data import PackedTrajectoryCache, query_indices_for_update | |
| from shared.modeling import load_model, load_processor, build_optimizer | |
| from shared.reference import FrozenReference | |
| def score_case(actor, reference, trajectory, chunk, device): | |
| actor.zero_grad(set_to_none=True) | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| torch.cuda.reset_peak_memory_stats(device) | |
| try: | |
| score = PolicyScorer(actor, reference, text_temperature=1., visual_temperature=1., | |
| text_chunk_size=chunk).score(trajectory, advantage=torch.tensor(1.)) | |
| loss = multimodal_grpo_loss([score.inputs], clip_low=.2, clip_high=.2, reference_beta=.01) | |
| loss.loss.backward() | |
| grad = torch.nn.utils.clip_grad_norm_(actor.parameters(), float("inf"), error_if_nonfinite=True) | |
| return {"ok": True, "loss": float(loss.loss.detach()), "grad_norm": float(grad), | |
| "peak_gib": torch.cuda.max_memory_allocated(device) / 1024**3} | |
| except torch.OutOfMemoryError: | |
| return {"ok": False, "error": traceback.format_exc(), | |
| "peak_gib": torch.cuda.max_memory_allocated(device) / 1024**3} | |
| 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) | |
| lock = None | |
| try: | |
| def prepare(): | |
| nonlocal lock | |
| claim_run_directory(run, resume=(run / ".run-claim.json").exists()) | |
| lock = acquire_run_writer(run) | |
| parallel.main_call(prepare) | |
| if parallel.main_call(lambda: (run / "probe.json").exists()): | |
| return | |
| checkpoint = Path(config.checkpoint.resume_from) | |
| verify_rl_checkpoint(checkpoint) | |
| cache = PackedTrajectoryCache(config.data.cache_root) | |
| indices = query_indices_for_update(cache.indices("train"), batch_size=8, | |
| seed=config.run.seed, update=1, shuffle=True) | |
| processor = load_processor(config.model.path) | |
| actor = load_model(config, device, for_training=True) | |
| optimizer, _, _ = build_optimizer(actor, config) | |
| state = torch.load(checkpoint / "training_state.pt", map_location="cpu", weights_only=False) | |
| optimizer.load_state_dict(state["optimizer"]) | |
| del state | |
| reference = FrozenReference(actor) | |
| load_reference_state(checkpoint.parent.parent, reference) | |
| slot, member = ((0, 0), (1, 7), (2, 7), (2, 3))[parallel.rank] | |
| case = run / f"rank-{parallel.rank}" | |
| case.mkdir(exist_ok=True) | |
| trace_path = case / "trajectory.pt" | |
| actor.eval() | |
| if trace_path.exists(): | |
| trajectory = torch.load(trace_path, map_location="cpu", weights_only=False) | |
| else: | |
| controller = GRPORolloutController(make_backend(actor, processor, config), protocol_limits(config)) | |
| trajectory, = controller.collect_group_members(indices[slot], cache.load(indices[slot]), | |
| group_size=8, group_indices=(member,), seed=config.run.seed, update=1, prompt_slot=slot) | |
| temporary = case / ".trajectory.tmp" | |
| torch.save(trajectory, temporary) | |
| temporary.rename(trace_path) | |
| lengths = {"sample_id": trajectory.sample_id, "slot": slot, "member": member, | |
| "sequence": len(trajectory.input_ids), "text_actions": len(trajectory.text_action_positions), | |
| "visual_steps": len(trajectory.visual_action_positions), "stop_reason": trajectory.stop_reason} | |
| print("RL_MEMORY_CASE " + json.dumps({"rank": parallel.rank, **lengths}), flush=True) | |
| actor.train() | |
| dense = score_case(actor, reference, trajectory, None, device) | |
| chunked = score_case(actor, reference, trajectory, 128, device) | |
| result = {"rank": parallel.rank, **lengths, "dense": dense, "chunked": chunked} | |
| write_json(case / "result.json", result) | |
| print("RL_MEMORY_RESULT " + json.dumps(result), flush=True) | |
| # Rank-local files remain inspectable if another worker is interrupted. | |
| results = parallel.all_gather_object(result) | |
| if not all(value["chunked"]["ok"] for value in results): | |
| raise RuntimeError("chunked replay memory probe failed") | |
| parallel.main_call(lambda: write_json(run / "probe.json", {"passed": True, "results": results})) | |
| finally: | |
| if lock: | |
| lock.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:
- 5.55 kB
- Xet hash:
- 26549a534ec7d10138e36b3fec80985a293206e0c8587f00edf06ff0adc4c858
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.