From b3195bba00fc8e1a34c79b42e74b20f7c84f83f8 Mon Sep 17 00:00:00 2001 From: Patrick Robert Miles Date: Wed, 8 Apr 2026 15:05:18 -0700 Subject: [PATCH 1/2] Rewrite sharded dice accumulation --- ScaFFold/utils/dice_score.py | 42 +++++++++++++++++++---------- ScaFFold/utils/evaluate.py | 12 ++++----- ScaFFold/utils/trainer.py | 31 +++++++++------------- ScaFFold/worker.py | 6 +++++ docs/perf/dice-path-rewrite.md | 48 ++++++++++++++++++++++++++++++++++ 5 files changed, 100 insertions(+), 39 deletions(-) create mode 100644 docs/perf/dice-path-rewrite.md diff --git a/ScaFFold/utils/dice_score.py b/ScaFFold/utils/dice_score.py index ed60fd4..e31fd37 100644 --- a/ScaFFold/utils/dice_score.py +++ b/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] diff --git a/ScaFFold/utils/evaluate.py b/ScaFFold/utils/evaluate.py index fbecf90..c57a4b9 100644 --- a/ScaFFold/utils/evaluate.py +++ b/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() diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 80a4d4e..e57a728 100644 --- a/ScaFFold/utils/trainer.py +++ b/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() diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index fbe6380..c548a8e 100644 --- a/ScaFFold/worker.py +++ b/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 diff --git a/docs/perf/dice-path-rewrite.md b/docs/perf/dice-path-rewrite.md new file mode 100644 index 0000000..f58975f --- /dev/null +++ b/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. From 46f0807b75bc63ab11a2305525b03c3ddfd06b29 Mon Sep 17 00:00:00 2001 From: Patrick Robert Miles Date: Wed, 8 Apr 2026 15:09:54 -0700 Subject: [PATCH 2/2] Sync benchmark testing config --- ScaFFold/configs/benchmark_testing.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) 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.