From a44efacf97616a393e337b2ddef83878d4b4ac3c Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 16 Jul 2026 10:35:48 -0700 Subject: [PATCH 1/7] Update benchmark_default.yml --- ScaFFold/configs/benchmark_default.yml | 29 +++++++++++++++----------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/ScaFFold/configs/benchmark_default.yml b/ScaFFold/configs/benchmark_default.yml index 687c785..2dc1671 100644 --- a/ScaFFold/configs/benchmark_default.yml +++ b/ScaFFold/configs/benchmark_default.yml @@ -2,28 +2,34 @@ 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. 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: 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 per rank. dataloader_num_workers: 1 # Number of DataLoader worker processes per rank. More workers will use more memory -optimizer: "ADAM" # "ADAM" is preferred option, otherwise training defautls to RMSProp. 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: -1 # Checkpoint every C epochs; set to -1 to disable checkpointing entirely. +# local_batch_size: 1 # Batch sizes for each vol size per rank. -# 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: 30 # In percent. -epochs: -1 # Number of training epochs. +# External problem specification +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: 7 # Determines dataset resolution and number of unet layers. Default is 7. +unet_bottleneck_dim: 3 # Power of 2 of the unet bottleneck layer dimension. Default of 3 -> bottleneck layer of size 8. +target_dice: 0.95 # Dice score required to reach convergence + +# External learning rate specification +batch_size: 1 # Batch sizes for each vol size per rank. +optimizer: "ADAM" # "ADAM" is preferred option, otherwise training defautls to RMSProp. starting_learning_rate: 0.001 # Initial learning rate for training. min_learning_rate: 0.0001 # Minimum learning rate for CosineAnnealingWarmRestarts. T_0: 100 # Epochs in the first cosine restart cycle. T_mult: 2 # Restart cycle growth factor. disable_scheduler: 0 # If 1, disable scheduler during training to use constant LR. + +# 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: 30 # In percent. +epochs: -1 # Number of training epochs. 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. @@ -37,4 +43,3 @@ group_norm_groups: 8 # Number of groups used by GroupNorm in the UN warmup_batches: 64 # How many warmup batches per rank to run before training. ce_weight_sample_fraction: 0.1 # Fraction of training masks to sample when estimating background vs foreground CE weights. dataset_reuse_enforce_commit_id: 0 # Enforce matching commit IDs for dataset reuse. -target_dice: 0.95 From c2fdf0a90fc624639382ccb0a3452016076a596c Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 16 Jul 2026 13:13:25 -0700 Subject: [PATCH 2/7] change to local_batch_size --- README.md | 2 +- ScaFFold/cli.py | 4 +++- ScaFFold/configs/benchmark_default.yml | 3 +-- ScaFFold/utils/config_utils.py | 2 +- ScaFFold/utils/trainer.py | 8 ++++---- ScaFFold/worker.py | 7 ++++--- 6 files changed, 14 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 9438b1d..be31404 100644 --- a/README.md +++ b/README.md @@ -75,7 +75,7 @@ n_instances_used_per_fractal: 145 # Number of unique instances to pull from eac problem_scale: 6 # 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. +local_batch_size: 1 # Batch size for each vol size per DDP rank. optimizer: "ADAM" # "ADAM" is preferred option, otherwise training defautls to RMSProp. # Internal/dev use only diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 9249aac..03888ee 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -143,7 +143,9 @@ def main(): ) benchmark_parser.add_argument("--seed", type=int, help="Random seed.") benchmark_parser.add_argument( - "--batch-size", type=int, help="Batch sizes for each volume size." + "--local-batch-size", + type=int, + help="Batch size for each vol size per DDP rank.", ) benchmark_parser.add_argument( "--warmup-batches", diff --git a/ScaFFold/configs/benchmark_default.yml b/ScaFFold/configs/benchmark_default.yml index 2dc1671..b6acff9 100644 --- a/ScaFFold/configs/benchmark_default.yml +++ b/ScaFFold/configs/benchmark_default.yml @@ -7,7 +7,6 @@ dataloader_num_workers: 1 # Number of DataLoader worker processes per r 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: -1 # Checkpoint every C epochs; set to -1 to disable checkpointing entirely. -# local_batch_size: 1 # Batch sizes for each vol size per rank. # External problem specification n_categories: 5 # Number of fractal categories present in the dataset. @@ -17,7 +16,7 @@ unet_bottleneck_dim: 3 # Power of 2 of the unet bottleneck layer dim target_dice: 0.95 # Dice score required to reach convergence # External learning rate specification -batch_size: 1 # Batch sizes for each vol size per rank. +local_batch_size: 1 # Batch size for each vol size per DDP rank. optimizer: "ADAM" # "ADAM" is preferred option, otherwise training defautls to RMSProp. starting_learning_rate: 0.001 # Initial learning rate for training. min_learning_rate: 0.0001 # Minimum learning rate for CosineAnnealingWarmRestarts. diff --git a/ScaFFold/utils/config_utils.py b/ScaFFold/utils/config_utils.py index 9cb632c..673d14c 100644 --- a/ScaFFold/utils/config_utils.py +++ b/ScaFFold/utils/config_utils.py @@ -52,7 +52,7 @@ def __init__(self, config_dict): self.n_fracts_per_vol = config_dict["n_fracts_per_vol"] self.n_instances_used_per_fractal = config_dict["n_instances_used_per_fractal"] self.scale = 1 - self.batch_size = config_dict["batch_size"] + self.local_batch_size = config_dict["local_batch_size"] self.dataloader_num_workers = config_dict["dataloader_num_workers"] self.epochs = config_dict["epochs"] self.optimizer = config_dict["optimizer"] diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index fc5af6e..9ea98c9 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -163,7 +163,7 @@ def create_dataloaders(self): num_workers = self.config.dataloader_num_workers loader_args = dict( - batch_size=self.config.batch_size, + batch_size=self.config.local_batch_size, num_workers=num_workers, pin_memory=True, ) @@ -182,9 +182,9 @@ def create_dataloaders(self): if len(self.val_loader) == 0: raise ValueError( "Validation DataLoader has zero batches. " - f"n_val={self.n_val}, batch_size={self.config.batch_size}, " + f"n_val={self.n_val}, local_batch_size={self.config.local_batch_size}, " f"data_num_replicas={self.data_num_replicas}. " - "Reduce batch_size or adjust validation sharding." + "Reduce local_batch_size or adjust validation sharding." ) def setup_training_components(self): @@ -675,7 +675,7 @@ def train(self): # We don't want to time partial batches, i.e. last batch (time will be lower than expected). time_minibatch = ( batch_idx - < len(self.train_sampler) // self.config.batch_size + < len(self.train_sampler) // self.config.local_batch_size ) if time_minibatch: minibatch_start_event = torch.cuda.Event(enable_timing=True) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index f0223d1..286eb33 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -220,7 +220,7 @@ def main(kwargs_dict: dict = {}): num_spatial_dims = len(ps.shard_dim) trainer.ddp_placements = [Shard(0)] + [Replicate()] * num_spatial_dims total_shards = math.prod(config.dc_num_shards) - global_batch_size = config.batch_size * (world_size // total_shards) + global_batch_size = config.local_batch_size * (world_size // total_shards) config.global_batch_size = global_batch_size ddp_ranks = world_size // total_shards adiak_value("global_batch_size", global_batch_size) @@ -230,7 +230,7 @@ def main(kwargs_dict: dict = {}): if rank == 0: log.info( f"Effective global batch size = {global_batch_size} " - f"(batch_size={config.batch_size} * " + f"(local_batch_size={config.local_batch_size} * " f"(world_size={world_size} / prod(dc_num_shards)={total_shards}))" ) log.info( @@ -247,7 +247,8 @@ def main(kwargs_dict: dict = {}): "Effective global batch size exceeds available samples: " f"global_batch_size={global_batch_size}, " f"{', '.join(too_small_splits)}, " - f"batch_size={config.batch_size}, world_size={world_size}, " + f"local_batch_size={config.local_batch_size}, " + f"world_size={world_size}, " f"dc_num_shards={config.dc_num_shards}" ) From 075a3565bd9a1921544bfda9a6b4bb1a1d977d01 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 16 Jul 2026 13:17:36 -0700 Subject: [PATCH 3/7] Refer to file in readme --- README.md | 37 +++---------------------------------- 1 file changed, 3 insertions(+), 34 deletions(-) diff --git a/README.md b/README.md index be31404..17cda9e 100644 --- a/README.md +++ b/README.md @@ -64,40 +64,9 @@ The model is trained from a random initialization until convergence, which is de After each run completes, statistics from the run are stored in `train_stats.csv`. Additionally, users can inspect plots of the training and validation losses over time in ` bottleneck layer of size 8. -seed: 42 # Random seed. -local_batch_size: 1 # Batch size for each vol size per DDP rank. -optimizer: "ADAM" # "ADAM" is preferred option, otherwise training defautls to RMSProp. - -# 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: 100 # Number of training epochs. -starting_learning_rate: .01 # Initial learning rate for training. -min_learning_rate: .001 # Minimum learning rate for CosineAnnealingWarmRestarts. -T_0: 100 # Epochs in the first cosine restart cycle. -T_mult: 2 # Restart cycle growth factor. -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. -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. -checkpoint_interval: 1 # Number of epochs between saving training checkpoints. -loss_freq: 1 # Number of epochs between logging the overall loss. -``` +Parameters are set in a `.yml` config file and can be modified by the user. See +[`ScaFFold/configs/benchmark_default.yml`](ScaFFold/configs/benchmark_default.yml) +for the default benchmark configuration. ## How the benchmark works From af79c8639b7ba3a7e93ecb2c41ecf660fa91a918 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 16 Jul 2026 13:31:17 -0700 Subject: [PATCH 4/7] More FOM params --- ScaFFold/worker.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index 286eb33..ffc61ab 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -290,7 +290,17 @@ def main(kwargs_dict: dict = {}): log.info( f"FOM = {fom} (1 / total_train_time={total_train_time:.6f} seconds). " f"This FOM is specific to problem_scale={config.problem_scale}, " - f"target_dice={config.target_dice}, seed={config.seed}." + f"target_dice={config.target_dice}, seed={config.seed}, " + f"local_batch_size={config.local_batch_size}, " + f"n_categories={config.n_categories}, " + f"n_instances_used_per_fractal={config.n_instances_used_per_fractal}, " + f"unet_bottleneck_dim={config.unet_bottleneck_dim}, " + f"optimizer={config.optimizer}, " + f"starting_learning_rate={config.starting_learning_rate}, " + f"min_learning_rate={config.min_learning_rate}, " + f"T_0={config.T_0}, T_mult={config.T_mult}, " + f"disable_scheduler={config.disable_scheduler}, " + f"dc_shard_dims={config.dc_shard_dims}." ) epochs = np.atleast_1d(train_data["epoch"]) total_epochs = int(epochs[-1]) From bf45e7751bf5597a36cfe2d646d646da94157382 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 16 Jul 2026 13:37:24 -0700 Subject: [PATCH 5/7] Update README --- README.md | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 17cda9e..ed909b4 100644 --- a/README.md +++ b/README.md @@ -29,9 +29,6 @@ The model is trained from a random initialization until convergence, which is de 1. Clone the repository: `git clone https://github.com/LBANN/ScaFFold.git && cd ScaFFold` -1. Build the ccl plugin (if not using WCI wheel) - `. scripts/install-rccl.sh` - 1. Create and activate a python venv for running the benchmark: `ml load python/3.11.5 && python3 -m venv .venvs/scaffoldvenv && source .venvs/scaffoldvenv/bin/activate && pip install --upgrade pip` @@ -52,6 +49,8 @@ The model is trained from a random initialization until convergence, which is de ## Running the benchmark +### Basic run configuration + 1. If running the benchmark for the first time, or running with different fractal parameters (`n_categories`, `variance_threshold`) than previously, generate fractal classes and instances: `scaffold generate_fractals -c ScaFFold/configs/benchmark_default.yml` @@ -68,6 +67,10 @@ Parameters are set in a `.yml` config file and can be modified by the user. See [`ScaFFold/configs/benchmark_default.yml`](ScaFFold/configs/benchmark_default.yml) for the default benchmark configuration. +### How to perform a scaling study + +See the specifications from [the benchmarking page.](https://software.llnl.gov/benchmarks/14_scaffold/scaffold.html) + ## How the benchmark works @@ -163,15 +166,17 @@ For n  in n_volumes: 3. Save volume and mask  to files ``` -### 1. Profiling with the PyTorch Profiler +### Performance Profiling + +#### 1. Profiling with the PyTorch Profiler Set `PROFILE_TORCH=ON` to generate a PyTorch profiling trace that can be read into [Perfetto](https://ui.perfetto.dev/). -### 2. Profiling with Caliper & Adiak +#### 2. Profiling with Caliper & Adiak -#### Building Caliper & Adiak with Python Bindings +##### Building Caliper & Adiak with Python Bindings -##### A. Using Benchpark (via Spack) +###### A. Using Benchpark (via Spack) 1. Initialize experiment with Caliper - `benchpark system init --dest tuolumne llnl-elcapitan cluster=tuolumne` @@ -179,7 +184,7 @@ Set `PROFILE_TORCH=ON` to generate a PyTorch profiling trace that can be read in 1. `benchpark setup scaffold wkp` 1. `# Follow ramble instructions ...` -##### B. Manually +###### B. Manually 1. Activate python environment used to run the benchmark - `$ source /usr/workspace/mckinsey/ScaFFold-profiling/.venvs/scaffoldvenv/bin/activate` @@ -219,7 +224,7 @@ export PYTHONPATH=/usr/workspace/mckinsey/Adiak/pybuild/lib/python3.11/site-pack export LD_LIBRARY_PATH=/usr/workspace/mckinsey/ScaFFold-profiling-manual/.venvs/scaffoldvenv/lib/python3.11/site-packages/torch/lib:$LD_LIBRARY_PATH ``` -#### Profiling ScaFFold with Caliper +##### Profiling ScaFFold with Caliper 1. Use the `CALI_CONFIG` environment variable to select a Caliper profiling configuration. If this variable is not defined, the annotated regions will not do anything, other than a function call and if check. - `$ CALI_CONFIG="spot(output=test.cali,profile.mpi)" scaffold benchmark -c ScaFFold/configs/benchmark_default.yml -j` From 998fbd1226dfd7038c048089b381947a1bf432f4 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 16 Jul 2026 13:54:31 -0700 Subject: [PATCH 6/7] rm seed --- ScaFFold/worker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index ffc61ab..ea879d5 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -290,7 +290,7 @@ def main(kwargs_dict: dict = {}): log.info( f"FOM = {fom} (1 / total_train_time={total_train_time:.6f} seconds). " f"This FOM is specific to problem_scale={config.problem_scale}, " - f"target_dice={config.target_dice}, seed={config.seed}, " + f"target_dice={config.target_dice}, " f"local_batch_size={config.local_batch_size}, " f"n_categories={config.n_categories}, " f"n_instances_used_per_fractal={config.n_instances_used_per_fractal}, " From 206102bf31c5ed6932d133a5eabdc06c5856ff78 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 16 Jul 2026 13:58:34 -0700 Subject: [PATCH 7/7] rm param --- ScaFFold/worker.py | 1 - 1 file changed, 1 deletion(-) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index ea879d5..3a89293 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -291,7 +291,6 @@ def main(kwargs_dict: dict = {}): f"FOM = {fom} (1 / total_train_time={total_train_time:.6f} seconds). " f"This FOM is specific to problem_scale={config.problem_scale}, " f"target_dice={config.target_dice}, " - f"local_batch_size={config.local_batch_size}, " f"n_categories={config.n_categories}, " f"n_instances_used_per_fractal={config.n_instances_used_per_fractal}, " f"unet_bottleneck_dim={config.unet_bottleneck_dim}, "