ScaFFold Benchmark

A scalable deep learning benchmark: UNet trained on procedurally-generated, 3D fractal data

About

ScaFFold is the Scale-free Fractal benchmark for deep learning.

Purpose

ScaFFold is a proxy application and benchmark representative of deep learning surrogate models that are trained on large, high-resolution, three-dimensional numerical simulations. It is meant to support benchmarking at a variety of system scales and be adaptable to future deep learning systems innovations. ScaFFold exercises much of the deep learning systems stack: I/O, compute, fine- and coarse-grained communication, and their integration in a framework.

Characteristics

ScaFFold trains a 3D U-Net to perform semantic segmentation on a synthetic dataset composed of 3D volumes containing different classes of fractals. The size of the problem is controlled by a scale parameter, which varies the size and complexity of the volumes and the depth of the U-Net. The scale parameter is exponential: each increase roughly doubles the problem size; e.g., a scale 7 problem has a volume size of :math:128^3 for each sample. Using fractals enables large datasets to be generated in-situ (rather than distributed) while ensuring a complex yet tractable semantic segmentation problem.

The model is trained from a random initialization until convergence, which is defined to be a validation Dice score of at least 0.95.

Setup

1. If running on an LLNL system, use the machine-specific install scripts in scripts/install-*.sh.

2. Clone the repository:
   git clone https://github.com/LBANN/ScaFFold.git && cd ScaFFold

3. Create and activate a python venv for running the benchmark:

   • Matrix: ml load python/3.13.2 && python3 -m venv .venvs/scaffoldvenv-matrix && source .venvs/scaffoldvenv-matrix/bin/activate && pip install --upgrade pip
   • Tuolumne: ml load python/3.13.2 && python3 -m venv .venvs/scaffoldvenv-tuo && source .venvs/scaffoldvenv-tuo/bin/activate && pip install --upgrade pip

4. Necessary LLNL settings:

   • CUDA (matrix):
     1. ml cuda/13.1.1 gcc/13.3.1 mvapich2/2.3.7
     2. export LD_LIBRARY_PATH=/usr/lib64:$LD_LIBRARY_PATH
   • ROCm (tuolumne):
     1. ml cce/21.0.2 cray-mpich/9.1.0 rocm/7.2.1 rccl/fast-env-slows-mpi

5. Install the benchmark in the python venv:

   • CUDA: pip install --no-binary=mpi4py -e .[cuda] --prefix=.venvs/scaffoldvenv-matrix --extra-index-url https://download.pytorch.org/whl/cu132 2>&1 | tee install.log
   • ROCm: pip install -e .[rocm] --find-links https://download.pytorch.org/whl/torch/ --find-links https://download.pytorch.org/whl/torchaudio/ --find-links https://download.pytorch.org/whl/torchvision/ --find-links https://download.pytorch.org/whl/triton-rocm/ 2>&1 | tee install.log

Running the benchmark

Basic run configuration

1. If running the benchmark for the first time, or running with different fractal parameters (n_categories, variance_threshold) than previously, generate fractal classes and instances:
   scaffold generate_fractals -c ScaFFold/configs/benchmark_default.yml

   Fractal category libraries are generated deterministically from the configured seed (under fract_base_dir/var<...>/seed<...>/) and reused by later runs with the same seed.

2. Once fractal generation completes, run the benchmark:
   torchrun-hpc -N 1 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark -c ScaFFold/configs/benchmark_default.yml

ScaFFold benchmark training always uses PyTorch distributed execution with DistConv spatial parallelism. For a singleton run, launch one distributed rank rather than disabling distributed execution.

Dataset cache and sharded datagen

benchmark creates or reuses datasets under dataset_dir. New datasets use the current physical-shard dataset format, which stores one volume and mask file per logical sample per DistConv shard. The physical layout is controlled by dc_num_shards and dc_shard_dims; for example, dc_num_shards: [1, 1, 2] writes two physical shards per logical volume, with filenames such as 120_shard000000.npy and 120_shard000001.npy. Datasets are generated with the same sharding configuration used for model training.

Unsharded runs use dc_num_shards: [1, 1, 1] and are still stored with _shard000000 files. Changing the DistConv shard layout changes the dataset cache key, so a run either reuses a dataset with matching sharding metadata or generates a new one. Older full-volume caches are ignored.

Each benchmark invocation performs exactly one benchmark run, in a run folder created under base_run_dir set in the config file. Every run parameter must be single-valued; a list (e.g. problem_scale: [6, 7]) is rejected by name, since parameter sweeps are not supported. To compare parameter settings, launch one benchmark run per setting. For reproducibility, the run folder holds a copy of the benchmark config yml as base_config.yaml plus the fully merged config.yaml for that run.

After the run completes, statistics from the run are stored in train_stats.csv. Additionally, users can inspect plots of the training and validation losses over time in <base_run_dir/figures.

Parameters are set in a .yml config file and can be modified by the user. See ScaFFold/configs/benchmark_default.yml for the default benchmark configuration.

How to perform a scaling study

See the specifications from the benchmarking page.

How the benchmark works

Overview

• 3D fractals, procedurally generated
• Train a UNet
  • Based on this work (link) which demonstrated that pretraining a UNet on synthetic fractals like this yielded signficant performance improvements

Generating the fractal classes

Fractals are generated via an Iterated Function System (IFS), which are composed of affine transformations. In our case, one affine transformation is determined by a set of 12 randomly generated parameters: the first 9 compose a 3x3 rotation matrix, and the last 3 compose a translation matrix:

$ F_{n+1} = \begin{pmatrix} a_n & b_n & c_n\ d_n & e_n & f_n\ g_n & h_n & i_n \end{pmatrix} F_n + \begin{pmatrix} j_n\ k_n \ l_n \end{pmatrix} $,
where $F_n$ is the $nth$ point in the fractal point cloud $F$.

One IFS, which defines a fractal category, is a set of 2-4 affine transformations, plus an associated probability for each affine transformation to be chosen during the fractal generation process. The process for generating an IFS looks like the following:

For n  in n_categories :

1. Choose random number of affine transformations to comprise this fractal category (between 2 and 8)

2. For each affine transformation:

    1. Generate 12 random values for the transformation parameters. First 9 are rotation matrix, last 3 are translation

    2. Calculate det(rotation matrix), which determines the (unnormalized) probability of selecting this transformation when generating a fractal from this IFS

3. Normalize affine transformation probabilities -- at this point, we have an IFS for a specific fractal category.


Once we have an IFS, we must determine if this fractal category meets our variance criteria:

4. Generate a fractal point cloud from this IFS

5. Process the fractal point cloud (normalize, center, check for NaNs) and calculate variance

6. If var(fractal_point_cloud) < criteria (default 0.15): this fractal category is valid, so we save the above IFS params to csv

Below is an example of an IFS for a fractal with four affine transformations:

Generating fractals from each class

With our fractal categories determined, we then generate several fractals from each class. Each fractal instance of a given class is created by scaling one of the 12 IFS parameters, then generating the fractal as usual. We scale each parameter by values in [0.4, 1.6] in intervals of 0.1. This gives 12 unique variations for each of the 12 IFS parameters, plus the unscaled/base fractal, for a total of 145 fractal instances for each class. See examples of fractal instances for a specific class below:

The weights we use to scale IFS parameters look like the following:

Dataset generation

Finally, we are ready to generate a dataset for training our model. Each sample in the dataset is composed of several fractal instances (default=3), randomly selected from any category, overlain with eachother in a 3D voxel grid. Each fractal instance in a sample is placed in a random, non-centered location in the voxel grid if the scale parameter set in the benchmark config file is <1; otherwise, each fractal instance is centered on the center of the voxel grid. Below is an outline of the data generation process:

For n  in n_volumes:

    1. Create volume  and mask  3D grids as matrix of 0s

    2. For fractal  in n_fractals_per_scene  (default=3):

        1. Pick a random fractal category from range(0, num_categories) 

        2. Pick a random instance of that fractal category and load that point cloud

        3. Project that fractal instance to a 3D voxel grid

        4. For xyz  in voxelgrid_coordinates :

            1. If instance  has points within voxel xyz:

                1. volume[xyz]   = [0, 0, 0.778] (make the voxel blue instead of black. the color is an arbitrary choice)

                2. mask[xyz]  = fractal_category 

    3. Save volume and mask  to files

