Transparent Checkpointing for Automatic Differentiation of Program Loops Through Expression Transformations

Michel Schanen, Sri Hari Krishna Narayanan, Sarah Williamson, Valentin Churavy, William S. Moses, Ludger Paehler

Research output: Chapter in Book/Report/Conference proceedingConference contribution

Abstract

Automatic differentiation (AutoDiff) in machine learning is largely restricted to expressions used for neural networks (NN), with the depth rarely exceeding a few tens of layers. Compared to NN, numerical simulations typically involve iterative algorithms like time steppers that lead to millions of iterations. Even for modest-sized models, this may yield infeasible memory requirements when applying the adjoint method, also called backpropagation, to time-dependent problems. In this situation, checkpointing algorithms provide a trade-off between recomputation and storage. This paper presents the package Checkpointing.jl that leverages expression transformations in the programming language Julia and the package ChainRules.jl to automatically and transparently transform loop iterations into differentiated loops. The user may choose between various checkpointing algorithm schemes and storage devices. We describe the unique design of Checkpointing.jl and demonstrate its features on an automatically differentiated MPI implementation of Burgers' equation on the Polaris cluster at the Argonne Leadership Computing Facility.
Original languageEnglish (US)
Title of host publicationComputational Science - ICCS 2023
Subtitle of host publication23rd International Conference, Prague, Czech Republic, July 3–5, 2023, Proceedings, Part III
EditorsJiří Mikyška, Clélia de Mulatier, Maciej Paszynski, Valeria V. Krzhizhanovskaya, Jack J. Dongarra, Peter M.A. Sloot
PublisherSpringer
Pages483-497
Number of pages15
ISBN (Electronic)9783031360244
ISBN (Print)9783031360237
StatePublished - Jun 26 2023
Externally publishedYes

Fingerprint

Dive into the research topics of 'Transparent Checkpointing for Automatic Differentiation of Program Loops Through Expression Transformations'. Together they form a unique fingerprint.

Cite this