Perf/remove host sync metrics

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.

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

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),

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.