EnzymeAD / EnzymeAD/Reactant.jl
How to add checkpointing scheme for reactant
- Dominant language
- Julia
- Stars
- 370
- Forks
- 74
- Avg merge
- 18h 47m
- Merged PRs (30d)
- 30
Description
Sorry if this isn't the right place to post, but in the process of adding reactant to my code I'm seeing two main issues and wanted to document here:
1) If I don't use Checkpointing.jl, the code increases in memory usage until my computer kills the process, and this is happening for a one day integration which shouldn't be enough to cause memory issues
2) I tried to change to the integration that uses checkpointing to circumvent (1) by adding a few lines:
```julia
S = Reactant.to_rarray(S)
dS = Reactant.to_rarray(dS)
revolve = Reactant.to_rarray(revolve) # first added line
compiled_outer = @compile outer(S, dS, revolve) # adjusted to take the checkpointing scheme as an argument
compiled_outer = outer
compiled_outer(S, dS, revolve)
return S, dS
```
inside the relevant functions `run_adjoint` and `outer`,
```julia
function outer(S, dS, revolve)
autodiff(Enzyme.ReverseWithPrimal, checkpointed_integration, Duplicated(S, dS), Const(revolve))
nothing
end
function run_adjoint(::Type{T}=Float32;
kwargs...
) where {T<:AbstractFloat}
P = ShallowWaters.Parameter(T=T;kwargs...)
@show typeof(P)
S = ShallowWaters.model_setup(P)
dS = Enzyme.Compiler.make_zero(S)
snaps = Int(floor(sqrt(S.grid.nt)))
revolve = Revolve{ShallowWaters.ModelSetup}(S.grid.nt,
snaps;
verbose=1,
gc=true,
write_checkpoints=false,
write_checkpoints_filename = "technicalpaper_checkingderivatives_30dayrun_withcheckpointing_120924",
write_checkpoints_period = 224
)
# autodiff(Enzyme.ReverseWithPrimal, checkpointed_integration, Duplicated(S, dS), Const(revolve))
# autodiff(Enzyme.ReverseWithPrimal, integration, Duplicated(S, dS))
# outer()
S = Reactant.to_rarray(S)
dS = Reactant.to_rarray(dS)
revolve = Reactant.to_rarray(revolve)
compiled_outer = @compile outer(S, dS, revolve)
compiled_outer = outer
compiled_outer(S, dS, revolve)
return S, dS
end
```
but this just resulted in the error message:
```julia
swilliamson@CRIOS-A66253 ~/D/G/S/eddy-stresses> julia --project=. eddy_paper.jl :) main!?#
┌ Warning: It's not recommended to use allowscalar([true]) to allow scalar indexing.
│ Instead, use `allowscalar() do end` or `@allowscalar` to denote exactly which operations can use scalar operations.
└ @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:184
typeof(P) = Parameter{Array{Float64, 3}, Vector{Float64}, Vector{Float64}}
[ Info: Revolve: Number of checkpoints: 15
[ Info: Revolve: Number of steps: 225
[ Info: Prediction:
[ Info: Forward steps : 522
[ Info: Overhead factor : 2.32
ERROR: LoadError: Abstract type Function does not have a definite size.
Stacktrace:
[1] sizeof
@ ./essentials.jl:631 [inlined]
[2] traced_type_inner(T::Type{<:Function}, seen::Dict{Type, Type}, mode::Reactant.TraceMode, track_numbers::Type)
@ Reactant ~/.julia/dev/Reactant/src/Tracing.jl:76
[3] traced_type_inner(T::Type, seen::Dict{Type, Type}, mode::Reactant.TraceMode, track_numbers::Type)
@ Reactant ~/.julia/dev/Reactant/src/Tracing.jl:415
[4] traced_type_inner(T::Type, seen::Dict{Type, Type}, mode::Reactant.TraceMode, track_numbers::Type)
@ Reactant ~/.julia/dev/Reactant/src/Tracing.jl:444
[5] traced_type(T::Type, ::Val{Reactant.ArrayToConcrete}, track_numbers::Type)
@ Reactant ~/.julia/dev/Reactant/src/Tracing.jl:611
[6] make_tracer(seen::Reactant.OrderedIdDict{Any, Any}, prev::Any, path::Any, mode::Reactant.TraceMode; toscalar::Bool, tobatch::Nothing, track_numbers::Type, kwargs::@Kwargs{})
@ Reactant ~/.julia/dev/Reactant/src/Tracing.jl:668
[7] make_tracer
@ ~/.julia/dev/Reactant/src/Tracing.jl:654 [inlined]
[8] to_rarray_internal
@ ~/.julia/dev/Reactant/src/Tracing.jl:1067 [inlined]
[9] #to_rarray#79
@ ~/.julia/dev/Reactant/src/Tracing.jl:1063 [inlined]
[10] to_rarray
@ ~/.julia/dev/Reactant/src/Tracing.jl:1061 [inlined]
[11] run_adjoint(::Type{Float32}; kwargs::@Kwargs{output::Bool, L_ratio::Int64, g::Float64, H::Int64, wind_forcing_x::String, Lx::Float64, seasonal_wind_x::Bool, topography::String, bc::String, bottom_drag::String, nn_forcing_dissipation::Bool, handwritten::Bool, α::Int64, nx::Int64, Ndays::Int64})
@ Main ~/Documents/GitHub/ShallowWaters_work/eddy-stresses/eddy_paper_experiment_functions.jl:39
[12] run_adjoint
@ ~/Documents/GitHub/ShallowWaters_work/eddy-stresses/eddy_paper_experiment_functions.jl:14 [inlined]
[13] top-level scope
@ ~/Documents/GitHub/ShallowWaters_work/eddy-stresses/eddy_paper_run_experiments.jl:3
[14] include(fname::String)
@ Base.MainInclude ./client.jl:494
[15] top-level scope
@ ~/Documents/GitHub/ShallowWaters_work/eddy-stresses/eddy_paper.jl:28
in expression starting at /Users/swilliamson/Documents/GitHub/ShallowWaters_work/eddy-stresses/eddy_paper_run_experiments.jl:3
in expression starting at /Users/swilliamson/Documents/GitHub/ShallowWaters_work/eddy-stresses/eddy_paper.jl:28
```
I tried to add the scheme in the same way that `S` and `dS` are treated, is this not the right idea? I mainly want to move to using checkpointing because, without it, I can't even see any error messages
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.