diff --git a/main/.documenter-siteinfo.json b/main/.documenter-siteinfo.json index c829872..e5b76df 100644 --- a/main/.documenter-siteinfo.json +++ b/main/.documenter-siteinfo.json @@ -1 +1 @@ -{"documenter":{"julia_version":"1.10.5","generation_timestamp":"2024-09-25T02:34:08","documenter_version":"1.7.0"}} \ No newline at end of file +{"documenter":{"julia_version":"1.10.5","generation_timestamp":"2024-09-29T01:15:35","documenter_version":"1.7.0"}} \ No newline at end of file diff --git a/main/index.html b/main/index.html index 222c510..4aec536 100644 --- a/main/index.html +++ b/main/index.html @@ -1,2 +1,2 @@ -Home · Checkpointing.jl

Checkpointing

Checkpointing.jl provides checkpointing schemes for adjoint computations using automatic differentiation (AD) of time-stepping loops. Currently, we support the macro @checkpoint_struct, which differentiates and checkpoints a struct used in a while or for the loop with a UnitRange.

Each loop iteration is differentiated using Enzyme.jl. We rely on ChainRulesCore.jl to integrate with AD tools applied to the code outside of the loop.

The schemes are agnostic to the AD tool being used and can be easily interfaced with any Julia AD tool. Currently, the package supports:

Scheme

  • Revolve/Binomial checkpointing [1]
  • Periodic checkpointing
  • Online r=2 checkpointing for a while loops with a priori unknown number of iterations [2]

Rules

Storage

  • ArrayStorage: Stores all checkpoints values in an array of type Array
  • HDF5Storage: Stores all checkpoints values in an HDF5 file

Limitations

  • Currently, the package only supports UnitRange ranges in for loops. We will add range types on a per-need basis. Please, open an issue if you need support for a specific range type.
  • We only support Enzyme as the differentiation tool of the loop body. This is due to our strict requirement for a mutation-enabled AD tool in our projects. However, there is no fundamental reason why we could not support other AD tools. Please, open an issue if you need support for a specific AD tool.
  • We don't support any activity analysis. This implies that loop iterators have to be part of the checkpointed struct if they are used in the loop body. Currently, we store the entire struct at each checkpoint. This is not necessary, and we will add support for storing only the required fields in the future.

Future

The following features are planned for development:

  • Support checkpoints on GPUs

Quick Start

API

References

[1] Andreas Griewank and Andrea Walther. 2000. Algorithm 799: revolve: an implementation of checkpointing for the reverse or adjoint mode of computational differentiation. ACM Trans. Math. Softw. 26, 1 (March 2000), 19–45. DOI:https://doi.org/10.1145/347837.347846

+Home · Checkpointing.jl

Checkpointing

Checkpointing.jl provides checkpointing schemes for adjoint computations using automatic differentiation (AD) of time-stepping loops. Currently, we support the macro @checkpoint_struct, which differentiates and checkpoints a struct used in a while or for the loop with a UnitRange.

Each loop iteration is differentiated using Enzyme.jl. We rely on ChainRulesCore.jl to integrate with AD tools applied to the code outside of the loop.

The schemes are agnostic to the AD tool being used and can be easily interfaced with any Julia AD tool. Currently, the package supports:

Scheme

  • Revolve/Binomial checkpointing [1]
  • Periodic checkpointing
  • Online r=2 checkpointing for a while loops with a priori unknown number of iterations [2]

Rules

Storage

  • ArrayStorage: Stores all checkpoints values in an array of type Array
  • HDF5Storage: Stores all checkpoints values in an HDF5 file

Limitations

  • Currently, the package only supports UnitRange ranges in for loops. We will add range types on a per-need basis. Please, open an issue if you need support for a specific range type.
  • We only support Enzyme as the differentiation tool of the loop body. This is due to our strict requirement for a mutation-enabled AD tool in our projects. However, there is no fundamental reason why we could not support other AD tools. Please, open an issue if you need support for a specific AD tool.
  • We don't support any activity analysis. This implies that loop iterators have to be part of the checkpointed struct if they are used in the loop body. Currently, we store the entire struct at each checkpoint. This is not necessary, and we will add support for storing only the required fields in the future.

Future

The following features are planned for development:

  • Support checkpoints on GPUs

Quick Start

API

References

[1] Andreas Griewank and Andrea Walther. 2000. Algorithm 799: revolve: an implementation of checkpointing for the reverse or adjoint mode of computational differentiation. ACM Trans. Math. Softw. 26, 1 (March 2000), 19–45. DOI:https://doi.org/10.1145/347837.347846

