Perf/optimizer step overhead

ScaFFold/cli.py

@@ -145,6 +145,11 @@ def main():
         type=int,
         help="Number of warmup batches to run per rank before training.",
     )
+    benchmark_parser.add_argument(
+        "--gradient-clip-max-norm",
+        type=float,
+        help="Clip gradients to this max norm. Values <= 0 disable clipping.",
+    )
     benchmark_parser.add_argument(
         "--optimizer",
         type=str,

ScaFFold/configs/benchmark_default.yml

@@ -25,6 +25,7 @@ datagen_from_scratch: 0            # If 1, delete existing fractals and instance
 train_from_scratch: 1              # If 1, delete existing train stats and checkpoint files. Keep 0 if want to restart runs where we left off.
 dist: 1                            # If 1, use torch DDP.
 torch_amp: 1                       # If 1, use mixed precision in training.
+gradient_clip_max_norm: 0.0        # Values > 0 enable gradient clipping. Set to 1.0 to match the previous default behavior.
 framework: "torch"                 # The DL framework to train with. Only valid option for now is "torch".
 checkpoint_dir: "checkpoints"      # Subfolder in which to save training checkpoints.
 loss_freq: 1                       # Number of epochs between logging the overall loss.

ScaFFold/configs/benchmark_testing.yml

@@ -1,30 +1,31 @@
 # External/user-facing
 base_run_dir: "benchmark_runs"     # Subfolder of $(pwd) in which to run jobs.
-dataset_dir: "datasets"            # Directory in which to store and query for datasets.
+dataset_dir: "/p/lustre5/miles30/benchmark_datasets/"            # Directory in which to store and query for datasets.
 fract_base_dir: "fractals"         # Base directory for fractal IFS and instances.
 n_categories: 5                    # Number of fractal categories present in the dataset.
 n_instances_used_per_fractal: 145  # Number of unique instances to pull from each fractal class. There are 145 unique; exceeding this number will reuse some instances.
-problem_scale: 6                   # Determines dataset resolution and number of unet layers. Default is 6.
+problem_scale: 7                   # Determines dataset resolution and number of unet layers. Default is 6.
 unet_bottleneck_dim: 3             # Power of 2 of the unet bottleneck layer dimension. Default of 3 -> bottleneck layer of size 8.
 seed: 42                           # Random seed.
 batch_size: 1                      # Batch sizes for each vol size.
 optimizer: "ADAM"                  # "ADAM" is preferred option, otherwise training defautls to RMSProp.
-num_shards: [1, 1, 1]              # DistConv param: number of shards to divide the tensor into. It's best to choose the fewest ranks needed to fit one sample in GPU memory, since that keeps communication at a minimum
-shard_dim: [2, 3, 4]               # DistConv param: dimension on which to shard
+dc_num_shards: [1, 1, 1]              # DistConv param: number of shards to divide the tensor into. It's best to choose the fewest ranks needed to fit one sample in GPU memory, since that keeps communication at a minimum
+dc_shard_dims: [2, 3, 4]               # DistConv param: dimension on which to shard
 checkpoint_interval: 100           # Checkpoint every C epochs. More frequent checkpointing can be very expensive on slow filesystems.
 
 # Internal/dev use only
 variance_threshold: 0.15           # Variance threshold for valid fractals. Default is 0.15.
 n_fracts_per_vol: 3                # Number of fractals overlaid in each volume. Default is 3.
 val_split: 25                      # In percent.
-epochs: 10                         # Number of training epochs.
+epochs: 3                         # Number of training epochs.
 learning_rate: .0001               # Learning rate for training.
 disable_scheduler: 1               # If 1, disable scheduler during training to use constant LR.
 more_determinism: 0                # If 1, improve model training determinism.
 datagen_from_scratch: 0            # If 1, delete existing fractals and instances, then regenerate from scratch.
 train_from_scratch: 1              # If 1, delete existing train stats and checkpoint files. Keep 0 if want to restart runs where we left off.
 dist: 1                            # If 1, use torch DDP.
 torch_amp: 1                       # If 1, use mixed precision in training.
+gradient_clip_max_norm: 0.0        # Values > 0 enable gradient clipping. Set to 1.0 to match the previous default behavior.
 framework: "torch"                 # The DL framework to train with. Only valid option for now is "torch".
 checkpoint_dir: "checkpoints"      # Subfolder in which to save training checkpoints.
 loss_freq: 1                       # Number of epochs between logging the overall loss.

ScaFFold/utils/config_utils.py

@@ -63,6 +63,9 @@ def __init__(self, config_dict):
         self.learning_rate = config_dict["learning_rate"]
         self.variance_threshold = config_dict["variance_threshold"]
         self.torch_amp = bool(config_dict["torch_amp"])
+        self.gradient_clip_max_norm = float(
+            config_dict.get("gradient_clip_max_norm", 0.0)
+        )
         self.loss_freq = config_dict["loss_freq"]
         self.checkpoint_dir = config_dict["checkpoint_dir"]
         self.normalize = config_dict["normalize"]

ScaFFold/utils/trainer.py

@@ -182,7 +182,7 @@ def setup_training_components(self):
         )
 
         self.log.info(
-            f"Optimizer: {self.optimizer}, Scheduler: {self.scheduler}, Gradient Scaler Enabled: {self.config.torch_amp}"
+            f"Optimizer: {self.optimizer}, Scheduler: {self.scheduler}, Gradient Scaler Enabled: {self.config.torch_amp}, Gradient Clip Max Norm: {self.config.gradient_clip_max_norm}"
         )
 
 
