Application-Driven Exascale: The JUPITER Benchmark Suite (2024)

GROMACS

GROMACS[Berendsen1995, pall_tackling_2014, Abraham2015, Pall2020] is a versatile package to perform molecular dynamics simulations, focusing on biochemical molecules and soft condensed matter systems. The application integrates Newton’s equations of motion for systems with hundreds to millions of particles and provides time-resolved trajectories.Two biological systems from the Unified European Applications Benchmark Suite (UEABS)[ueabs2024] are used, test cases A and C. Test case A simulates a GluCl ion channel embedded in a membrane. Test case C contains 27 replicas of the STMV with about $28\,000\,000$ atoms and allows testing the scalability of system-supplied Fast Fourier Transform (FFT) libraries.

ICON

The ICOsahedral Non-hydrostatic model (ICON)[Zaengl2015] is a modelling framework for weather, climate, and environmental prediction used for operational weather forecasting at the German Weather Service. ICON also provides an Earth System Model for climate simulations, i.e., a general circulation model of the atmosphere, including a land module[Giorgetta2018] and an ocean model[Korn2022]. While the atmosphere part has been ported to GPUs[Giorgetta2022], the ocean component is still running on CPUs only. ICON is available under a permissive open source licence.The JUPITER benchmark case is based on the atmosphere component, with global forecast simulation in two resolutions, resulting in two sub-benchmarks: R02B09 (\qty5\kilo grid point distance) and R02B10 (\qty2.5\kilo grid point distance)[Hohenegger2023]. The coarser resolution is targeted for execution on 120 nodes, and the finer resolution is for 300 nodes. While reasonable scaling to $2\times$ the node count (240 nodes and 600 nodes, respectively) is possible, this is not the usual mode of operation for ICON.These simulations are crucial for ICON’s development towards a storm-resolving climate model with \qty1\kilo resolution or even less[Stevens2020]. A unique aspect of the ICON benchmark is its large input dataset: R02B09 requires\qty1.8\tera of data, R02B10 needs \qty4.5\tera. Therefore, the ICON benchmark also tests the performance of I/O operations on a system.

Megatron-LM

Megatron-LM[shoeybi2020megatronlm] is a prominent codebase in Natural Language Processing (NLP), known for its vast scale and performance capabilities. It employs the transformer model architecture[vaswani2023attention] and leverages various parallelization techniques and optimizations[narayanan2021efficient, korthikanti2022reducing, dao2023flashattention2, rajbhandari2020zero] through PyTorch to achieve high hardware utilization with excellent efficiency. The training of various recent open source GPT-like models was carried out with Megatron-LM[John2023OGX], making this benchmark crucial to assess a system’s capability to handle disruptive generative AI workloads.The benchmark trains a 175 billion parameter model, converting the usual throughput metric (tokens per time) to a hypothetical time-to-solution FOM by training 20 million tokens.

MMoCLIP

Contrastive Language-Image Pre-training (CLIP)[radford2021learning] is a method that conducts self-supervised learning on weakly-aligned image-text pairs with open vocabulary language, resulting in language-vision foundation models. The approach enables usage of substantially increased datasets, like web-scale datasets[Schuhmann2021, Schuhmann2022]. The OpenCLIP[ilharco_gabriel_2021_5143773, Cherti2023] codebase, an open source implementation of CLIP, enables efficiently distributed training of CLIP models by using multiple data parallelism schemes through PyTorch, scaling to more than a thousand GPUs.Due to its strong transferability and robustness, OpenCLIP is used in diverse multi-modal learning approaches and downstream applications[ldm, liu2023visual, sun2023emu], and efficient training is crucial for the machine learning community dealing with open, fully reproducible foundation models.

The MMoCLIP benchmark is curated from OpenCLIP. It trains an ViT-L-14 model on a synthetic dataset of $3\,200\,000$ image-text pairs and records the total training time as a FOM.

Quantum ESPRESSO

