From 42dafef9c95b7b78610b55dc1d01e525c72d4278 Mon Sep 17 00:00:00 2001 From: Patrick Robert Miles Date: Wed, 8 Apr 2026 15:02:24 -0700 Subject: [PATCH 1/2] Make gradient clipping configurable --- ScaFFold/cli.py | 5 +++ ScaFFold/configs/benchmark_default.yml | 1 + ScaFFold/configs/benchmark_testing.yml | 1 + ScaFFold/utils/config_utils.py | 3 ++ ScaFFold/utils/trainer.py | 26 +++++++----- docs/perf/optimizer-step-overhead.md | 59 ++++++++++++++++++++++++++ 6 files changed, 85 insertions(+), 10 deletions(-) create mode 100644 docs/perf/optimizer-step-overhead.md diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 5b76c87..fda5ed0 100644 --- a/ScaFFold/cli.py +++ b/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, diff --git a/ScaFFold/configs/benchmark_default.yml b/ScaFFold/configs/benchmark_default.yml index 180a0dc..9fbb9da 100644 --- a/ScaFFold/configs/benchmark_default.yml +++ b/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. diff --git a/ScaFFold/configs/benchmark_testing.yml b/ScaFFold/configs/benchmark_testing.yml index 6b8c30a..8bfda6e 100644 --- a/ScaFFold/configs/benchmark_testing.yml +++ b/ScaFFold/configs/benchmark_testing.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. diff --git a/ScaFFold/utils/config_utils.py b/ScaFFold/utils/config_utils.py index 640ad19..91729a0 100644 --- a/ScaFFold/utils/config_utils.py +++ b/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"] diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 80a4d4e..8d3ed10 100644 --- a/ScaFFold/utils/trainer.py +++ b/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 diff --git a/docs/perf/optimizer-step-overhead.md b/docs/perf/optimizer-step-overhead.md new file mode 100644 index 0000000..4027c6f --- /dev/null +++ b/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. From fb80e879bee3e02404e8fe7529ae21030a0e0f61 Mon Sep 17 00:00:00 2001 From: Patrick Robert Miles Date: Wed, 8 Apr 2026 15:09:54 -0700 Subject: [PATCH 2/2] Sync benchmark testing config --- ScaFFold/configs/benchmark_testing.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/ScaFFold/configs/benchmark_testing.yml b/ScaFFold/configs/benchmark_testing.yml index 8bfda6e..40e674a 100644 --- a/ScaFFold/configs/benchmark_testing.yml +++ b/ScaFFold/configs/benchmark_testing.yml @@ -1,23 +1,23 @@ # 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.