Add per-epoch profiling to benchmark, push trigger for iteration

- Track epoch times in `EpochLossCallback` (train epoch start → end)
- Print epoch timing breakdown in both CLI (`--verbose`) and Modal benchmark
- Shows overhead = wallclock - sum(epoch times) to isolate setup/teardown
- Temporary `push: branches: [modal-benchmark-v2]` for GHA iteration

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: ryan-williams

.github/workflows/gpu-benchmark-modal.yml

@@ -1,6 +1,8 @@
 name: GPU Benchmark (Modal)
 
 on:
+  push:
+    branches: [modal-benchmark-v2]
   workflow_dispatch:
     inputs:
       gpu:

modal/benchmark.py

@@ -220,6 +220,15 @@ def run_benchmark(
     log.info("train_loss: %.6f", results["final_train_loss"])
     log.info("Wallclock: %.1fs", results["wallclock_s"])
     log.info("GPU: %s (Modal)", gpu_type)
+    epoch_times = results.get("epoch_times", [])
+    if epoch_times:
+        for i, t in enumerate(epoch_times):
+            log.info("  Epoch %d: %.1fs", i, t)
+        log.info("  Mean epoch: %.1fs", sum(epoch_times) / len(epoch_times))
+        log.info(
+            "  Overhead (wallclock - sum epochs): %.1fs",
+            results["wallclock_s"] - sum(epoch_times),
+        )
 
     return results
 

scripts/e2e_train.py

@@ -202,8 +202,19 @@ class EpochLossCallback(Callback):
         def __init__(self):
             self.epoch_train_losses: list[float] = []
             self.epoch_val_losses: list[float] = []
+            self.epoch_times: list[float] = []
+            self._epoch_start: float = 0
+
+        def on_train_epoch_start(self, _trainer, _pl_module):
+            from time import monotonic
+
+            self._epoch_start = monotonic()
 
         def on_train_epoch_end(self, trainer, _pl_module):
+            from time import monotonic
+
+            if self._epoch_start:
+                self.epoch_times.append(monotonic() - self._epoch_start)
             train_loss = trainer.callback_metrics.get("train_loss_epoch")
             if train_loss is not None:
                 self.epoch_train_losses.append(float(train_loss))
@@ -316,6 +327,7 @@ def on_validation_epoch_end(self, trainer, _pl_module):
         "final_train_loss": final_train_loss,
         "epoch_val_losses": loss_callback.epoch_val_losses,
         "epoch_train_losses": loss_callback.epoch_train_losses,
+        "epoch_times": loss_callback.epoch_times,
         "wallclock_s": train_wallclock,
         "wandb_run_url": wandb_run_url,
     }
@@ -377,6 +389,11 @@ def main(
             echo(f"  Final train_loss: {results['final_train_loss']:.6f}")
         echo(f"  Epoch val_losses: {[f'{v:.6f}' for v in results['epoch_val_losses']]}")
         echo(f"  Epoch train_losses: {[f'{v:.6f}' for v in results['epoch_train_losses']]}")
+        epoch_times = results.get("epoch_times", [])
+        if epoch_times:
+            echo(f"  Epoch times: {[f'{t:.1f}s' for t in epoch_times]}")
+            echo(f"  Mean epoch: {sum(epoch_times) / len(epoch_times):.1f}s")
+            echo(f"  Overhead (wallclock - sum epochs): {results['wallclock_s'] - sum(epoch_times):.1f}s")
 
     echo(f"Final val_loss: {final_val_loss:.6f}")