Miles30/performance

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

@@ -17,6 +17,9 @@
 
 CALI_PERF_ENV_VAR = "CALI_CONFIG"
 TORCH_PERF_ENV_VAR = "PROFILE_TORCH"
+TORCH_PROFILE_WAIT_STEPS = 5
+TORCH_PROFILE_WARMUP_STEPS = 2
+TORCH_PROFILE_ACTIVE_STEPS = 8
 
 _CALI_PERF_ENABLED = False
 TORCH_PERF_ENABLED = False
@@ -37,6 +40,7 @@
     try:
         from torch.profiler import ProfilerActivity
         from torch.profiler import profile as torchprofile
+        from torch.profiler import schedule as torchschedule
 
         TORCH_PERF_ENABLED = True
     except Exception:
@@ -92,12 +96,19 @@ def adiak_fini():
 
 def get_torch_context(ranks_per_node, rank):
     if TORCH_PERF_ENABLED:
-        TORCH_PERF_LOCAL = TORCH_PERF_ENABLED and (rank % ranks_per_node == 0)
+        # Restrict profiling to global rank 0 to minimize benchmark distortion.
+        TORCH_PERF_LOCAL = TORCH_PERF_ENABLED and rank == 0
         prof_ctx = (
             torchprofile(
                 activities=[ProfilerActivity.CUDA, ProfilerActivity.CPU],
-                record_shapes=True,
-                with_stack=True,
+                schedule=torchschedule(
+                    wait=TORCH_PROFILE_WAIT_STEPS,
+                    warmup=TORCH_PROFILE_WARMUP_STEPS,
+                    active=TORCH_PROFILE_ACTIVE_STEPS,
+                    repeat=1,
+                ),
+                record_shapes=False,
+                with_stack=False,
             )
             if TORCH_PERF_LOCAL
             else nullcontext()

ScaFFold/utils/trainer.py

@@ -76,6 +76,7 @@ def __init__(self, model, config, device, log):
         self.ps = None  # DistConv ParallelStrategy
         self.spatial_mesh = None  # Spatial mesh for use w/ DistConv
         self.ddp_placements = None  # DDP placements for use w/ DistConv
+        self.profiler = None
 
         self.checkpoint_path_absolute = str(
             self.config.run_dir + "/" + self.config.checkpoint_dir
@@ -658,6 +659,8 @@ def train(self):
                         self.global_step += 1
                         # Stay on GPU
                         epoch_loss += loss.detach()
+                        if self.profiler is not None:
+                            self.profiler.step()
                         end_code_region("update_loss")
                     end_code_region("batch_loop")
 

ScaFFold/worker.py

@@ -228,14 +228,15 @@ def main(kwargs_dict: dict = {}):
     # Run the training
     #
     ranks_per_node = get_local_size()
+    begin_code_region("cleanup_or_resume")
+    trainer.cleanup_or_resume()
+    end_code_region("cleanup_or_resume")
+    begin_code_region("warmup")
+    trainer.warmup()
+    end_code_region("warmup")
     prof_ctx, TORCH_PERF_LOCAL = get_torch_context(ranks_per_node, rank)
     with prof_ctx as prof:
-        begin_code_region("cleanup_or_resume")
-        trainer.cleanup_or_resume()
-        end_code_region("cleanup_or_resume")
-        begin_code_region("warmup")
-        trainer.warmup()
-        end_code_region("warmup")
+        trainer.profiler = prof if TORCH_PERF_LOCAL else None
         begin_code_region("train")
         trainer.train()
         end_code_region("train")

docs/perf/base-measurement-baseline.md

@@ -0,0 +1,55 @@
+# Performance Measurement Baseline
+
+This branch establishes the shared evaluation harness for the isolated performance branches.
+
+## What Changed
+
+- Throughput runs no longer enable torch profiling by default.
+- Torch profiling is now an explicit opt-in mode using `PROFILE_TORCH=ON`.
+- Profiling captures a fixed short training window on global rank `0` only.
+- Warmup is excluded from the profiled region.
+- The standardized comparison matrix is:
+  - `problem_scale=7` on `16` GPUs with `--dc-num-shards 1 2 2`
+  - `problem_scale=8` on `32` GPUs with `--dc-num-shards 1 2 2`
+
+## Profiling Behavior
+
+When `PROFILE_TORCH=ON` is set:
+
+- only global rank `0` is profiled
+- cleanup and warmup run outside the profiler
+- the training loop uses a fixed profiler schedule:
+  - wait `5` batches
+  - warmup `2` batches
+  - record `8` batches
+
+This keeps trace size and profiler overhead bounded while still capturing steady-state work.
+
+## Evaluation Commands
+
+Use these throughput runs for branch-to-branch comparisons:
+
+```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
+```
+
+Use this dedicated profiling mode only when a branch needs a 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
+```
+
+## Comparison Rules
+
+- Compare `train_stats.csv` epoch durations from epoch `2` onward.
+- Keep dataset, job shape, and sharding fixed across branches.
+- Use profiler traces only as attribution support, not as the throughput baseline.
+- Treat `problem_scale=9` as exploratory and outside the required comparison matrix.

docs/perf/distconv-input-layout-proposal.md

@@ -0,0 +1,67 @@
+# DistConv Input-Layout Proposal
+
+This document captures the deferred DistConv/input-layout optimization work so it can be picked up in a separate implementation chat.
+
+## Current Hot Path
+
+The warmup, training, and evaluation loops all repeat the same input-layout sequence for every batch:
+
+1. Load dense local tensors from the `DataLoader`.
+2. Move tensors to device.
+3. Build a DTensor from the local tensor with batch sharding semantics.
+4. Immediately call `.to_local()` on that DTensor.
+5. Redistribute the resulting tensor with `DCTensor.distribute(...)`.
+
+That sequence currently appears in all three paths:
+
+- `PyTorchTrainer.warmup`
+- `PyTorchTrainer.train`
+- `evaluate`
+
+The code also reconstructs logically identical placement metadata in the hot path, especially in evaluation.
+
+## Why This Is Suspect
+
+The saved traces show material host time in:
+
+- `distconv.__torch_function__`
+- `torch.overrides.handle_torch_function`
+- DTensor mesh scatter/shard helpers
+- repeated per-batch Python dispatch around layout conversion
+
+That does not prove the layout conversion is the dominant bottleneck at all scales, but it is substantial enough to justify isolating in its own branch.
+
+## Proposed Investigation Direction
+
+The future branch should focus only on reducing repeated layout-conversion overhead without changing model math or sharding semantics.
+
+Recommended implementation direction:
+
+- Precompute and cache placement metadata once during trainer setup.
+- Factor the shared dense-to-distconv conversion path into one helper used by warmup, train, and eval.
+- Minimize repeated DTensor construction and redundant `.to_local()` transitions where DistConv can consume a cheaper equivalent form.
+- Cache or precompute the flattened spatial process-group handle if the revised implementation still needs frequent collectives tied to the same spatial mesh.
+
+## Non-Goals for That Branch
+
+- No Dice math changes.
+- No DataLoader changes.
+- No optimizer or AMP changes.
+- No change to the user-visible sharding interface.
+
+## Evidence To Collect In The Future Branch
+
+Before and after traces should specifically compare:
+
+- total time in DTensor distribution helpers
+- total time in DistConv tensor conversion helpers
+- CPU gaps between `DataLoader` completion and the first major forward kernels
+- steady-state epoch duration for:
+  - `problem_scale=7`, `16` GPUs, `1 2 2`
+  - `problem_scale=8`, `32` GPUs, `1 2 2`
+
+## Risks
+
+- DistConv may rely on the current handoff pattern for correctness even if it looks redundant.
+- A cleaner helper extraction may reduce code duplication without producing a measurable speedup.
+- The right optimization may depend on DistConv internals that are not obvious from the application code alone.

scripts/scaffold-tuolumne_scale7.job

@@ -0,0 +1,34 @@
+#!/bin/bash
+
+# flux: --exclusive
+# flux: -N 4
+# flux: -g=1
+# flux: -t 60m
+# flux: -qpdebug
+
+ml cce/21.0.0 cray-mpich/9.1.0 rocm/7.1.1 rccl/fast-env-slows-mpi
+
+. .venvs/scaffoldvenv-tuo/bin/activate
+
+# (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"
+
+# Disable direct convolution benchmarking (should speedup warmup by a significant amount, does the below three options together)
+# export MIOPEN_DEBUG_CONV_DIRECT=0
+# Disable direct naive convolution benchmarking (naive_conv_ab_nonpacked_fwd_ndhwc_half_double_half.kd)
+export MIOPEN_DEBUG_CONV_DIRECT_NAIVE_CONV_FWD=0
+# Disable naive_conv_ab_nonpacked_bwd_ndhwc_half_double_half.kd
+# export MIOPEN_DEBUG_CONV_DIRECT_NAIVE_CONV_BWD=0
+# Disable naive_conv_ab_nonpacked_wrw_ndhwc_half_double_half.kd
+export MIOPEN_DEBUG_CONV_DIRECT_NAIVE_CONV_WRW=0
+
+# torchrun-hpc -N 1 -n 1 $(which scaffold) generate_fractals -c $(pwd)/ScaFFold/configs/benchmark_default.yml
+
+# Uncomment for a short post-warmup torch trace on rank 0 only.
+# export 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 1 1
+# 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 1 2
+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 4 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark -c $(pwd)/ScaFFold/configs/benchmark_testing.yml --problem-scale 7 --dc-num-shards 2 2 2

scripts/scaffold-tuolumne_scale8.job

@@ -0,0 +1,33 @@
+#!/bin/bash
+
+# flux: --exclusive
+# flux: -N 8
+# flux: -g=1
+# flux: -t 60m
+# flux: -qpdebug
+
+ml cce/21.0.0 cray-mpich/9.1.0 rocm/7.1.1 rccl/fast-env-slows-mpi
+
+. .venvs/scaffoldvenv-tuo/bin/activate
+
+# (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"
+
+# Disable direct convolution benchmarking (should speedup warmup by a significant amount, does the below three options together)
+# export MIOPEN_DEBUG_CONV_DIRECT=0
+# Disable direct naive convolution benchmarking (naive_conv_ab_nonpacked_fwd_ndhwc_half_double_half.kd)
+export MIOPEN_DEBUG_CONV_DIRECT_NAIVE_CONV_FWD=0
+# Disable naive_conv_ab_nonpacked_bwd_ndhwc_half_double_half.kd
+# export MIOPEN_DEBUG_CONV_DIRECT_NAIVE_CONV_BWD=0
+# Disable naive_conv_ab_nonpacked_wrw_ndhwc_half_double_half.kd
+export MIOPEN_DEBUG_CONV_DIRECT_NAIVE_CONV_WRW=0
+
+# torchrun-hpc -N 1 -n 1 $(which scaffold) generate_fractals -c $(pwd)/ScaFFold/configs/benchmark_default.yml
+
+# Uncomment for a short post-warmup torch trace on rank 0 only.
+# export 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 1 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
+# 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 2 2 2

scripts/scaffold-tuolumne_scale9.job

@@ -0,0 +1,34 @@
+#!/bin/bash
+
+# flux: --exclusive
+# flux: -N 16
+# flux: -g=1
+# flux: -t 60m
+# flux: -qpdebug
+
+ml cce/21.0.0 cray-mpich/9.1.0 rocm/7.1.1 rccl/fast-env-slows-mpi
+
+. .venvs/scaffoldvenv-tuo/bin/activate
+
+# (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"
+
+# Disable direct convolution benchmarking (should speedup warmup by a significant amount, does the below three options together)
+# export MIOPEN_DEBUG_CONV_DIRECT=0
+# Disable direct naive convolution benchmarking (naive_conv_ab_nonpacked_fwd_ndhwc_half_double_half.kd)
+export MIOPEN_DEBUG_CONV_DIRECT_NAIVE_CONV_FWD=0
+# Disable naive_conv_ab_nonpacked_bwd_ndhwc_half_double_half.kd
+# export MIOPEN_DEBUG_CONV_DIRECT_NAIVE_CONV_BWD=0
+# Disable naive_conv_ab_nonpacked_wrw_ndhwc_half_double_half.kd
+export MIOPEN_DEBUG_CONV_DIRECT_NAIVE_CONV_WRW=0
+
+# torchrun-hpc -N 1 -n 1 $(which scaffold) generate_fractals -c $(pwd)/ScaFFold/configs/benchmark_default.yml
+
+# Uncomment for a short post-warmup torch trace on rank 0 only.
+# export PROFILE_TORCH=ON
+
+torchrun-hpc -N 16 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark -c $(pwd)/ScaFFold/configs/benchmark_testing.yml --problem-scale 9 --dc-num-shards 2 2 4 --unet-bottleneck-dim 5
+# 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 1 2
+# 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 4 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark -c $(pwd)/ScaFFold/configs/benchmark_testing.yml --problem-scale 7 --dc-num-shards 2 2 2

scripts/scaffold-tuolumne_test.job

@@ -0,0 +1,34 @@
+#!/bin/bash
+
+# flux: --exclusive
+# flux: -N 1
+# flux: -g=1
+# flux: -t 20m
+# flux: -qpdebug
+
+ml cce/21.0.0 cray-mpich/9.1.0 rocm/7.1.1 rccl/fast-env-slows-mpi
+
+. .venvs/scaffoldvenv-tuo/bin/activate
+
+# (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"
+
+# Disable direct convolution benchmarking (should speedup warmup by a significant amount, does the below three options together)
+# export MIOPEN_DEBUG_CONV_DIRECT=0
+# Disable direct naive convolution benchmarking (naive_conv_ab_nonpacked_fwd_ndhwc_half_double_half.kd)
+export MIOPEN_DEBUG_CONV_DIRECT_NAIVE_CONV_FWD=0
+# Disable naive_conv_ab_nonpacked_bwd_ndhwc_half_double_half.kd
+# export MIOPEN_DEBUG_CONV_DIRECT_NAIVE_CONV_BWD=0
+# Disable naive_conv_ab_nonpacked_wrw_ndhwc_half_double_half.kd
+export MIOPEN_DEBUG_CONV_DIRECT_NAIVE_CONV_WRW=0
+
+# torchrun-hpc -N 1 -n 1 $(which scaffold) generate_fractals -c $(pwd)/ScaFFold/configs/benchmark_default.yml
+
+# Uncomment for a short post-warmup torch trace on rank 0 only.
+# export PROFILE_TORCH=ON
+
+torchrun-hpc -N 1 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark -c $(pwd)/ScaFFold/configs/benchmark_testing.yml --problem-scale 7 --dc-num-shards 1 1 1
+# 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 1 2
+# 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 4 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark -c $(pwd)/ScaFFold/configs/benchmark_testing.yml --problem-scale 7 --dc-num-shards 2 2 2