Source-linked AI summary
Exascale Deep Learning for Climate Analytics
Thorsten Kurth, Sean Treichler, Joshua Romero, Mayur Mudigonda, Nathan Luehr, Everett Phillips, Ankur Mahesh, Michael Matheson, Jack Deslippe, Massimiliano Fatica, Prabhat, Michael Houston
TL;DR
Extreme-weather analysis requires automated pixel-level segmentation of large climate-simulation outputs. The paper adapts Tiramisu and DeepLabv3+ and develops data, communication, and optimization techniques for distributed training. The resulting DeepLabv3+ system reaches 999.0 PF/s sustained performance and 90.7% parallel efficiency on 27360 GPUs in half precision.
Problem
Climate simulations generate tens of terabytes of output, while existing extreme-weather analysis has relied on heuristic threshold conditions rather than pixel-level masks.
Method
The paper adapts Tiramisu and DeepLabv3+ for climate-data segmentation and combines system-level and training-algorithm innovations for distributed scaling.
Results
999.0 PF/s sustained performance and 90.7% parallel efficiency are achieved on 27360 GPUs in half precision.
Takeaways & Limitations
The resulting methods support exascale-class deep learning for high-resolution extreme-weather segmentation and are described as applicable to broader deep learning workloads at scale.
Abstract
from arXiv · showhide
We extract pixel-level masks of extreme weather patterns using variants of Tiramisu and DeepLabv3+ neural networks. We describe improvements to the software frameworks, input pipeline, and the network training algorithms necessary to efficiently scale deep learning on the Piz Daint and Summit systems. The Tiramisu network scales to 5300 P100 GPUs with a sustained throughput of 21.0 PF/s and parallel efficiency of 79.0%. DeepLabv3+ scales up to 27360 V100 GPUs with a sustained throughput of 325.8 PF/s and a parallel efficiency of 90.7% in single precision. By taking advantage of the FP16 Tensor Cores, a half-precision version of the DeepLabv3+ network achieves a peak and sustained throughput of 1.13 EF/s and 999.0 PF/s respectively.
I. JUSTIFICATION
The paper addresses automated, pixel-level segmentation of extreme weather patterns in large climate datasets using deep learning, while targeting leadership-scale GPU performance. It combines segmentation architectures with system and algorithmic improvements for scalable analysis.
- 27360 GPUs deliver 1.13 EF/s peak and 1.0 EF/s sustained FP16 performance for the segmentation application.
- Climate simulations produce tens of terabytes of output, motivating automated analysis of extreme weather patterns.
- Deep learning is used to extract high-quality, pixel-level masks rather than relying only on multivariate threshold conditions.
- The modified DeepLabv3+ uses an encoder with ASPP blocks adapted for larger inputs and a full-resolution decoder for precise boundaries.
1) Segmentation Architecture:
The study evaluates modified Tiramisu and DeepLabv3+ architectures on high-resolution, multivariate climate data. The models produce three-class pixel masks for tropical cyclones, atmospheric rivers, and background.
- Tiramisu uses dense blocks, concatenative skip connections, a down path, an up path, and ADAM training for pixel-level segmentation.
- DeepLabv3+ uses a ResNet-50 encoder, atrous convolutions, and a simplified decoder to preserve information while detecting larger features.
- The dataset contains 0.25-degree CAM5 output on an 1152×768 grid with 3-hour temporal resolution and over 100 years of simulations.
- About 63K samples use all 16 variables and labels for tropical cyclone, atmospheric river, and background classes.
B. Contributions
The paper contributes adaptations of segmentation networks, system-level scaling methods, and algorithmic changes for distributed training. It demonstrates high-efficiency half-precision execution at extreme GPU concurrency.
- The authors adapt Tiramisu and DeepLabv3+ for high-resolution, multivariate scientific segmentation datasets.
- System-level innovations address data staging, parallel I/O, and networking collectives for large GPU-based HPC systems.
- Algorithmic innovations enable the deep learning networks to converge at scale.
- 999.0 PF/s sustained performance and 90.7% parallel efficiency are achieved on 27360 GPUs in half precision.
- The TensorFlow and Horovod optimizations are described as broadly applicable to deep learning workloads at scale.
A. State-of-the-art in Scientific Deep Learning
Scientific deep learning is expanding, but scaling representative architectures to leadership-class systems requires advances beyond prior GPU demonstrations. The paper addresses computational, data-movement, and storage constraints in this setting.
- Scientific deep learning is an emerging field increasingly incorporating scientific principles and domain assumptions.
- Prior demonstrations scaled ResNet-50 to 1024 P100 GPUs and another system to 2048 P40 GPUs, motivating substantially larger scales.
- 1) High speed parallel data staging:: A single GPU consumes 189 MB/s of training data, implying 5.23 TB/s for the full Summit system.
- 1) High speed parallel data staging:: Summit provides 800 GB of node-local SSD storage, while the 3.5 TB climate dataset can be distributed across nodes for training.
- 1) High speed parallel data staging:: Naive staging required 10–20 minutes and heavily loaded the global file system, motivating distributed data staging.
- 1) High speed parallel data staging:: Piz Daint lacks local SSDs, so Linux tmpfs is the only sufficiently fast node-local storage, but its capacity is limited.
2) Optimized data ingestion pipeline:
The pipeline combines staged local data with optimized distributed training to prevent input and communication bottlenecks from reducing GPU utilization at scale.
- Input pipeline: Placing input operations in the network graph can idle GPUs while CPUs read and transform data, motivating separate input-pipeline optimization.The staging approach removes global-file-system variability, but TensorFlow input operations still require optimization.
- Distributed training: Horovod distributes training by replicating the TensorFlow graph across MPI ranks and averaging gradients synchronously with all-reduce operations.Each Summit node uses six ranks, one per GPU, allowing the same code on Summit and Piz Daint.
- Hierarchical all-reduce: A centralized Horovod controller can become a bottleneck because ranks must coordinate all-reduce ordering to avoid deadlock.The controller receives readiness messages and dynamically reorders operations consistently across ranks.
- Hierarchical all-reduce: Hierarchical all-reduce uses local GPU reductions, MPI reductions across selected ranks, and NCCL broadcasts to replicate the final result within each node.Four local ranks each reduce a quarter of the data across nodes before NCCL broadcasts the complete result to all six GPUs.
B. Deep Learning Innovations
The paper adapts segmentation losses, optimization, input channels, and network designs to the severe imbalance and scientific structure of climate imagery.
- Segmentation challenge: 98.2% of pixels are background, 1.7% are atmospheric rivers, and fewer than 0.1% are tropical cyclones, allowing unweighted training to reach 98.2% accuracy by predicting background.The class distribution makes overall accuracy insufficient for evaluating rare-event masks.
- Optimization: LARC assigns independent layer-wise learning rates based on layer weight norms and is designed to improve large-network accuracy with large batch sizes.It controls update magnitude relative to each layer’s weight norm rather than assigning separate rates per weight.
- Input channels: Using all 16 CAM5 input channels on Summit dramatically improves model accuracy compared with the initial four-channel configuration.The channels include climate fields such as temperature, wind speed, pressure, and humidity at different altitudes.
- Input channels: The optimal channel subset likely lies between four and 16 channels, and the authors plan to tune it using rapid large-scale training.Summit’s additional computational capability enabled the 16-channel experiments.
4) Gradient lag:
The training design addresses large-scale performance through gradient lag, network-specific architectural choices, and graph-based throughput measurement.
- 4) Gradient lag:: Lagging gradient updates uses the previous step’s gradients, improving communication-computation overlap and enabling more efficient all-reduce tensor batching.The authors note that optimizer changes may require hyperparameter changes to preserve convergence.
- Network design: A Tiramisu growth rate of 32 was selected after profiling showed it was more efficient than the initial growth rate of 16 while maintaining roughly similar network size.The architecture compensates by reducing the number of layers.
- Network design: DeepLabv3+ operates at the dataset’s native resolution because quarter-resolution segmentation would reduce mask fidelity for irregular, fine-scale labels.Native resolution is more computationally expensive than the standard 288 × 192 design.
- Performance measurement: The application computes FLOP counts by traversing the TensorFlow graph, accounting for shared nodes and convolution algorithm choices selected dynamically by cuDNN.cuDNN API tracing identifies whether convolutions use implicit GEMMs or direct convolution.
- Performance measurement: Sustained throughput is reported as the median-over-time mean samples per step across ranks after normalizing the FLOP count by images processed.The procedure also computes asymmetric error bars from the step time series.
- Performance measurement: Keeping epoch steps constant during scaling makes validation overhead negligible after amortization across the training steps.Validation runs after each full pass through the training data.
A. HPC Systems and Environment
The evaluation spans Piz Daint and Summit, whose GPU, memory, storage, interconnect, and software configurations shape the scalable training environment.
- Piz Daint: Piz Daint’s XC50 portion has 5,320 nodes, each with one P100 GPU, 64 GB DDR memory, and Aries interconnect connectivity.The P100 provides 16 GB of HBM2 memory and 32 GB/s bidirectional PCIe bandwidth.
- Piz Daint: Piz Daint uses TensorFlow 1.6, CUDA 8.0, cuDNN 7.1.1, improved Horovod, Cray MPICH, and CUDA-aware collectives.The software stack includes a hierarchical Horovod control plane.
- Summit: Summit has 4,608 nodes, each with two Power9 CPUs, six V100 GPUs, 512 GB system memory, and 1.6 TB NVMe storage.NVLink connects each CPU to three GPUs, while dual-rail InfiniBand connects nodes through a nonblocking fat-tree.
- Summit: Summit’s Volta Tensor Cores support mixed-precision operations, with each GPU capable of 125 TFLOP/s and each node reaching 750 TF/s peak performance.Tensor Cores accept FP16 inputs and produce FP16 or FP32 outputs.
- Performance analysis: Figure 3 compares Tiramisu and DeepLabv3+ by kernel category, showing time fractions alongside achieved fractions of peak math and memory performance for FP32 and FP16.Values include measurement uncertainty, and asterisks denote fractions below 0.1%.
VII. PERFORMANCE RESULTS
The performance analysis attributes efficiency differences to GPU kernel behavior, especially memory demands in FP16, and identifies TensorFlow-level optimizations as the practical path forward.
- Tiramisu convolutions become memory limited in FP16 because its layers have small filter sizes, whereas DeepLabv3+ has higher computational intensity from larger channel counts.The DeepLabv3+ design reduces overall memory demand and improves datapath utilization.
- FP32 GPU utilization is already complete, so further gains require optimizing or eliminating kernels rather than simply increasing GPU usage.The analysis groups thousands of kernels into eight categories to assess computational and memory needs.
- A cuDNN upgrade produced a 5% performance improvement without application changes, while rewriting the networks with direct cuDNN calls was rejected because maintenance costs outweighed a 5-10% gain.The final strategy keeps TensorFlow while improving memory management and fusing point-wise operations.
- FP16 makes memory-bound kernels a larger share of step time, increasing the value of eliminating copies and fusing point-wise operations.These optimizations are expected to help FP16 more than FP32.
- FP16 profiling reveals periods when the GPU runs out of work, implicating the CPU-side input pipeline or TensorFlow scheduler as additional optimization targets.
B. Scaling Experiments
Scaling experiments on Piz Daint and Summit demonstrate high-throughput training across thousands of GPUs, with gradient lag and input staging materially affecting scalability.
- Scaling Experiments: 21.0 PF/s and 79.0% parallel efficiency are achieved by Tiramisu on all 5300 Piz Daint nodes in FP32.
- Scaling Experiments: 325.8 PF/s and 90.7% parallel efficiency are achieved by DeepLabv3+ on 27360 Summit GPUs in FP32.
- Scaling Experiments: 1.13 EF/s peak, 999.0 PF/s sustained, and 90.7% parallel efficiency are reached by DeepLabv3+ in FP16 at 27360 GPUs.
- Scaling Experiments: 176.8 PF/s in FP32 and 492.2 PF/s in FP16 are sustained by Tiramisu at 4096 Summit nodes, with parallel efficiencies above 90% in both cases.
- Scaling Experiments: Gradient lag improves overall application scalability, with the highest-performing Summit results corresponding to the “lag 1” configuration.
- Scaling Experiments: 75.8% parallel efficiency at 2048 GPUs represents a 9.5% penalty when input data is read from the global filesystem instead of staged in local tmpfs storage.
- Convergence at Scale: All tested FP16 and FP32 configurations converge at large scale, while DeepLabv3+ generally converges faster than Tiramisu.The experiments use up to 1024 Summit nodes and compare with smaller 384- and 1536-GPU runs.
D. Climate Science Results
The modified DeepLabv3+ achieves higher validation IoU than Tiramisu and produces detailed masks, while large-scale training converges across precisions and concurrencies. The analysis also exposes trade-offs from weighted loss and identifies future needs for temporal modeling and improved ground truth.
- 73% IoU for modified DeepLabv3+ exceeds Tiramisu’s 59% IoU on the validation dataset.
- DeepLabv3+ masks capture atmospheric features and approximate their exact boundaries, with some predictions appearing superior to heuristic labels.
- Weighted loss penalizes false-negative tropical-cyclone pixels roughly 37× more than false positives, contributing to an observed overprediction case.
- Pixel-level masks enable conditional precipitation, wind-velocity, and power-dissipation metrics for individual storms, supporting higher-fidelity impact characterization.
- Future work includes temporal storm architectures and high-quality ground-truth datasets that avoid reliance on heuristics.
B. Future Systems
Future exascale systems will require model parallelism and stronger data-movement support beyond purely data-parallel training. The paper connects these needs to communication, storage, and system-level investments for efficient scientific deep learning.
- Purely data-parallel scaling will become numerically difficult on future exascale machines, making model-parallel approaches indispensable in the foreseeable future.
- Domain decomposition across processors is a natural extension, but it requires more complex collectives and optimizations in software libraries and network switches.
- Further training optimizations will intensify parallel-I/O pressure by increasing the rate at which networks must receive input data.
- Node-local nonvolatile memory may reduce pressure on the global file system, while compression trades processor utilization for lower data movement.
- Efficient exascale deep learning requires balancing CPUs, accelerators, memory, storage, I/O, and networks to maintain high scaling efficiency.
- The demonstrated environment scales to 27360 Volta GPUs with 999.0 PF/s sustained performance and 90.7% parallel efficiency.