@@ -341,7 +341,7 @@ def warmup(self):
 
         # Match the main training path as closely as possible.
         self.model.train()
-        self.optimizer.zero_grad(set_to_none=False)
+        self.optimizer.zero_grad(set_to_none=True)
         start_warmup = time.time()
         max_batches = min(warmup_batches, len(self.train_loader))
         self.log.info(f"Running {max_batches} warmup batch(es) per rank")
@@ -442,8 +442,12 @@ def warmup(self):
 
             # Backward pass
             self.grad_scaler.scale(loss).backward()
-            self.grad_scaler.unscale_(self.optimizer)
-            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+            if self.config.gradient_clip_max_norm > 0:
+                self.grad_scaler.unscale_(self.optimizer)
+                torch.nn.utils.clip_grad_norm_(
+                    self.model.parameters(),
+                    max_norm=self.config.gradient_clip_max_norm,
+                )
             self.log.debug(f"  warmup: backward pass complete. Stepping optimizer")
 
             self.grad_scaler.step(self.optimizer)
@@ -505,7 +509,7 @@ def train(self):
                     self.train_loader.sampler.set_epoch(epoch)
                     self.val_loader.sampler.set_epoch(epoch)
                 self.model.train()
-                self.optimizer.zero_grad(set_to_none=False)
+                self.optimizer.zero_grad(set_to_none=True)
 
                 estr = (
                     f"{epoch}"
@@ -643,14 +647,16 @@ def train(self):
                         gather_and_print_mem(self.log, "post_backward")
 
                         begin_code_region("step_and_update")
-                        self.grad_scaler.unscale_(self.optimizer)
-                        torch.nn.utils.clip_grad_norm_(
-                            self.model.parameters(), max_norm=1.0
-                        )
+                        if self.config.gradient_clip_max_norm > 0:
+                            self.grad_scaler.unscale_(self.optimizer)
+                            torch.nn.utils.clip_grad_norm_(
+                                self.model.parameters(),
+                                max_norm=self.config.gradient_clip_max_norm,
+                            )
                         self.grad_scaler.step(self.optimizer)
                         gather_and_print_mem(self.log, "after_optim_step")
                         self.grad_scaler.update()
-                        self.optimizer.zero_grad(set_to_none=False)
+                        self.optimizer.zero_grad(set_to_none=True)
                         end_code_region("step_and_update")
 
                         # Update the loss

docs/perf/optimizer-step-overhead.md

@@ -0,0 +1,59 @@
+# Optimizer Step Overhead
+
+This branch reduces optimizer-step overhead relative to `miles30/performance`.
+
+## What Changed
+
+- All training-path `optimizer.zero_grad(...)` calls now use `set_to_none=True`.
+- Gradient clipping is now controlled by `gradient_clip_max_norm`.
+- Clipping is disabled by default with `gradient_clip_max_norm: 0.0`.
+- When clipping is disabled, the training loop skips both `grad_scaler.unscale_(optimizer)` and `clip_grad_norm_`.
+
+## Public Interface
+
+- New config key: `gradient_clip_max_norm`
+- New CLI override: `--gradient-clip-max-norm`
+
+Set `gradient_clip_max_norm: 1.0` to reproduce the previous clipping threshold.
+
+## What Did Not Change
+
+- Optimizer stepping is still once per batch.
+- Loss math and Dice math are unchanged.
+- Data loading and DistConv behavior are unchanged.
+
+## Expected Effect
+
+- Lower per-step optimizer overhead.
+- Fewer full-parameter traversals in the benchmark default configuration.
+- Cleaner separation between numerical stability tuning and throughput measurement.
+
+## Evaluation
+
+Default benchmark behavior now runs with clipping disabled:
+
+```bash
+torchrun-hpc -N 4 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark \
+  -c $(pwd)/ScaFFold/configs/benchmark_testing.yml \
+  --problem-scale 7 --dc-num-shards 1 2 2
+
+torchrun-hpc -N 8 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark \
+  -c $(pwd)/ScaFFold/configs/benchmark_testing.yml \
+  --problem-scale 8 --dc-num-shards 1 2 2
+```
+
+To compare against the old clipping behavior:
+
+```bash
+torchrun-hpc -N 4 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark \
+  -c $(pwd)/ScaFFold/configs/benchmark_testing.yml \
+  --problem-scale 7 --dc-num-shards 1 2 2 \
+  --gradient-clip-max-norm 1.0
+```
+
+## Acceptance Criteria
+
+- Default config disables clipping.
+- Enabling `--gradient-clip-max-norm 1.0` restores the old threshold.
+- Throughput improves or stays flat in epoch `2+` timing.
+- Training still runs correctly with and without clipping enabled.