diff options
Diffstat (limited to 'experiments/run.py')
| -rw-r--r-- | experiments/run.py | 35 |
1 files changed, 35 insertions, 0 deletions
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: |