New config categories

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`
 
@@ -66,45 +65,13 @@ ScaFFold benchmark training always uses PyTorch distributed execution with DistC
 
 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 `<base_run_dir/figures`.
 
-Parameters are set in a `.yml` config file and can be modified by the user:
-
-```yml
-# 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 generated datasets.
-fract_base_dir: "fractals"         # Base directory for fractal IFS and instances.
-n_categories: 5                    # Positive 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.
-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 size per data-parallel rank.
-dataloader_num_workers: 1          # Number of DataLoader worker processes per rank.
-optimizer: "ADAM"                  # "ADAM" is preferred option, otherwise training defaults to RMSProp.
-dc_num_shards: [1, 1, 1]           # DistConv spatial shard counts.
-dc_shard_dims: [2, 3, 4]           # Tensor dimensions sharded by DistConv.
-checkpoint_interval: -1            # Checkpoint every C epochs; set to -1 to disable checkpointing entirely.
-
-# 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.
-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.
-loss_freq: 1                       # Number of epochs between logging the overall loss.
-warmup_batches: 64                 # How many warmup batches per rank to run before training.
-target_dice: 0.95                  # Validation Dice score threshold for convergence when epochs is -1.
-```
+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
 
@@ -201,23 +168,25 @@ 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`
     - `benchpark experiment init --dest scaffold --system tuolumne scaffold+rocm package_manager=spack-pip caliper=mpi,time,rocm`
 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`
@@ -257,7 +226,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`

ScaFFold/cli.py

@@ -142,7 +142,9 @@ def main():
     )
     benchmark_parser.add_argument("--seed", type=int, help="Random seed.")
     benchmark_parser.add_argument(
-        "--batch-size", type=int, help="Batch size per data-parallel rank."
+        "--local-batch-size",
+        type=int,
+        help="Batch size for each vol size per DDP rank.",
     )
     benchmark_parser.add_argument(
         "--warmup-batches",

ScaFFold/configs/benchmark_default.yml

@@ -2,28 +2,33 @@
 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                    # Positive 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 defaults 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.
 
-# 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                    # Positive 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
+local_batch_size: 1                # Batch size for each vol size per DDP rank.
+optimizer: "ADAM"                  # "ADAM" is preferred option, otherwise training defaults 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.
@@ -36,4 +41,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

ScaFFold/utils/config_utils.py

@@ -60,7 +60,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"]

ScaFFold/utils/trainer.py

@@ -160,7 +160,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,
         )
@@ -179,9 +179,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):
@@ -661,7 +661,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)

ScaFFold/worker.py

@@ -195,7 +195,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)
@@ -205,7 +205,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(
@@ -222,7 +222,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}"
             )
 
@@ -277,7 +278,16 @@ 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"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}."
             )
             extra_msg = f"Trained to >= {config.target_dice} validation dice score in {total_train_time:.2f} seconds, {total_epochs} epochs, {total_optimizer_steps} optimizer steps."
         else: