Perf/dice path rewrite

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/dice_score.py

@@ -64,11 +64,9 @@ def dice_loss(input: Tensor, target: Tensor, multiclass: bool = False):
 
 class SpatialAllReduce(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, input, spatial_mesh):
+    def forward(ctx, input, reduce_group):
         output = input.clone()
-        for mesh_dim in range(spatial_mesh.ndim):
-            pg = spatial_mesh.get_group(mesh_dim)
-            dist.all_reduce(output, op=dist.ReduceOp.SUM, group=pg)
+        dist.all_reduce(output, op=dist.ReduceOp.SUM, group=reduce_group)
         return output
 
     @staticmethod
@@ -79,28 +77,44 @@ def backward(ctx, grad_output):
 @annotate()
 def compute_sharded_dice(
     preds: torch.Tensor,
-    targets: torch.Tensor,
-    spatial_mesh,
+    labels: torch.Tensor,
+    reduce_group,
+    num_classes: int,
     epsilon: float = 1e-6,
 ):
     """
     Computes the globally sharded Dice score.
     Returns the raw score tensor of shape [Batch, Channels].
     """
-    assert preds.size() == targets.size(), (
-        f"Shape mismatch: {preds.size()} vs {targets.size()}"
-    )
     assert preds.dim() == 5, f"Expected 5D tensor, got {preds.dim()}D"
+    assert labels.dim() == 4, f"Expected 4D labels tensor, got {labels.dim()}D"
+    assert preds.size(0) == labels.size(0), (
+        f"Batch mismatch: {preds.size(0)} vs {labels.size(0)}"
+    )
+    assert preds.shape[2:] == labels.shape[1:], (
+        f"Spatial mismatch: {preds.shape} vs {labels.shape}"
+    )
 
-    sum_dim = (-1, -2, -3)  # D, H, W
+    batch_size = preds.size(0)
+    preds_flat = preds.reshape(batch_size, num_classes, -1)
+    labels_flat = labels.reshape(batch_size, -1).long()
 
-    local_inter = 2.0 * (preds * targets).sum(dim=sum_dim)
-    local_sets_sum_raw = preds.sum(dim=sum_dim) + targets.sum(dim=sum_dim)
+    pred_sums = preds_flat.sum(dim=2)
+    true_class_probs = preds_flat.gather(1, labels_flat.unsqueeze(1)).squeeze(1)
+
+    intersections = torch.zeros_like(pred_sums)
+    intersections.scatter_add_(1, labels_flat, true_class_probs)
+    intersections.mul_(2.0)
+
+    target_sums = torch.zeros_like(pred_sums)
+    target_sums.scatter_add_(
+        1, labels_flat, torch.ones_like(true_class_probs, dtype=preds.dtype)
+    )
 
-    packed = torch.stack([local_inter, local_sets_sum_raw])
+    packed = torch.stack([intersections, pred_sums + target_sums])
 
     # Global reduce across spatial mesh
-    packed_global = SpatialAllReduce.apply(packed, spatial_mesh)
+    packed_global = SpatialAllReduce.apply(packed, reduce_group)
 
     global_inter = packed_global[0]
     global_sets_sum_raw = packed_global[1]

ScaFFold/utils/evaluate.py

@@ -41,7 +41,7 @@ def evaluate(
     total_dice_score = 0.0
     processed_batches = 0
 
-    spatial_mesh = parallel_strategy.device_mesh[parallel_strategy.distconv_dim_names]
+    spatial_reduce_group = getattr(parallel_strategy, "spatial_reduce_group")
 
     if primary:
         print(
@@ -99,7 +99,7 @@ def evaluate(
 
             # --- 1. Sharded CE Loss ---
             local_ce_sum = F.cross_entropy(local_preds, local_labels, reduction="sum")
-            global_ce_sum = SpatialAllReduce.apply(local_ce_sum, spatial_mesh)
+            global_ce_sum = SpatialAllReduce.apply(local_ce_sum, spatial_reduce_group)
 
             # Divide by total global voxels to get the mean CE Loss
             global_total_voxels = local_labels.numel() * math.prod(
@@ -109,13 +109,13 @@ def evaluate(
 
             # --- 2. Format Predictions & Labels (Strictly Multiclass) ---
             mask_pred_probs = F.softmax(local_preds, dim=1).float()
-            mask_true_onehot = (
-                F.one_hot(local_labels, n_categories + 1).permute(0, 4, 1, 2, 3).float()
-            )
 
             # Dice loss uses probabilities
             dice_score_probs = compute_sharded_dice(
-                mask_pred_probs, mask_true_onehot, spatial_mesh
+                mask_pred_probs,
+                local_labels,
+                spatial_reduce_group,
+                num_classes=n_categories + 1,
             )
             dice_loss_curr = 1.0 - dice_score_probs.mean()
 

ScaFFold/utils/trainer.py

@@ -75,6 +75,7 @@ def __init__(self, model, config, device, log):
         self.start_epoch = -1
         self.ps = None  # DistConv ParallelStrategy
         self.spatial_mesh = None  # Spatial mesh for use w/ DistConv
+        self.spatial_reduce_group = None  # Flattened spatial group for collectives
         self.ddp_placements = None  # DDP placements for use w/ DistConv
         self.profiler = None
 
@@ -414,7 +415,9 @@ def warmup(self):
                 )
 
                 # Pass the spatial_mesh directly
-                global_ce_sum = SpatialAllReduce.apply(local_ce_sum, self.spatial_mesh)
+                global_ce_sum = SpatialAllReduce.apply(
+                    local_ce_sum, self.spatial_reduce_group
+                )
 
                 global_total_voxels = local_labels.numel() * math.prod(
                     self.config.dc_num_shards
@@ -423,13 +426,11 @@ def warmup(self):
 
                 # 2. Sharded Dice Loss
                 local_preds_softmax = F.softmax(local_preds, dim=1).float()
-                local_labels_one_hot = (
-                    F.one_hot(local_labels, num_classes=self.config.n_categories + 1)
-                    .permute(0, 4, 1, 2, 3)
-                    .float()
-                )
                 dice_scores = compute_sharded_dice(
-                    local_preds_softmax, local_labels_one_hot, self.spatial_mesh
+                    local_preds_softmax,
+                    local_labels,
+                    self.spatial_reduce_group,
+                    num_classes=self.config.n_categories + 1,
                 )
                 loss_dice = 1.0 - dice_scores.mean()
 
@@ -455,7 +456,6 @@ def warmup(self):
                 local_preds,
                 local_labels,
                 local_preds_softmax,
-                local_labels_one_hot,
             )
             del loss_ce, loss_dice, loss, images_dp, true_masks_dp
 
@@ -603,7 +603,7 @@ def train(self):
 
                             # Pass the spatial_mesh directly
                             global_ce_sum = SpatialAllReduce.apply(
-                                local_ce_sum, self.spatial_mesh
+                                local_ce_sum, self.spatial_reduce_group
                             )
 
                             global_total_voxels = local_labels.numel() * math.prod(
@@ -613,20 +613,13 @@ def train(self):
 
                             # 2. Sharded Dice Loss
                             local_preds_softmax = F.softmax(local_preds, dim=1).float()
-                            local_labels_one_hot = (
-                                F.one_hot(
-                                    local_labels,
-                                    num_classes=self.config.n_categories + 1,
-                                )
-                                .permute(0, 4, 1, 2, 3)
-                                .float()
-                            )
 
                             # Compute sharded dice using new function
                             dice_scores = compute_sharded_dice(
                                 local_preds_softmax,
-                                local_labels_one_hot,
-                                self.spatial_mesh,
+                                local_labels,
+                                self.spatial_reduce_group,
+                                num_classes=self.config.n_categories + 1,
                             )
                             loss_dice = 1.0 - dice_scores.mean()
 

ScaFFold/worker.py

@@ -216,6 +216,12 @@ def main(kwargs_dict: dict = {}):
         trainer = PyTorchTrainer(model, config, device, log)
         trainer.ps = ps
         trainer.spatial_mesh = ps.device_mesh[ps.distconv_dim_names]
+        trainer.spatial_reduce_group = (
+            trainer.spatial_mesh.get_group()
+            if trainer.spatial_mesh.ndim == 1
+            else trainer.spatial_mesh._flatten().get_group()
+        )
+        trainer.ps.spatial_reduce_group = trainer.spatial_reduce_group
         num_spatial_dims = len(ps.shard_dim)
         trainer.ddp_placements = [Shard(0)] + [Replicate()] * num_spatial_dims
 

docs/perf/dice-path-rewrite.md

@@ -0,0 +1,48 @@
+# Dice Path Rewrite
+
+This branch rewrites the Dice path relative to `miles30/performance` while leaving cross-entropy semantics unchanged.
+
+## What Changed
+
+- Removed dense one-hot expansion of the full 3D label volume from train, warmup, and validation.
+- Replaced the Dice statistic construction with classwise scatter-based accumulation from integer labels.
+- Switched the spatial reduction from one all-reduce per mesh dimension to one all-reduce on a flattened spatial group.
+
+## What Did Not Change
+
+- Cross-entropy loss semantics are unchanged.
+- Reported Dice still excludes background class `0`.
+- Optimizer stepping, DataLoader behavior, and user-visible sharding arguments are unchanged.
+
+## Expected Effect
+
+- Less temporary memory traffic in the Dice path.
+- Lower communication overhead for Dice and sharded CE reductions.
+- Better GPU saturation at larger spatial scales where one-hot tensors are expensive.
+
+## Evaluation
+
+```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 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
+```
+
+## Acceptance Criteria
+
+- No dense label one-hot allocation remains in train, warmup, or eval.
+- Reported Dice stays finite and continues to exclude background.
+- CE loss semantics remain unchanged.
+- Epoch `2+` duration improves or stays flat for the required `ps7`/`ps8` comparison matrix.