diff --git a/ScaFFold/configs/benchmark_testing.yml b/ScaFFold/configs/benchmark_testing.yml index 6b8c30a..956a9a0 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. diff --git a/ScaFFold/utils/evaluate.py b/ScaFFold/utils/evaluate.py index fbecf90..9dc2459 100644 --- a/ScaFFold/utils/evaluate.py +++ b/ScaFFold/utils/evaluate.py @@ -38,8 +38,8 @@ def evaluate( ): net.eval() num_val_batches = len(dataloader) - total_dice_score = 0.0 - processed_batches = 0 + total_dice_score = torch.zeros((), device=device, dtype=torch.float64) + processed_batches = torch.zeros((), device=device, dtype=torch.float64) spatial_mesh = parallel_strategy.device_mesh[parallel_strategy.distconv_dim_names] @@ -49,7 +49,7 @@ def evaluate( ) with torch.autocast(device.type if device.type != "mps" else "cpu", enabled=amp): - val_loss_epoch = 0.0 + val_loss_epoch = torch.zeros((), device=device, dtype=torch.float64) for batch in tqdm( dataloader, total=num_val_batches, @@ -125,15 +125,15 @@ def evaluate( # --- Combine and Accumulate --- loss = CE_loss + dice_loss_curr - val_loss_epoch += loss.item() - total_dice_score += batch_dice_score.item() - processed_batches += 1 + val_loss_epoch += loss.detach().to(torch.float64) + total_dice_score += batch_dice_score.detach().to(torch.float64) + processed_batches += 1.0 net.train() - val_loss_avg = val_loss_epoch / max(processed_batches, 1) + val_loss_avg = val_loss_epoch / processed_batches.clamp_min(1.0) if primary: print( - f"evaluate.py: dice_score={total_dice_score}, val_loss_epoch={val_loss_epoch}, val_loss_avg={val_loss_avg}, num_val_batches={processed_batches}" + f"evaluate.py: dice_score={total_dice_score.item()}, val_loss_epoch={val_loss_epoch.item()}, val_loss_avg={val_loss_avg.item()}, num_val_batches={int(processed_batches.item())}" ) return total_dice_score, val_loss_epoch, val_loss_avg, processed_batches diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 80a4d4e..4044f7b 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -485,7 +485,7 @@ def train(self): """ epoch = 1 - dice_score_train = 0 + dice_score_train = 0.0 with open(self.outfile_path, "a", newline="") as outfile: start = time.time() while dice_score_train < self.config.target_dice: @@ -497,8 +497,12 @@ def train(self): # Timer and tracking variables epoch_start_time = time.time() - train_dice_total = 0 - epoch_loss = 0 # Accumulator for per-batch losses + train_dice_total = torch.zeros( + (), device=self.device, dtype=torch.float64 + ) + epoch_loss = torch.zeros( + (), device=self.device, dtype=torch.float64 + ) # Accumulator for per-batch losses # Set necessary modes/states if self.config.dist: @@ -632,7 +636,9 @@ def train(self): # 3. Combine Loss loss = loss_ce + loss_dice - train_dice_total += dice_scores[:, 1:].mean().item() + train_dice_total += ( + dice_scores[:, 1:].mean().detach().to(torch.float64) + ) end_code_region("calculate_loss") @@ -658,14 +664,14 @@ def train(self): pbar.update(images_dc.shape[0]) self.global_step += 1 # Stay on GPU - epoch_loss += loss.detach() + epoch_loss += loss.detach().to(torch.float64) if self.profiler is not None: self.profiler.step() end_code_region("update_loss") end_code_region("batch_loop") # Calculate overall loss as average of per-batch loss - overall_loss = epoch_loss.item() / len(self.train_loader) + overall_loss = (epoch_loss / len(self.train_loader)).item() # # Evaluate model on validation set, update LR if necessary @@ -680,13 +686,14 @@ def train(self): self.config.n_categories, self.config._parallel_strategy, ) - dice_info = torch.tensor([dice_sum, numbatch]) + dice_info = torch.stack( + [dice_sum, numbatch.to(dtype=dice_sum.dtype)], dim=0 + ) if self.config.dist: - dice_info = dice_info.to(device=self.device) torch.distributed.all_reduce( dice_info, op=torch.distributed.ReduceOp.SUM ) - val_score = dice_info[0].item() / max(dice_info[1].item(), 1) + val_score = (dice_info[0] / dice_info[1].clamp_min(1.0)).item() if not self.config.disable_scheduler: # The following is true when trying to overfit, # in which case we only care about train loss @@ -706,7 +713,7 @@ def train(self): # # Write out data for this epoch to train stats csv # - train_dice = float(train_dice_total / len(self.train_loader)) + train_dice = (train_dice_total / len(self.train_loader)).item() self.log.info( f" epoch {epoch} \ | train_dice_loss {train_dice:.6f} (type {type(train_dice)}) \ @@ -720,8 +727,8 @@ def train(self): str(epoch), str(epoch_loss.item()), str(overall_loss), - str(val_loss_epoch), - str(val_loss_avg), + str(val_loss_epoch.item()), + str(val_loss_avg.item()), str(train_dice), str(val_score), str(epoch_duration), diff --git a/docs/perf/remove-host-sync-metrics.md b/docs/perf/remove-host-sync-metrics.md new file mode 100644 index 0000000..7f8d800 --- /dev/null +++ b/docs/perf/remove-host-sync-metrics.md @@ -0,0 +1,51 @@ +# Remove Host Sync From Metric Accumulation + +This branch removes avoidable host synchronization from training and validation metric accumulation relative to `miles30/performance`. + +## What Changed + +- Training epoch loss and training Dice totals stay on device during the batch loop. +- Validation loss, Dice totals, and processed-batch counts stay on device during the validation loop. +- Conversion to Python scalars now happens once per epoch or validation pass instead of once per batch. + +## What Did Not Change + +- Optimizer stepping is still once per batch. +- Loss math is unchanged. +- DistConv and DataLoader behavior are unchanged. +- CSV columns and printed summary fields are unchanged. + +## Expected Effect + +- Fewer per-batch host/device synchronization points. +- Lower CPU-side idle gaps around metric logging. +- Cleaner profiler timelines without changing convergence behavior. + +## Evaluation + +Throughput: + +```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 +``` + +Optional trace: + +```bash +PROFILE_TORCH=ON 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 +``` + +## Acceptance Criteria + +- `train_stats.csv` format is unchanged. +- No per-batch `.item()` remains in the train/validation accumulation path. +- Optimizer step cadence remains unchanged. +- Epoch `2+` duration improves or stays flat with cleaner CPU gaps in the trace.