Quantum ESPRESSO (QE)[QE1, QE2, QEGPU] is an open source, density-functional-theory-based electronic structure software used both in academia and industry. QE calculates different material properties using a plane wave basis set and pseudo-potentials and can exploit novel accelerators well[QEexa].The dominant kernel in QE performs a three-dimensional FFT, which is usually a memory-bound kernel and is communication-bound for large systems[QEexa].

For the benchmark suite, the Car-Parrinello Molecular Dynamics model was chosen. The benchmark is based on a use case created in the MaX project[MaXrepo, MaX] and does calculations for a slab of ZrO2 with 792 atoms.

NAStJA

The Neoteric Autonomous Stencil code for Jolly Algorithms (NAStJA) is a massively-parallel simulation framework of biological tissues using a Cellular Potts Model[BerghoffNastja, Graner1992]. This model relies on nearest neighbour interactions and is parallelized by dividing the overall workload into multiple sub-regions, called blocks. Each block is treated independently by an MPI process, with boundaries being exchanged. Using NAStJA, tissues composed of thousands to millions of cells can be simulated at subcellular resolution[berghoff2020cells].As a test case, adhesion-driven cell sorting is used, a common process in tissue development and segregation[Steinberg1962].The benchmark investigates the first 5050 Monte Carlo (MC) steps of a system of size $720\times 720\times 1152\,\unit{\micro\cubed}$, containing roughly $600\,000$ cells. NAStJA utilizes MPI for parallelization/distribution and is one of the few CPU-only benchmarks in the suite. The application exhibits an irregular memory access pattern at each iteration, which is not suitable for GPU execution.

Arbor

Arbor is a library for simulating biophysically-realistic neural networks, bridging the gap between point and nanoscale models[arbor].Developed in the HBP[hbp], it aims at efficient use of modern HPC hardware behind an intuitive interface.Neurons are modelled by morphology, ion channels, and connections.Arbor is written in C++ with a Python interface and available under a permissive open source license.The user-centric description is discretized and aggregated to optimize data layouts for individual hardware.At runtime, the cable equation is integrated alternating with a system of ODEs for the channels.Users model channels via a domain-specific language that must be compiled for the target hardware.Communication is performed, concurrently with time evolution, every $n$ steps, determined by neural delays.The benchmark is parameterized to fill the GPU memory in the variants T, S, M, L.To differentiate from point models, it is weighted heavily towards computation, emphasized by a sparse connectivity.A complex cell from the Allen Institute was selected and adapted to random morphologies of fixed depth[allen].Cells are organized into rings propagating a single spike.Rings are interconnected to place load on the network without altering dynamics, yielding a deterministic, scalable workload.Profiling shows two cost centers: \qty52 ion channels and \qty33 cable equation; hiding communication completely.The Base version runs on 32 JUWELS Booster nodes, filling all 4 GPUs’ memory.This was scaled up to the full Booster to verify efficient resource usage and extrapolated to \qty1\exa.The number of generated spikes is used for validation.

Chroma

Chroma[Edwards:2004sx] is an all-purpose application for LQCD computations. It is compiled on top of the USQCD software stack[usqcd], which provides LQCD-specific libraries for communication, data-parallelism, I/O, and, importantly, sparse linear solver libraries optimized for different architectures.Key libraries used in this benchmark include QMP for MPI wrapping, QDP-JIT for data-parallelism and parallel-I/O via QIO, and the GPU-targeted QUDA solver library[Clark:2009wm].Chroma and the USQCD stack are open source and community-developed. Chroma is one of the most widely used LQCD suites and is representative of LQCD codes in general.

LQCD calculations generally depend heavily on solving very large, regular, sparse linear systems (dimension $10^{6}-10^{9}$ generally). Due to the regularity of the data and the calculation, LQCD lends itself to many levels of concurrency.

