From e583d85ff932e4f3a9c59c141a3a812746772147 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Wed, 6 May 2026 12:51:51 -0700 Subject: [PATCH 01/19] fix dtypes for torch --- ScaFFold/utils/data_types.py | 5 ++++- ScaFFold/utils/trainer.py | 10 ++++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/ScaFFold/utils/data_types.py b/ScaFFold/utils/data_types.py index b555811..ef1515d 100644 --- a/ScaFFold/utils/data_types.py +++ b/ScaFFold/utils/data_types.py @@ -19,7 +19,10 @@ # Masks are values 0 <= x <= n_categories MASK_DTYPE = np.uint16 # Volumes/img are 0 <= x <= 1 -VOLUME_DTYPE = np.float32 +VOLUME_DTYPE_NAME = "float32" +VOLUME_NP_DTYPE = getattr(np, VOLUME_DTYPE_NAME) +VOLUME_TORCH_DTYPE = getattr(torch, VOLUME_DTYPE_NAME) +VOLUME_DTYPE = VOLUME_NP_DTYPE # Shared AMP dtype selection for torch.autocast. AMP_DTYPE = torch.bfloat16 diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 1a1d2e0..3add912 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -30,7 +30,7 @@ from ScaFFold.utils.checkpointing import CheckpointManager from ScaFFold.utils.data_loading import FractalDataset, SpatialShardSpec -from ScaFFold.utils.data_types import AMP_DTYPE, VOLUME_DTYPE +from ScaFFold.utils.data_types import AMP_DTYPE, VOLUME_TORCH_DTYPE from ScaFFold.utils.dice_score import compute_sharded_dice from ScaFFold.utils.distributed import get_local_rank, get_world_rank, get_world_size @@ -436,7 +436,7 @@ def warmup(self): images = images.to( device=self.device, - dtype=VOLUME_DTYPE, + dtype=VOLUME_TORCH_DTYPE, memory_format=torch.channels_last_3d, non_blocking=True, ) @@ -611,7 +611,7 @@ def train(self): begin_code_region("image_to_device") images = images.to( device=self.device, - dtype=VOLUME_DTYPE, + dtype=VOLUME_TORCH_DTYPE, memory_format=torch.channels_last_3d, # NDHWC (channels last) vs NCDHW (channels first) non_blocking=True, ) @@ -749,7 +749,9 @@ def train(self): self.config.n_categories, self.config._parallel_strategy, ) - dice_info = torch.tensor([dice_sum, numsamples], dtype=VOLUME_DTYPE) + dice_info = torch.tensor( + [dice_sum, numsamples], dtype=VOLUME_TORCH_DTYPE + ) if self.config.dist: dice_info = dice_info.to(device=self.device) torch.distributed.all_reduce( From 3dfbd138977624bcc33acc7bdcfca7ebd6e98b6e Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 7 May 2026 10:39:04 -0700 Subject: [PATCH 02/19] Add per minibatch timer --- ScaFFold/utils/trainer.py | 50 +++++++++++++++++++++++++++++++++------ 1 file changed, 43 insertions(+), 7 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 1a1d2e0..1aa399a 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -267,6 +267,20 @@ def _current_learning_rate(self): return self.config.starting_learning_rate return self.optimizer.param_groups[0]["lr"] + def _timing_ddp_rank(self): + if self.ps is None: + return self.world_rank + return self.ps.ddp_ind + + def _timing_shard_label(self): + if self.ps is None: + return "replicated" + return "x".join(str(shard_index) for shard_index in self.ps.shard_ind) + + def _sync_device_for_timing(self): + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + class PyTorchTrainer(BaseTrainer): """ @@ -349,18 +363,18 @@ def cleanup_or_resume(self): with open(self.outfile_path, "a", newline="") as outfile: outfile.write(",".join(headers) + "\n") - def _truncate_stats_file(self, start_epoch): + def _truncate_stats_file(self, start_epoch, path=None): """ Scans the stats file and truncates it at the first occurrence of an epoch >= start_epoch. This is O(1) memory and safe for large logs. """ - self.log.info( - f"Truncating {self.outfile_path} to remove epochs >= {start_epoch}" - ) + if path is None: + path = self.outfile_path + self.log.info(f"Truncating {path} to remove epochs >= {start_epoch}") try: # Open in read+update mode ('r+') to allow seeking and truncating - with open(self.outfile_path, "r+") as f: + with open(path, "r+") as f: header = f.readline() if not header: return @@ -401,7 +415,7 @@ def _truncate_stats_file(self, start_epoch): pass except Exception as e: - self.log.warning(f"Failed to truncate stats file: {e}") + self.log.warning(f"Failed to truncate stats file {path}: {e}") def _get_memsize(self, tensor, tensor_label: str, verbosity: int = 0): """Log size of tensor in memory""" @@ -604,7 +618,12 @@ def train(self): disable=True if self.world_rank != 0 else False, ) as pbar: begin_code_region("batch_loop") - for batch in self.train_loader: + for batch_idx, batch in enumerate(self.train_loader): + time_minibatch = batch_idx == 0 and self.world_rank == 0 + if time_minibatch: + self._sync_device_for_timing() + minibatch_start_time = time.perf_counter() + # Load initial samples and labels images, true_masks = batch["image"], batch["mask"] @@ -724,6 +743,23 @@ def train(self): self.global_step += 1 # Stay on GPU epoch_loss += loss.detach() + if time_minibatch: + self._sync_device_for_timing() + minibatch_time_s = ( + time.perf_counter() - minibatch_start_time + ) + print( + "MINIBATCH_TIMER " + f"epoch={epoch} " + f"batch_idx={batch_idx} " + f"global_step={self.global_step} " + f"ddp_rank={self._timing_ddp_rank()} " + f"world_rank={self.world_rank} " + f"local_rank={self.local_rank} " + f"shard_index={self._timing_shard_label()} " + f"batch_size={images_dc.shape[0]} " + f"minibatch_time_s={minibatch_time_s:.6f}", + ) end_code_region("update_loss") end_code_region("batch_loop") From c9ef075c1f786510692e05b36867b5956b5ccef6 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 7 May 2026 13:11:13 -0700 Subject: [PATCH 03/19] cleanup --- ScaFFold/utils/trainer.py | 33 ++++----------------------------- 1 file changed, 4 insertions(+), 29 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 77db2bf..e981e8a 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -267,20 +267,6 @@ def _current_learning_rate(self): return self.config.starting_learning_rate return self.optimizer.param_groups[0]["lr"] - def _timing_ddp_rank(self): - if self.ps is None: - return self.world_rank - return self.ps.ddp_ind - - def _timing_shard_label(self): - if self.ps is None: - return "replicated" - return "x".join(str(shard_index) for shard_index in self.ps.shard_ind) - - def _sync_device_for_timing(self): - if self.device.type == "cuda": - torch.cuda.synchronize(self.device) - class PyTorchTrainer(BaseTrainer): """ @@ -621,7 +607,6 @@ def train(self): for batch_idx, batch in enumerate(self.train_loader): time_minibatch = batch_idx == 0 and self.world_rank == 0 if time_minibatch: - self._sync_device_for_timing() minibatch_start_time = time.perf_counter() # Load initial samples and labels @@ -744,22 +729,12 @@ def train(self): # Stay on GPU epoch_loss += loss.detach() if time_minibatch: - self._sync_device_for_timing() + # This sync has some potential performance impact + # TODO: Would be better to measure this with Caliper, which uses CUDA events. + torch.cuda.synchronize(self.device) minibatch_time_s = ( time.perf_counter() - minibatch_start_time ) - print( - "MINIBATCH_TIMER " - f"epoch={epoch} " - f"batch_idx={batch_idx} " - f"global_step={self.global_step} " - f"ddp_rank={self._timing_ddp_rank()} " - f"world_rank={self.world_rank} " - f"local_rank={self.local_rank} " - f"shard_index={self._timing_shard_label()} " - f"batch_size={images_dc.shape[0]} " - f"minibatch_time_s={minibatch_time_s:.6f}", - ) end_code_region("update_loss") end_code_region("batch_loop") @@ -827,7 +802,7 @@ def train(self): ) outfile.flush() print( - f"Epoch {epoch} completed in {epoch_duration} seconds. Total train time so far: {time.time() - start}" + f"Epoch {epoch} completed in {epoch_duration} seconds. Total train time so far: {time.time() - start}. Rank 0 first batch minibatch_time_s={minibatch_time_s:.6f}." ) # From c0ba9c90af0516197f6583969aad572a838156ae Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 21 May 2026 18:41:26 -0700 Subject: [PATCH 04/19] Add adiak metadata --- ScaFFold/worker.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index 168bd85..e933d85 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -39,6 +39,7 @@ ) from ScaFFold.utils.perf_measure import ( annotate, + adiak_value, begin_code_region, end_code_region, get_torch_context, @@ -220,6 +221,10 @@ def main(kwargs_dict: dict = {}): total_shards = math.prod(config.dc_num_shards) global_batch_size = config.batch_size * (world_size // total_shards) ddp_ranks = world_size // total_shards + adiak_value("global_batch_size", global_batch_size) + adiak_value("ddp_ranks", ddp_ranks) + adiak_value("total_shards", total_shards) + adiak_value("num_spatial_dims", num_spatial_dims) if rank == 0: log.info( f"Effective global batch size = {global_batch_size} " From 1f7af3220cbd5c2a510cae23bb678dc1418fa33a Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Fri, 22 May 2026 12:41:50 -0700 Subject: [PATCH 05/19] Update worker.py --- ScaFFold/worker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index e933d85..fde4582 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -38,8 +38,8 @@ initialize_dist, ) from ScaFFold.utils.perf_measure import ( - annotate, adiak_value, + annotate, begin_code_region, end_code_region, get_torch_context, From fa113f12cafcfd5e757e428161980c16a12b86ab Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Tue, 26 May 2026 17:44:35 -0700 Subject: [PATCH 06/19] Define FOM --- ScaFFold/utils/trainer.py | 102 ++++++++++++++++++++++++++++++++++---- ScaFFold/worker.py | 1 + 2 files changed, 93 insertions(+), 10 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index e981e8a..ce356bc 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -15,6 +15,7 @@ # Standard library import math import os +import statistics import shutil import time from pathlib import Path @@ -267,6 +268,29 @@ def _current_learning_rate(self): return self.config.starting_learning_rate return self.optimizer.param_groups[0]["lr"] + def _publish_fom(self, minibatch_time_s): + """Publish ScaFFold throughput FOM from the representative minibatch time.""" + global_batch_size = getattr(self.config, "global_batch_size", None) + if global_batch_size is None: + self.log.warning("Skipping FOM reporting: global_batch_size is not set") + return + if minibatch_time_s <= 0: + self.log.warning( + f"Skipping FOM reporting: invalid minibatch_time_s={minibatch_time_s}" + ) + return + + fom = global_batch_size * self.config.problem_scale / minibatch_time_s + adiak_value("minibatch_time_s", minibatch_time_s) + adiak_value("FOM", fom) + if self.world_rank == 0: + self.log.info( + f"FOM = {fom:.6f} " + f"(global_batch_size={global_batch_size} * " + f"problem_scale={self.config.problem_scale} / " + f"minibatch_time_s={minibatch_time_s:.6f})" + ) + class PyTorchTrainer(BaseTrainer): """ @@ -570,6 +594,8 @@ def train(self): epoch = 1 dice_score_train = 0 + epoch_minibatch_times_s = [] + warned_no_full_minibatches = False with open(self.outfile_path, "a", newline="") as outfile: start = time.time() while dice_score_train < self.config.target_dice: @@ -583,6 +609,8 @@ def train(self): epoch_start_time = time.time() train_dice_total = 0 epoch_loss = 0 # Accumulator for per-batch losses + minibatch_time_s = None + minibatch_events = [] # Set necessary modes/states if self.config.dist: @@ -603,11 +631,31 @@ def train(self): unit="img", disable=True if self.world_rank != 0 else False, ) as pbar: + full_train_batches = ( + len(self.train_sampler) // self.config.batch_size + ) + time_minibatches = full_train_batches > 0 + if full_train_batches == 0 and not warned_no_full_minibatches: + self.log.warning( + "Skipping FOM reporting: no full training minibatches" + ) + warned_no_full_minibatches = True begin_code_region("batch_loop") for batch_idx, batch in enumerate(self.train_loader): - time_minibatch = batch_idx == 0 and self.world_rank == 0 + time_minibatch = ( + time_minibatches and batch_idx < full_train_batches + ) if time_minibatch: - minibatch_start_time = time.perf_counter() + minibatch_start_event = torch.cuda.Event( + enable_timing=True + ) + minibatch_end_event = torch.cuda.Event( + enable_timing=True + ) + minibatch_start_event.record() + minibatch_events.append( + (minibatch_start_event, minibatch_end_event) + ) # Load initial samples and labels images, true_masks = batch["image"], batch["mask"] @@ -728,16 +776,43 @@ def train(self): self.global_step += 1 # Stay on GPU epoch_loss += loss.detach() - if time_minibatch: - # This sync has some potential performance impact - # TODO: Would be better to measure this with Caliper, which uses CUDA events. - torch.cuda.synchronize(self.device) - minibatch_time_s = ( - time.perf_counter() - minibatch_start_time - ) end_code_region("update_loss") + + begin_code_region("record_minibatch_time") + if time_minibatch: + minibatch_end_event.record() + end_code_region("record_minibatch_time") end_code_region("batch_loop") + # Sync for batch time happens once after epoch is already done (low overhead) + if minibatch_events: + minibatch_events[-1][1].synchronize() + local_minibatch_times = torch.tensor( + [ + start_event.elapsed_time(end_event) / 1000.0 + for start_event, end_event in minibatch_events + ], + device=self.device, + ) + if self.config.dist: + gathered_minibatch_times = [ + torch.empty_like(local_minibatch_times) + for _ in range(self.world_size) + ] + torch.distributed.all_gather( + gathered_minibatch_times, local_minibatch_times + ) + minibatch_times = torch.stack(gathered_minibatch_times) + minibatch_times = torch.max( + minibatch_times, dim=0 + ).values + else: + minibatch_times = local_minibatch_times + minibatch_time_s = statistics.median( + minibatch_times.cpu().tolist() + ) + epoch_minibatch_times_s.append(minibatch_time_s) + # Calculate overall loss as average of per-batch loss overall_loss = epoch_loss.item() / len(self.train_loader) @@ -801,8 +876,13 @@ def train(self): + "\n" ) outfile.flush() + minibatch_time_msg = ( + f" Median full batch minibatch_time_s={minibatch_time_s:.6f}." + if minibatch_time_s is not None + else "" + ) print( - f"Epoch {epoch} completed in {epoch_duration} seconds. Total train time so far: {time.time() - start}. Rank 0 first batch minibatch_time_s={minibatch_time_s:.6f}." + f"Epoch {epoch} completed in {epoch_duration} seconds. Total train time so far: {time.time() - start}.{minibatch_time_msg}" ) # @@ -829,4 +909,6 @@ def train(self): "Invalid value (NaN) encountered in dice score computation" ) + if epoch_minibatch_times_s: + self._publish_fom(statistics.median(epoch_minibatch_times_s)) adiak_value("final_epochs", epoch) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index fde4582..ece8c62 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -220,6 +220,7 @@ def main(kwargs_dict: dict = {}): 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) + config.global_batch_size = global_batch_size ddp_ranks = world_size // total_shards adiak_value("global_batch_size", global_batch_size) adiak_value("ddp_ranks", ddp_ranks) From 1e9df1985c5a0e427aaef910e011829052195732 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Wed, 27 May 2026 11:42:06 -0700 Subject: [PATCH 07/19] rm gbs, divide by epochs --- ScaFFold/utils/trainer.py | 41 ++++++++++++++++++++------------------- 1 file changed, 21 insertions(+), 20 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index ce356bc..3698044 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -268,27 +268,31 @@ def _current_learning_rate(self): return self.config.starting_learning_rate return self.optimizer.param_groups[0]["lr"] - def _publish_fom(self, minibatch_time_s): + def _publish_fom(self, minibatch_time_s, completed_epochs): """Publish ScaFFold throughput FOM from the representative minibatch time.""" - global_batch_size = getattr(self.config, "global_batch_size", None) - if global_batch_size is None: - self.log.warning("Skipping FOM reporting: global_batch_size is not set") - return if minibatch_time_s <= 0: self.log.warning( f"Skipping FOM reporting: invalid minibatch_time_s={minibatch_time_s}" ) return + fom_epochs = ( + completed_epochs if self.config.epochs == -1 else self.config.epochs + ) + if fom_epochs <= 0: + self.log.warning(f"Skipping FOM reporting: invalid epochs={fom_epochs}") + return - fom = global_batch_size * self.config.problem_scale / minibatch_time_s + problem_size = (2**3) ** self.config.problem_scale + fom = problem_size / minibatch_time_s / fom_epochs adiak_value("minibatch_time_s", minibatch_time_s) + adiak_value("FOM_epochs", fom_epochs) adiak_value("FOM", fom) if self.world_rank == 0: self.log.info( f"FOM = {fom:.6f} " - f"(global_batch_size={global_batch_size} * " - f"problem_scale={self.config.problem_scale} / " - f"minibatch_time_s={minibatch_time_s:.6f})" + f"(((2^3)^problem_scale={problem_size}) / " + f"minibatch_time_s={minibatch_time_s:.6f} / " + f"epochs={fom_epochs})" ) @@ -646,12 +650,8 @@ def train(self): time_minibatches and batch_idx < full_train_batches ) if time_minibatch: - minibatch_start_event = torch.cuda.Event( - enable_timing=True - ) - minibatch_end_event = torch.cuda.Event( - enable_timing=True - ) + minibatch_start_event = torch.cuda.Event(enable_timing=True) + minibatch_end_event = torch.cuda.Event(enable_timing=True) minibatch_start_event.record() minibatch_events.append( (minibatch_start_event, minibatch_end_event) @@ -803,9 +803,7 @@ def train(self): gathered_minibatch_times, local_minibatch_times ) minibatch_times = torch.stack(gathered_minibatch_times) - minibatch_times = torch.max( - minibatch_times, dim=0 - ).values + minibatch_times = torch.max(minibatch_times, dim=0).values else: minibatch_times = local_minibatch_times minibatch_time_s = statistics.median( @@ -909,6 +907,9 @@ def train(self): "Invalid value (NaN) encountered in dice score computation" ) + completed_epochs = epoch - 1 if epoch_minibatch_times_s: - self._publish_fom(statistics.median(epoch_minibatch_times_s)) - adiak_value("final_epochs", epoch) + self._publish_fom( + statistics.median(epoch_minibatch_times_s), completed_epochs + ) + adiak_value("final_epochs", completed_epochs) From cc7961cd6dd4e244d9588185d0d2e39bf7e90403 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 28 May 2026 12:00:44 -0700 Subject: [PATCH 08/19] lint --- ScaFFold/utils/trainer.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 238f4e6..73b74e0 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -15,8 +15,8 @@ # Standard library import math import os -import statistics import shutil +import statistics import time from pathlib import Path @@ -776,13 +776,6 @@ def train(self): self.global_step += 1 # Stay on GPU epoch_loss += loss.detach() - if time_minibatch: - # This sync has some potential performance impact - # TODO: Would be better to measure this with Caliper, which uses CUDA events. - torch.cuda.synchronize(self.device) - minibatch_time_s = ( - time.perf_counter() - minibatch_start_time - ) end_code_region("update_loss") begin_code_region("record_minibatch_time") From d0a61f47d8ec4e8a1a4c235eec91dc0dc5295e1f Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 28 May 2026 15:01:59 -0700 Subject: [PATCH 09/19] Redefine FOM --- ScaFFold/utils/trainer.py | 28 ++++++---------------------- ScaFFold/worker.py | 13 +++++++++++++ 2 files changed, 19 insertions(+), 22 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 73b74e0..2935827 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -268,31 +268,17 @@ def _current_learning_rate(self): return self.config.starting_learning_rate return self.optimizer.param_groups[0]["lr"] - def _publish_fom(self, minibatch_time_s, completed_epochs): - """Publish ScaFFold throughput FOM from the representative minibatch time.""" + def _publish_minibatch_time(self, minibatch_time_s): + """Publish representative full minibatch time.""" if minibatch_time_s <= 0: self.log.warning( - f"Skipping FOM reporting: invalid minibatch_time_s={minibatch_time_s}" + f"Skipping minibatch timing reporting: invalid minibatch_time_s={minibatch_time_s}" ) return - fom_epochs = ( - completed_epochs if self.config.epochs == -1 else self.config.epochs - ) - if fom_epochs <= 0: - self.log.warning(f"Skipping FOM reporting: invalid epochs={fom_epochs}") - return - - problem_size = (2**3) ** self.config.problem_scale - fom = problem_size / minibatch_time_s / fom_epochs adiak_value("minibatch_time_s", minibatch_time_s) - adiak_value("FOM_epochs", fom_epochs) - adiak_value("FOM", fom) if self.world_rank == 0: self.log.info( - f"FOM = {fom:.6f} " - f"(((2^3)^problem_scale={problem_size}) / " - f"minibatch_time_s={minibatch_time_s:.6f} / " - f"epochs={fom_epochs})" + f"Representative full batch minibatch_time_s={minibatch_time_s:.6f}" ) @@ -641,7 +627,7 @@ def train(self): time_minibatches = full_train_batches > 0 if full_train_batches == 0 and not warned_no_full_minibatches: self.log.warning( - "Skipping FOM reporting: no full training minibatches" + "Skipping minibatch timing reporting: no full training minibatches" ) warned_no_full_minibatches = True begin_code_region("batch_loop") @@ -909,7 +895,5 @@ def train(self): completed_epochs = epoch - 1 if epoch_minibatch_times_s: - self._publish_fom( - statistics.median(epoch_minibatch_times_s), completed_epochs - ) + self._publish_minibatch_time(statistics.median(epoch_minibatch_times_s)) adiak_value("final_epochs", completed_epochs) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index ece8c62..7ca2488 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -282,6 +282,19 @@ def main(kwargs_dict: dict = {}): outfile_path = trainer.outfile_path train_data = np.genfromtxt(outfile_path, dtype=float, delimiter=",", names=True) total_train_time = train_data["epoch_duration"].sum() + if total_train_time > 0: + fom = 1.0 / total_train_time + adiak_value("FOM", fom) + if rank == 0: + log.info( + f"FOM = {fom:.6f} (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}." + ) + else: + log.warning( + f"Skipping FOM reporting: invalid total_train_time={total_train_time}" + ) epochs = np.atleast_1d(train_data["epoch"]) total_epochs = int(epochs[-1]) if config.epochs == -1: From e1a7defa3750a3ddf5a6b0cfeb354090995bb1d6 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 28 May 2026 15:04:28 -0700 Subject: [PATCH 10/19] Cleanup --- ScaFFold/utils/trainer.py | 20 ++++++-------------- 1 file changed, 6 insertions(+), 14 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 2935827..96c3c00 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -268,19 +268,6 @@ def _current_learning_rate(self): return self.config.starting_learning_rate return self.optimizer.param_groups[0]["lr"] - def _publish_minibatch_time(self, minibatch_time_s): - """Publish representative full minibatch time.""" - if minibatch_time_s <= 0: - self.log.warning( - f"Skipping minibatch timing reporting: invalid minibatch_time_s={minibatch_time_s}" - ) - return - adiak_value("minibatch_time_s", minibatch_time_s) - if self.world_rank == 0: - self.log.info( - f"Representative full batch minibatch_time_s={minibatch_time_s:.6f}" - ) - class PyTorchTrainer(BaseTrainer): """ @@ -895,5 +882,10 @@ def train(self): completed_epochs = epoch - 1 if epoch_minibatch_times_s: - self._publish_minibatch_time(statistics.median(epoch_minibatch_times_s)) + minibatch_time_s = statistics.median(epoch_minibatch_times_s) + adiak_value("minibatch_time_s", minibatch_time_s) + if self.world_rank == 0: + self.log.info( + f"Representative full batch minibatch_time_s={minibatch_time_s:.6f}" + ) adiak_value("final_epochs", completed_epochs) From 8b3429545501cae57699df4fc0a5b710727063b2 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Wed, 3 Jun 2026 18:31:08 -0700 Subject: [PATCH 11/19] Update worker.py --- ScaFFold/worker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index 7ca2488..33b43c8 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -287,7 +287,7 @@ def main(kwargs_dict: dict = {}): adiak_value("FOM", fom) if rank == 0: log.info( - f"FOM = {fom:.6f} (1 / total_train_time={total_train_time:.6f} seconds). " + 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}." ) From b3a4848a7abcdaf0ff334551c54cdd84d2cfa2b0 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 4 Jun 2026 15:33:05 -0700 Subject: [PATCH 12/19] Fix merge artifact --- ScaFFold/utils/trainer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 26ee573..7e96169 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -719,7 +719,8 @@ def train(self): # Sync for batch time happens once after epoch is already done (low overhead) if time_minibatch: - epoch_minibatch_times_s.append(self._sync_gather_minibatch_timer(minibatch_events)) + minibatch_time_s = self._sync_gather_minibatch_timer(minibatch_events) + epoch_minibatch_times_s.append(minibatch_time_s) # Calculate overall loss as average of per-batch loss overall_loss = epoch_loss.item() / len(self.train_loader) From 4aa9238ccd234f82d547a07801013c46ff8c10f6 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 4 Jun 2026 15:34:03 -0700 Subject: [PATCH 13/19] lint --- ScaFFold/utils/trainer.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 7e96169..a4e041b 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -555,8 +555,7 @@ def _sync_gather_minibatch_timer(self, minibatch_events): ) if self.config.dist: gathered_minibatch_times = [ - torch.empty_like(local_minibatch_times) - for _ in range(self.world_size) + torch.empty_like(local_minibatch_times) for _ in range(self.world_size) ] torch.distributed.all_gather( gathered_minibatch_times, local_minibatch_times @@ -565,9 +564,7 @@ def _sync_gather_minibatch_timer(self, minibatch_events): minibatch_times = torch.max(minibatch_times, dim=0).values else: minibatch_times = local_minibatch_times - minibatch_time_s = statistics.median( - minibatch_times.cpu().tolist() - ) + minibatch_time_s = statistics.median(minibatch_times.cpu().tolist()) return minibatch_time_s def warmup(self): @@ -719,7 +716,9 @@ def train(self): # Sync for batch time happens once after epoch is already done (low overhead) if time_minibatch: - minibatch_time_s = self._sync_gather_minibatch_timer(minibatch_events) + minibatch_time_s = self._sync_gather_minibatch_timer( + minibatch_events + ) epoch_minibatch_times_s.append(minibatch_time_s) # Calculate overall loss as average of per-batch loss From 6014356699ed8174d73ce06310c725849c8501f4 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 4 Jun 2026 15:39:31 -0700 Subject: [PATCH 14/19] Cleanup --- ScaFFold/utils/trainer.py | 23 ++++++----------------- 1 file changed, 6 insertions(+), 17 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index a4e041b..6642532 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -635,7 +635,6 @@ def train(self): epoch = self.start_epoch dice_score_train = 0 epoch_minibatch_times_s = [] - warned_no_full_minibatches = False with open(self.outfile_path, "a", newline="") as outfile: start = time.time() while dice_score_train < self.config.target_dice: @@ -671,19 +670,12 @@ def train(self): unit="img", disable=True if self.world_rank != 0 else False, ) as pbar: - full_train_batches = ( - len(self.train_sampler) // self.config.batch_size - ) - time_minibatches = full_train_batches > 0 - if full_train_batches == 0 and not warned_no_full_minibatches: - self.log.warning( - "Skipping minibatch timing reporting: no full training minibatches" - ) - warned_no_full_minibatches = True begin_code_region("batch_loop") for batch_idx, batch in enumerate(self.train_loader): + # We don't want to time partial batches, i.e. last batch (time will be lower than expected). time_minibatch = ( - time_minibatches and batch_idx < full_train_batches + batch_idx + < len(self.train_sampler) // self.config.batch_size ) if time_minibatch: minibatch_start_event = torch.cuda.Event(enable_timing=True) @@ -692,6 +684,7 @@ def train(self): minibatch_events.append( (minibatch_start_event, minibatch_end_event) ) + begin_code_region("minibatch_time") begin_code_region("run_training_batch") batch_size, batch_loss, batch_dice_score = ( self._run_training_batch( @@ -709,6 +702,7 @@ def train(self): # Stay on GPU epoch_loss += batch_loss end_code_region("update_loss") + end_code_region("minibatch_time") if time_minibatch: minibatch_end_event.record() @@ -784,13 +778,8 @@ def train(self): + "\n" ) outfile.flush() - minibatch_time_msg = ( - f" Median full batch minibatch_time_s={minibatch_time_s:.6f}." - if minibatch_time_s is not None - else "" - ) print( - f"Epoch {epoch} completed in {epoch_duration} seconds. Total train time so far: {time.time() - start}.{minibatch_time_msg}" + f"Epoch {epoch} completed in {epoch_duration} seconds. Total train time so far: {time.time() - start}. Median full batch minibatch_time_s={minibatch_time_s:.6f}." ) # From 54c40dea4eda5ad6f3c17db5f90c1b1153355056 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 4 Jun 2026 15:43:34 -0700 Subject: [PATCH 15/19] Cleanup --- ScaFFold/utils/trainer.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 6642532..8baf2a8 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -708,13 +708,6 @@ def train(self): minibatch_end_event.record() end_code_region("batch_loop") - # Sync for batch time happens once after epoch is already done (low overhead) - if time_minibatch: - minibatch_time_s = self._sync_gather_minibatch_timer( - minibatch_events - ) - epoch_minibatch_times_s.append(minibatch_time_s) - # Calculate overall loss as average of per-batch loss overall_loss = epoch_loss.item() / len(self.train_loader) @@ -753,6 +746,14 @@ def train(self): epoch_end_time = time.time() epoch_duration = epoch_end_time - epoch_start_time + + # Sync for batch time happens once after epoch is already done (low overhead) + if len(minibatch_events) > 0: + minibatch_time_s = self._sync_gather_minibatch_timer( + minibatch_events + ) + epoch_minibatch_times_s.append(minibatch_time_s) + # # Write out data for this epoch to train stats csv # From 881996531af0536b49ab2a6d3274a96530aeebb2 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 4 Jun 2026 15:48:17 -0700 Subject: [PATCH 16/19] Cleanup --- ScaFFold/utils/trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 8baf2a8..06c1247 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -780,7 +780,7 @@ def train(self): ) outfile.flush() print( - f"Epoch {epoch} completed in {epoch_duration} seconds. Total train time so far: {time.time() - start}. Median full batch minibatch_time_s={minibatch_time_s:.6f}." + f"Epoch {epoch} completed in {epoch_duration:.6f} seconds. Total train time so far: {time.time() - start:.6f} seconds. Median of minibatch times: {minibatch_time_s:.6f} seconds." ) # @@ -813,6 +813,6 @@ def train(self): adiak_value("minibatch_time_s", minibatch_time_s) if self.world_rank == 0: self.log.info( - f"Representative full batch minibatch_time_s={minibatch_time_s:.6f}" + f"Median of epoch minibatch time medians: {minibatch_time_s:.6f} seconds." ) adiak_value("final_epochs", completed_epochs) From 2de4c55b8adc5094618e15584ce4a04f9a5dd876 Mon Sep 17 00:00:00 2001 From: Patrick R Miles <78748866+PatrickRMiles@users.noreply.github.com> Date: Thu, 28 May 2026 15:08:41 -0700 Subject: [PATCH 17/19] fix restart epoch bug in trainer (#72) Co-authored-by: Patrick Miles --- ScaFFold/utils/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 96c3c00..d6a802c 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -569,7 +569,7 @@ def train(self): Execute model training """ - epoch = 1 + epoch = self.start_epoch dice_score_train = 0 epoch_minibatch_times_s = [] warned_no_full_minibatches = False From fb69a72bcd55cc3c2fcbc4cbeaeb72c77dc070f4 Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Wed, 3 Jun 2026 12:10:41 -0700 Subject: [PATCH 18/19] Update scaffold-tuolumne.job --- scripts/scaffold-tuolumne.job | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/scripts/scaffold-tuolumne.job b/scripts/scaffold-tuolumne.job index bbbba33..a22d8c6 100644 --- a/scripts/scaffold-tuolumne.job +++ b/scripts/scaffold-tuolumne.job @@ -13,7 +13,8 @@ ml cce/21.0.0 cray-mpich/9.1.0 rocm/7.1.1 rccl/fast-env-slows-mpi # (1) Avoid libmagma error # (2) Removing libmpi may cause segfault on mpi4py import -export LD_PRELOAD="/opt/rocm-7.1.1/llvm/lib/libomp.so /opt/cray/pe/mpich/9.1.0/ofi/gnu/11.2/lib/libmpi_gnu.so.12" +# (3-5) undefined symbol: cblas_gemm_f16f16f32 +export LD_PRELOAD="/opt/rocm-7.1.1/llvm/lib/libomp.so /opt/cray/pe/mpich/9.1.0/ofi/gnu/11.2/lib/libmpi_gnu.so.12 /opt/intel/oneapi/mkl/2024.2/lib/libmkl_core.so.2 /opt/intel/oneapi/mkl/2024.2/lib/libmkl_gnu_thread.so.2 /opt/intel/oneapi/mkl/2024.2/lib/libmkl_intel_lp64.so.2" # Disable direct convolution benchmarking (should speedup warmup by a significant amount, does the below three options together) # export MIOPEN_DEBUG_CONV_DIRECT=0 From 3528cbcba65a719e3bf7c7ee95357fe29ef7152b Mon Sep 17 00:00:00 2001 From: Michael McKinsey Date: Thu, 11 Jun 2026 13:30:47 -0700 Subject: [PATCH 19/19] fix region. Make file reading scope to 1 proc --- ScaFFold/utils/trainer.py | 2 +- ScaFFold/worker.py | 46 +++++++++++++++++---------------------- 2 files changed, 21 insertions(+), 27 deletions(-) diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 06c1247..65e55e7 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -693,7 +693,7 @@ def train(self): ) ) train_dice_total += batch_dice_score - begin_code_region("run_training_batch") + end_code_region("run_training_batch") # Update the loss begin_code_region("update_loss") diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index 33b43c8..02c9118 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -279,37 +279,31 @@ def main(kwargs_dict: dict = {}): # # Calculate benchmark score # - outfile_path = trainer.outfile_path - train_data = np.genfromtxt(outfile_path, dtype=float, delimiter=",", names=True) - total_train_time = train_data["epoch_duration"].sum() - if total_train_time > 0: + if rank == 0: + outfile_path = trainer.outfile_path + train_data = np.genfromtxt(outfile_path, dtype=float, delimiter=",", names=True) + total_train_time = train_data["epoch_duration"].sum() fom = 1.0 / total_train_time adiak_value("FOM", fom) - if rank == 0: - 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}." - ) - else: - log.warning( - f"Skipping FOM reporting: invalid total_train_time={total_train_time}" - ) - epochs = np.atleast_1d(train_data["epoch"]) - total_epochs = int(epochs[-1]) - if config.epochs == -1: - extra_msg = f"Trained to >= {config.target_dice} validation dice score in {total_train_time:.2f} seconds, {total_epochs} epochs." - else: - extra_msg = ( - f"Completed in {total_train_time:.2f} seconds, {total_epochs} epochs." + 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}." ) + epochs = np.atleast_1d(train_data["epoch"]) + total_epochs = int(epochs[-1]) + if config.epochs == -1: + extra_msg = f"Trained to >= {config.target_dice} validation dice score in {total_train_time:.2f} seconds, {total_epochs} epochs." + else: + extra_msg = ( + f"Completed in {total_train_time:.2f} seconds, {total_epochs} epochs." + ) - log.info(f"Benchmark run at scale {config.problem_scale} complete. \n{extra_msg}") + log.info( + f"Benchmark run at scale {config.problem_scale} complete. \n{extra_msg}" + ) - # - # Generate plots - # - if rank == 0: + # Generate plots log.info("Generating figures on rank 0...") begin_code_region("generate_figures") standard_viz.main(config)