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. diff --git a/ScaFFold/utils/perf_measure.py b/ScaFFold/utils/perf_measure.py index 5af8d5b..26feeec 100644 --- a/ScaFFold/utils/perf_measure.py +++ b/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() diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index ec5fc7b..80a4d4e 100644 --- a/ScaFFold/utils/trainer.py +++ b/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") diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index ab20c4e..fbe6380 100644 --- a/ScaFFold/worker.py +++ b/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") diff --git a/docs/perf/base-measurement-baseline.md b/docs/perf/base-measurement-baseline.md new file mode 100644 index 0000000..dfe1dec --- /dev/null +++ b/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. diff --git a/docs/perf/distconv-input-layout-proposal.md b/docs/perf/distconv-input-layout-proposal.md new file mode 100644 index 0000000..d90587e --- /dev/null +++ b/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. diff --git a/scripts/scaffold-tuolumne_scale7.job b/scripts/scaffold-tuolumne_scale7.job new file mode 100644 index 0000000..2f4c5f9 --- /dev/null +++ b/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 diff --git a/scripts/scaffold-tuolumne_scale8.job b/scripts/scaffold-tuolumne_scale8.job new file mode 100644 index 0000000..66f4096 --- /dev/null +++ b/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 diff --git a/scripts/scaffold-tuolumne_scale9.job b/scripts/scaffold-tuolumne_scale9.job new file mode 100644 index 0000000..b585235 --- /dev/null +++ b/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 diff --git a/scripts/scaffold-tuolumne_test.job b/scripts/scaffold-tuolumne_test.job new file mode 100644 index 0000000..5153e24 --- /dev/null +++ b/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