| Home > Workflow collections > In process > The Road to Exascale: Lessons Learned from Scaling a Scientific AI Workflow to 16,384 GPUs |
| Conference Presentation (After Call) | FZJ-2026-03288 |
; ; ; ; ;
2026
Abstract: Foundation models have progressed by scaling parameters and training data, driving rapidly increasing computational demands. This trend is especially pronounced in scientific imaging, where datasets can span terabytes to petabytes. Exascale systems provide the compute to train models at this scale. However, using these machines efficiently is non-trivial. At extreme scale, bottlenecks shift from GPU throughput to end-to-end workflow behavior, including startup overheads, storage access, communication, and synchronization.We study these effects and derive practical scaling lessons from a real training pipeline. This work was carried out on the JUPITER system at Jülich Supercomputing Centre within the JUPITER Research and Early Access Program (JUREAP) and the GCS Exascale Pioneer project brainfm. We adapt the execution environment and I/O path to reduce indirect I/O and metadata pressure caused by containers, runtime-generated artifacts, and logging. In parallel, we evaluate model- and loss-level choices that reduce synchronization and collective communication. To quantify data access performance, we compare HDF5 and Zarr for highly concurrent random access across file layouts and backends. We demonstrate the approach by training a neuroscience vision foundation model with contrastive learning on terabyte-scale microscopic images of histological human brain sections (CytoNet, https://arxiv.org/abs/2511.01870). We scale the workflow up to 16,384 NVIDIA GH200 superchips across 4,096 compute nodes.Across large runs, indirect I/O emerges as a primary scalability limiter, driven by container image access, startup scripts, bytecode generation, temporary-directory traffic, and uncontrolled logging. Staging container images into node-local memory and redirecting runtime-generated files away from shared storage reduces filesystem metadata storms and improves startup robustness. On the algorithmic side, synchronization-heavy components constrain scaling, motivating architecture choices that avoid batch-level collectives (e.g., batch normalization) and a contrastive-loss implementation that reduces redundant per-rank compute while limiting collective communication. For highly concurrent data access, we find that Zarr with the TensorStore backend provides the lowest and most stable access times.We distill these findings into practical guidelines that link workflow engineering, model design, and storage choices for training scientific foundation models at extreme scale.
|
The record appears in these collections: |