diff --git a/main/lib/checkpointing/index.html b/main/lib/checkpointing/index.html index 296e69f..f4c4e36 100644 --- a/main/lib/checkpointing/index.html +++ b/main/lib/checkpointing/index.html @@ -3,4 +3,4 @@ alg, model, loop, -)

This macro is supposed to be only used in conjunction with ChainRules. It does not initialize the shadowcopy. Apply the checkpointing scheme alg on the loop loop expression. model is the primal struct. shadowmodel contains the adjoints and is created here. It is supposed to be initialized by ChainRules.

source

Supported Schemes

Checkpointing.Revolve — Type
Revolve

This is a Julia adaptation of the functionality of Revolve; see Alg. 799 published by Griewank et al. A minor extension is the optional bundle parameter that allows to treat as many loop iterations in one tape/adjoint sweep. If bundle is 1, the default, then the behavior is that of Alg. 799.

source
Checkpointing.Periodic — Type
Periodic

Periodic checkpointing scheme.

source

Supported Storages

Checkpointing.ArrayStorage — Type
ArrayStorage

Array (RAM) storage for checkpointing.

source
Checkpointing.HDF5Storage — Type
HDF5Storage

A storage type for checkpointing that uses HDF5 files to store the checkpoints.

source

Developer variables for implementing new schemes

Checkpointing.Scheme — Type
Scheme

Abstract type from which all checkpointing schemes are derived.

source
Checkpointing.ActionFlag — Type
ActionFlag

Each checkpointing algorithm currently uses the same ActionFlag type for setting the next action in the checkpointing scheme none: no action store: store a checkpoint now equivalent to TAKESHOT in Alg. 79 restore: restore a checkpoint now equivalent to RESTORE in Alg. 79 forward: execute iteration(s) forward equivalent to ADVANCE in Alg. 79 firstuturn: tape iteration(s); optionally leave to return later; and (upon return) do the adjoint(s) equivalent to FIRSTTURN in Alg. 799 uturn: tape iteration(s) and do the adjoint(s) equivalent to YOUTURN in Alg. 79 done: we are done with adjoining the loop equivalent to the terminate enum value in Alg. 79

source
Checkpointing.Action — Type
Action

Stores the state of the checkpointing scheme after an action is taken. * actionflag is the next action * iteration is number of iterations for move forward * startiteration is the loop step to start from * cpnum is the checkpoint index number

source
+)

This macro is supposed to be only used in conjunction with ChainRules. It does not initialize the shadowcopy. Apply the checkpointing scheme alg on the loop loop expression. model is the primal struct. shadowmodel contains the adjoints and is created here. It is supposed to be initialized by ChainRules.

source

Supported Schemes

Checkpointing.Revolve — Type
Revolve

This is a Julia adaptation of the functionality of Revolve; see Alg. 799 published by Griewank et al. A minor extension is the optional bundle parameter that allows to treat as many loop iterations in one tape/adjoint sweep. If bundle is 1, the default, then the behavior is that of Alg. 799.

source
Checkpointing.Periodic — Type
Periodic

Periodic checkpointing scheme.

source

Supported Storages

Checkpointing.ArrayStorage — Type
ArrayStorage

Array (RAM) storage for checkpointing.

source
Checkpointing.HDF5Storage — Type
HDF5Storage

A storage type for checkpointing that uses HDF5 files to store the checkpoints.

source

Developer variables for implementing new schemes

Checkpointing.Scheme — Type
Scheme

Abstract type from which all checkpointing schemes are derived.

source
Checkpointing.ActionFlag — Type
ActionFlag

Each checkpointing algorithm currently uses the same ActionFlag type for setting the next action in the checkpointing scheme none: no action store: store a checkpoint now equivalent to TAKESHOT in Alg. 79 restore: restore a checkpoint now equivalent to RESTORE in Alg. 79 forward: execute iteration(s) forward equivalent to ADVANCE in Alg. 79 firstuturn: tape iteration(s); optionally leave to return later; and (upon return) do the adjoint(s) equivalent to FIRSTTURN in Alg. 799 uturn: tape iteration(s) and do the adjoint(s) equivalent to YOUTURN in Alg. 79 done: we are done with adjoining the loop equivalent to the terminate enum value in Alg. 79

source
Checkpointing.Action — Type
Action

Stores the state of the checkpointing scheme after an action is taken. * actionflag is the next action * iteration is number of iterations for move forward * startiteration is the loop step to start from * cpnum is the checkpoint index number

