Differentiable, mixed-precision multigrid for weather and climate applications

This project is part of the ExaGEO Centre for Doctoral Training (CDT), more details can be found here.

Primary Supervisor: Professor Michèle Weiland with the Met Office

Further Description

Climate and weather prediction are undergoing rapid changes, with key questions such as “how can traditional algorithms be combined with AI approaches?” and “how can the resulting methods be executed efficiently on state-of-the-art computing hardware?”. Modern hardware (primarily GPUs) is evolving to make extensive use of floating-point precisions lower than 64-bit, the historical norm for simulation algorithms. This is partly because Machine Learning remains efficient at lower precision, but also because lower-precision computational units occupy less silicon, and lower precision improves performance through better vector-unit utilisation and reduced demand on memory and network bandwidth. The theoretical peak performance of state-of-the art GPUs increases in line with the reduction in precision (4x from 64-bit to 16-bit).

This project will investigate applying low-precision computation to weather and climate simulation codes, such as the Met Office’s LFRic model. The focus will be on the semi-implicit solver & multigrid preconditioner. Differentiability, in the sense of being able to back-propagate gradients through the solver, is required to integrate the algorithm in hybrid physical/ML models. The project will explore how making multigrid methods differentiable can be combined with mixed- and low-precision approaches. While this offers potentially significant performance gains, it raises new questions about gradient accuracy.

Next-generation numerical weather prediction (NWP) applications such as the Met Office’s LFRic atmosphere model exploit techniques such as geometric multigrid preconditioning (to improve convergence and cut the number of global communication calls between processing units) and mixed precision, where model components such as the linear solver and transport scheme routinely run as 32-bit “bubbles” within a higher-precision code, the success of this approach has led to the target configuration running fully at 32-bit. Yet these codes do not yet take advantage of the smallest precisions (16-bit, 8-bit, 4-bit), nor of the emerging capability to differentiate the dynamical core itself.

Overview of research area

The student will develop mixed- and low-precision, differentiable multigrid implementations and perform rigorous evaluation of both the numerical accuracy of the resulting gradients and the performance benefits, with a particular focus on Exascale hardware architectures (GPUs and the memory hierarchy they expose).

• Teaser Project 1: How accurate are the gradients of a low-precision multigrid solver? Reducing floating-point precision has well-known implications for numerical accuracy. In multigrid those implications are subtle because coarsening and refining interchange between geometric and precision scales, and numerical error from reduced precision during the coarsening steps may be acceptable. When the solver is also made differentiable, a second source of error appears: the accuracy of the gradient itself. The student will investigate, on a simple 2D problem, how far precision can be aggressively lowered (looking at both geometric and precision-based refinement) before the gradient — not just the forward solution — breaks down.
• Teaser Project 2: The costs and benefits of differentiability at (Exa)scale. Modern GPUs accelerate lower-precision arithmetic in hardware, though the exact performance impact is hardware-specific and problem dependent; lower precision also reduces demand on the memory subsystem and improves vector-unit utilisation. When a multigrid solver is made differentiable, the performance story is set by the whole pipeline — the forward solve plus the computation of gradients — and the way gradients are obtained matters enormously. Back-propagation “replays” the forward operations in reverse and may retain intermediate state, imposing very different memory traffic and synchronisation demands from an adjoint-state solve, which is forward-like in structure.

Student Requirements

A UK 2:1 honours degree, or its international equivalent, in a relevant subject such as computer science and informatics, physics, mathematics, engineering, biology, chemistry and geosciences.

You must be a competent programmer in at least one of C, C++, Python, Fortran, or Java and should be familiar with mathematical concepts such as algebra, linear algebra and probability and statistics.

English Language requirements as set by University of Edinburgh.