diff --git a/README.md b/README.md index 0652b93..f2827e8 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` @@ -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 ` 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,15 +168,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` @@ -217,7 +186,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` @@ -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` diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 14a689e..130b3ef 100644 --- a/ScaFFold/cli.py +++ b/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", diff --git a/ScaFFold/configs/benchmark_default.yml b/ScaFFold/configs/benchmark_default.yml index 14fcbf9..c41f4ed 100644 --- a/ScaFFold/configs/benchmark_default.yml +++ b/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 diff --git a/ScaFFold/utils/config_utils.py b/ScaFFold/utils/config_utils.py index 064f295..c13d611 100644 --- a/ScaFFold/utils/config_utils.py +++ b/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"] diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index af255b7..d8d6547 100644 --- a/ScaFFold/utils/trainer.py +++ b/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) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index 245778d..938206c 100644 --- a/ScaFFold/worker.py +++ b/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: