H-R evolve

Abstract

We study the problem of checkpointing strategies for adjoint computation on synchronous hierarchical platforms, specifically computational platforms with several levels of storage with different writing and reading costs. When reversing a large adjoint chain, choosing which data to checkpoint and where is a critical decision for the overall performance of the computation. We introduce H-Revolve, an optimal algorithm for this problem. We make it available in a public Python library along with the implementation of several state-ofthe-art algorithms for the variant of the problem with two levels of storage. We provide a detailed description of how one can use this library in an adjoint computation software in the field of automatic differentiation or backpropagation. Finally, we evaluate the performance of H-Revolve and other checkpointing heuristics though an extensive campaign of simulation.

checkpoint_schedules: schedules for incremental checkpointing of adjoint simulations

“…This can be used in conjunction with an adjoint framework such as tlm_adjoint or pyadjoint and a compatible PDE framework, such as Firedrake (Ham et al, 2023) or FEniCS (Alnaes et al, 2015), to enable users to create adjoint solvers for their choice of PDE, numerical methods, and checkpointing algorithm, all without recoding the underlying algorithms. Some of the algorithms supported by checkpointing_schedules have been implemented many times, while for others, such as H-Revolve the only previously published implementation was a simple proof of concept in the original paper (Herrmann & Pallez, 2020). The checkpoint schedule API provided by checkpoint_schedules further provides a toolkit for the implementation of further checkpoint schedules, thereby providing a direct route from algorithm developers to users.…”

“…Currently, checkpoint_schedules is able to generate schedules for the following checkpointing schemes: Revolve (Stumm & Walther, 2009), disk-revolve (Aupy et al, 2016), periodic-disk revolve (Aupy & Herrmann, 2017), two-level (Pringle et al, 2016), H-Revolve (Herrmann & Pallez, 2020), and mixed storage checkpointing (James R. Maddison, 2023). It also contains trivial schedules that store the entire forward state.…”

“…GP and JH wrote the original reference implementation of H-Revolve and related schedules originally published in Herrmann & Pallez (2020), and contributed to the fixed and enhanced version of that code that is included in checkpoint_schedules. DH and JM designed the original checkpoint_schedules API, which was implemented by DH, JM and DD.…”

“…Griewank & Walther (2000) proposed a checkpointing algorithm, which is optimal under certain assumptions, including that the number of steps is known in advance, and that all the storage has equal access cost. Subsequent authors have produced checkpointing algorithms that relax these requirements in various ways, such as by accounting for different types of storage (e.g., memory and disk) or by not requiring the number of steps to be known in advance, for example Stumm & Walther (2009), Aupy et al (2016), Schanen et al (2016), Aupy & Herrmann (2017), Herrmann & Pallez (2020), James R. Maddison (2023), and Zhang & Constantinescu (2023).…”