From d4477d2d9a04c8e941a35173154e89517f7e5998 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Wed, 22 Jul 2026 02:20:20 -0500 Subject: experiments: log peak memory and total training work --- experiments/run.py | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/experiments/run.py b/experiments/run.py index a44e72e..bc66bae 100644 --- a/experiments/run.py +++ b/experiments/run.py @@ -53,6 +53,34 @@ def device_sync(device): torch.cuda.synchronize() +def reset_peak_memory(device): + if str(device).startswith("cuda") and torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats(torch.device(device)) + + +def hardware_report(device): + report = {"device": str(device), "torch_version": torch.__version__} + if str(device).startswith("cuda") and torch.cuda.is_available(): + cuda_device = torch.device(device) + props = torch.cuda.get_device_properties(cuda_device) + report.update({ + "cuda_device_name": props.name, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "device_total_memory_bytes": props.total_memory, + "peak_memory_allocated_bytes": torch.cuda.max_memory_allocated(cuda_device), + "peak_memory_reserved_bytes": torch.cuda.max_memory_reserved(cuda_device), + }) + else: + report.update({ + "cuda_device_name": None, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "device_total_memory_bytes": None, + "peak_memory_allocated_bytes": None, + "peak_memory_reserved_bytes": None, + }) + return report + + def calibration_work_per_event(net, cfg): """Hardware-independent causal-calibration work for one minibatch. @@ -192,6 +220,9 @@ def train(args): train_loader, eval_loader, n_in, n_out, split = load_task(args, device) args.n_in, args.n_out = n_in, n_out net, cfg = build(args, device) + # Reset after persistent data/model allocation so the peak includes the + # resident training state plus transient perturbation-batch expansion. + reset_peak_memory(device) # A fixed probe batch is loaded only when diagnostics are requested. In a # frozen test-only run with diagnostics disabled, the test set is therefore @@ -384,8 +415,12 @@ def train(args): "calibration_batch_loss_evaluations": calibration_batch_loss_evaluations, "calibration_example_loss_evaluations": calibration_example_loss_evaluations, "calibration_forward_equivalent_examples": calibration_forward_equivalent_examples, + "training_forward_equivalent_examples": ( + ordinary_forward_examples + warmup_examples + + calibration_forward_equivalent_examples), "max_perturbation_batch_expansion": perturbation_batch_expansion, } + log["hardware"] = hardware_report(device) os.makedirs(args.outdir, exist_ok=True) outpath = os.path.join(args.outdir, f"{args.tag}.json") with open(outpath, "w") as f: -- cgit v1.2.3