experiments: log peak memory and total training work

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: