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

doi:10.1007/978-3-031-36024-4_37

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 proceeding › Conference 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 language	English (US)
Title of host publication	Computational Science - ICCS 2023 - 23rd International Conference, Proceedings
Subtitle of host publication	23rd International Conference, Prague, Czech Republic, July 3–5, 2023, Proceedings, Part III
Editors	Jiří Mikyška, Clélia de Mulatier, Maciej Paszynski, Valeria V. Krzhizhanovskaya, Jack J. Dongarra, Peter M.A. Sloot
Publisher	Springer
Pages	483-497
Number of pages	15
ISBN (Electronic)	9783031360244
ISBN (Print)	9783031360237
DOIs	https://doi.org/10.1007/978-3-031-36024-4_37
State	Published - Jun 26 2023
Externally published	Yes

Publication series

Name	Lecture Notes in Computer Science (including subseries Lecture Notes in Artificial Intelligence and Lecture Notes in Bioinformatics)
Volume	14075 LNCS
ISSN (Print)	0302-9743
ISSN (Electronic)	1611-3349

Keywords

Automatic differentiation
Checkpointing
Julia

ASJC Scopus subject areas

Theoretical Computer Science
General Computer Science

Online availability

10.1007/978-3-031-36024-4_37

Library availability

Discover UIUC Full Text

Cite this

Schanen, M., Narayanan, S. H. K., Williamson, S., Churavy, V., Moses, W. S., & Paehler, L. (2023). Transparent Checkpointing for Automatic Differentiation of Program Loops Through Expression Transformations. In J. Mikyška, C. de Mulatier, M. Paszynski, V. V. Krzhizhanovskaya, J. J. Dongarra, & P. M. A. Sloot (Eds.), Computational Science - ICCS 2023 - 23rd International Conference, Proceedings: 23rd International Conference, Prague, Czech Republic, July 3–5, 2023, Proceedings, Part III (pp. 483-497). (Lecture Notes in Computer Science (including subseries Lecture Notes in Artificial Intelligence and Lecture Notes in Bioinformatics); Vol. 14075 LNCS). Springer. https://doi.org/10.1007/978-3-031-36024-4_37

Transparent Checkpointing for Automatic Differentiation of Program Loops Through Expression Transformations. / Schanen, Michel; Narayanan, Sri Hari Krishna; Williamson, Sarah et al.
Computational Science - ICCS 2023 - 23rd International Conference, Proceedings: 23rd International Conference, Prague, Czech Republic, July 3–5, 2023, Proceedings, Part III. ed. / Jiří Mikyška; Clélia de Mulatier; Maciej Paszynski; Valeria V. Krzhizhanovskaya; Jack J. Dongarra; Peter M.A. Sloot. Springer, 2023. p. 483-497 (Lecture Notes in Computer Science (including subseries Lecture Notes in Artificial Intelligence and Lecture Notes in Bioinformatics); Vol. 14075 LNCS).

Research output: Chapter in Book/Report/Conference proceeding › Conference contribution

Schanen, M, Narayanan, SHK, Williamson, S, Churavy, V, Moses, WS & Paehler, L 2023, Transparent Checkpointing for Automatic Differentiation of Program Loops Through Expression Transformations. in J Mikyška, C de Mulatier, M Paszynski, VV Krzhizhanovskaya, JJ Dongarra & PMA Sloot (eds), Computational Science - ICCS 2023 - 23rd International Conference, Proceedings: 23rd International Conference, Prague, Czech Republic, July 3–5, 2023, Proceedings, Part III. Lecture Notes in Computer Science (including subseries Lecture Notes in Artificial Intelligence and Lecture Notes in Bioinformatics), vol. 14075 LNCS, Springer, pp. 483-497. https://doi.org/10.1007/978-3-031-36024-4_37