The Chroma LQCD benchmark in the JUPITER Benchmark Suitecontains the representative benchmarks for the Hybrid Monte-Carlo (HMC) component of the LQCD simulations. In the benchmark, a number of HMC update trajectories are performed using the $3+1$ flavours of Clover Wilson fermions — three light quark flavours with identical mass, and a fourth flavour with heavier mass — and the Lüscher-Weisz gauge action. The 4D lattice is initialized with a random $SU(3)$ element on each link. Checkpointing is disabled by a source patch to remove the I/O overhead for the calculations.The benchmark also contains a fix to Chroma, allowing simulation of 4D lattice volumes greater than $2^{31}$ and among the largest LQCD simulations anywhere to date. The benchmark performance is sensitive to the decomposition configuration used for distributing the 4D lattice to different tasks and to the affinity between the CPU, GPU, NUMA domains, and the network controller for each task.

The benchmark is validated by comparing the output with a reference solution with a tolerance of $10^{-10}$ for the Base benchmark and $10^{-8}$ for High-Scaling benchmarks.

The relevant metric (FOM) is the total time spent in HMC updates, excluding the first update, which includes overhead for tuning QUDA parameters. So a minimum of two updates must be prescribed.

JUQCS

JUQCS is a massively parallel simulator for universal gate-based Quantum Computers (QCs)[DeRaedt2007JUQCS] written in Fortran90 using MPI, OpenMP, and CUDA. During the past decades, JUQCS has been used to benchmark some of the largest supercomputers worldwide, including the Sunway TaihuLight and the K computer[DeRaedt2019JUQCS] as well as JUWELS Booster[Willsch2022JUQCSG], and it was part of Google’s quantum supremacy demonstration[Google2019QuantumSupremacy]. Various versions of JUQCS are available in binary form as part of a container[JUQCSDocker]; a light version with available sources was created specifically for this benchmark suite[JUQCSLight].JUQCS simulates an $n$-qubit gate-based QC by iteratively updating a rank-$n$ tensor of $2^{n}$ complex numbers (state vector) stored in double precision and distributed over the supercomputer’s memory. The total available memory determines the size of the largest QC that can be simulated. For instance, a universal simulation of $n=45$ qubits requires a little over $16\times 2^{45}\,\unit{}=\qty{0.5}{\pebi}$. Many operations require the transfer of half of all memory, i.e.,$2^{n}/2$ complex double-precision numbers, across the network, which can help to assess the performance of a supercomputer’s communication network[DeRaedt2019JUQCS, Willsch2022JUQCSG].As Figure3 shows, the deviation of JUQCS w.r.t.the theoretically expected linear scaling (green triangles) reveals both a drop in performance from intra-node to inter-node GPU communication (from 1 to 2 nodes) and another drop when communication enters the large-scale regime at 256 nodes.

All present JUQCS benchmarks simulate successive applications of a single-qubit quantum gate that requires large memory transfers. The Base benchmark simulates $n=36$ qubits requiring $\qty{1}{\tebi}$ of GPU memory. The High-Scaling benchmark contains two memory variants: a large memory variant with $n=42$ qubits requiring $\qty{64}{\tebi}$ (L) and a small memory variant with $n=41$ qubits requiring $\qty{32}{\tebi}$ of GPU memory (S). Rules are given for extrapolation to an Exascale setup using $n=46$ (L) or $n=45$ qubits (S). For all test cases, verification is done using the theoretically known results[DeRaedt2019JUQCS].

In addition, an MSA version of the JUQCS benchmark simulates $n=34$ qubits on both JUWELS Cluster and Booster simultaneously. The total amount of memory is split into two parts, with $\qty{128}{\gibi}$ residing on the CPU nodes and $\qty{128}{\gibi}$ residing on the GPU nodes. MPI is used for communication between the Cluster and the Booster, and the number of MPI tasks is similarly split into two. On the Cluster, each MPI task launches 12 OpenMP threads, with one thread per CPU core. On the Booster, each MPI task controls one of the GPUs.

nekRS

