Retool AD specification

[rsenne]

Originally done as a workaround but the logic is a bit silly

[penelopeysm]

I guess, maybe part of the question is where / how do you want the adtype to be specified? There are several places where the user can specify an adtype

• pass it to the DensityModel constructor (in the case of DynamicPPLExt)
• prepare a LogDensityProblem with the ad prep (in the case of LogDensityProblemsExt; the example there uses DynamicPPL but there are other possibilities, e.g. LogDensityProblemsAD.jl)
• in the sampler

and several places where this gets used

• gradient of logdensity
• hvp

Right now it's not immediately clear how the former maps to the latter (one has to read the source code to find out).

There are several options here, but I think that the best interface places adtype specification on the model, not the sampler.* Furthermore, in principle the adtypes for gradient and hvp could (and indeed should) be able to be separately specified. (That's already the case now of course!) So, something like this:

function DensityModel(
    logdensity::Any,          # callable
    grad_logdensity::Any,     # ADType OR callable
    hvp::Any,                 # ADType OR callable OR nothing if not needed
    ...
)

At the point of constructing DensityModel, you can't actually do the gradient prep etc yet, because you don't have a sample input to do it on.

So this places some extra load on the sampler to do some prep work before starting sampling. Just considering the gradient for now (HVP is the same idea, just more of a faff):

dm = DensityModel(sum, AutoForwardDiff())

Then the sampler would have to do a kind of 'upconversion' before starting any sampling:

x0 = ...
prep, actual_gradlogp = if dm.grad_logdensity isa AbstractADType
    prep = DI.prepare_gradient(...)
    grad = x -> DI.gradient(...)
    (prep, grad)
else
    # assume it's a callable of the right form already
    nothing, dm.grad_logdensity
end

On the bright side, this upconversion depends only on the DensityModel plus some sample input so it could be nicely encapsulated inside a function which returns a, say, PreppedDensityModel.

If you wanted to retain backwards compat and allow specifying adtypes via the sampler, then you could also inject the sampler's adtype into this function at this stage.

[rsenne]

Thank you for the thoughtful input! I largely agree with a lot of this. AD design for this package has been a bit tricky because i) I'm trying to support both Turing and non-Turing use cases (especially relevant given the next point), and ii) I'm supporting both CPU and GPU models, with GPU being the largest target since that's where you see the biggest algorithmic gains--which is further complicated by the newfound knowledge that GPU support isn't considered a priority by the Turing team (probably should have checked in on this, but oh well).

This has caused me a bit of struggle, because in reality the current AD stack itself is a bit hacky: we aren't really computing "real" HVPs in the sense of true second-order AD, due to the many difficulties I've encountered with the AD + GPU + SecondOrder stacks, and one still isn't able to supply only a log posterior (or the unnormalized target). Though initial experiments with Reactant are very promising--it's just further complicated by my not totally understanding how it should exist with the current DI infrastructure as to my understanding the don't interoperate at the moment

guess all of this being said, I like the structure of your proposed design and would like to see something like it implemented--anything that makes AD specification simpler overall is a plus. Some of these points are useful constraints, though: for example, we currently don't support supplying only the log posterior, at least not explicitly, which this suggestion would help address--but we're then blocked on the second-order part of the GPU stack (in the case we are supplying only the log posterior). Unless we can get Reactant to land, which frankly feels inevitable: it's the only true second-order approach I've been able to use that works right out of the box. (The forward/reverse-on-grad paths we currently support should stick around regardless, since they can speed things up when you can write the gradient down.)

Some of these points above will necessarily be useful design constraints, others are a bit more orthogonal, but I think this is much needed discussion so thanks for that!

Also unrelated--but are you interested in being a more formal contributor? You've already helped a good bit so I'm more than happy to have you join on, but no pressure, just thought I'd ask :)

[penelopeysm]

Hmmm - so you can effectively only differentiate one order on GPU (so, e.g., getting gradlogp from analytical logp, or getting HVP from analytical gradlogp) but any second-order differentiation fails?

a more formal contributor