Schanen M, Narayanan SHK, Williamson S, Churavy V, Moses WS, Paehler L. Transparent Checkpointing for Automatic Differentiation of Program Loops Through Expression Transformations. In Mikyška J, de Mulatier C, Paszynski M, Krzhizhanovskaya VV, Dongarra JJ, Sloot PMA, editors, Computational Science - ICCS 2023 - 23rd International Conference, Proceedings: 23rd International Conference, Prague, Czech Republic, July 3–5, 2023, Proceedings, Part III. Springer. 2023. p. 483-497. (Lecture Notes in Computer Science (including subseries Lecture Notes in Artificial Intelligence and Lecture Notes in Bioinformatics)). doi: 10.1007/978-3-031-36024-4_37

Schanen, Michel ; Narayanan, Sri Hari Krishna ; Williamson, Sarah et al. / Transparent Checkpointing for Automatic Differentiation of Program Loops Through Expression Transformations. Computational Science - ICCS 2023 - 23rd International Conference, Proceedings: 23rd International Conference, Prague, Czech Republic, July 3–5, 2023, Proceedings, Part III. editor / Jiří Mikyška ; Clélia de Mulatier ; Maciej Paszynski ; Valeria V. Krzhizhanovskaya ; Jack J. Dongarra ; Peter M.A. Sloot. Springer, 2023. pp. 483-497 (Lecture Notes in Computer Science (including subseries Lecture Notes in Artificial Intelligence and Lecture Notes in Bioinformatics)).

@inproceedings{98d555af1b3a4b79a6b9c4912d536a1b,

title = "Transparent Checkpointing for Automatic Differentiation of Program Loops Through Expression Transformations",

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.",

keywords = "Automatic differentiation, Checkpointing, Julia",

author = "Michel Schanen and Narayanan, {Sri Hari Krishna} and Sarah Williamson and Valentin Churavy and Moses, {William S.} and Ludger Paehler",

year = "2023",

month = jun,

day = "26",

doi = "10.1007/978-3-031-36024-4_37",

language = "English (US)",

isbn = "9783031360237",

series = "Lecture Notes in Computer Science (including subseries Lecture Notes in Artificial Intelligence and Lecture Notes in Bioinformatics)",

publisher = "Springer",

pages = "483--497",

editor = "Ji{\v r}{\'i} Miky{\v s}ka and {de Mulatier}, Cl{\'e}lia and Maciej Paszynski and Krzhizhanovskaya, {Valeria V.} and Dongarra, {Jack J.} and Sloot, {Peter M.A.}",

booktitle = "Computational Science - ICCS 2023 - 23rd International Conference, Proceedings",

address = "Germany",

}

TY - GEN

T1 - Transparent Checkpointing for Automatic Differentiation of Program Loops Through Expression Transformations

AU - Schanen, Michel

AU - Narayanan, Sri Hari Krishna

AU - Williamson, Sarah

AU - Churavy, Valentin

AU - Moses, William S.

AU - Paehler, Ludger

PY - 2023/6/26

Y1 - 2023/6/26

N2 - 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.

AB - 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.

KW - Automatic differentiation

KW - Checkpointing

KW - Julia

UR - http://www.scopus.com/inward/record.url?scp=85179433713&partnerID=8YFLogxK

UR - http://www.scopus.com/inward/citedby.url?scp=85179433713&partnerID=8YFLogxK

U2 - 10.1007/978-3-031-36024-4_37

DO - 10.1007/978-3-031-36024-4_37

M3 - Conference contribution

SN - 9783031360237

T3 - Lecture Notes in Computer Science (including subseries Lecture Notes in Artificial Intelligence and Lecture Notes in Bioinformatics)

SP - 483

EP - 497

BT - Computational Science - ICCS 2023 - 23rd International Conference, Proceedings

A2 - Mikyška, Jiří

A2 - de Mulatier, Clélia

A2 - Paszynski, Maciej

A2 - Krzhizhanovskaya, Valeria V.

A2 - Dongarra, Jack J.

A2 - Sloot, Peter M.A.

PB - Springer

ER -