In the current physical-shard dataset format, this save step writes each logical sample as one or more shard files, matching the requested dc_num_shards layout. The dataloader reads only the shard file needed by the current DistConv rank instead of loading a full volume and slicing it locally.

U-Net architecture

ScaFFold trains a configurable 3D U-Net for semantic segmentation. Each encoder block applies two 3D convolutions with GroupNorm and ReLU, preceded by 2x downsampling with MaxPool3d after the first block. Decoder blocks upsample with ConvTranspose3d, concatenate the matching encoder skip connection, and apply the same double-convolution block. The final OutConv maps the last feature volume to per-voxel class logits.

The schematic below shows the scale 8 configuration, with a 256^3 input volume, five downsampling stages, five skip-connected upsampling stages, and a 6-class output mask.

Performance Profiling

1. Profiling with the PyTorch Profiler

Set PROFILE_TORCH=ON to generate a PyTorch profiling trace that can be read into Perfetto. The trace is written into the run directory. 1, true, on and yes (any case) enable profiling; every other value, including 0, false, no and off, leaves it disabled.

2. Profiling with Caliper & Adiak

Building Caliper & Adiak with Python Bindings

A. Using Benchpark (via Spack)

1. Initialize experiment with Caliper
   • benchpark system init --dest tuolumne llnl-elcapitan cluster=tuolumne
   • benchpark experiment init --dest scaffold --system tuolumne scaffold+rocm package_manager=spack-pip caliper=mpi,time,rocm
2. benchpark setup scaffold wkp
3. # Follow ramble instructions ...

B. Manually

1. Activate python environment used to run the benchmark
   • $ source /usr/workspace/mckinsey/ScaFFold-profiling/.venvs/scaffoldvenv/bin/activate
2. Install required packages for profiling
   • $ pip install ScaFFold[profiling]
3. Build Adiak for metadata with -DWITH_PYTHON_BINDINGS=ON:

git clone https://github.com/LLNL/Adiak.git
cd Adiak && git submodule init && git submodule update
mkdir pybuild && cd pybuild
cmake -DENABLE_PYTHON_BINDINGS=ON -DENABLE_TESTS=OFF -DENABLE_MPI=ON -DCMAKE_INSTALL_PREFIX=. -Dpybind11_DIR=$(pybind11-config --cmakedir) ..
make && make install

4. Build Caliper with -DWITH_PYTHON_BINDINGS=ON (ROCm/CUDA flags depend on arch):

git clone https://github.com/LLNL/Caliper.git
cd Caliper
mkdir pybuild && cd pybuild
ml rocm/7.2.1
ml cuda/13.1.1
cmake -DWITH_PYTHON_BINDINGS=ON \
   -DWITH_ROCPROFILER=ON \
   -DWITH_CUPTI=ON \
   -DWITH_NVTX=ON \
   -DWITH_MPI=ON \
   -DWITH_ADIAK=ON \
   -Dadiak_DIR=/usr/workspace/mckinsey/Adiak/pybuild/lib/cmake/adiak/ \
   -Dpybind11_DIR=$(pybind11-config --cmakedir) \
   -DCMAKE_INSTALL_PREFIX=. ..
make && make install

5. Note: manual build requires manually exporting pycaliper and pyadiak into PYTHONPATH at runtime (using a spack environment would avoid this)

export PYTHONPATH=/usr/workspace/mckinsey/Caliper/pybuild-adiak/lib/python3.11/site-packages/:$PYTHONPATH
export PYTHONPATH=/usr/workspace/mckinsey/Adiak/pybuild/lib/python3.11/site-packages:$PYTHONPATH
# Avoid error with MPI profiling service
export LD_LIBRARY_PATH=/usr/workspace/mckinsey/ScaFFold-profiling-manual/.venvs/scaffoldvenv/lib/python3.11/site-packages/torch/lib:$LD_LIBRARY_PATH

Profiling ScaFFold with Caliper

1. Use the CALI_CONFIG environment variable to select a Caliper profiling configuration. If this variable is not defined, the annotated regions will not do anything, other than a function call and if check.
   • $ CALI_CONFIG="spot(output=test.cali,profile.mpi)" scaffold benchmark -c ScaFFold/configs/benchmark_default.yml -j