No! I think you should have the credit for making it. I don't mind contributing PRs here and there, that is the whole point of open source :)

[rsenne]

so you can effectively only differentiate one order on GPU

unfortunately so :/ Enzyme can work but its incredibly brittle (unless of course its paired with Reactant) and Mooncake also explodes pretty much immediately on almost all examples I've given it.

I don't mind contributing PRs here and there, that is the whole point of open source :)

Fair enough, well the help is much appreciated!

[rsenne]

@gdalle tagging you here because this discussion touched on most aspects of AD in this package so curious to hear your thoughts

[penelopeysm]

I see -- well, it's up to you how you want to handle it. My 2p is that the API should reflect the ideal scenario and if things don't work they can be documented or you can insert warnings in the codebase. It's a bit like how there are almost certainly AD backends that don't work even with first-order stuff, but in principle DI lets you supply any AbstractADType, including those that won't work. The point being that if AD doesn't work, it's not really your library's job to handle it or to restrict the API.

[rsenne]

Yes I'm very much in agreement, on this point. I think the plan as you've constructed seems pretty close to ideal. Most of my points were about just thinking how we can maybe nicely support the multiplicity of cases nicely / talking out loud, but I think your proposed design pretty much handles them (excl. maybe reactant as that i'm a lot more unclear on). Is this something you'd like to work on?

[penelopeysm]

Not particularly, I'm afraid 😅. I could handle the first order stuff but I'm kind of lost with the HVPs. In principle I could figure it out but that might take too long (also a bit busy with other stuff these few days).

[rsenne]

Lmao fair enough! FWIW i have full faith you could have pulled it off--nonetheless thanks for the ideas! I'll probably ask for a review once I have a prototype working

[gdalle]

Gonna try to react to a few remarks but since I'm not familiar with this model-sampler architecture I may not be the best person to evaluate design here:

Originally done as a workaround but the logic is a bit silly

Why do you want to specify a default ADType in the first place? If you don't, it seems you won't need ForwardDiff as a trigger to ensure AutoForwardDiff is available?

we aren't really computing "real" HVPs in the sense of true second-order AD, due to the many difficulties I've encountered with the AD + GPU + SecondOrder stacks

For future references, are the MWEs provided in #32 for second-order failures still the right ones?

Though initial experiments with Reactant are very promising--it's just further complicated by my not totally understanding how it should exist with the current DI infrastructure as to my understanding the don't interoperate at the moment

DI's support of Reactant is a work in progress (emphasis on "work" more than on "progress"). It will probably look something like JuliaDiff/DifferentiationInterface.jl#918, where the user can wrap AutoEnzyme inside AutoReactant to have the gradient operator compile at preparation time. Even once that exists however, it will only be able to speed up each gradient individually, as opposed to the whole function where those gradients are used. Depending on your use case, gradient compilation alone may either be lifesaving as it is, or quite insufficient compared to full-program optimization.

If this is helpful for you, I'll try to revive the Reactant PR alongside the batched DI API we've been discussing.

[rsenne]

Why do you want to specify a default ADType in the first place? If you don't, it seems you won't need ForwardDiff as a trigger

I don't really--it originally was a workaround for #25 but, i realized that it doesn't really make sense to do it this way and would be better to just document. So plan is no default and remove

are the MWEs provided in #32 for second-order failures still the right ones

They should be--though I can probably condense / put in one comment here as they are strewn across a few threads / minimized if that would be helpful

RE: Reactant, my experience thus far with Enzyme + GPU has been that Enzyme without Reactant seems to be relatively brittle. Part of this may be fixed with newer versions, as it seems I was able to uncover a fairly pervasive bug that seems to be now solved though I'm unsure if this fix has been integrated into the Julia side of Enzyme. So I guess having Reactant is probably better than not?

RE batching: i think this will be really useful, but as of now I've still had no luck with any second order AD backend (except prelim experiments with Reactant + Enzyme though need to try out Zygote, has been on my TODO list, but its now stalled I hear?)

[rsenne]

@gdalle

Enzyme MWE

using Enzyme: Enzyme
using CUDA: CUDA

logp(β) = -0.5f0 * sum(abs2, β)

const D = 8
β = CUDA.CuArray(Float32.(1:D) ./ D)
v = CUDA.CuArray(ones(Float32, D))
res = CUDA.CuArray(zeros(Float32, D))

Enzyme.hvp!(res, logp, β, v)

Enzyme Error

unsupported tag gc-transition for   %13 = call i32 @cuMemcpyDtoHAsync_v2(ptr %"dst::Ptr", i64 %"src::CuPtr", i64 4, ptr %.unbox) #929 [ "jl_roots"(ptr addrspace(10) %value_phi), "gc-transition"() ], !dbg !911
unsupported tag
UNREACHABLE executed at /workspace/srcdir/Enzyme/enzyme/Enzyme/GradientUtils.cpp:309!

Mooncake MWE

using Mooncake
using CUDA
using Random

logp(x) = -0.5f0 * sum(abs2, x)
x_gpu = CUDA.CuArray(randn(Float32, 8))
v_gpu = CUDA.CuArray(ones(Float32, 8))

cache = prepare_hvp_cache(logp, x_gpu)
value_and_hvp!!(cache, logp, v_gpu, x_gpu)

Mooncake Error

ERROR: Mooncake.IntrinsicsWrappers.MissingIntrinsicWrapperException("Unable to translate the intrinsic Val{Core.Intrinsics.atomic_pointerref}() into a regular Julia function. Please see github.com/chalk-lab/Mooncake.jl/issues/208 for more discussion.")
Stacktrace:
  [1] translate(f::Val{Core.Intrinsics.atomic_pointerref})
    @ Mooncake.IntrinsicsWrappers /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/rules/builtins.jl:128
  [2] lift_intrinsic(::Core.IntrinsicFunction, ::Core.SSAValue, ::QuoteNode, ::Vararg{Any})
    @ Mooncake /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/ir_normalisation.jl:309
  [3] intrinsic_to_function
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/ir_normalisation.jl:300 [inlined]
  [4] normalise!(ir::Compiler.IRCode, spnames::Vector{Symbol})
    @ Mooncake /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/ir_normalisation.jl:35
  [5] generate_dual_ir(interp::Mooncake.MooncakeInterpreter{Mooncake.DefaultCtx, Mooncake.ForwardMode}, sig_or_mi::Core.MethodInstance; debug_mode::Bool, do_inline::Bool, do_optimize::Bool)
    @ Mooncake /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/forward_mode.jl:196
  [6] build_frule(interp::Mooncake.MooncakeInterpreter{Mooncake.DefaultCtx, Mooncake.ForwardMode}, sig_or_mi::Any; debug_mode::Bool, silence_debug_messages::Bool, skip_world_age_check::Bool)
    @ Mooncake /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/forward_mode.jl:98
  [7] _build_rule!(rule::Mooncake.LazyFRule{Tuple{…}, Mooncake.DerivedFRule{…}}, args::Tuple{Mooncake.Dual{…}, Mooncake.Dual{…}, Mooncake.Dual{…}, Mooncake.Dual{…}, Mooncake.Dual{…}})
    @ Mooncake /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/forward_mode.jl:532
  [8] LazyFRule
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/forward_mode.jl:527 [inlined]
  [9] cufunction
    @ /projectnb/ssmsvi/rsenne/.julia/packages/CUDA/bjncr/src/compiler/execution.jl:365 [inlined]
 [10] (::Tuple{…})(_2::Mooncake.Dual{…}, _3::Mooncake.Dual{…}, _4::Mooncake.Dual{…}, _5::Mooncake.Dual{…}, _6::Mooncake.Dual{…})
    @ Base.Experimental ./<missing>:0
 [11] (::MistyClosures.MistyClosure{…})(::Mooncake.Dual{…}, ::Mooncake.Dual{…}, ::Mooncake.Dual{…}, ::Mooncake.Dual{…}, ::Mooncake.Dual{…})
    @ MistyClosures /projectnb/ssmsvi/rsenne/.julia/packages/MistyClosures/2vtLL/src/MistyClosures.jl:22
 [12] (::Mooncake.DerivedFRule{…})(::Mooncake.Dual{…}, ::Mooncake.Dual{…}, ::Mooncake.Dual{…}, ::Mooncake.Dual{…}, ::Mooncake.Dual{…})
    @ Mooncake /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/forward_mode.jl:122
 [13] __call_rule
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/utils.jl:639 [inlined]
 [14] (::Mooncake.DynamicFRule{…})(::Mooncake.Dual{…}, ::Mooncake.Dual{…}, ::Mooncake.Dual{…}, ::Mooncake.Dual{…}, ::Mooncake.Dual{…})
    @ Mooncake /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/forward_mode.jl:582
 [15] macro expansion
    @ /projectnb/ssmsvi/rsenne/.julia/packages/CUDA/bjncr/src/compiler/execution.jl:112 [inlined]
 [16] #_#6
    @ /projectnb/ssmsvi/rsenne/.julia/packages/CUDA/bjncr/src/CUDAKernels.jl:129 [inlined]
 [17] (::Tuple{…})(_2::Mooncake.Dual{…}, _3::Mooncake.Dual{…}, _4::Mooncake.Dual{…}, _5::Mooncake.Dual{…}, _6::Mooncake.Dual)
    @ Base.Experimental ./<missing>:0
 [18] (::MistyClosures.MistyClosure{…})(::Mooncake.Dual{…}, ::Mooncake.Dual{…}, ::Mooncake.Dual{…}, ::Mooncake.Dual{…}, ::Mooncake.Dual{…})
    @ MistyClosures /projectnb/ssmsvi/rsenne/.julia/packages/MistyClosures/2vtLL/src/MistyClosures.jl:22
 [19] DerivedFRule
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/forward_mode.jl:122 [inlined]
 [20] __call_rule
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/utils.jl:639 [inlined]
 [21] _build_rule!(rule::Mooncake.LazyFRule{Tuple{…}, Mooncake.DerivedFRule{…}}, args::Tuple{Mooncake.Dual{…}, Mooncake.Dual{…}, Mooncake.Dual{…}, Mooncake.Dual{…}, Mooncake.Dual{…}, Mooncake.Dual{…}})
    @ Mooncake /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/forward_mode.jl:533
 [22] LazyFRule
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/forward_mode.jl:527 [inlined]
 [23] #412
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interface.jl:2240 [inlined]
 [24] (::Tuple{Mooncake.LazyFRule{…}, Mooncake.LazyFRule{…}, Mooncake.LazyFRule{…}})(_2::Mooncake.Dual{Mooncake.var"#412#413"{…}, Mooncake.Tangent{…}}, _3::Mooncake.Dual{CuArray{…}, CuArray{…}})
    @ Base.Experimental ./<missing>:0
 [25] (::MistyClosures.MistyClosure{Core.OpaqueClosure{…}})(::Mooncake.Dual{Mooncake.var"#412#413"{…}, Mooncake.Tangent{…}}, ::Mooncake.Dual{CuArray{…}, CuArray{…}})
    @ MistyClosures /projectnb/ssmsvi/rsenne/.julia/packages/MistyClosures/2vtLL/src/MistyClosures.jl:22
 [26] DerivedFRule
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interpreter/forward_mode.jl:122 [inlined]
 [27] __call_rule
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/utils.jl:639 [inlined]
 [28] value_and_derivative!!
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interface.jl:2128 [inlined]
 [29] value_and_hvp!!(cache::Mooncake.HVPCache{…}, f::typeof(logp), v::CuArray{…}, x1::CuArray{…})
    @ Mooncake /projectnb/ssmsvi/rsenne/.julia/packages/Mooncake/O5brV/src/interface.jl:2302
 [30] top-level scope
    @ REPL[9]:1
Some type information was truncated. Use `show(err)` to see complete types.

[gdalle]

Have you tried opening an Enzyme issue for that one?