AERIS: Argonne-Erdsystemmodell für zuverlässige und kompetente Vorhersagen

papers.abstract

Generatives maschinelles Lernen bietet neue Möglichkeiten, um die komplexe Dynamik des Erdsystems besser zu verstehen. Neuere, auf Diffusion basierende Methoden adressieren spektrale Verzerrungen und verbessern die Ensemble-Kalibrierung in der Wettervorhersage im Vergleich zu deterministischen Methoden, haben sich jedoch bisher als schwierig erwiesen, stabil in hohen Auflösungen zu skalieren. Wir stellen AERIS vor, einen pixelbasierten Swin-Diffusion-Transformer mit 1,3 bis 80 Milliarden Parametern, um diese Lücke zu schließen, sowie SWiPe, eine generalisierbare Technik, die Fensterparallelismus mit Sequenz- und Pipeline-Parallelismus kombiniert, um fensterbasierte Transformer zu partitionieren, ohne zusätzliche Kommunikationskosten oder eine erhöhte globale Batch-Größe zu verursachen. Auf Aurora (10.080 Knoten) erreicht AERIS 10,21 ExaFLOPS (gemischte Präzision) und eine Spitzenleistung von 11,21 ExaFLOPS mit einer Patch-Größe von 1x1 auf dem 0,25° ERA5-Datensatz, wobei eine schwache Skalierungseffizienz von 95,5 % und eine starke Skalierungseffizienz von 81,6 % erzielt wird. AERIS übertrifft das IFS ENS und bleibt auf saisonalen Skalen bis zu 90 Tagen stabil, was das Potenzial von milliardenparametrigen Diffusionsmodellen für die Wetter- und Klimavorhersage unterstreicht.

English

Generative machine learning offers new opportunities to better understand complex Earth system dynamics. Recent diffusion-based methods address spectral biases and improve ensemble calibration in weather forecasting compared to deterministic methods, yet have so far proven difficult to scale stably at high resolutions. We introduce AERIS, a 1.3 to 80B parameter pixel-level Swin diffusion transformer to address this gap, and SWiPe, a generalizable technique that composes window parallelism with sequence and pipeline parallelism to shard window-based transformers without added communication cost or increased global batch size. On Aurora (10,080 nodes), AERIS sustains 10.21 ExaFLOPS (mixed precision) and a peak performance of 11.21 ExaFLOPS with 1 times 1 patch size on the 0.25{\deg} ERA5 dataset, achieving 95.5% weak scaling efficiency, and 81.6% strong scaling efficiency. AERIS outperforms the IFS ENS and remains stable on seasonal scales to 90 days, highlighting the potential of billion-parameter diffusion models for weather and climate prediction.

AERIS: Argonne-Erdsystemmodell für zuverlässige und kompetente Vorhersagen

AERIS: Argonne Earth Systems Model for Reliable and Skillful Predictions

papers.abstract

Support