nekRS[Fischer2022] is a fast CFD solver designed for GPUs that solves the low-Mach Navier-Stokes equations (NSEs), potentially coupled with multiphysics effects. nekRS has been run at scale on many large supercomputers, featuring excellenttime-to-solution due to its high GPU throughput, and was nominated for the 2023 Gordon-Bell Award[Merzari2023]. nekRS uses high-order spectral elements[Patera1984] in which the solution, data, and test functions are represented as locally structured $N^{\mathrm{th}}$-order tensor product polynomials on a set of $E$ globally unstructured curvilinear hexahedral brick elements.There are two key benefits to this strategy. First, high-order polynomial expansions significantly reduce the number of unknowns ($n\approx EN^{3}$) needed to achieve engineering tolerances. Second, the locally structured forms allow tensor product sum factorization, which yields low $\mathcal{O}(n)$ storage cost and $\mathcal{O}(nN)$ work complexity[Orszag1980].The leading order $\mathcal{O}(nN)$ work terms can be cast as small dense matrix-matrix products with good computational intensity[Deville2002]. nekRS is written in C++ and the kernels are implemented using the portable Open Concurrent Compute Abstraction (OCCA) library [occa] for abstraction between different parallel languages/hardware architectures.

The benchmark case is derived from a Rayleigh-Bénard convection (RBC) application[Samuel2024, Bode2024] which simulates turbulence induced by a temperature gradient — a typical case executed at scale.The simulation domain is a sheet. It is much more extended in the periodic directions than in the wall-bounded direction. The chosen polynomial order is $9$ with 600 time steps per run. Verification is based on pre-computed results and derived tolerances.The High-Scaling benchmark variants use between $28\,836\,900$ (small, $\sim 11\,229$ per GPU) and $57\,760\,000$ (large, $\sim 22.492$ per GPU) elements, which is more than the minimum number of elements required for the "strong scaling limit" of $$7000$-$8000$$ elements per GPU. The Base benchmark case uses $719\,104$ elements resulting in $22\,472$ elements per GPU.

PIConGPU

PIConGPU is an open source, fully relativistic particle-in-cell (PIC) code designed for studying laser-plasma interactions and astrophysical phenomena. It uses the PIC algorithm with several key components, namely, particle initialization, charge calculations using grid interpolation, field calculations using densities, and time-marching due to Lorentz force. This approach allows particles to interact via fields on the grid rather than direct pairwise interactions, reducing computational steps from $N^{2}$ to $N$ for $N$ particles.PIConGPU employs a unique data model with asynchronous data transfers to handle the computational challenge. It can simulate complex plasma systems with billions of particles on GPU clusters[PIConGPU2013]. PIConGPU is developed with a hardware-agnostic approach using the Alpaka library[ZenkerAsHES2016, Worpitz2015], providing outstanding performance across all supported platforms, like CPUs, AMD and NVIDIA GPUs, and FPGAs[PIConGPU_alpaka].The benchmark suite uses a 3D test-case simulating the Kelvin-Helmholtz Instability (KHI), a non-relativistic shear-flow instability, utilizing a pre-ionized hydrogen plasma with periodic boundary conditions.While relevant for various research communities, the nature of shear-flows and the use of periodic boundary conditions does not impose a significant load imbalance throughout the simulation. Therefore, the performance of the code is based on its structure rather than the physics of the problem. In the KHI use case, the number of particles per cell is kept constant to 25, using as many cells as the GPU memory allows. A grid size of $\vec{x}=(4096,2048,1024)$ is chosen for the small memory variant, and extended to $(4096,2048,2048)$ (M) and $(4096,4096,2560)$ (L) for the larger variants. PIConGPU employs domain decomposition for distribution, dividing the computational domain into smaller subdomains along three dimensions. To distribute along these three dimensions, the maximum number of nodes that can be utilized is limited to $640$, rather than $642$.

V Lessons Learned

We developed the JUPITER Benchmark Suite building upon our experience from previous HPC system procurements.The suite constitutes a substantial expansion from those earlier endeavours, and should be considered as a living object that will continue to evolve over time.In the following, we summarize the lessons learned, covering first the perspective of application developers, then the benchmark suite creation, and finally the overall procurement process.