source
diff --git a/main/quickstart/2bd1b6ee.svg b/main/quickstart/f1c215a2.svg similarity index 83% rename from main/quickstart/2bd1b6ee.svg rename to main/quickstart/f1c215a2.svg index b0a7914..58ae01e 100644 --- a/main/quickstart/2bd1b6ee.svg +++ b/main/quickstart/f1c215a2.svg @@ -1,44 +1,44 @@ - + - + - + - + - + - - - - - - - - - - - - - - - - - - - - - - - - - + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/main/quickstart/index.html b/main/quickstart/index.html index a7e56d4..9ea2cb1 100644 --- a/main/quickstart/index.html +++ b/main/quickstart/index.html @@ -54,4 +54,4 @@ return heat.Tnext, dheat.Tnext[2:end-1] end
heat (generic function with 1 method)

Plot function values:

tsteps = 500
-T, dT = heat(Revolve{Heat}(tsteps,4), tsteps)
([20.0, 19.28670708893302, 18.5762559822926, 17.865804875652184, 17.163811520281293, 16.461818164910405, 15.773698196009292, 15.085578227108181, 14.416422194357889, 13.747266161607596  …  0.0010844643164675824, 0.0008681767992163476, 0.0007135124080400669, 0.0005588480168637862, 0.000445160866081522, 0.00033147371529925775, 0.0002427191223882328, 0.00015396452947720788, 7.698226473860394e-5, 0.0], [0.0356607964401121, 0.07117950265889618, 0.10669762026127153, 0.14179285030017058, 0.17688683371117586, 0.2112871477986924, 0.24568541302418895, 0.27913548144214495, 0.31258246870379675, 0.3448485923782427  …  0.3448485923782427, 0.31258246870379675, 0.27913548144214495, 0.24568541302418895, 0.2112871477986924, 0.17688683371117586, 0.14179285030017058, 0.10669762026127153, 0.07117950265889618, 0.0356607964401121])

Plot gradient with respect to sum(T):

plot(dT)
Example block output +T, dT = heat(Revolve{Heat}(tsteps,4), tsteps)
([20.0, 19.28670708893302, 18.5762559822926, 17.865804875652184, 17.163811520281293, 16.461818164910405, 15.773698196009292, 15.085578227108181, 14.416422194357889, 13.747266161607596  …  0.0010844643164675824, 0.0008681767992163476, 0.0007135124080400669, 0.0005588480168637862, 0.000445160866081522, 0.00033147371529925775, 0.0002427191223882328, 0.00015396452947720788, 7.698226473860394e-5, 0.0], [0.0356607964401121, 0.07117950265889618, 0.10669762026127153, 0.14179285030017058, 0.17688683371117586, 0.2112871477986924, 0.24568541302418895, 0.27913548144214495, 0.31258246870379675, 0.3448485923782427  …  0.3448485923782427, 0.31258246870379675, 0.27913548144214495, 0.24568541302418895, 0.2112871477986924, 0.17688683371117586, 0.14179285030017058, 0.10669762026127153, 0.07117950265889618, 0.0356607964401121])

Plot gradient with respect to sum(T):

plot(dT)
Example block output diff --git a/main/rules/index.html b/main/rules/index.html index 9d2dc7f..5af5dfd 100644 --- a/main/rules/index.html +++ b/main/rules/index.html @@ -13,4 +13,4 @@ model_input::MT, shadowmodel::MT, condition::Function -) where {MT} end +) where {MT} end diff --git a/main/schemes/index.html b/main/schemes/index.html index 0615eaa..1d49133 100644 --- a/main/schemes/index.html +++ b/main/schemes/index.html @@ -13,4 +13,4 @@ model_input::MT, shadowmodel::MT, condition::Function -) where {MT}

The reverse pass for the checkpointing scheme will be called with the loop body body, the scheme alg, the model input model_input, the initialized shadow model shadowmodel, and the loop range range or the loop condition condition.

+) where {MT}

The reverse pass for the checkpointing scheme will be called with the loop body body, the scheme alg, the model input model_input, the initialized shadow model shadowmodel, and the loop range range or the loop condition condition.

diff --git a/main/storage/index.html b/main/storage/index.html index e552b0a..257b342 100644 --- a/main/storage/index.html +++ b/main/storage/index.html @@ -2,4 +2,4 @@ Storage · Checkpointing.jl

Storage

The checkpoint storage types are derived from AbstractStorage 

using Checkpointing
 struct MyStorage{MT} <: AbstractStorage where {MT} end

and have to implement the following functions.

  • A constructor invoked by the user
function MyStorage{MT}(n::Int) where {MT} end
  • Set and get functions for the storage type
Base.getindex(storage::MyStorage{MT}, i::Int) where {MT}
 Base.setindex!(storage::MyStorage{MT}, value, i::Int) where {MT}
  • Size and dimension functions
Base.size(storage::MyStorage{MT}) where {MT}
-Base.ndims(storage::MyStorage{MT}) where {MT}
+Base.ndims(storage::MyStorage{MT}) where {MT}