From b2291e6c095bc9df583da41253635ca8621b09a6 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Wed, 5 Aug 2026 09:38:38 -0400 Subject: [PATCH 01/11] Some initial changes e.g., remove gating --- CHANGELOG.md | 33 ++++++++++++++---- ext/EnzymeExt.jl | 54 +++++++++-------------------- src/DEER/DEER.jl | 73 ++++++++++++++++++++------------------- test/test-HVP-Strategy.jl | 73 +++++++++++++++++++++++++++++++-------- 4 files changed, 139 insertions(+), 94 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8b29749..d2d204a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -50,12 +50,33 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 its strategy was routed on `DI.outer(backend)`, so a `DifferentiationInterface.SecondOrder` ran the wrong half of the pair. A `SecondOrder` now goes to the true second-order path instead of either - strategy, and the half-selecting helpers it still uses agree: normalization - applies to the outer, the pass whose mode the backend extensions care about. - Unwrapping to that half happens before the normalization hooks are dispatched - on, so a `SecondOrder(AutoEnzyme(), ...)` still reaches `EnzymeExt` and gets - its mode and function annotation pinned rather than running as a bare - `AutoEnzyme()` (which aborts on GPU). + strategy, and the half-selecting helper it still uses agrees: normalization + applies to the outer pass. Unwrapping to that half happens before the + normalization hook is dispatched on, so a `SecondOrder(AutoEnzyme(), ...)` still + reaches `EnzymeExt` and gets its function annotation filled in rather than + running as a bare `AutoEnzyme()`. +- Backend normalization no longer picks a differentiation mode on the user's + behalf (#62). `EnzymeExt` pinned `mode=Enzyme.Forward` (with + `set_runtime_activity`) onto an `AutoEnzyme()` left mode-agnostic, on the + grounds that reverse mode hit a gc-transition abort on GPU and that composed + `pmcmc_matmul` calls needed runtime activity. The `pmcmc_*` Enzyme rules keep + Enzyme off both paths on their own now, so the pin bought nothing — and it cost + correctness, because it silently rewrote the direction of a `SecondOrder`'s + outer half. `SecondOrder(AutoEnzyme(), AutoForwardDiff())` is + reverse-over-forward to `hvp_mode`, its inner half being forward-only, and came + out forward-over-forward. Normalization now fills in only + `function_annotation=Enzyme.Const`, which is about this package's own read-only + HVP wrappers rather than about Enzyme's mode, and leaves `mode` exactly as given + — unset included, for DI to resolve from the operator it runs. `hvp_mode` is + therefore identical before and after normalization for every backend pair. + + A mode set explicitly was never overridden, so only mode-agnostic backends were + affected, and the HVP was a correct HVP either way; what changes is that the + composition asked for is the one that runs. The two normalization hooks + (`_hvp_forward_backend`, `_hvp_closure_backend`) collapse into a single + `_normalized_backend`, since without a mode to choose they no longer differ. + Users relying on a plain `AutoEnzyme()` being run forward should now pass + `AutoEnzyme(; mode=Enzyme.Forward)` explicitly. ### Changed diff --git a/ext/EnzymeExt.jl b/ext/EnzymeExt.jl index 5c0859f..82e63f8 100644 --- a/ext/EnzymeExt.jl +++ b/ext/EnzymeExt.jl @@ -20,46 +20,26 @@ using Enzyme.EnzymeCore.EnzymeRules: # TODO: Implement matmul overloads upstream in Enzyme. See: https://github.com/EnzymeAD/Enzyme.jl/issues/3122 #= -Tell DEER's forward-on-grad HVP path how to normalize a plain `AutoEnzyme()`: -pin `mode=Enzyme.Forward` and `function_annotation=Enzyme.Const`. Pinning -Forward is load-bearing on GPU — without it DI defaults to reverse mode, which -hits the gc-transition abort documented on the Enzyme rules below. +Normalization of the user's `AutoEnzyme` for DEER's AD-HVP paths: fill in +`function_annotation=Enzyme.Const` when they left it open, so Enzyme doesn't +throw `EnzymeMutabilityException` on the read-only `_HvpReverseClosure` / +`_BatchHvpReverseClosure` wrappers, which capture `gradlogp`. Those wrapper types +belong to this package, so declaring them constant is this package's business. -`mode` and `function_annotation` are normalized independently — a user who -sets one keeps that choice, but still gets the default for the other. +`mode` is passed through exactly as given, unset included. Choosing a direction +on the user's behalf is not our call: a mode they set is a decision, and an unset +one is DI's to resolve from the operator it runs. -`set_runtime_activity` is load-bearing for composed `pmcmc_matmul` calls -(e.g. `pmcmc_matmul(transpose(X), pmcmc_matmul(X, β))`). Static activity -analysis can't prove the outer call's `transpose(X)` shadow is safe to -reuse, and Enzyme aborts with `EnzymeRuntimeActivityError`. With runtime -activity, the shadow is tracked dynamically. +An earlier version pinned `mode=Enzyme.Forward` here (with +`set_runtime_activity`) against the gc-transition abort on GPU and +`EnzymeRuntimeActivityError` on composed `pmcmc_matmul` calls. The rules below +keep Enzyme off both paths on their own, so the pin bought nothing and cost +correctness: it silently rewrote the direction of a `SecondOrder`'s outer half +(see `DEER._normalized_backend`). =# -function DEER._hvp_forward_backend(backend::ADTypes.AutoEnzyme{M,A}) where {M,A} - mode = if backend.mode === nothing - Enzyme.set_runtime_activity(Enzyme.Forward) - else - backend.mode - end - annotation = A === Nothing ? Enzyme.Const : A - return ADTypes.AutoEnzyme(; mode=mode, function_annotation=annotation) -end - -#= -Tell DEER's reverse-on-grad HVP path how to normalize a plain `AutoEnzyme()`: -fill in `function_annotation=Enzyme.Const` so Enzyme doesn't throw -`EnzymeMutabilityException` on the read-only `_HvpReverseClosure` / -`_BatchHvpReverseClosure` wrappers, and default `mode` to reverse with -runtime activity. As in `_hvp_forward_backend`, the two fields are -normalized independently. -=# -function DEER._hvp_closure_backend(backend::ADTypes.AutoEnzyme{M,A}) where {M,A} - mode = if backend.mode === nothing - Enzyme.set_runtime_activity(Enzyme.Reverse) - else - backend.mode - end - annotation = A === Nothing ? Enzyme.Const : A - return ADTypes.AutoEnzyme(; mode=mode, function_annotation=annotation) +function DEER._normalized_backend(backend::ADTypes.AutoEnzyme{M,A}) where {M,A} + A === Nothing || return backend + return ADTypes.AutoEnzyme(; mode=backend.mode, function_annotation=Enzyme.Const) end #= diff --git a/src/DEER/DEER.jl b/src/DEER/DEER.jl index 03896ef..f16e8e9 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -93,7 +93,7 @@ function _prepare_hvp(f, backend::AbstractADType, x_template::AbstractVector) v_template = similar(x_template) fill!(v_template, zero(eltype(x_template))) return DI.prepare_pushforward( - f, _hvp_forward_backend(backend), x_template, (v_template,) + f, _normalized_backend(backend), x_template, (v_template,) ) end @@ -116,14 +116,14 @@ function _hvp_prepared( ) x_exec = _materialize_ad_vector(x) v_exec = _tangent_like(x_exec, v) - res = DI.pushforward(f, prep, _hvp_forward_backend(backend), x_exec, (v_exec,)) + res = DI.pushforward(f, prep, _normalized_backend(backend), x_exec, (v_exec,)) return res isa Tuple ? first(res) : res end function _hvp_nopre(f, backend::AbstractADType, x::AbstractVector, v::AbstractVector) x_exec = _materialize_ad_vector(x) v_exec = _tangent_like(x_exec, v) - res = DI.pushforward(f, _hvp_forward_backend(backend), x_exec, (v_exec,)) + res = DI.pushforward(f, _normalized_backend(backend), x_exec, (v_exec,)) return res isa Tuple ? first(res) : res end @@ -133,7 +133,7 @@ function _prepare_batch_hvp_from_grad( V_template = similar(X_template) fill!(V_template, zero(eltype(X_template))) return DI.prepare_pushforward( - grad_batch, _hvp_forward_backend(backend), X_template, (V_template,) + grad_batch, _normalized_backend(backend), X_template, (V_template,) ) end @@ -142,7 +142,7 @@ function _batch_hvp_from_grad_prepared( ) X_exec = _materialize_ad_matrix(X) V_exec = _tangent_like(X_exec, V) - res = DI.pushforward(grad_batch, prep, _hvp_forward_backend(backend), X_exec, (V_exec,)) + res = DI.pushforward(grad_batch, prep, _normalized_backend(backend), X_exec, (V_exec,)) return res isa Tuple ? first(res) : res end @@ -212,33 +212,30 @@ function _hvp_strategy(backend::Union{AbstractADType,DI.SecondOrder}) end #= -Hooks for backend-specific normalization of the user's `backend`. - -`_hvp_forward_backend` is for the forward-on-grad pushforward path -(differentiates the user's `gradlogp` directly). EnzymeExt specializes it -to pin `mode=Enzyme.Forward` and `function_annotation=Enzyme.Const` when -the user passed plain `AutoEnzyme()`, without pinning Forward, DI lowers -through reverse mode and hits the gc-transition abort on GPU (see -`ext/EnzymeExt.jl`). - -`_hvp_closure_backend` is for the reverse-on-grad gradient path on the -read-only `_HvpReverseClosure` / `_BatchHvpReverseClosure` wrappers. -EnzymeExt specializes it to set `function_annotation=Enzyme.Const` so -Enzyme doesn't throw `EnzymeMutabilityException` on a closure that captures -`gradlogp`. - -A `SecondOrder` is unwrapped to its outer half first, since the pass we are -about to run is the outer one i.e., the inner derivative is whatever -`gradlogp` already is. That keeps the half `_hvp_strategy` routed on, so the -strategy and the backend that carries it out can't end up disagreeing. The -unwrapping recurses rather than calling `DI.outer` in the generic method, so -that a wrapped backend still reaches its own normalization: dispatch happens -on what comes out of `DI.outer`, not on the `SecondOrder` around it. +Hook for backend-specific normalization of the user's `backend`, applied on every +AD-HVP path before the backend reaches DI. + +It supplies what the wrappers DEER differentiates need, and nothing else. Those +wrapper types are ours, so annotating them is ours to do: EnzymeExt specializes +this to fill `function_annotation=Enzyme.Const`, without which Enzyme throws +`EnzymeMutabilityException` on the read-only `_HvpReverseClosure` / +`_BatchHvpReverseClosure`, which capture `gradlogp`. + +It deliberately does not choose a differentiation mode. Which direction a pass +runs is the user's call when they state one and DI's to resolve from the operator +when they don't; this package is not an AD package and has no business overriding +either. Picking one here also used to corrupt a `SecondOrder`, whose halves carry +directions of their own (see `_normalized_second_order`). + +A `SecondOrder` normalizes to its outer half, because the paths that call this +differentiate the already-built `gradlogp`: the inner derivative has run, so the +outer is the only pass left. The unwrapping recurses rather than calling +`DI.outer` in the generic method, so that a wrapped backend still reaches its own +specialization: dispatch happens on what comes out of `DI.outer`, not on the +`SecondOrder` around it. =# -_hvp_forward_backend(backend::DI.SecondOrder) = _hvp_forward_backend(DI.outer(backend)) -_hvp_closure_backend(backend::DI.SecondOrder) = _hvp_closure_backend(DI.outer(backend)) -_hvp_forward_backend(backend::AbstractADType) = backend -_hvp_closure_backend(backend::AbstractADType) = backend +_normalized_backend(backend::DI.SecondOrder) = _normalized_backend(DI.outer(backend)) +_normalized_backend(backend::AbstractADType) = backend function _prepare_hvp_via_grad_reverse( gradlogp, backend::AbstractADType, x_template::AbstractVector @@ -246,7 +243,7 @@ function _prepare_hvp_via_grad_reverse( v_template = similar(x_template) fill!(v_template, zero(eltype(x_template))) f = _HvpReverseClosure(gradlogp) - eff_backend = _hvp_closure_backend(backend) + eff_backend = _normalized_backend(backend) prep = DI.prepare_gradient(f, eff_backend, x_template, DI.Constant(v_template)) return (f, prep, eff_backend) end @@ -262,7 +259,7 @@ function _prepare_batch_hvp_via_grad_reverse( V_template = similar(X_template) fill!(V_template, zero(eltype(X_template))) f = _BatchHvpReverseClosure(grad_batch) - eff_backend = _hvp_closure_backend(backend) + eff_backend = _normalized_backend(backend) prep = DI.prepare_gradient(f, eff_backend, X_template, DI.Constant(V_template)) return (f, prep, eff_backend) end @@ -315,8 +312,12 @@ takes both passes over the log-density, so these never touch the gradient slot. Preferred over pushing tangents through a prepared DI gradient, which drops out of its preparation once the outer pass hands it an unexpected tangent type. -Only the outer half is normalized: `_hvp_forward_backend` selects it out of the -pair. The inner is a plain first-order gradient and needs no pinning. +Both halves are passed to `DI.hvp` as the user composed them, so the direction +each one runs in is theirs and DI's, not ours. Normalization touches only the +outer half, and only to fill in annotations for the wrappers being +differentiated; because it never substitutes a mode, `DI.hvp_mode` of the pair is +the same before and after. The inner half is a plain first-order gradient over +the user's own `logdensity` and is passed straight through. The batched form differentiates `sum(logdensity_batch(X))`, whose Hessian is block-diagonal by column independence, so its HVP along `V` is the columnwise @@ -324,7 +325,7 @@ HVP. Same argument the batched gradient rests on. --------------------------------------------------------------------------- =# function _normalized_second_order(backend::DI.SecondOrder) - return DI.SecondOrder(_hvp_forward_backend(backend), DI.inner(backend)) + return DI.SecondOrder(_normalized_backend(DI.outer(backend)), DI.inner(backend)) end function _make_hvp_fn_second_order( diff --git a/test/test-HVP-Strategy.jl b/test/test-HVP-Strategy.jl index 8bfea83..b69e1bd 100644 --- a/test/test-HVP-Strategy.jl +++ b/test/test-HVP-Strategy.jl @@ -36,26 +36,69 @@ const DI_STRAT = ParallelMCMC.DEER.DI @test DEER_STRAT._hvp_strategy(so_agnostic_outer) isa DEER_STRAT.ReverseOnGrad end - @testset "the backend that runs is the one routed on" begin - # both paths differentiate the already-built gradlogp, so both take the outer + @testset "normalization unwraps a SecondOrder to its outer half" begin + # The strategy paths differentiate the already-built gradlogp, so the + # outer pass is the only one left to normalize. so_fwd = DI_STRAT.SecondOrder(AutoForwardDiff(), AutoZygote()) - @test DEER_STRAT._hvp_forward_backend(so_fwd) === AutoForwardDiff() + @test DEER_STRAT._normalized_backend(so_fwd) === AutoForwardDiff() so_rev = DI_STRAT.SecondOrder(AutoZygote(), AutoForwardDiff()) - @test DEER_STRAT._hvp_closure_backend(so_rev) === AutoZygote() + @test DEER_STRAT._normalized_backend(so_rev) === AutoZygote() + + #= Unwrapping has to happen before the backend-specific hook is dispatched + on, or a wrapped `AutoEnzyme()` comes out bare and misses EnzymeExt. =# + so_enz = DI_STRAT.SecondOrder(AutoEnzyme(), AutoForwardDiff()) + @test DEER_STRAT._normalized_backend(so_enz) === + DEER_STRAT._normalized_backend(AutoEnzyme()) + end + + @testset "normalization supplies Const but never a mode" begin + #= The wrappers DEER differentiates are its own types, so annotating them + `Const` is its business. The mode is not: one the user set is a decision, + and an unset one is DI's to resolve from the operator it runs. =# + bare = DEER_STRAT._normalized_backend(AutoEnzyme()) + @test bare isa AutoEnzyme{<:Any,Enzyme.Const} + @test bare.mode === nothing + + for mode in (Enzyme.Forward, Enzyme.Reverse) + normalized = DEER_STRAT._normalized_backend(AutoEnzyme(; mode=mode)) + @test normalized.mode === mode + @test normalized isa AutoEnzyme{<:Any,Enzyme.Const} + end + + # An annotation the user chose is left alone. + annotated = AutoEnzyme(; function_annotation=Enzyme.Duplicated) + @test DEER_STRAT._normalized_backend(annotated) === annotated + + # Backends with no specialization pass straight through. + @test DEER_STRAT._normalized_backend(AutoForwardDiff()) === AutoForwardDiff() + @test DEER_STRAT._normalized_backend(AutoZygote()) === AutoZygote() end - @testset "unwrapping a SecondOrder still reaches backend normalization" begin - #= The outer half has to be taken before the backend-specific hook is - dispatched on, or a wrapped `AutoEnzyme()` comes out bare: unnormalized, - it lowers through reverse mode and aborts on GPU (see ext/EnzymeExt.jl). =# - so = DI_STRAT.SecondOrder(AutoEnzyme(), AutoForwardDiff()) - @test DEER_STRAT._hvp_forward_backend(so) === - DEER_STRAT._hvp_forward_backend(AutoEnzyme()) - @test DEER_STRAT._hvp_closure_backend(so) === - DEER_STRAT._hvp_closure_backend(AutoEnzyme()) - @test DEER_STRAT._hvp_forward_backend(so).mode isa Enzyme.ForwardMode - @test DEER_STRAT._hvp_closure_backend(so) isa AutoEnzyme{<:Any,Enzyme.Const} + @testset "normalizing a SecondOrder keeps the composition DI resolved" begin + #= Regression for #62. Normalization used to route the outer half through + a forward-only hook, which pinned `Enzyme.Forward` onto it. For a pair + `hvp_mode` resolves to reverse -- `SecondOrder(AutoEnzyme(), + AutoForwardDiff())` is reverse-over-forward, its inner half being + forward-only -- that silently made it forward-over-forward. =# + for so in ( + DI_STRAT.SecondOrder(AutoEnzyme(), AutoForwardDiff()), + DI_STRAT.SecondOrder(AutoEnzyme(), AutoZygote()), + DI_STRAT.SecondOrder(AutoEnzyme(; mode=Enzyme.Reverse), AutoForwardDiff()), + DI_STRAT.SecondOrder(AutoEnzyme(; mode=Enzyme.Forward), AutoZygote()), + DI_STRAT.SecondOrder(AutoForwardDiff(), AutoZygote()), + DI_STRAT.SecondOrder(AutoZygote(), AutoForwardDiff()), + ) + normalized = DEER_STRAT._normalized_second_order(so) + @test DI_STRAT.hvp_mode(normalized) == DI_STRAT.hvp_mode(so) + # The inner half is the user's own first-order gradient, untouched. + @test DI_STRAT.inner(normalized) === DI_STRAT.inner(so) + # And no mode is invented for the outer half either. Only `AutoEnzyme` + # carries a `mode`, so this is the case that could regress. + if DI_STRAT.outer(so) isa AutoEnzyme + @test DI_STRAT.outer(normalized).mode === DI_STRAT.outer(so).mode + end + end end @testset "strategy resolution is type-stable" begin From af84abd7c93f7353dfee302c62fb07b0f05b7cc3 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Thu, 6 Aug 2026 09:43:24 -0400 Subject: [PATCH 02/11] Fix wording. --- src/interface.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/interface.jl b/src/interface.jl index 9568b0a..69ede25 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -27,8 +27,8 @@ DI's second-order operator. Passing a `SecondOrder` yourself always means the latter, and bypasses the gradient slot even when you wrote it by hand. - `logdensity(x::AbstractVector) -> Real` -- `grad_logdensity` — callable `x -> AbstractVector`, or a backend applied to - `logdensity`. +- `grad_logdensity` — callable `x -> AbstractVector`, or a backend to + differentiate `logdensity` with. - `hvp` — optional callable `(x, v) -> AbstractVector`, or a backend. If `nothing`, DEER builds the HVP from the sampler's `backend`. - `logdensity_batch(X::AbstractMatrix) -> AbstractVector` — optional batched @@ -37,7 +37,7 @@ latter, and bypasses the gradient slot even when you wrote it by hand. A batched gradient derived from this is one gradient of its sum, so coupling between columns would go unnoticed and give wrong derivatives. - `grad_logdensity_batch` — optional callable `X -> AbstractMatrix`, or a - backend applied to `logdensity_batch`. Left out alongside a + backend to differentiate `logdensity_batch` with. Left out alongside a `logdensity_batch`, it is derived when `grad_logdensity` is a backend. - `hvp_batch` — optional callable `(X, V) -> AbstractMatrix`, or a backend, resolved against `grad_logdensity_batch` the same way `hvp` is against From ad1f64c550757ab1bf3368f33685431ff84ec76e Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Thu, 6 Aug 2026 09:59:08 -0400 Subject: [PATCH 03/11] delete `_normalized_backend(::DI.SecondOrder)` --- src/DEER/DEER.jl | 21 ++++++++------------- test/test-HVP-Strategy.jl | 19 +++---------------- 2 files changed, 11 insertions(+), 29 deletions(-) diff --git a/src/DEER/DEER.jl b/src/DEER/DEER.jl index f16e8e9..f0d677b 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -196,9 +196,8 @@ relying on constant propagation through `===`. The routing follows DI's `hvp_mode`: a forward outer pass (`DI.ForwardOverAnything`) takes `ForwardOnGrad`, anything else `ReverseOnGrad`. Only the outer direction matters since we differentiate -the already-built `gradlogp`. Plain `AutoEnzyme()` lands on `ForwardOnGrad`, -which we need: Enzyme reverse hits the gc-transition abort on GPU (see -`ext/EnzymeExt.jl`). +the already-built `gradlogp`. The mode is whatever the user's backend carries; +nothing here substitutes one (see `_normalized_backend`). =# abstract type HVPStrategy end struct ForwardOnGrad <: HVPStrategy end @@ -207,9 +206,7 @@ struct ReverseOnGrad <: HVPStrategy end _strategy_from(::DI.ForwardOverAnything) = ForwardOnGrad() _strategy_from(::DI.HVPMode) = ReverseOnGrad() -function _hvp_strategy(backend::Union{AbstractADType,DI.SecondOrder}) - return _strategy_from(DI.hvp_mode(backend)) -end +_hvp_strategy(backend::AbstractADType) = _strategy_from(DI.hvp_mode(backend)) #= Hook for backend-specific normalization of the user's `backend`, applied on every @@ -227,14 +224,12 @@ when they don't; this package is not an AD package and has no business overridin either. Picking one here also used to corrupt a `SecondOrder`, whose halves carry directions of their own (see `_normalized_second_order`). -A `SecondOrder` normalizes to its outer half, because the paths that call this -differentiate the already-built `gradlogp`: the inner derivative has run, so the -outer is the only pass left. The unwrapping recurses rather than calling -`DI.outer` in the generic method, so that a wrapped backend still reaches its own -specialization: dispatch happens on what comes out of `DI.outer`, not on the -`SecondOrder` around it. +Callers hand this a single pass, never a `SecondOrder`: the strategy paths below +run one AD pass over a hand-written `gradlogp`, and `_resolve_hvp` sends every +`SecondOrder` to `_make_hvp_fn_second_order` before they are reached. +`_normalized_second_order` is the one caller that starts from a pair, and it +selects the outer half itself. =# -_normalized_backend(backend::DI.SecondOrder) = _normalized_backend(DI.outer(backend)) _normalized_backend(backend::AbstractADType) = backend function _prepare_hvp_via_grad_reverse( diff --git a/test/test-HVP-Strategy.jl b/test/test-HVP-Strategy.jl index b69e1bd..e56997e 100644 --- a/test/test-HVP-Strategy.jl +++ b/test/test-HVP-Strategy.jl @@ -36,22 +36,6 @@ const DI_STRAT = ParallelMCMC.DEER.DI @test DEER_STRAT._hvp_strategy(so_agnostic_outer) isa DEER_STRAT.ReverseOnGrad end - @testset "normalization unwraps a SecondOrder to its outer half" begin - # The strategy paths differentiate the already-built gradlogp, so the - # outer pass is the only one left to normalize. - so_fwd = DI_STRAT.SecondOrder(AutoForwardDiff(), AutoZygote()) - @test DEER_STRAT._normalized_backend(so_fwd) === AutoForwardDiff() - - so_rev = DI_STRAT.SecondOrder(AutoZygote(), AutoForwardDiff()) - @test DEER_STRAT._normalized_backend(so_rev) === AutoZygote() - - #= Unwrapping has to happen before the backend-specific hook is dispatched - on, or a wrapped `AutoEnzyme()` comes out bare and misses EnzymeExt. =# - so_enz = DI_STRAT.SecondOrder(AutoEnzyme(), AutoForwardDiff()) - @test DEER_STRAT._normalized_backend(so_enz) === - DEER_STRAT._normalized_backend(AutoEnzyme()) - end - @testset "normalization supplies Const but never a mode" begin #= The wrappers DEER differentiates are its own types, so annotating them `Const` is its business. The mode is not: one the user set is a decision, @@ -97,6 +81,9 @@ const DI_STRAT = ParallelMCMC.DEER.DI # carries a `mode`, so this is the case that could regress. if DI_STRAT.outer(so) isa AutoEnzyme @test DI_STRAT.outer(normalized).mode === DI_STRAT.outer(so).mode + #= The outer half still reaches EnzymeExt's specialization rather + than passing through as a bare backend. =# + @test DI_STRAT.outer(normalized) isa AutoEnzyme{<:Any,Enzyme.Const} end end end From 4c3e12b2c27033c86bd18c5baef0c4919258b098 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Sat, 8 Aug 2026 22:40:21 -0400 Subject: [PATCH 04/11] Initial reactant attempt. --- CHANGELOG.md | 8 ++ Project.toml | 5 +- docs/src/10-getting-started.md | 2 +- docs/src/15-gpu.md | 23 ++++- ext/ReactantExt.jl | 152 ++++++++++++++++++++++++++++++++ src/DEER/DEER.jl | 33 +++++++ src/ParallelMCMC.jl | 2 +- src/interface.jl | 77 +++++++++++++++- test/Project.toml | 1 + test/test-HVP-Strategy.jl | 8 ++ test/test-Reactant-HVP.jl | 156 +++++++++++++++++++++++++++++++++ 11 files changed, 461 insertions(+), 6 deletions(-) create mode 100644 ext/ReactantExt.jl create mode 100644 test/test-Reactant-HVP.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index d2d204a..851d2e1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 batched gradient derived for it when `grad_logdensity` is a backend, rather than leaving the batched DEER path switched off (#52). `hvp_batch` can be a backend in that case too, and differentiates the derived gradient. +- New `ReactantExt`: `ADTypes.AutoReactant()` in any derivative slot, or as the + sampler `backend`, traces the derivative with Enzyme-MLIR and compiles it to + an XLA executable via Reactant.jl. This bypasses Enzyme's LLVM pipeline and + is currently the only path that computes a genuine second-order HVP on GPU, + so `ParallelMALASampler` now works on CUDA for log-density-only models + (#37, #52). Requires `using Reactant` and a Reactant-traceable log-density; + `AutoReactant` cannot be paired with a DifferentiationInterface backend + across the two passes of an HVP, which raises an `ArgumentError`. - An HVP backend over an AD-derived gradient is now taken as true second-order AD, `DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` handed to `DI.hvp`, instead of an outer AD pass over the prepared DI gradient (#37). diff --git a/Project.toml b/Project.toml index 82d250a..7f7f77a 100644 --- a/Project.toml +++ b/Project.toml @@ -20,15 +20,17 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [extensions] DynamicPPLExt = ["DynamicPPL", "FlexiChains", "LogDensityProblems"] EnzymeExt = "Enzyme" LogDensityProblemsExt = "LogDensityProblems" +ReactantExt = ["Reactant", "Enzyme"] [compat] -ADTypes = "1.21.0" +ADTypes = "1.22.0" AbstractMCMC = "5.10.0" CUDA = "5.11.0" CUDA_Runtime_jll = "0.21" @@ -42,6 +44,7 @@ LogDensityProblems = "2" Mooncake = "0.5.26" OrderedCollections = "1" Random = "1" +Reactant = "0.2.278" Statistics = "1" Zygote = "0.7.10" julia = "1.10" diff --git a/docs/src/10-getting-started.md b/docs/src/10-getting-started.md index b24e81e..b257131 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -49,7 +49,7 @@ Backends become prepared [DifferentiationInterface](https://github.com/JuliaDiff Naming both passes yourself is the one route that ignores the gradient slot, hand-written or not. It is also the only AD route to an HVP for a Turing or LogDensityProblems model, whose gradient arrives already prepared and cannot be differentiated again. -Which pairs work is up to the backends. On CPU, ForwardDiff, ReverseDiff, Zygote and Enzyme all serve a log-density-only model. `AutoMooncake` serves neither direction: it has no reverse-over-reverse, and its gradient rejects an outer pass's tangents. Give Mooncake a hand-written `grad_logdensity` instead. No second-order pair works on GPU yet (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37) and the [GPU page](15-gpu.md)). +Which pairs work is up to the backends. On CPU, ForwardDiff, ReverseDiff, Zygote and Enzyme all serve a log-density-only model. `AutoMooncake` serves neither direction: it has no reverse-over-reverse, and its gradient rejects an outer pass's tangents. Give Mooncake a hand-written `grad_logdensity` instead. On GPU, no DI-driven second-order pair works yet (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)); `AutoReactant()` in both slots is the exception, and compiles the pair to a single XLA program. See the [GPU page](15-gpu.md). The batched pair works the same way, on `sum(logdensity_batch(X))`. That sum's gradient is the stacked per-column gradients only because the columns are independent, so `logdensity_batch` must not couple them. Omitting `grad_logdensity_batch` derives one when `grad_logdensity` is a backend; with a hand-written gradient the batched path stays off and the unbatched update covers it. Both batched derivative slots require `logdensity_batch`, which is also useful on its own for scoring a whole trajectory at once. diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index ea99f64..1f8c5a2 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -233,11 +233,32 @@ DEER needs a Hessian–vector product $H v$ at every Newton step. `DensityModel - **You only supply `gradlogp` / `grad_logdensity_batch`.** The sampler builds the HVP by differentiating your gradient — either a forward-mode pushforward of `gradlogp` ([`ForwardOnGrad`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/src/DEER/DEER.jl), the default for most backends) or a reverse-mode gradient of `x -> dot(gradlogp(x), v)` ([`ReverseOnGrad`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/src/DEER/DEER.jl), used for `AutoMooncake` and `AutoZygote`). This is the **AD-HVP fallback**, and it is what the logistic-regression example above uses. !!! warning "Log-density-only models on GPU" - `grad_logdensity` can itself be an AD backend (`DensityModel(logp, AutoEnzyme(), dim)`, see [Getting started](10-getting-started.md)), but don't do that with `ParallelMALASampler` on GPU. The HVP becomes `SecondOrder(hvp_backend, grad_backend)` on your log-density, which currently fails on GPU with both Enzyme and Mooncake (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)). The same goes for passing a `SecondOrder` explicitly. Write `gradlogp` out by hand for DEER, so the HVP is a single pass over it. The sequential samplers only need the gradient, so log-density-only models work there. + `grad_logdensity` can itself be an AD backend (`DensityModel(logp, AutoEnzyme(), dim)`, see [Getting started](10-getting-started.md)), but don't do that with `ParallelMALASampler` on GPU for the DI-driven backends. The HVP becomes `SecondOrder(hvp_backend, grad_backend)` on your log-density, which currently fails on GPU with both Enzyme and Mooncake (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)). The same goes for passing a `SecondOrder` explicitly. Write `gradlogp` out by hand for DEER, so the HVP is a single pass over it — or use `AutoReactant()` (below), the one backend where log-density-only DEER works on GPU. The sequential samplers only need the gradient, so log-density-only models work there. !!! note "A backend in `grad_logdensity` reaches `logdensity_batch` too" The batched path needs a batched gradient, and derives one from `logdensity_batch` when `grad_logdensity` is a backend. That puts `logdensity_batch` under the same restrictions as the rest of your AD-visible code. Supply `grad_logdensity_batch` to avoid it. +### Reactant: genuine second-order HVPs on GPU + +`ADTypes.AutoReactant()` (requires `using Reactant`) takes a different route entirely. The derivative is traced with Enzyme-MLIR and compiled to an XLA executable by [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl), bypassing Enzyme's LLVM pipeline — and with it the gc-transition abort and the `pmcmc_*` wrapper requirements above. It is currently the only path that computes a genuine second-order HVP on GPU, so log-density-only models work with DEER: + +```julia +using Reactant, ADTypes + +model = DensityModel(logp, AutoReactant(), D) # no hand-written gradient +sampler = ParallelMALASampler(0.005f0; T=16, backend=AutoReactant()) +``` + +When both the gradient slot and the HVP source are `AutoReactant()`, the HVP compiles as explicit forward-over-reverse from the raw log-density, in a single fused XLA program. With a hand-written `gradlogp`, it compiles a forward pushforward of your gradient instead. + +Three caveats: + +- The traced function must be **Reactant-traceable**: plain array operations. DynamicPPL-built log-densities do not trace as-is. +- **Reactant does not mix with the DI backends across the two passes of an HVP.** Both `grad_logdensity` and the HVP source take `AutoReactant()`, or neither does; a hand-written gradient pairs with either. Mixing raises an `ArgumentError` at preparation time, since Reactant cannot trace a DI-prepared gradient and DI cannot differentiate a compiled XLA executable. +- Calls cross a marshalling boundary (package arrays ↔ Reactant's XLA device memory) on every invocation, and executables are shape-specialized at preparation time. Correct everywhere, but it leaves fusion on the table relative to a future end-to-end Reactant pipeline. + +DifferentiationInterface cannot drive Reactant yet; when that support lands ([DI#918](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/918)), this path folds into the standard `hvp_mode` routing with no user-facing change. + ### When the fallback is the right call - **Complex or composed models.** Bayesian neural nets, hierarchical models with many transformations, mixtures, or anything where the Hessian has no convenient closed form. Deriving and maintaining `hvp` by hand for these is error-prone; AD removes a whole class of bugs. diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl new file mode 100644 index 0000000..21df47a --- /dev/null +++ b/ext/ReactantExt.jl @@ -0,0 +1,152 @@ +module ReactantExt + +#= +Reactant-compiled derivative paths, selected by `ADTypes.AutoReactant()` in +any derivative slot of `DensityModel` (or as the sampler `backend`). +Derivatives are traced with Enzyme-MLIR and compiled to XLA executables via +`Reactant.@compile`, bypassing Enzyme's LLVM pipeline — and with it the GPU +`cuMemcpyDtoHAsync_v2` gc-transition abort. This is the only path that +computes a genuine second-order HVP on GPU. + +Requirements / limitations: + - The traced function (`logdensity` / `gradlogp` / batched forms) must be + Reactant-traceable: plain array ops. DynamicPPL-built log-densities do + NOT trace as-is. + - Executables are shape-specialized to the preparation templates, and + arguments are marshalled to/from Reactant's own (XLA) device memory on + every call. A boundary copy, not a fused in-place path — a target for + later optimization. + - HVPs are explicit forward-over-(gradient) compositions; NEVER + `Enzyme.hvp`, which silently returns zeros under `@compile`. + - `AutoReactant.mode` (the wrapped `AutoEnzyme`) is currently ignored: + gradients always trace as Enzyme reverse, HVPs as forward-over-that. +=# + +using ParallelMCMC: ParallelMCMC +using ParallelMCMC.DEER: DEER +using ADTypes: AutoReactant +using Reactant: Reactant, @compile +using Enzyme: Enzyme + +# Host-materialize for the Reactant boundary: Reactant manages its own (XLA) +# device memory, so we round-trip through a plain host Array regardless of +# where the package's arrays live (Vector / CuArray). +_host(x::Array) = x +_host(x::AbstractArray) = Array(x) + +function _from_host(template::AbstractArray, out) + res = similar(template, size(out)) + copyto!(res, Array(out)) + return res +end + +#= +Compile `core` for the template shapes once and return a closure that +marshals package arrays <-> Reactant arrays. +=# +function _compiled(core, t1::AbstractArray) + r1 = Reactant.to_rarray(_host(t1)) + compiled = @compile core(r1) + return x -> _from_host(x, compiled(Reactant.to_rarray(_host(x)))) +end + +function _compiled(core, t1::AbstractArray, t2::AbstractArray) + r1 = Reactant.to_rarray(_host(t1)) + r2 = Reactant.to_rarray(_host(t2)) + compiled = @compile core(r1, r2) + return function (x, v) + out = compiled(Reactant.to_rarray(_host(x)), Reactant.to_rarray(_host(v))) + return _from_host(x, out) + end +end + +# Forward-mode JVP of `g` in direction `v`, i.e. J(g)·v. For g = gradlogp +# this is the HVP H·v. `Const(g)` so Enzyme doesn't treat captures as active. +function _jvp(g, x, v) + return only(Enzyme.autodiff(Enzyme.Forward, Enzyme.Const(g), Enzyme.Duplicated(x, v))) +end + +_rev_gradient(f, x) = Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(f), x)[1] + +# Columns of X are independent samples, so ∇_X sum(logp_batch(X)) stacks the +# per-column gradients. +_sumbatch(f, X) = sum(f(X)) + +#= +Gradient slots. The wrappers keep the raw log-density so the HVP factories +below can re-trace forward-over-reverse from it, instead of trying to trace +through an already-compiled executable. +=# +struct _ReactantGradient{F,C} + logdensity::F + compiled::C +end +(g::_ReactantGradient)(x) = g.compiled(x) + +function ParallelMCMC._reactant_resolve_gradient( + logdensity, backend::AutoReactant, x_template::AbstractVector +) + core = Base.Fix1(_rev_gradient, logdensity) + return _ReactantGradient(logdensity, _compiled(core, x_template)) +end + +struct _ReactantGradientBatch{F,C} + logdensity_batch::F + compiled::C +end +(g::_ReactantGradientBatch)(X) = g.compiled(X) + +function ParallelMCMC._reactant_resolve_gradient_batch( + logdensity_batch, backend::AutoReactant, X_template::AbstractMatrix +) + core = Base.Fix1(_rev_gradient, Base.Fix1(_sumbatch, logdensity_batch)) + return _ReactantGradientBatch(logdensity_batch, _compiled(core, X_template)) +end + +#= +HVP factories. Two tracings depending on where the gradient came from: + + - `_ReactantGradient` (the gradient slot was itself AutoReactant): + re-trace from the raw log-density as explicit forward-over-reverse — + genuine second-order AD fused into one XLA program. + - any other callable (hand-written gradient): forward JVP over it, + provided it is traceable. +=# +function DEER._make_hvp_fn( + ::DEER.ReactantHVP, + gradlogp::_ReactantGradient, + backend::AutoReactant, + x_template::AbstractVector, +) + inner = Base.Fix1(_rev_gradient, gradlogp.logdensity) + core(x, v) = _jvp(inner, x, v) + return _compiled(core, x_template, x_template) +end + +function DEER._make_hvp_fn( + ::DEER.ReactantHVP, gradlogp, backend::AutoReactant, x_template::AbstractVector +) + core(x, v) = _jvp(gradlogp, x, v) + return _compiled(core, x_template, x_template) +end + +function DEER._make_hvp_batch_fn( + ::DEER.ReactantHVP, + grad_batch::_ReactantGradientBatch, + backend::AutoReactant, + X_template::AbstractMatrix, +) + inner = Base.Fix1(_rev_gradient, Base.Fix1(_sumbatch, grad_batch.logdensity_batch)) + core(X, V) = _jvp(inner, X, V) + return _compiled(core, X_template, X_template) +end + +function DEER._make_hvp_batch_fn( + ::DEER.ReactantHVP, grad_batch, backend::AutoReactant, X_template::AbstractMatrix +) + # Column-independent batched gradient ⇒ forward JVP is the columnwise HVP. + core(X, V) = _jvp(grad_batch, X, V) + return _compiled(core, X_template, X_template) +end + +end # module diff --git a/src/DEER/DEER.jl b/src/DEER/DEER.jl index f0d677b..3abe628 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -203,10 +203,23 @@ abstract type HVPStrategy end struct ForwardOnGrad <: HVPStrategy end struct ReverseOnGrad <: HVPStrategy end +#= +ReactantHVP — trace the HVP with Enzyme-MLIR and compile it to an XLA +executable via Reactant.jl (see `ext/ReactantExt.jl`). Selected by +`ADTypes.AutoReactant()`, which DI cannot drive yet, so it short-circuits ahead +of the `hvp_mode` routing above; once DI gains Reactant support the +`AutoReactant` specializations can be deleted. Reactant bypasses Enzyme's LLVM +pipeline, avoiding the GPU gc-transition abort — the only path that computes a +genuine second-order HVP on GPU. The traced functions must be Reactant-traceable +(plain array ops). +=# +struct ReactantHVP <: HVPStrategy end + _strategy_from(::DI.ForwardOverAnything) = ForwardOnGrad() _strategy_from(::DI.HVPMode) = ReverseOnGrad() _hvp_strategy(backend::AbstractADType) = _strategy_from(DI.hvp_mode(backend)) +_hvp_strategy(::ADTypes.AutoReactant) = ReactantHVP() #= Hook for backend-specific normalization of the user's `backend`, applied on every @@ -300,6 +313,26 @@ function _make_hvp_batch_fn( return (X, V) -> _batch_hvp_via_grad_reverse_prepared(prep, X, V) end +#= +`ReactantHVP` fallbacks. `ReactantExt` adds methods with `backend` pinned to +`ADTypes.AutoReactant` (strictly more specific — no method overwriting, which +precompilation forbids); without Reactant loaded these give a clear error +instead of a `MethodError`. +=# +const _REACTANT_LOAD_HINT = "AutoReactant requires Reactant.jl: add `using Reactant` to load ParallelMCMC's ReactantExt." + +function _make_hvp_fn( + ::ReactantHVP, gradlogp, backend::AbstractADType, x_template::AbstractVector +) + return error(_REACTANT_LOAD_HINT) +end + +function _make_hvp_batch_fn( + ::ReactantHVP, grad_batch, backend::AbstractADType, X_template::AbstractMatrix +) + return error(_REACTANT_LOAD_HINT) +end + #= --------------------------------------------------------------------------- Second-order HVP, for a model whose gradient is itself AD-derived. `DI.hvp` diff --git a/src/ParallelMCMC.jl b/src/ParallelMCMC.jl index e98eebb..7041e90 100644 --- a/src/ParallelMCMC.jl +++ b/src/ParallelMCMC.jl @@ -1,7 +1,7 @@ module ParallelMCMC using AbstractMCMC -using ADTypes: AbstractADType +using ADTypes: ADTypes, AbstractADType using CUDA using DifferentiationInterface: DifferentiationInterface using FlexiChains diff --git a/src/interface.jl b/src/interface.jl index 69ede25..4e89cda 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -211,6 +211,29 @@ function _resolve_gradient_batch( ) end +#= +`AutoReactant` gradients bypass DI (which cannot drive Reactant yet) and go +through hook functions that `ReactantExt` fills in with strictly more specific +methods; the untyped fallbacks give a clear load-order error. +=# +function _resolve_gradient( + logdensity, backend::ADTypes.AutoReactant, x_template::AbstractVector +) + return _reactant_resolve_gradient(logdensity, backend, x_template) +end +function _reactant_resolve_gradient(logdensity, backend, x_template) + return error(DEER._REACTANT_LOAD_HINT) +end + +function _resolve_gradient_batch( + logdensity_batch, backend::ADTypes.AutoReactant, X_template::AbstractMatrix +) + return _reactant_resolve_gradient_batch(logdensity_batch, backend, X_template) +end +function _reactant_resolve_gradient_batch(logdensity_batch, backend, X_template) + return error(DEER._REACTANT_LOAD_HINT) +end + #= Resolve an HVP slot given as a backend. `grad_backend` is the backend that produced `grad`, or nothing when the gradient slot held a callable. Dispatch is @@ -220,9 +243,18 @@ statically known. A `SecondOrder` bypasses the gradient slot even when that slot is hand-written: naming both passes asks for two derivatives of `logdensity`. The slot is still the drift term the MALA step uses. + +`AutoReactant` short-circuits ahead of both: DI cannot drive Reactant, so it can +neither build the `SecondOrder` nor route through `hvp_mode`. `ReactantExt` takes +over the second-order case by dispatching on `grad` instead — a Reactant-resolved +gradient carries the raw `logdensity` with it and re-traces forward-over-reverse +from there. =# function _resolve_hvp(logdensity, grad, grad_backend, hvp_backend, x_template) - if hvp_backend isa DI.SecondOrder + _check_reactant_pair(grad_backend, hvp_backend) + if hvp_backend isa ADTypes.AutoReactant + return DEER._make_hvp_fn(DEER.ReactantHVP(), grad, hvp_backend, x_template) + elseif hvp_backend isa DI.SecondOrder return DEER._make_hvp_fn_second_order(logdensity, hvp_backend, x_template) elseif grad_backend !== nothing return DEER._make_hvp_fn_second_order( @@ -239,7 +271,12 @@ end function _resolve_hvp_batch( logdensity_batch, grad_batch, grad_batch_backend, hvp_backend, X_template ) - if hvp_backend isa DI.SecondOrder + _check_reactant_pair(grad_batch_backend, hvp_backend) + if hvp_backend isa ADTypes.AutoReactant + return DEER._make_hvp_batch_fn( + DEER.ReactantHVP(), grad_batch, hvp_backend, X_template + ) + elseif hvp_backend isa DI.SecondOrder return DEER._make_hvp_batch_fn_second_order( _BatchLogdensitySum(logdensity_batch), hvp_backend, X_template ) @@ -256,6 +293,42 @@ function _resolve_hvp_batch( end end +#= +Reactant does not pair with a DI backend across the two passes of an HVP: the +compiled gradient is an opaque XLA executable DI cannot differentiate, and a +DI-prepared gradient is not Reactant-traceable. Both slots take `AutoReactant` +or neither does; a hand-written gradient pairs with either. + +Dispatch rather than a runtime `isa` chain, so the check folds away with the rest +of `_resolve_hvp`'s branching. +=# +_check_reactant_pair(grad_backend, hvp_backend) = nothing +_check_reactant_pair(::ADTypes.AutoReactant, ::ADTypes.AutoReactant) = nothing +_check_reactant_pair(::Nothing, ::ADTypes.AutoReactant) = nothing + +function _check_reactant_pair(grad_backend::ADTypes.AutoReactant, hvp_backend) + return throw( + ArgumentError( + "an AutoReactant gradient needs an AutoReactant Hessian-vector product: " * + "got hvp backend $(hvp_backend). Reactant compiles the gradient to an XLA " * + "executable, which DifferentiationInterface cannot differentiate. Set the " * + "model's `hvp` (or the sampler's `backend`) to AutoReactant() as well.", + ), + ) +end + +function _check_reactant_pair(grad_backend, hvp_backend::ADTypes.AutoReactant) + return throw( + ArgumentError( + "an AutoReactant Hessian-vector product needs an AutoReactant or " * + "hand-written gradient: got gradient backend $(grad_backend). Reactant " * + "traces the HVP from the log-density (or from your gradient) and cannot " * + "trace a DifferentiationInterface-prepared gradient. Set " * + "`grad_logdensity` to AutoReactant() or supply a callable.", + ), + ) +end + """ _prepare_model(model, x_template) -> PreppedDensityModel _prepare_model(model, x_template, T::Int, backend) -> PreppedDensityModel diff --git a/test/Project.toml b/test/Project.toml index 7dbcce5..2892146 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -15,6 +15,7 @@ LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" ParallelMCMC = "1a970f40-4406-51c9-a967-cb3143c111e8" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" diff --git a/test/test-HVP-Strategy.jl b/test/test-HVP-Strategy.jl index e56997e..d28a02b 100644 --- a/test/test-HVP-Strategy.jl +++ b/test/test-HVP-Strategy.jl @@ -24,6 +24,14 @@ const DI_STRAT = ParallelMCMC.DEER.DI DEER_STRAT.ReverseOnGrad end + @testset "AutoReactant short-circuits hvp_mode" begin + #= DI cannot drive Reactant, so `AutoReactant` never reaches `DI.hvp_mode` + and the strategy is picked by dispatch instead. The method lives in DEER + rather than in ReactantExt, so this holds with Reactant unloaded. =# + @test DEER_STRAT._hvp_strategy(AutoReactant()) isa DEER_STRAT.ReactantHVP + @test @inferred(DEER_STRAT._hvp_strategy(AutoReactant())) isa DEER_STRAT.ReactantHVP + end + @testset "SecondOrder follows hvp_mode's composition" begin so_fwd_outer = DI_STRAT.SecondOrder(AutoForwardDiff(), AutoZygote()) @test DEER_STRAT._hvp_strategy(so_fwd_outer) isa DEER_STRAT.ForwardOnGrad diff --git a/test/test-Reactant-HVP.jl b/test/test-Reactant-HVP.jl new file mode 100644 index 0000000..c2b03eb --- /dev/null +++ b/test/test-Reactant-HVP.jl @@ -0,0 +1,156 @@ +using Test +using Random +using LinearAlgebra +using FlexiChains + +using ParallelMCMC +using ADTypes +# Stands in for "some backend that is not Reactant" in the pairing tests below. +using ForwardDiff: ForwardDiff + +#= +Reactant-compiled derivative paths (`AutoReactant`), see ext/ReactantExt.jl. + +The quartic target keeps second-order structure honest: logp = -0.25‖x‖⁴ has +H = -(‖x‖² I + 2 x xᵀ), so an HVP that silently drops the second-order term +(the failure mode of `Enzyme.hvp` under `@compile`) is caught, unlike a +Gaussian where H is constant. +=# +logp_r(x) = -0.25 * sum(abs2, x)^2 +gradlogp_r(x) = -sum(abs2, x) .* x +hvp_r(x, v) = -(sum(abs2, x) .* v .+ 2 .* dot(x, v) .* x) +logp_batch_r(X) = vec(-0.25 .* sum(abs2, X; dims=1) .^ 2) +gradlogp_batch_r(X) = -X .* sum(abs2, X; dims=1) + +const D_R = 4 +const CT_R = FlexiChains.FlexiChain{Symbol} + +#= The pairing rule lives in `_resolve_hvp`, not in the extension, so its +dispatch table is checked whether or not Reactant loads. =# +@testset "Reactant does not pair with a DI backend" begin + # Both slots Reactant, or a hand-written gradient (`nothing`), are accepted. + @test ParallelMCMC._check_reactant_pair(AutoReactant(), AutoReactant()) === nothing + @test ParallelMCMC._check_reactant_pair(nothing, AutoReactant()) === nothing + @test ParallelMCMC._check_reactant_pair(AutoForwardDiff(), AutoForwardDiff()) === + nothing + + # One of each is refused in both directions. + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + AutoReactant(), AutoForwardDiff() + ) + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + AutoForwardDiff(), AutoReactant() + ) +end + +reactant_ok = try + using Reactant: Reactant + using Enzyme: Enzyme + true +catch err + @warn "Reactant not available — skipping Reactant HVP tests" err + false +end + +if reactant_ok + @testset "extension is loaded" begin + @test Base.get_extension(ParallelMCMC, :ReactantExt) !== nothing + end + + #= The same rule reached through `_prepare_model`, where the gradient slot + resolves first: a mixed pair must still surface as an ArgumentError and not + as whatever DI or Reactant would say downstream. =# + @testset "mixed pairs are refused at preparation" begin + x = zeros(D_R) + + reactant_grad = DensityModel(logp_r, AutoReactant(), D_R; hvp=AutoForwardDiff()) + @test_throws ArgumentError ParallelMCMC._prepare_model(reactant_grad, x, 8, nothing) + + reactant_hvp = DensityModel(logp_r, AutoForwardDiff(), D_R; hvp=AutoReactant()) + @test_throws ArgumentError ParallelMCMC._prepare_model(reactant_hvp, x, 8, nothing) + end + + @testset "HVP matches analytic" begin + rng = MersenneTwister(71) + x = randn(rng, D_R) + v = randn(rng, D_R) + + @testset "forward over user gradient (sampler backend)" begin + model = DensityModel(logp_r, gradlogp_r, D_R) + m_p = ParallelMCMC._prepare_model(model, x, 8, AutoReactant()) + @test m_p.hvp(x, v) ≈ hvp_r(x, v) + end + + @testset "forward-over-reverse from logp alone (both slots AutoReactant)" begin + model = DensityModel(logp_r, AutoReactant(), D_R; hvp=AutoReactant()) + m_p = ParallelMCMC._prepare_model(model, x, 8, nothing) + @test m_p.grad_logdensity(x) ≈ gradlogp_r(x) + @test m_p.hvp(x, v) ≈ hvp_r(x, v) + end + + @testset "batched slots" begin + T = 8 + X = randn(rng, D_R, T) + V = randn(rng, D_R, T) + Hv_cols = reduce(hcat, [hvp_r(X[:, t], V[:, t]) for t in 1:T]) + + model = DensityModel( + logp_r, + AutoReactant(), + D_R; + logdensity_batch=logp_batch_r, + grad_logdensity_batch=AutoReactant(), + hvp=AutoReactant(), + hvp_batch=AutoReactant(), + ) + m_p = ParallelMCMC._prepare_model(model, X[:, 1], T, nothing) + @test m_p.grad_logdensity_batch(X) ≈ gradlogp_batch_r(X) + @test m_p.hvp_batch(X, V) ≈ Hv_cols + end + end + + @testset "end-to-end sampling with AutoReactant" begin + model = DensityModel(logp_r, AutoReactant(), D_R) + s = ParallelMALASampler(0.02; T=16, backend=AutoReactant()) + chain = sample(MersenneTwister(72), model, s, 64; chain_type=CT_R, progress=false) + @test size(chain) == (64, 1) + @test all(x -> all(isfinite, x), chain[:x]) + end + + reactant_gpu_ok = try + using CUDA: CUDA + CUDA.functional() && (CUDA.CuArray([1.0f0]); true) + catch + false + end + + if !reactant_gpu_ok + @info "Reactant HVP test: CUDA not functional — skipping CuArray boundary" + else + #= + The compiled executable lives in Reactant's own (XLA) device memory; + what's checked here is the CuArray <-> Reactant marshalling boundary: + CuArray in, CuArray out, values matching the analytic HVP. + =# + logp_r32(x) = -0.25f0 * sum(abs2, x)^2 + + @testset "CuArray boundary" begin + rng = MersenneTwister(73) + x_h = randn(rng, Float32, D_R) + v_h = randn(rng, Float32, D_R) + x_d = CUDA.CuArray(x_h) + v_d = CUDA.CuArray(v_h) + + model = DensityModel(logp_r32, AutoReactant(), D_R; hvp=AutoReactant()) + m_p = ParallelMCMC._prepare_model(model, x_d, 8, nothing) + + g = m_p.grad_logdensity(x_d) + @test g isa CUDA.CuArray + @test Array(g) ≈ gradlogp_r(x_h) + + Hv = m_p.hvp(x_d, v_d) + @test Hv isa CUDA.CuArray + @test Array(Hv) ≈ hvp_r(x_h, v_h) + end + end +end From 07a72188af42fd023224ef20138467add4d5b46e Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Wed, 12 Aug 2026 17:36:06 -0400 Subject: [PATCH 05/11] Some rewrites to text --- CHANGELOG.md | 184 +++++++++++---------- Project.toml | 2 +- docs/src/10-getting-started.md | 4 +- docs/src/15-gpu.md | 22 ++- docs/src/95-reference.md | 9 ++ ext/DynamicPPLExt.jl | 20 +-- ext/EnzymeExt.jl | 19 +-- ext/LogDensityProblemsExt.jl | 6 +- ext/ReactantExt.jl | 167 +++++++++++++------ src/DEER/DEER.jl | 120 ++++++++------ src/ParallelMCMC.jl | 4 + src/interface.jl | 259 ++++++++++++++++++++--------- test/Project.toml | 8 +- test/runtests.jl | 11 ++ test/test-HVP-Strategy.jl | 7 +- test/test-Reactant-HVP.jl | 288 ++++++++++++++++++++++++++++++--- 16 files changed, 796 insertions(+), 334 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 851d2e1..bde0fe4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,115 +10,123 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - The derivative slots of `DensityModel` (`grad_logdensity`, `hvp`, - `grad_logdensity_batch`, `hvp_batch`) now take an `ADTypes.AbstractADType` - in place of a callable, so `DensityModel(logp, AutoForwardDiff(), dim)` - builds a model from the log-density alone. Backends are turned into - prepared DifferentiationInterface callables when sampling starts, and that - preparation is reused across steps (#40, #52). The prepared model rides - along in the sampler state to get that reuse; a state handed to `step` for a - different model, which `initial_state` allows, is re-prepared from the model - passed rather than reused, so the model given to `sample` is the one sampled. -- `ParallelMALASampler`'s `backend` keyword is now optional. It is only the - fallback source of Hessian-vector products, so a `DensityModel` carrying - its own `hvp` / `hvp_batch` does not need it (#52). With no sampler - `backend`, a batched HVP is derived from the model's own `hvp` backend. + `grad_logdensity_batch`, `hvp_batch`) now accept an `ADTypes.AbstractADType` + in place of a callable, so `DensityModel(logp, AutoForwardDiff(), dim)` builds + a model from the log-density alone (#40, #52). Backends become prepared + DifferentiationInterface callables when sampling starts, and the prepared + model rides along in the sampler state so the preparation is reused. A state + handed to `step` for a different model — which `initial_state` allows — is + re-prepared, so the model given to `sample` is the one sampled. +- `ParallelMALASampler`'s `backend` keyword is now optional: it is only the + fallback source of Hessian-vector products, so a `DensityModel` carrying its + own `hvp` / `hvp_batch` does not need one. Without it, a batched HVP comes + from the model's own `hvp` backend (#52). - A `logdensity_batch` given without a `grad_logdensity_batch` now has the - batched gradient derived for it when `grad_logdensity` is a backend, rather - than leaving the batched DEER path switched off (#52). `hvp_batch` can be a - backend in that case too, and differentiates the derived gradient. -- New `ReactantExt`: `ADTypes.AutoReactant()` in any derivative slot, or as the + batched gradient derived from it when `grad_logdensity` is a backend, instead + of leaving the batched DEER path switched off. `hvp_batch` can be a backend in + that case too, and differentiates the derived gradient (#52). +- An HVP backend over an AD-derived gradient is now real second-order AD, + `DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` handed to + `DI.hvp`, rather than an outer AD pass over the prepared DI gradient — which + dropped out of its preparation as soon as tangents were pushed through it + (#37). Around 10x fewer allocations for a `logdensity`-only model. A backend + over a hand-written gradient still differentiates that gradient once. +- `hvp` / `hvp_batch` accept a `SecondOrder` with both halves honoured, so the + log-density is differentiated twice and the gradient slot is not the inner + pass. The inner half used to be discarded. This is the one AD route to an HVP + for a Turing or LogDensityProblems model, whose gradient arrives already + prepared and cannot be differentiated again. +- New `ReactantExt`. `ADTypes.AutoReactant()` in a derivative slot, or as the sampler `backend`, traces the derivative with Enzyme-MLIR and compiles it to - an XLA executable via Reactant.jl. This bypasses Enzyme's LLVM pipeline and - is currently the only path that computes a genuine second-order HVP on GPU, - so `ParallelMALASampler` now works on CUDA for log-density-only models - (#37, #52). Requires `using Reactant` and a Reactant-traceable log-density; - `AutoReactant` cannot be paired with a DifferentiationInterface backend - across the two passes of an HVP, which raises an `ArgumentError`. -- An HVP backend over an AD-derived gradient is now taken as true second-order - AD, `DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` handed - to `DI.hvp`, instead of an outer AD pass over the prepared DI gradient (#37). - That nesting dropped out of its preparation as soon as the outer pass pushed - tangents in, so the composed operator is both what was asked for and cheaper: - around 10x fewer allocations for a `logdensity`-only model. A backend over a - hand-written gradient still differentiates that gradient once, as before. -- `hvp` / `hvp_batch` accept a `SecondOrder` with both halves honoured, meaning - the log-density is differentiated twice and the gradient slot is not the inner - pass. Previously the inner half was silently discarded and only the outer used. - This is the one AD route to an HVP for a Turing or LogDensityProblems model, - whose gradient arrives already prepared and so cannot be differentiated again. + an XLA executable via Reactant.jl, off Enzyme's LLVM pipeline and off + DifferentiationInterface entirely, which yields a genuine second-order HVP for + a log-density-only model (#37, #52). Requires `using Reactant` and a + Reactant-traceable log-density. `AutoReactant` does not pair with a + DifferentiationInterface backend, a `LogDensityProblems` gradient, or a + `SecondOrder` across the two passes of an HVP, and doing so raises an + `ArgumentError` before any AD runs, as does a non-default + `AutoReactant(; mode=...)`. Two caveats, spelled out in `ext/ReactantExt.jl`'s + module docstring and the Reactant section of `docs/src/15-gpu.md`: a traced + log-density must be pure with respect to the data it captures, since a captured + array mutated after preparation stays frozen at its old value in every + derivative compiled from it, and the compiled derivative runs on whatever + device Reactant's XLA client targets, which need not be the GPU the model's + arrays live on. - The `DensityModel` constructors in `DynamicPPLExt` and `LogDensityProblemsExt` forward `logdensity_batch`, `grad_logdensity_batch` and `hvp_batch`, so a Turing or LogDensityProblems model can reach the batched DEER path. Neither provides a batched log-density, so `logdensity_batch` has to be written by hand. - Adds `JuliaFormatter` testing which was forgotten (#60). -### Fixed - -- The reverse-on-grad HVP path differentiated with `DI.inner(backend)` while - its strategy was routed on `DI.outer(backend)`, so a - `DifferentiationInterface.SecondOrder` ran the wrong half of the pair. A - `SecondOrder` now goes to the true second-order path instead of either - strategy, and the half-selecting helper it still uses agrees: normalization - applies to the outer pass. Unwrapping to that half happens before the - normalization hook is dispatched on, so a `SecondOrder(AutoEnzyme(), ...)` still - reaches `EnzymeExt` and gets its function annotation filled in rather than - running as a bare `AutoEnzyme()`. -- Backend normalization no longer picks a differentiation mode on the user's - behalf (#62). `EnzymeExt` pinned `mode=Enzyme.Forward` (with - `set_runtime_activity`) onto an `AutoEnzyme()` left mode-agnostic, on the - grounds that reverse mode hit a gc-transition abort on GPU and that composed - `pmcmc_matmul` calls needed runtime activity. The `pmcmc_*` Enzyme rules keep - Enzyme off both paths on their own now, so the pin bought nothing — and it cost - correctness, because it silently rewrote the direction of a `SecondOrder`'s - outer half. `SecondOrder(AutoEnzyme(), AutoForwardDiff())` is - reverse-over-forward to `hvp_mode`, its inner half being forward-only, and came - out forward-over-forward. Normalization now fills in only - `function_annotation=Enzyme.Const`, which is about this package's own read-only - HVP wrappers rather than about Enzyme's mode, and leaves `mode` exactly as given - — unset included, for DI to resolve from the operator it runs. `hvp_mode` is - therefore identical before and after normalization for every backend pair. - - A mode set explicitly was never overridden, so only mode-agnostic backends were - affected, and the HVP was a correct HVP either way; what changes is that the - composition asked for is the one that runs. The two normalization hooks - (`_hvp_forward_backend`, `_hvp_closure_backend`) collapse into a single - `_normalized_backend`, since without a mode to choose they no longer differ. - Users relying on a plain `AutoEnzyme()` being run forward should now pass - `AutoEnzyme(; mode=Enzyme.Forward)` explicitly. - ### Changed - The AD-HVP fallback strategy (forward-on-grad vs reverse-on-grad) now comes from DifferentiationInterface's `hvp_mode` trait rather than a hardcoded per-backend list, so `AutoEnzyme(mode=Enzyme.Reverse)` routes to the reverse-on-grad path (#38). -- Because a `logdensity_batch` without a `grad_logdensity_batch` now has the - batched gradient derived rather than switching the batched DEER path off, a - model whose `grad_logdensity` is a backend runs the batched update where it - used to run the unbatched one, and AD is applied to its `logdensity_batch`. On - GPU that subjects a function nothing was differentiating before to the - backend's restrictions (`pmcmc_*` wrappers for Enzyme). Supply - `grad_logdensity_batch` to keep AD out of the batched path. A model with a - hand-written `grad_logdensity` is unaffected: nothing derives a batched - gradient for it, so the batched path stays off as before. +- A model whose `grad_logdensity` is a backend now runs the batched update where + it used to run the unbatched one, since a `logdensity_batch` without a + `grad_logdensity_batch` has one derived for it. On GPU that puts the backend's + restrictions (`pmcmc_*` wrappers for Enzyme) on a `logdensity_batch` nothing + was differentiating before; supply `grad_logdensity_batch` to keep AD out of + it. A hand-written `grad_logdensity` is unaffected. - `ParallelMALASampler`'s `backend` no longer derives a batched gradient, only Hessian-vector products. It could previously switch the batched DEER path on - for a model with a hand-written gradient, which made a keyword that reads as - an HVP fallback decide which update path ran and put AD on a - `logdensity_batch` the user had not opted into differentiating. Models that - relied on that should pass `grad_logdensity_batch` explicitly, or a backend in - `grad_logdensity` for one to be derived from. -- Both batched derivative slots now require `logdensity_batch`, which the - batched update evaluates directly, and the constructor rejects them without - one. A callable `grad_logdensity_batch` or `hvp_batch` supplied on its own - used to be accepted and then silently ignored. `logdensity_batch` alone is - still valid and still used to score whole trajectories. + for a model with a hand-written gradient, which let a keyword that reads as an + HVP fallback decide which update path ran. Pass `grad_logdensity_batch` + explicitly, or a backend in `grad_logdensity` to derive one from. +- Both batched derivative slots now require `logdensity_batch`, which the batched + update evaluates directly, and the constructor rejects them without one. A + callable `grad_logdensity_batch` or `hvp_batch` on its own used to be accepted + and then ignored. `logdensity_batch` alone still scores whole trajectories. - An `hvp_batch` that reaches sampling with no batched gradient to pair it with - now raises rather than silently falling back to the unbatched update. + now raises rather than falling back to the unbatched update. + +### Fixed + +- The reverse-on-grad HVP path differentiated with `DI.inner(backend)` while its + strategy was routed on `DI.outer(backend)`, so a + `DifferentiationInterface.SecondOrder` ran the wrong half of the pair. A + `SecondOrder` now goes to the second-order path instead of either strategy, and + normalization applies to the outer half after unwrapping, so + `SecondOrder(AutoEnzyme(), ...)` still reaches `EnzymeExt` rather than running + as a bare `AutoEnzyme()`. +- Backend normalization no longer picks a differentiation mode on the user's + behalf (#62). `EnzymeExt` pinned `mode=Enzyme.Forward`, with + `set_runtime_activity`, onto a mode-agnostic `AutoEnzyme()`, against a + gc-transition abort on GPU and an `EnzymeRuntimeActivityError` on composed + `pmcmc_matmul` calls. The `pmcmc_*` Enzyme rules keep Enzyme off both paths on + their own, so the pin bought nothing and cost correctness: it rewrote the + direction of a `SecondOrder`'s outer half, turning + `SecondOrder(AutoEnzyme(), AutoForwardDiff())` — reverse-over-forward to + `hvp_mode`, its inner half being forward-only — into forward-over-forward. + Normalization now fills in `function_annotation=Enzyme.Const` and leaves `mode` + exactly as given, unset included, so `hvp_mode` reads the same before and after + it for every pair. Only mode-agnostic backends were affected and the HVP was + correct either way; what changes is that the composition asked for is the one + that runs. The two hooks `_hvp_forward_backend` and `_hvp_closure_backend` + collapse into one `_normalized_backend`. Pass + `AutoEnzyme(; mode=Enzyme.Forward)` to keep the old direction. +- `_check_reactant_pair` now runs before `_prepare_model` resolves a gradient, so + a mismatched `AutoReactant` pairing is reported immediately instead of after a + full XLA compile, and `AutoReactant` nested inside a `SecondOrder` is rejected + there too rather than reaching `DI.prepare_hvp`. An `AutoReactant` gradient + composes through the same second-order branch as any other AD-derived gradient, + and `_hvp_strategy(::AutoReactant)` is now live rather than a dead branch. +- `test/test-Reactant-HVP.jl` asserts a posterior mean against a target with a + known mean, and cross-checks against the analytic HVP on the same noise tape. + `size(chain)` and `all(isfinite, ...)` pass even for a badly wrong HVP, since + DEER's Newton iteration then just fails to converge rather than producing + `NaN`s. ### Removed - `DynamicPPLExt` no longer requires `ForwardDiff` as a triggering library to load. +- `Reactant` moved from `test/Project.toml`'s `[deps]` to `[extras]`, and + `test/test-Reactant-HVP.jl` is skipped unless `PARALLELMCMC_TEST_REACTANT` is + set. `Reactant_jll` ships a prebuilt XLA and is a large download most CI runs + and most local `Pkg.test()` calls should not have to pay for. Opt in by setting + the env var and adding `Reactant` to the test environment. ## [0.2.0] - 2026-06-29 diff --git a/Project.toml b/Project.toml index 7f7f77a..40b34e9 100644 --- a/Project.toml +++ b/Project.toml @@ -30,7 +30,7 @@ LogDensityProblemsExt = "LogDensityProblems" ReactantExt = ["Reactant", "Enzyme"] [compat] -ADTypes = "1.22.0" +ADTypes = "1.21.0" AbstractMCMC = "5.10.0" CUDA = "5.11.0" CUDA_Runtime_jll = "0.21" diff --git a/docs/src/10-getting-started.md b/docs/src/10-getting-started.md index b257131..50f4d4a 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -49,11 +49,11 @@ Backends become prepared [DifferentiationInterface](https://github.com/JuliaDiff Naming both passes yourself is the one route that ignores the gradient slot, hand-written or not. It is also the only AD route to an HVP for a Turing or LogDensityProblems model, whose gradient arrives already prepared and cannot be differentiated again. -Which pairs work is up to the backends. On CPU, ForwardDiff, ReverseDiff, Zygote and Enzyme all serve a log-density-only model. `AutoMooncake` serves neither direction: it has no reverse-over-reverse, and its gradient rejects an outer pass's tangents. Give Mooncake a hand-written `grad_logdensity` instead. On GPU, no DI-driven second-order pair works yet (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)); `AutoReactant()` in both slots is the exception, and compiles the pair to a single XLA program. See the [GPU page](15-gpu.md). +Which pairs work is up to the backends. On CPU, ForwardDiff, ReverseDiff, Zygote and Enzyme all serve a log-density-only model. `AutoMooncake` serves neither direction: it has no reverse-over-reverse, and its gradient rejects an outer pass's tangents, so give Mooncake a hand-written `grad_logdensity`. On GPU no DI-driven second-order pair works yet (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)). `AutoReactant()` in both slots is the exception: it skips DI and traces forward-over-reverse straight from the log-density as one compiled XLA program. See the [GPU page](15-gpu.md). The batched pair works the same way, on `sum(logdensity_batch(X))`. That sum's gradient is the stacked per-column gradients only because the columns are independent, so `logdensity_batch` must not couple them. Omitting `grad_logdensity_batch` derives one when `grad_logdensity` is a backend; with a hand-written gradient the batched path stays off and the unbatched update covers it. Both batched derivative slots require `logdensity_batch`, which is also useful on its own for scoring a whole trajectory at once. -`backend` on [`ParallelMALASampler`](@ref) supplies Hessian-vector products for a model that brings no `hvp` / `hvp_batch` of its own, and nothing else. A model carrying its own can leave it out. +`backend` on [`ParallelMALASampler`](@ref) supplies Hessian-vector products for a model that brings no `hvp` / `hvp_batch` of its own, and nothing else, so a model carrying its own can leave it out. --- diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index 1f8c5a2..318477a 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -233,14 +233,14 @@ DEER needs a Hessian–vector product $H v$ at every Newton step. `DensityModel - **You only supply `gradlogp` / `grad_logdensity_batch`.** The sampler builds the HVP by differentiating your gradient — either a forward-mode pushforward of `gradlogp` ([`ForwardOnGrad`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/src/DEER/DEER.jl), the default for most backends) or a reverse-mode gradient of `x -> dot(gradlogp(x), v)` ([`ReverseOnGrad`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/src/DEER/DEER.jl), used for `AutoMooncake` and `AutoZygote`). This is the **AD-HVP fallback**, and it is what the logistic-regression example above uses. !!! warning "Log-density-only models on GPU" - `grad_logdensity` can itself be an AD backend (`DensityModel(logp, AutoEnzyme(), dim)`, see [Getting started](10-getting-started.md)), but don't do that with `ParallelMALASampler` on GPU for the DI-driven backends. The HVP becomes `SecondOrder(hvp_backend, grad_backend)` on your log-density, which currently fails on GPU with both Enzyme and Mooncake (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)). The same goes for passing a `SecondOrder` explicitly. Write `gradlogp` out by hand for DEER, so the HVP is a single pass over it — or use `AutoReactant()` (below), the one backend where log-density-only DEER works on GPU. The sequential samplers only need the gradient, so log-density-only models work there. + `grad_logdensity` can itself be an AD backend (`DensityModel(logp, AutoEnzyme(), dim)`, see [Getting started](10-getting-started.md)), but not with `ParallelMALASampler` on GPU for the DI-driven backends. The HVP then becomes `SecondOrder(hvp_backend, grad_backend)` on your log-density, which currently fails on GPU with both Enzyme and Mooncake (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)); passing a `SecondOrder` explicitly hits the same wall. Write `gradlogp` out by hand so the HVP is a single pass over it, or use `AutoReactant()` below. The sequential samplers only need the gradient, so log-density-only models are fine there. !!! note "A backend in `grad_logdensity` reaches `logdensity_batch` too" The batched path needs a batched gradient, and derives one from `logdensity_batch` when `grad_logdensity` is a backend. That puts `logdensity_batch` under the same restrictions as the rest of your AD-visible code. Supply `grad_logdensity_batch` to avoid it. -### Reactant: genuine second-order HVPs on GPU +### Reactant HVPs, off the DI path -`ADTypes.AutoReactant()` (requires `using Reactant`) takes a different route entirely. The derivative is traced with Enzyme-MLIR and compiled to an XLA executable by [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl), bypassing Enzyme's LLVM pipeline — and with it the gc-transition abort and the `pmcmc_*` wrapper requirements above. It is currently the only path that computes a genuine second-order HVP on GPU, so log-density-only models work with DEER: +`ADTypes.AutoReactant()` (requires `using Reactant`) goes another way. The derivative is traced with Enzyme-MLIR and compiled to an XLA executable by [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl), never touching Enzyme's LLVM pipeline, so neither the gc-transition abort nor the `pmcmc_*` wrappers above apply. It gives a genuine second-order HVP for a log-density-only model: ```julia using Reactant, ADTypes @@ -249,15 +249,21 @@ model = DensityModel(logp, AutoReactant(), D) # no hand-written gradient sampler = ParallelMALASampler(0.005f0; T=16, backend=AutoReactant()) ``` -When both the gradient slot and the HVP source are `AutoReactant()`, the HVP compiles as explicit forward-over-reverse from the raw log-density, in a single fused XLA program. With a hand-written `gradlogp`, it compiles a forward pushforward of your gradient instead. +With both the gradient slot and the HVP source `AutoReactant()`, the HVP compiles as explicit forward-over-reverse from the raw log-density, as its own XLA program. That is a separate executable from the plain gradient's: what gets fused is the HVP's two AD passes, not the gradient and the HVP. With a hand-written `gradlogp` it compiles a forward pushforward of your gradient instead. -Three caveats: +!!! warning "Two silent failure modes" + **Captured data is frozen at compile time.** `Reactant.@compile` bakes any plain `Array` or `Ref` reached through the closure into the executable as a constant. A `logdensity` written as `x -> f(x, data)` whose `data` you later mutate (`data .= new_values`) keeps returning the pre-mutation derivative from every executable compiled before the mutation, with no error and no warning, which biases the samples. Traced functions must be pure with respect to what they capture; pass data that can change in as an argument. + + **"GPU" here describes your array type, not the device Reactant runs on.** Which XLA client Reactant compiles for is a Reactant/`Reactant_jll`-wide setting (`Reactant.set_default_backend`) that nothing in this package controls, and it is `"cpu"` unless a GPU client was selected explicitly. Preparing a model whose `x_template` is a `CuArray` against a CPU client still compiles and still gives correct answers, but every call round-trips `CuArray → host → XLA-CPU → host → CuArray`. `_prepare_model` warns on that combination and can do nothing else about it; point Reactant at a GPU client yourself. + +Other caveats: - The traced function must be **Reactant-traceable**: plain array operations. DynamicPPL-built log-densities do not trace as-is. -- **Reactant does not mix with the DI backends across the two passes of an HVP.** Both `grad_logdensity` and the HVP source take `AutoReactant()`, or neither does; a hand-written gradient pairs with either. Mixing raises an `ArgumentError` at preparation time, since Reactant cannot trace a DI-prepared gradient and DI cannot differentiate a compiled XLA executable. -- Calls cross a marshalling boundary (package arrays ↔ Reactant's XLA device memory) on every invocation, and executables are shape-specialized at preparation time. Correct everywhere, but it leaves fusion on the table relative to a future end-to-end Reactant pipeline. +- **Reactant does not mix with the DI backends across the two passes of an HVP.** Both `grad_logdensity` and the HVP source take `AutoReactant()`, or neither does; a hand-written gradient pairs with either. Mixing raises an `ArgumentError` at preparation time, since Reactant cannot trace a DI-prepared gradient and DI cannot differentiate a compiled XLA executable. Same for a `LogDensityProblems`/Turing gradient, and for nesting `AutoReactant()` inside a `DifferentiationInterface.SecondOrder`. +- **`AutoReactant(; mode=...)` raises** rather than being ignored. Derivatives always trace as Enzyme reverse-mode with the HVP as forward-over-that, whatever `mode` says. +- **Every call crosses a marshalling boundary** between package arrays and Reactant's XLA device memory: two host round-trips and a handful of allocations, not a fused in-place path. Executables are shape-specialized at preparation time, and `@compile` does not memoize across preparations, so every `sample()` call recompiles every `AutoReactant` slot rather than only the first one in the process. -DifferentiationInterface cannot drive Reactant yet; when that support lands ([DI#918](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/918)), this path folds into the standard `hvp_mode` routing with no user-facing change. +DifferentiationInterface cannot interoperate with Reactant. When that support lands ([DI#918](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/918)) this path folds into the ordinary `hvp_mode` routing with no user-facing change. ### When the fallback is the right call diff --git a/docs/src/95-reference.md b/docs/src/95-reference.md index b7f2fb9..aac80ed 100644 --- a/docs/src/95-reference.md +++ b/docs/src/95-reference.md @@ -6,6 +6,15 @@ CurrentModule = ParallelMCMC This page documents all public types and functions exported by ParallelMCMC.jl. +## Reactant + +Loading `Reactant` (`using Reactant`) enables `ADTypes.AutoReactant()` as a +derivative-slot backend on `DensityModel` and as `ParallelMALASampler`'s +`backend`. It is not a `DifferentiationInterface` backend like the others. +[GPU Execution](15-gpu.md) covers what it does, its pairing rules, and its +caveats, in particular that traced log-densities must be pure with respect +to any data they capture. + ## Extension constructors `DensityModel` also has extension constructors for common probabilistic-programming interfaces: diff --git a/ext/DynamicPPLExt.jl b/ext/DynamicPPLExt.jl index 064c6d6..f824f8c 100644 --- a/ext/DynamicPPLExt.jl +++ b/ext/DynamicPPLExt.jl @@ -18,17 +18,19 @@ computation via DynamicPPL's `adtype` interface. Requires `DynamicPPL` and `LogDensityProblems` to be loaded (these are the weak-dependency triggers for this extension), plus any AD backend that is used. -`ad_backend` is DynamicPPL's own `adtype`, not a `DensityModel` slot: it goes to -the `LogDensityFunction` that fills the log-density and gradient slots, which is -why it takes a backend only and never a callable. The rest are `DensityModel` -slots forwarded unchanged. +`ad_backend` is DynamicPPL's own `adtype` rather than a `DensityModel` slot: it +goes to the `LogDensityFunction` that fills the log-density and gradient slots, +which is why it takes a backend and never a callable. The remaining keywords are +`DensityModel` slots, forwarded unchanged. -`ParallelMALASampler` also needs an HVP. Give it a callable or a +`ParallelMALASampler` also wants an HVP. Give it a callable or a `DifferentiationInterface.SecondOrder`, which differentiates the log-density and -so bypasses DynamicPPL's gradient. A plain backend fails, since it would -differentiate the gradient `ad_backend` produced and that preparation rejects an -outer pass's tangents. DynamicPPL supplies no batched log-density either, so -reaching the batched DEER path means writing `logdensity_batch` by hand. +so bypasses DynamicPPL's gradient. A plain backend fails: it would differentiate +the gradient `ad_backend` produced, whose preparation rejects an outer pass's +tangents. `ADTypes.AutoReactant()` fails as well, bypassing DI or not, since +Reactant cannot trace DynamicPPL's model evaluation. DynamicPPL supplies no +batched log-density either, so reaching the batched DEER path means writing +`logdensity_batch` by hand. # Example ```julia diff --git a/ext/EnzymeExt.jl b/ext/EnzymeExt.jl index 82e63f8..33899fd 100644 --- a/ext/EnzymeExt.jl +++ b/ext/EnzymeExt.jl @@ -23,19 +23,14 @@ using Enzyme.EnzymeCore.EnzymeRules: Normalization of the user's `AutoEnzyme` for DEER's AD-HVP paths: fill in `function_annotation=Enzyme.Const` when they left it open, so Enzyme doesn't throw `EnzymeMutabilityException` on the read-only `_HvpReverseClosure` / -`_BatchHvpReverseClosure` wrappers, which capture `gradlogp`. Those wrapper types -belong to this package, so declaring them constant is this package's business. +`_BatchHvpReverseClosure` wrappers that capture `gradlogp`. -`mode` is passed through exactly as given, unset included. Choosing a direction -on the user's behalf is not our call: a mode they set is a decision, and an unset -one is DI's to resolve from the operator it runs. - -An earlier version pinned `mode=Enzyme.Forward` here (with -`set_runtime_activity`) against the gc-transition abort on GPU and -`EnzymeRuntimeActivityError` on composed `pmcmc_matmul` calls. The rules below -keep Enzyme off both paths on their own, so the pin bought nothing and cost -correctness: it silently rewrote the direction of a `SecondOrder`'s outer half -(see `DEER._normalized_backend`). +`mode` passes through exactly as given, unset included. An earlier version pinned +`mode=Enzyme.Forward` here, with `set_runtime_activity`, against the GPU +gc-transition abort and the `EnzymeRuntimeActivityError` on composed +`pmcmc_matmul` calls. The rules below keep Enzyme off both paths by themselves, +so the pin bought nothing and cost correctness: it rewrote the direction of a +`SecondOrder`'s outer half (see `DEER._normalized_backend`). =# function DEER._normalized_backend(backend::ADTypes.AutoEnzyme{M,A}) where {M,A} A === Nothing || return backend diff --git a/ext/LogDensityProblemsExt.jl b/ext/LogDensityProblemsExt.jl index 6f59578..b146b2b 100644 --- a/ext/LogDensityProblemsExt.jl +++ b/ext/LogDensityProblemsExt.jl @@ -24,8 +24,10 @@ parameter named `:x` will be chosen, unless you also pass `param_names` to `samp and keep their meaning there, with one caveat. `ld` fills the gradient slot, and a gradient `ld` computes by AD carries a preparation tied to its input type, so it rejects the tangents a plain `hvp` backend would push through it. Use a -callable, or a `DifferentiationInterface.SecondOrder` which differentiates the -log-density instead. Same for `hvp_batch`. +callable, or a `DifferentiationInterface.SecondOrder`, which differentiates the +log-density instead. Same for `hvp_batch`. `ADTypes.AutoReactant()` is out here +too, DI or no DI: Reactant cannot trace the DynamicPPL/LogDensityProblems +machinery `ld`'s gradient dispatches into. `ld` supplies no batched log-density, so reaching the batched DEER path means writing `logdensity_batch` by hand. diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl index 21df47a..6c14c01 100644 --- a/ext/ReactantExt.jl +++ b/ext/ReactantExt.jl @@ -1,12 +1,30 @@ module ReactantExt #= -Reactant-compiled derivative paths, selected by `ADTypes.AutoReactant()` in -any derivative slot of `DensityModel` (or as the sampler `backend`). -Derivatives are traced with Enzyme-MLIR and compiled to XLA executables via -`Reactant.@compile`, bypassing Enzyme's LLVM pipeline — and with it the GPU -`cuMemcpyDtoHAsync_v2` gc-transition abort. This is the only path that -computes a genuine second-order HVP on GPU. +Reactant-compiled derivative paths, selected by `ADTypes.AutoReactant()` in any +derivative slot of `DensityModel`, or as the sampler `backend`. Derivatives are +traced with Enzyme-MLIR and compiled to XLA executables by `Reactant.@compile`, +which keeps them off Enzyme's LLVM pipeline and so off the GPU +`cuMemcpyDtoHAsync_v2` gc-transition abort, and off DifferentiationInterface +entirely. For a log-density-only model it yields a genuine second-order HVP. + +Two silent failure modes, ahead of the ordinary limitations: + +Captured data is frozen at compile time. `@compile` bakes any plain `Array` / +`Ref` reached through the traced closure into the executable as a constant. A +`logdensity` written as `x -> f(x, data)` whose `data` is mutated afterwards +keeps handing back the pre-mutation derivative from every executable compiled +before the mutation, with no error and no warning. Traced functions must be pure +with respect to what they capture; pass data that can change in as an argument. + +The compiled program runs wherever Reactant's XLA client points, which need not +be a GPU. That client is a Reactant/`Reactant_jll`-wide setting +(`Reactant.set_default_backend`) and has nothing to do with where the package's +own `Vector`s / `CuArray`s live. Preparing on `CuArray` parameters while the +client targets `"cpu"` still compiles and still gives correct answers, but every +call round-trips `CuArray -> host -> XLA-CPU -> host -> CuArray`. +`_warn_reactant_host_roundtrip` below warns once per compiled slot on that +combination; it cannot fix it. Requirements / limitations: - The traced function (`logdensity` / `gradlogp` / batched forms) must be @@ -15,28 +33,83 @@ Requirements / limitations: - Executables are shape-specialized to the preparation templates, and arguments are marshalled to/from Reactant's own (XLA) device memory on every call. A boundary copy, not a fused in-place path — a target for - later optimization. + later optimization. `@compile` also does not memoize across preparations: + every `sample()` call recompiles every `AutoReactant` slot the model uses. - HVPs are explicit forward-over-(gradient) compositions; NEVER `Enzyme.hvp`, which silently returns zeros under `@compile`. - - `AutoReactant.mode` (the wrapped `AutoEnzyme`) is currently ignored: - gradients always trace as Enzyme reverse, HVPs as forward-over-that. + - `AutoReactant.mode` (the wrapped `AutoEnzyme`) is not honoured: gradients + always trace as Enzyme reverse, HVPs as forward-over-that. A non-default + `mode` is rejected outright (`_check_reactant_mode`) rather than ignored. =# +# Both `Reactant` and `Enzyme` trigger this extension (see Project.toml): the +# traced derivatives call `Enzyme.autodiff` / `Enzyme.gradient` themselves +# rather than reaching Enzyme-MLIR through Reactant. using ParallelMCMC: ParallelMCMC using ParallelMCMC.DEER: DEER -using ADTypes: AutoReactant +using ADTypes: ADTypes, AutoReactant, AutoEnzyme using Reactant: Reactant, @compile using Enzyme: Enzyme +#= +`AutoReactant()` defaults to `AutoReactant(; mode=AutoEnzyme())`, i.e. +`AutoReactant{AutoEnzyme{Nothing,Nothing}}` — the exact type matched below. +Anything else names an Enzyme mode or annotation, and this extension honours +neither. +=# +_check_reactant_mode(::AutoReactant{AutoEnzyme{Nothing,Nothing}}) = nothing +function _check_reactant_mode(backend::AutoReactant) + return throw( + ArgumentError( + "AutoReactant(; mode=$(backend.mode)) is not supported: gradients always " * + "trace as Enzyme reverse-mode and HVPs as forward-over-that, whatever " * + "`mode` says. Use the default AutoReactant().", + ), + ) +end + +#= +Warn once per compiled slot when the template is not a plain `Array` (so looks +like it lives on a GPU) while Reactant's default XLA client targets "cpu": every +call then pays a host round trip on top of the usual marshalling. Guarded, so a +Reactant version without `XLA.platform_name` / `XLA.default_backend` degrades to +no warning instead of erroring out of `_prepare_model`. +=# +function _reactant_client_platform() + return try + string(Reactant.XLA.platform_name(Reactant.XLA.default_backend())) + catch + nothing + end +end + +_warn_reactant_host_roundtrip(::Array) = nothing +function _warn_reactant_host_roundtrip(x::AbstractArray) + if _reactant_client_platform() == "cpu" + @warn "AutoReactant: preparing on a $(typeof(x)), but Reactant's default XLA " * + "client targets \"cpu\". Every call will round-trip to the host and back " * + "instead of running where the array lives, which is likely slower than not " * + "using Reactant at all. Point Reactant at a GPU client with " * + "`Reactant.set_default_backend(\"gpu\")` if one is available." maxlog = 1 + end + return nothing +end + # Host-materialize for the Reactant boundary: Reactant manages its own (XLA) # device memory, so we round-trip through a plain host Array regardless of -# where the package's arrays live (Vector / CuArray). +# where the package's arrays live (Vector / CuArray / SubArray). _host(x::Array) = x _host(x::AbstractArray) = Array(x) +#= +Eltype comes from `out`, what Reactant actually computed, not from `template`: a +traced computation that promotes internally would otherwise be narrowed back to +the template's eltype on the way out. +=# function _from_host(template::AbstractArray, out) - res = similar(template, size(out)) - copyto!(res, Array(out)) + out_h = Array(out) + res = similar(template, eltype(out_h), size(out_h)) + copyto!(res, out_h) return res end @@ -45,12 +118,14 @@ Compile `core` for the template shapes once and return a closure that marshals package arrays <-> Reactant arrays. =# function _compiled(core, t1::AbstractArray) + _warn_reactant_host_roundtrip(t1) r1 = Reactant.to_rarray(_host(t1)) compiled = @compile core(r1) return x -> _from_host(x, compiled(Reactant.to_rarray(_host(x)))) end function _compiled(core, t1::AbstractArray, t2::AbstractArray) + _warn_reactant_host_roundtrip(t1) r1 = Reactant.to_rarray(_host(t1)) r2 = Reactant.to_rarray(_host(t2)) compiled = @compile core(r1, r2) @@ -68,17 +143,11 @@ end _rev_gradient(f, x) = Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(f), x)[1] -# Columns of X are independent samples, so ∇_X sum(logp_batch(X)) stacks the -# per-column gradients. -_sumbatch(f, X) = sum(f(X)) - #= -Gradient slots. The wrappers keep the raw log-density so the HVP factories -below can re-trace forward-over-reverse from it, instead of trying to trace -through an already-compiled executable. +Gradient slots. They hold only the compiled callable; the HVP factories below +get `logdensity` from `_resolve_hvp`, which already has it. =# -struct _ReactantGradient{F,C} - logdensity::F +struct _ReactantGradient{C} compiled::C end (g::_ReactantGradient)(x) = g.compiled(x) @@ -86,12 +155,12 @@ end function ParallelMCMC._reactant_resolve_gradient( logdensity, backend::AutoReactant, x_template::AbstractVector ) + _check_reactant_mode(backend) core = Base.Fix1(_rev_gradient, logdensity) - return _ReactantGradient(logdensity, _compiled(core, x_template)) + return _ReactantGradient(_compiled(core, x_template)) end -struct _ReactantGradientBatch{F,C} - logdensity_batch::F +struct _ReactantGradientBatch{C} compiled::C end (g::_ReactantGradientBatch)(X) = g.compiled(X) @@ -99,26 +168,32 @@ end function ParallelMCMC._reactant_resolve_gradient_batch( logdensity_batch, backend::AutoReactant, X_template::AbstractMatrix ) - core = Base.Fix1(_rev_gradient, Base.Fix1(_sumbatch, logdensity_batch)) - return _ReactantGradientBatch(logdensity_batch, _compiled(core, X_template)) + _check_reactant_mode(backend) + # `_BatchLogdensitySum` is the same sum-over-columns the DI-driven batched + # gradient uses (src/interface.jl), so both batched paths differentiate the + # same thing. + core = Base.Fix1(_rev_gradient, ParallelMCMC._BatchLogdensitySum(logdensity_batch)) + return _ReactantGradientBatch(_compiled(core, X_template)) end #= -HVP factories. Two tracings depending on where the gradient came from: - - - `_ReactantGradient` (the gradient slot was itself AutoReactant): - re-trace from the raw log-density as explicit forward-over-reverse — - genuine second-order AD fused into one XLA program. - - any other callable (hand-written gradient): forward JVP over it, - provided it is traceable. +HVP factories. `_resolve_hvp` / `_resolve_hvp_batch` (src/interface.jl) route to +one of two shapes, the same two every other backend gets. + + - Both `grad_logdensity` and the HVP source `AutoReactant`: the AD-derived + gradient case, with `_second_order` collapsing the pair to one + `AutoReactant()` since DI cannot form a `SecondOrder` from it. + `_make_hvp_fn_second_order` here traces forward-over-reverse from + `logdensity` as a single XLA program. + - A hand-written `gradlogp` with an `AutoReactant` HVP source: routed by + `_hvp_strategy(::AutoReactant) = ReactantHVP()` in `DEER.jl` to + `_make_hvp_fn` below, a forward JVP over that callable. =# -function DEER._make_hvp_fn( - ::DEER.ReactantHVP, - gradlogp::_ReactantGradient, - backend::AutoReactant, - x_template::AbstractVector, +function DEER._make_hvp_fn_second_order( + logdensity, backend::AutoReactant, x_template::AbstractVector ) - inner = Base.Fix1(_rev_gradient, gradlogp.logdensity) + _check_reactant_mode(backend) + inner = Base.Fix1(_rev_gradient, logdensity) core(x, v) = _jvp(inner, x, v) return _compiled(core, x_template, x_template) end @@ -126,17 +201,16 @@ end function DEER._make_hvp_fn( ::DEER.ReactantHVP, gradlogp, backend::AutoReactant, x_template::AbstractVector ) + _check_reactant_mode(backend) core(x, v) = _jvp(gradlogp, x, v) return _compiled(core, x_template, x_template) end -function DEER._make_hvp_batch_fn( - ::DEER.ReactantHVP, - grad_batch::_ReactantGradientBatch, - backend::AutoReactant, - X_template::AbstractMatrix, +function DEER._make_hvp_batch_fn_second_order( + logdensity_batch_sum, backend::AutoReactant, X_template::AbstractMatrix ) - inner = Base.Fix1(_rev_gradient, Base.Fix1(_sumbatch, grad_batch.logdensity_batch)) + _check_reactant_mode(backend) + inner = Base.Fix1(_rev_gradient, logdensity_batch_sum) core(X, V) = _jvp(inner, X, V) return _compiled(core, X_template, X_template) end @@ -145,6 +219,7 @@ function DEER._make_hvp_batch_fn( ::DEER.ReactantHVP, grad_batch, backend::AutoReactant, X_template::AbstractMatrix ) # Column-independent batched gradient ⇒ forward JVP is the columnwise HVP. + _check_reactant_mode(backend) core(X, V) = _jvp(grad_batch, X, V) return _compiled(core, X_template, X_template) end diff --git a/src/DEER/DEER.jl b/src/DEER/DEER.jl index 3abe628..07919ea 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -164,7 +164,7 @@ We bundle the closure with the prep so `prepare_gradient` and `gradient` see the same function instance (DI keys preparations on function identity). --------------------------------------------------------------------------- =# -import ..ParallelMCMC: pmcmc_dot, pmcmc_dotsum +import ..ParallelMCMC: pmcmc_dot, pmcmc_dotsum, _REACTANT_LOAD_HINT struct _HvpReverseClosure{F} grad::F @@ -177,17 +177,15 @@ end (c::_BatchHvpReverseClosure)(X, V) = pmcmc_dotsum(c.grad_batch(X), V) #= -Pick the AD-HVP fallback strategy from the user's backend. These two apply when -the HVP is one AD pass over a gradient we already have, i.e., a hand-written -`gradlogp`, which neither of them differentiates twice: +Pick the AD-HVP fallback strategy from the user's backend. Both are one AD pass +over a hand-written `gradlogp`; an AD-derived gradient goes to +`_make_hvp_fn_second_order` instead. ForwardOnGrad() — `pushforward(gradlogp, x, v)`. Routes through the `pmcmc_matmul` frule. ReverseOnGrad() — `gradient(x -> pmcmc_dot(gradlogp(x), v))`. Routes through the matmul and dot/sum rrules. -An AD-derived gradient takes neither and goes to `_make_hvp_fn_second_order`. - These are singleton types rather than symbols so the choice dispatches statically — `_make_hvp_fn(_hvp_strategy(backend), ...)` resolves to one concrete method (and one concrete return type) at compile time, without @@ -204,14 +202,17 @@ struct ForwardOnGrad <: HVPStrategy end struct ReverseOnGrad <: HVPStrategy end #= -ReactantHVP — trace the HVP with Enzyme-MLIR and compile it to an XLA -executable via Reactant.jl (see `ext/ReactantExt.jl`). Selected by -`ADTypes.AutoReactant()`, which DI cannot drive yet, so it short-circuits ahead -of the `hvp_mode` routing above; once DI gains Reactant support the -`AutoReactant` specializations can be deleted. Reactant bypasses Enzyme's LLVM -pipeline, avoiding the GPU gc-transition abort — the only path that computes a -genuine second-order HVP on GPU. The traced functions must be Reactant-traceable -(plain array ops). +ReactantHVP traces the HVP with Enzyme-MLIR and compiles it to an XLA +executable, off Enzyme's LLVM pipeline and so off the GPU gc-transition abort. +DI cannot drive Reactant, so `AutoReactant` short-circuits the `hvp_mode` +routing above; drop these specializations once it can. What the traced function +has to look like, and which device the compiled program actually runs on, are in +`ext/ReactantExt.jl`'s module docstring. + +Only reached over a hand-written `gradlogp`. An `AutoReactant` gradient slot is +an AD-derived gradient like any other and goes to `_make_hvp_fn_second_order` / +`_make_hvp_batch_fn_second_order`, which dispatch on the backend rather than on +the resolved gradient's type (see `_resolve_hvp`). =# struct ReactantHVP <: HVPStrategy end @@ -222,26 +223,21 @@ _hvp_strategy(backend::AbstractADType) = _strategy_from(DI.hvp_mode(backend)) _hvp_strategy(::ADTypes.AutoReactant) = ReactantHVP() #= -Hook for backend-specific normalization of the user's `backend`, applied on every -AD-HVP path before the backend reaches DI. - -It supplies what the wrappers DEER differentiates need, and nothing else. Those -wrapper types are ours, so annotating them is ours to do: EnzymeExt specializes -this to fill `function_annotation=Enzyme.Const`, without which Enzyme throws +Hook for backend-specific normalization, applied on every AD-HVP path before the +backend reaches DI. It fills in what the wrappers DEER differentiates need and +nothing else: those wrapper types are ours, so EnzymeExt sets +`function_annotation=Enzyme.Const` on them, without which Enzyme throws `EnzymeMutabilityException` on the read-only `_HvpReverseClosure` / -`_BatchHvpReverseClosure`, which capture `gradlogp`. - -It deliberately does not choose a differentiation mode. Which direction a pass -runs is the user's call when they state one and DI's to resolve from the operator -when they don't; this package is not an AD package and has no business overriding -either. Picking one here also used to corrupt a `SecondOrder`, whose halves carry -directions of their own (see `_normalized_second_order`). - -Callers hand this a single pass, never a `SecondOrder`: the strategy paths below -run one AD pass over a hand-written `gradlogp`, and `_resolve_hvp` sends every -`SecondOrder` to `_make_hvp_fn_second_order` before they are reached. -`_normalized_second_order` is the one caller that starts from a pair, and it -selects the outer half itself. +`_BatchHvpReverseClosure` that capture `gradlogp`. + +It does not choose a differentiation mode. A mode the user set is a decision, an +unset one is DI's to resolve from the operator it runs, and substituting one here +used to rewrite the outer half of a `SecondOrder` out from under `hvp_mode` (see +`_normalized_second_order`). + +Only ever handed a single pass. `_resolve_hvp` sends every `SecondOrder` to +`_make_hvp_fn_second_order` before the strategy paths below are reached, and +`_normalized_second_order` picks the outer half itself. =# _normalized_backend(backend::AbstractADType) = backend @@ -314,13 +310,11 @@ function _make_hvp_batch_fn( end #= -`ReactantHVP` fallbacks. `ReactantExt` adds methods with `backend` pinned to -`ADTypes.AutoReactant` (strictly more specific — no method overwriting, which -precompilation forbids); without Reactant loaded these give a clear error -instead of a `MethodError`. +`ReactantHVP` fallbacks, so a missing `using Reactant` gives the load hint +rather than a `MethodError`. `ReactantExt` pins `backend` to +`ADTypes.AutoReactant`, which is strictly more specific, so nothing is +overwritten — precompilation forbids that. =# -const _REACTANT_LOAD_HINT = "AutoReactant requires Reactant.jl: add `using Reactant` to load ParallelMCMC's ReactantExt." - function _make_hvp_fn( ::ReactantHVP, gradlogp, backend::AbstractADType, x_template::AbstractVector ) @@ -333,23 +327,43 @@ function _make_hvp_batch_fn( return error(_REACTANT_LOAD_HINT) end +#= +Fallback for the "both slots `AutoReactant`" second-order path, which +`_second_order` in `interface.jl` routes here with `backend::AutoReactant` +rather than a `DI.SecondOrder`. Signature is `AbstractADType` and not +`AutoReactant` because `ReactantExt`'s method is `AutoReactant` exactly, and +precompilation refuses an identical signature; less specific still loses to it. +JET needs a method here too, since it cannot see a conditionally-loaded +extension and would otherwise flag `_resolve_hvp`'s `AutoReactant` branch as +having none. +=# +function _make_hvp_fn_second_order( + logdensity, backend::AbstractADType, x_template::AbstractVector +) + return error(_REACTANT_LOAD_HINT) +end + +function _make_hvp_batch_fn_second_order( + logdensity_batch_sum, backend::AbstractADType, X_template::AbstractMatrix +) + return error(_REACTANT_LOAD_HINT) +end + #= --------------------------------------------------------------------------- Second-order HVP, for a model whose gradient is itself AD-derived. `DI.hvp` -takes both passes over the log-density, so these never touch the gradient slot. -Preferred over pushing tangents through a prepared DI gradient, which drops out -of its preparation once the outer pass hands it an unexpected tangent type. - -Both halves are passed to `DI.hvp` as the user composed them, so the direction -each one runs in is theirs and DI's, not ours. Normalization touches only the -outer half, and only to fill in annotations for the wrappers being -differentiated; because it never substitutes a mode, `DI.hvp_mode` of the pair is -the same before and after. The inner half is a plain first-order gradient over -the user's own `logdensity` and is passed straight through. - -The batched form differentiates `sum(logdensity_batch(X))`, whose Hessian is -block-diagonal by column independence, so its HVP along `V` is the columnwise -HVP. Same argument the batched gradient rests on. +takes both passes over the log-density, so the gradient slot is never touched. +The alternative — pushing tangents through the prepared DI gradient — drops out +of that preparation the moment the outer pass hands it an unexpected tangent +type. + +Both halves reach `DI.hvp` as the user composed them. Normalization touches the +outer one, and only to fill in annotations, so `DI.hvp_mode` of the pair reads +the same before and after; the inner half is a first-order gradient over the +user's own `logdensity` and goes through untouched. + +The batched form differentiates `sum(logdensity_batch(X))`. Column independence +makes that Hessian block-diagonal, so its HVP along `V` is the columnwise HVP. --------------------------------------------------------------------------- =# function _normalized_second_order(backend::DI.SecondOrder) diff --git a/src/ParallelMCMC.jl b/src/ParallelMCMC.jl index 7041e90..126e64f 100644 --- a/src/ParallelMCMC.jl +++ b/src/ParallelMCMC.jl @@ -28,6 +28,10 @@ pmcmc_matmul(A::AbstractVecOrMat, B::AbstractVecOrMat) = A * B pmcmc_dot(a::AbstractVector, b::AbstractVector) = dot(a, b) pmcmc_dotsum(A::AbstractVecOrMat, B::AbstractVecOrMat) = sum(A .* B) +#= Lives here rather than in `DEER` because both DEER's `ReactantHVP` fallbacks +and `interface.jl`'s gradient hooks report it. =# +const _REACTANT_LOAD_HINT = "AutoReactant requires Reactant.jl: add `using Reactant` to load ParallelMCMC's ReactantExt." + include("MALA/MALA.jl") include("DEER/DEERScan.jl") include("DEER/DEER.jl") diff --git a/src/interface.jl b/src/interface.jl index 4e89cda..b664b5b 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -11,20 +11,28 @@ Defines model/sampler/state/transition types and implements Wraps a log-density function, its gradient, and optional Hessian-vector product helpers for use with ParallelMCMC samplers. -The derivative slots (`grad_logdensity`, `hvp`, `grad_logdensity_batch`, -`hvp_batch`) take either a callable or an `ADTypes.AbstractADType`. Backends -are turned into prepared DifferentiationInterface callables when sampling -starts, and any AD failure surfaces there. So a model needs nothing beyond -the log-density: +Each derivative slot (`grad_logdensity`, `hvp`, `grad_logdensity_batch`, +`hvp_batch`) takes a callable or an `ADTypes.AbstractADType`, so a model can be +built from the log-density alone: DensityModel(logp, AutoForwardDiff(), dim) -How a backend in `hvp` / `hvp_batch` gets its second derivative depends on the -gradient slot. Over a hand-written gradient it is a single AD pass across your -own code. Over an AD-derived one it is +Backends become prepared DifferentiationInterface callables when sampling starts, +and an AD failure surfaces there. `ADTypes.AutoReactant()` is the one backend DI +cannot drive; it is traced with Enzyme-MLIR and compiled to an XLA executable by +Reactant.jl instead, which needs `using Reactant` and brings requirements of its +own — see the GPU guide, `docs/src/15-gpu.md`, and `ext/ReactantExt.jl`'s module +docstring. + +A backend in `hvp` / `hvp_batch` over a hand-written gradient is a single AD pass +across your own code. Over an AD-derived one it is `DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)`, taken through DI's second-order operator. Passing a `SecondOrder` yourself always means the latter, and bypasses the gradient slot even when you wrote it by hand. +`AutoReactant` sits outside both: it cannot go inside a `SecondOrder`, and +`grad_logdensity` and `hvp`/`hvp_batch` must either both be `AutoReactant()` or +neither, since mixing it with a DifferentiationInterface backend across the two +passes of an HVP raises an `ArgumentError` at preparation time. - `logdensity(x::AbstractVector) -> Real` - `grad_logdensity` — callable `x -> AbstractVector`, or a backend to @@ -49,12 +57,11 @@ latter, and bypasses the gradient slot even when you wrote it by hand. information. Both batched derivative slots require `logdensity_batch`, which the batched -update evaluates directly. `logdensity_batch` alone is allowed and scores whole -trajectories at once without switching the batched update on. - -`ParallelMALASampler` runs the batched DEER update once it has a -`logdensity_batch` and a batched gradient: either `grad_logdensity_batch`, or one -derived from `logdensity_batch` when `grad_logdensity` is a backend. +update evaluates directly. `ParallelMALASampler` runs that update once it has a +`logdensity_batch` and a batched gradient — `grad_logdensity_batch`, or one +derived from `logdensity_batch` when `grad_logdensity` is a backend. A +`logdensity_batch` on its own is fine too: it scores whole trajectories at once +and leaves the batched update off. """ struct DensityModel{F,G,H,FB,GB,HB,PN} <: AbstractMCMC.AbstractModel logdensity::F @@ -95,15 +102,14 @@ function DensityModel( "grad_logdensity must be a callable or an ADTypes.AbstractADType backend" ), ) - #= The batched DEER update evaluates `logdensity_batch` itself, so neither - batched derivative is usable without one: a backend would have nothing to - differentiate, and a callable would never be reached. Rejected here rather - than silently ignored. A `logdensity_batch` on its own is allowed, and - `_prepare_model` decides from the gradient slot whether the path can run. =# + #= The batched DEER update evaluates `logdensity_batch` itself, so without one + a batched derivative slot is unusable either way: a backend has nothing to + differentiate and a callable is never reached. Rejected rather than ignored. + A `logdensity_batch` on its own is fine; `_prepare_model` decides from the + gradient slot whether the batched path can run. =# _batch_needs_logp(name) = throw( ArgumentError( - "$name requires logdensity_batch: the batched DEER path evaluates the " * - "batched log-density, so it cannot run without one", + "$name requires logdensity_batch, which the batched DEER update evaluates" ), ) if logdensity_batch === nothing @@ -123,12 +129,13 @@ function DensityModel( end """ -A [`DensityModel`](@ref) with its backend slots resolved to prepared -DifferentiationInterface callables. `_prepare_model` builds these, and the -sampler internals take them rather than a `DensityModel`, so no slot of one -of these ever holds an `AbstractADType`. Which slots are filled depends on -the sampler: the sequential samplers only need `grad_logdensity` and get -`nothing` for the DEER-only slots, DEER fills the rest. +A [`DensityModel`](@ref) with its backend slots resolved to prepared callables — +DifferentiationInterface ones, or a compiled Reactant/XLA executable for +`AutoReactant()`. `_prepare_model` builds these, and the sampler internals take +them rather than a `DensityModel`, so no slot here ever holds an +`AbstractADType`. Which slots are filled depends on the sampler: the sequential +samplers only need `grad_logdensity` and get `nothing` for the DEER-only slots, +DEER fills the rest. `source` is the `DensityModel` this was prepared from. A sampler state carries a prepped model so the preparation is reused across steps, and `initial_state` @@ -158,8 +165,7 @@ _prepped_for(prepped::PreppedDensityModel, model::DensityModel) = prepped.source #= Resolved gradient wrappers. Structs rather than anonymous closures since DI keys preparations on function identity. `TX` is the input type the prep was made for; -anything else falls back to unprepared `DI.gradient` rather than failing. In a -normal run the prepared branch is the one that fires. +anything else falls back to an unprepared `DI.gradient` rather than failing. =# struct _ADGradient{F,B<:AbstractADType,P,TX} logdensity::F @@ -222,7 +228,7 @@ function _resolve_gradient( return _reactant_resolve_gradient(logdensity, backend, x_template) end function _reactant_resolve_gradient(logdensity, backend, x_template) - return error(DEER._REACTANT_LOAD_HINT) + return error(_REACTANT_LOAD_HINT) end function _resolve_gradient_batch( @@ -231,7 +237,7 @@ function _resolve_gradient_batch( return _reactant_resolve_gradient_batch(logdensity_batch, backend, X_template) end function _reactant_resolve_gradient_batch(logdensity_batch, backend, X_template) - return error(DEER._REACTANT_LOAD_HINT) + return error(_REACTANT_LOAD_HINT) end #= @@ -244,21 +250,21 @@ A `SecondOrder` bypasses the gradient slot even when that slot is hand-written: naming both passes asks for two derivatives of `logdensity`. The slot is still the drift term the MALA step uses. -`AutoReactant` short-circuits ahead of both: DI cannot drive Reactant, so it can -neither build the `SecondOrder` nor route through `hvp_mode`. `ReactantExt` takes -over the second-order case by dispatching on `grad` instead — a Reactant-resolved -gradient carries the raw `logdensity` with it and re-traces forward-over-reverse -from there. +`AutoReactant` takes the same three branches as anything else. `_second_order` +below collapses an `AutoReactant` pair to a single `AutoReactant()` rather than a +`DI.SecondOrder`, and `_hvp_strategy(::AutoReactant)` sends the +hand-written-gradient case to `ReactantHVP`; `ReactantExt` supplies both matching +methods. `_check_reactant_pair` runs first, so a mismatched pairing is reported +before either path pays for a gradient resolution or an HVP compile. =# function _resolve_hvp(logdensity, grad, grad_backend, hvp_backend, x_template) _check_reactant_pair(grad_backend, hvp_backend) - if hvp_backend isa ADTypes.AutoReactant - return DEER._make_hvp_fn(DEER.ReactantHVP(), grad, hvp_backend, x_template) - elseif hvp_backend isa DI.SecondOrder + _check_reactant_hvp_source(grad, hvp_backend) + if hvp_backend isa DI.SecondOrder return DEER._make_hvp_fn_second_order(logdensity, hvp_backend, x_template) elseif grad_backend !== nothing return DEER._make_hvp_fn_second_order( - logdensity, DI.SecondOrder(hvp_backend, grad_backend), x_template + logdensity, _second_order(hvp_backend, grad_backend), x_template ) else return DEER._make_hvp_fn( @@ -272,18 +278,14 @@ function _resolve_hvp_batch( logdensity_batch, grad_batch, grad_batch_backend, hvp_backend, X_template ) _check_reactant_pair(grad_batch_backend, hvp_backend) - if hvp_backend isa ADTypes.AutoReactant - return DEER._make_hvp_batch_fn( - DEER.ReactantHVP(), grad_batch, hvp_backend, X_template - ) - elseif hvp_backend isa DI.SecondOrder + if hvp_backend isa DI.SecondOrder return DEER._make_hvp_batch_fn_second_order( _BatchLogdensitySum(logdensity_batch), hvp_backend, X_template ) elseif grad_batch_backend !== nothing return DEER._make_hvp_batch_fn_second_order( _BatchLogdensitySum(logdensity_batch), - DI.SecondOrder(hvp_backend, grad_batch_backend), + _second_order(hvp_backend, grad_batch_backend), X_template, ) else @@ -293,6 +295,16 @@ function _resolve_hvp_batch( end end +#= The HVP backend composed with the backend that produced the gradient under it. +For a DI-driven pair that is literally `DI.SecondOrder(hvp_backend, +grad_backend)`, taken through `DI.hvp`. Two `AutoReactant`s are not a pair DI +could run at all, so they collapse to the one backend that traces +forward-over-reverse from `logdensity` itself (`ReactantExt`'s +`_make_hvp_fn_second_order`). A mixed pair is already out by the time this runs, +via `_check_reactant_pair`. =# +_second_order(hvp_backend, grad_backend) = DI.SecondOrder(hvp_backend, grad_backend) +_second_order(::ADTypes.AutoReactant, ::ADTypes.AutoReactant) = ADTypes.AutoReactant() + #= Reactant does not pair with a DI backend across the two passes of an HVP: the compiled gradient is an opaque XLA executable DI cannot differentiate, and a @@ -329,6 +341,50 @@ function _check_reactant_pair(grad_backend, hvp_backend::ADTypes.AutoReactant) ) end +#= Would otherwise be ambiguous between the two methods above, and wants its own +message anyway: "an AutoReactant Hessian-vector product" does not describe a +`SecondOrder`. =# +function _check_reactant_pair(::ADTypes.AutoReactant, hvp_backend::DI.SecondOrder) + return throw( + ArgumentError( + "an AutoReactant gradient needs a bare AutoReactant Hessian-vector " * + "product: got hvp backend $(hvp_backend). An AutoReactant gradient " * + "compiles to an XLA executable, which DifferentiationInterface's " * + "SecondOrder cannot drive. Set the model's `hvp` (or the sampler's " * + "`backend`) to AutoReactant(), unwrapped.", + ), + ) +end + +#= `SecondOrder(AutoReactant(), AutoReactant())` is a natural thing to try, given +the pairing table in `10-getting-started.md`, and would otherwise land in +`DI.prepare_hvp` several frames deep with no Reactant support. Checked whatever +the gradient slot holds, since a `SecondOrder` bypasses it anyway (see +`_resolve_hvp`). =# +function _check_reactant_pair(grad_backend, hvp_backend::DI.SecondOrder) + _second_order_has_reactant(hvp_backend) || return nothing + return throw( + ArgumentError( + "AutoReactant cannot go inside a DifferentiationInterface SecondOrder: " * + "got hvp backend $(hvp_backend). DifferentiationInterface has no Reactant " * + "support at all. Set `hvp` (or the sampler's `backend`) to a bare " * + "AutoReactant() instead.", + ), + ) +end + +function _second_order_has_reactant(so::DI.SecondOrder) + return DI.outer(so) isa ADTypes.AutoReactant || DI.inner(so) isa ADTypes.AutoReactant +end + +#= A `LogDensityProblemGradient` (defined below) is a callable, so it clears the +`grad_backend === nothing` test that otherwise means "hand-written gradient" — +but it dispatches into DynamicPPL/LogDensityProblems and is not +Reactant-traceable. `_check_reactant_pair` only sees the backend, which is +`nothing` for both, so this checks `grad`'s type instead. The specialization has +to wait for `LogDensityProblemGradient` to exist and sits further down. =# +_check_reactant_hvp_source(grad, hvp_backend) = nothing + """ _prepare_model(model, x_template) -> PreppedDensityModel _prepare_model(model, x_template, T::Int, backend) -> PreppedDensityModel @@ -354,10 +410,9 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector) else model.grad_logdensity end - #= The derivative slots DEER alone reads are dropped rather than passed - along: a sequential sampler never looks at them, and carrying an - unresolved backend would break the invariant that no slot here holds an - `AbstractADType`. =# + #= The DEER-only slots are dropped rather than passed along: a sequential + sampler never reads them, and an unresolved backend sitting in one would + break the invariant that no slot here holds an `AbstractADType`. =# return PreppedDensityModel( model.logdensity, grad, @@ -374,31 +429,48 @@ end function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, backend) grad_backend = model.grad_logdensity isa AbstractADType ? model.grad_logdensity : nothing - grad = if grad_backend !== nothing - _resolve_gradient(model.logdensity, grad_backend, x_template) - else - model.grad_logdensity - end - hvp = if model.hvp === nothing || model.hvp isa AbstractADType - hvp_backend = model.hvp === nothing ? backend : model.hvp - hvp_backend === nothing && throw( + #= Settle the HVP backend and check its pairing with `grad_backend` before + resolving the gradient. Otherwise a mismatched `AutoReactant` pair, or the + LogDensityProblems-gradient case, surfaces only once `_resolve_gradient` has + paid for an XLA compile: 18+ seconds to report a config error. Neither check + needs the resolved gradient. `grad_backend` is known already, and + `_check_reactant_hvp_source` only looks at `model.grad_logdensity`'s type, + which is `grad` unchanged whenever `grad_backend` is `nothing`. =# + needs_hvp = model.hvp === nothing || model.hvp isa AbstractADType + hvp_backend = if needs_hvp + hb = model.hvp === nothing ? backend : model.hvp + hb === nothing && throw( ArgumentError( "ParallelMALASampler needs a Hessian-vector product: supply `hvp` " * "on the DensityModel (callable or AD backend), or pass `backend=` " * "to ParallelMALASampler", ), ) + _check_reactant_pair(grad_backend, hb) + _check_reactant_hvp_source(model.grad_logdensity, hb) + hb + else + nothing + end + + grad = if grad_backend !== nothing + _resolve_gradient(model.logdensity, grad_backend, x_template) + else + model.grad_logdensity + end + + hvp = if needs_hvp _resolve_hvp(model.logdensity, grad, grad_backend, hvp_backend, x_template) else model.hvp end #= A batched log-density with no batched gradient gets one from the model's - own gradient backend, never the sampler's: a model with a hand-written - gradient has not opted into AD, and deriving one anyway would let `backend=` - decide which update path runs. Failing to derive leaves the path off rather - than raising, since `_trajectory_logps` uses `logdensity_batch` regardless. =# + own gradient backend, never the sampler's: a hand-written gradient has not + opted into AD, and deriving one anyway would let `backend=` decide which + update path runs. Not deriving one leaves the batched path off rather than + raising, since `_trajectory_logps` uses `logdensity_batch` either way. =# grad_batch = model.grad_logdensity_batch if grad_batch === nothing && model.logdensity_batch !== nothing grad_batch = grad_backend @@ -413,26 +485,35 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, X_template = similar(x_template, length(x_template), T) X_template .= x_template - if grad_batch_backend !== nothing - grad_batch = _resolve_gradient_batch( - model.logdensity_batch, grad_batch_backend, X_template - ) - end - - if hvp_batch === nothing || hvp_batch isa AbstractADType + # Same reasoning as the unbatched case above: check before compiling. + needs_hvp_batch = hvp_batch === nothing || hvp_batch isa AbstractADType + hvp_batch_backend = if needs_hvp_batch # The model's own HVP backend if it has one, else the sampler's. - hvp_batch_backend = if hvp_batch === nothing + hbb = if hvp_batch === nothing model.hvp isa AbstractADType ? model.hvp : backend else hvp_batch end - hvp_batch_backend === nothing && throw( + hbb === nothing && throw( ArgumentError( "the batched DEER path needs a batched Hessian-vector product: " * "supply `hvp_batch` on the DensityModel (callable or AD backend), " * "or pass `backend=` to ParallelMALASampler", ), ) + _check_reactant_pair(grad_batch_backend, hbb) + hbb + else + nothing + end + + if grad_batch_backend !== nothing + grad_batch = _resolve_gradient_batch( + model.logdensity_batch, grad_batch_backend, X_template + ) + end + + if needs_hvp_batch hvp_batch = _resolve_hvp_batch( model.logdensity_batch, grad_batch, @@ -467,10 +548,10 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, ) end -#= Callable structs that allow us to dispatch on the type of the LogDensityProblems object in -the postprocessing stage. Ideally these would be defined in the LogDensityProblemsExt. -However, structs defined in extensions are hard to get hold of so we define them here. -The callable behaviour itself is implemented in LogDensityProblemsExt =# +# Callable structs that allow us to dispatch on the type of the LogDensityProblems object in +# the postprocessing stage. Ideally these would be defined in the LogDensityProblemsExt. +# However, structs defined in extensions are hard to get hold of so we define them here. +# The callable behaviour itself is implemented in LogDensityProblemsExt. struct LogDensityProblemPrimal{L} ld::L end @@ -478,6 +559,20 @@ struct LogDensityProblemGradient{L} ld::L end +# The `_check_reactant_hvp_source` specialization promised further up. +function _check_reactant_hvp_source( + ::LogDensityProblemGradient, hvp_backend::ADTypes.AutoReactant +) + return throw( + ArgumentError( + "an AutoReactant Hessian-vector product needs an AutoReactant or " * + "hand-written gradient: got a LogDensityProblems-derived gradient. Reactant " * + "cannot trace DynamicPPL/LogDensityProblems machinery. Set `grad_logdensity` " * + "to AutoReactant(), or supply a Reactant-traceable callable.", + ), + ) +end + """ MALASampler(epsilon; cholM=nothing) @@ -654,11 +749,15 @@ DEER-parallelized MALA sampler. Supported Jacobian modes are `:stoch_diag` (the default Hutchinson diagonal estimator) and `:diag` (exact diagonal via `D` JVPs). -`backend` is the fallback source of Hessian-vector products, used when the -`DensityModel` brings no `hvp` / `hvp_batch` of its own. That is all it does: it -never supplies a gradient, so it cannot change which update path runs or put AD -on a function the model did not already have a backend for. A model carrying its -own HVPs does not need it. +`backend` supplies Hessian-vector products when the `DensityModel` brings no +`hvp` / `hvp_batch` of its own, and does nothing else. It never supplies a +gradient, so it cannot decide which update path runs, nor put AD on a function +the model had no backend for. Leave it out for a model that carries its own HVPs. + +`backend = ADTypes.AutoReactant()` constrains the model too, since +`AutoReactant` does not pair with a DifferentiationInterface backend across the +two passes of an HVP: `grad_logdensity` must then be `AutoReactant()` as well, or +a hand-written callable. Mixing raises an `ArgumentError` at preparation time. """ struct ParallelMALASampler{FP<:AbstractFloat,CM,AD} <: AbstractMCMC.AbstractSampler epsilon::FP diff --git a/test/Project.toml b/test/Project.toml index 2892146..a2370e2 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -15,7 +15,6 @@ LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" ParallelMCMC = "1a970f40-4406-51c9-a967-cb3143c111e8" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" -Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" @@ -23,8 +22,9 @@ TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" +[compat] +JuliaFormatter = "2.12" + [extras] CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" - -[compat] -JuliaFormatter = "2.12" \ No newline at end of file +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" diff --git a/test/runtests.jl b/test/runtests.jl index 45cecaf..4d5b1d8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,13 @@ using ParallelMCMC using Test +#= `Reactant_jll` ships a prebuilt XLA and is a large download, so `Reactant` sits +in test/Project.toml's `[extras]` and `test-Reactant-HVP.jl` runs only when this +is set. Add Reactant to the test environment as well; the file's own +`try ... using Reactant ... catch` skips its testsets if it still won't load. =# +const _RUN_REACTANT_TESTS = + lowercase(get(ENV, "PARALLELMCMC_TEST_REACTANT", "false")) in ("1", "true", "yes") + @testset verbose=true "ParallelMCMC" begin #= Don't add your tests to runtests.jl. Instead, create files named @@ -14,6 +21,10 @@ using Test if isnothing(match(r"^test-.*\.jl$", file)) continue end + if file == "test-Reactant-HVP.jl" && !_RUN_REACTANT_TESTS + @info "Skipping $file (set PARALLELMCMC_TEST_REACTANT=true to opt in)" + continue + end title = titlecase(replace(splitext(file[6:end])[1], "-" => " ")) @testset verbose=true "$title" begin include(joinpath(root, file)) # robust if walkdir recurses diff --git a/test/test-HVP-Strategy.jl b/test/test-HVP-Strategy.jl index d28a02b..44305c1 100644 --- a/test/test-HVP-Strategy.jl +++ b/test/test-HVP-Strategy.jl @@ -46,8 +46,7 @@ const DI_STRAT = ParallelMCMC.DEER.DI @testset "normalization supplies Const but never a mode" begin #= The wrappers DEER differentiates are its own types, so annotating them - `Const` is its business. The mode is not: one the user set is a decision, - and an unset one is DI's to resolve from the operator it runs. =# + `Const` is its business. The mode is not. =# bare = DEER_STRAT._normalized_backend(AutoEnzyme()) @test bare isa AutoEnzyme{<:Any,Enzyme.Const} @test bare.mode === nothing @@ -70,9 +69,9 @@ const DI_STRAT = ParallelMCMC.DEER.DI @testset "normalizing a SecondOrder keeps the composition DI resolved" begin #= Regression for #62. Normalization used to route the outer half through a forward-only hook, which pinned `Enzyme.Forward` onto it. For a pair - `hvp_mode` resolves to reverse -- `SecondOrder(AutoEnzyme(), + `hvp_mode` resolves to reverse — `SecondOrder(AutoEnzyme(), AutoForwardDiff())` is reverse-over-forward, its inner half being - forward-only -- that silently made it forward-over-forward. =# + forward-only — that made it forward-over-forward. =# for so in ( DI_STRAT.SecondOrder(AutoEnzyme(), AutoForwardDiff()), DI_STRAT.SecondOrder(AutoEnzyme(), AutoZygote()), diff --git a/test/test-Reactant-HVP.jl b/test/test-Reactant-HVP.jl index c2b03eb..562e33a 100644 --- a/test/test-Reactant-HVP.jl +++ b/test/test-Reactant-HVP.jl @@ -1,32 +1,46 @@ using Test using Random using LinearAlgebra +using Statistics using FlexiChains using ParallelMCMC using ADTypes # Stands in for "some backend that is not Reactant" in the pairing tests below. using ForwardDiff: ForwardDiff +using LogDensityProblems: LogDensityProblems + +const DI_R = ParallelMCMC.DEER.DI #= Reactant-compiled derivative paths (`AutoReactant`), see ext/ReactantExt.jl. -The quartic target keeps second-order structure honest: logp = -0.25‖x‖⁴ has -H = -(‖x‖² I + 2 x xᵀ), so an HVP that silently drops the second-order term -(the failure mode of `Enzyme.hvp` under `@compile`) is caught, unlike a -Gaussian where H is constant. +The quartic target keeps the second-order structure honest: logp = -0.25‖x‖⁴ has +H = -(‖x‖² I + 2 x xᵀ), so an HVP that drops the second-order term — the failure +mode of `Enzyme.hvp` under `@compile` — is caught, where a Gaussian's constant H +would hide it. Every derivative-accuracy testset below uses it; do not swap in a +Gaussian. =# logp_r(x) = -0.25 * sum(abs2, x)^2 gradlogp_r(x) = -sum(abs2, x) .* x hvp_r(x, v) = -(sum(abs2, x) .* v .+ 2 .* dot(x, v) .* x) logp_batch_r(X) = vec(-0.25 .* sum(abs2, X; dims=1) .^ 2) gradlogp_batch_r(X) = -X .* sum(abs2, X; dims=1) +logp_r32(x) = -0.25f0 * sum(abs2, x)^2 + +#= Standard Gaussian, constant Hessian (H = -I). Only for the end-to-end sampling +tests, which ask whether the sampler converged to the right posterior; whether +the HVP is second-order-correct is the quartic target's job. =# +logp_gauss(x) = -0.5 * sum(abs2, x) +gradlogp_gauss(x) = -x +hvp_gauss(x, v) = -v const D_R = 4 const CT_R = FlexiChains.FlexiChain{Symbol} -#= The pairing rule lives in `_resolve_hvp`, not in the extension, so its -dispatch table is checked whether or not Reactant loads. =# +#= The pairing rule lives in `_resolve_hvp` / `_prepare_model`, not in the +extension, so its dispatch table is checked whether or not Reactant loads — +none of these calls resolve a gradient or compile anything. =# @testset "Reactant does not pair with a DI backend" begin # Both slots Reactant, or a hand-written gradient (`nothing`), are accepted. @test ParallelMCMC._check_reactant_pair(AutoReactant(), AutoReactant()) === nothing @@ -43,6 +57,67 @@ dispatch table is checked whether or not Reactant loads. =# ) end +#= `SecondOrder(AutoReactant(), AutoReactant())`, or `AutoReactant()` paired with +a `SecondOrder` at all, is a natural thing to try given the pairing table in +10-getting-started.md, and would otherwise land in `DI.prepare_hvp` several +frames deep with no Reactant support. A dispatch-table property, so this runs +whether or not Reactant is loaded. =# +@testset "AutoReactant cannot appear inside a SecondOrder" begin + # Sanity: an ordinary SecondOrder is unaffected. + @test ParallelMCMC._check_reactant_pair( + nothing, DI_R.SecondOrder(AutoForwardDiff(), AutoForwardDiff()) + ) === nothing + + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + nothing, DI_R.SecondOrder(AutoReactant(), AutoReactant()) + ) + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + AutoForwardDiff(), DI_R.SecondOrder(AutoReactant(), AutoForwardDiff()) + ) + # (AutoReactant, SecondOrder) exercises the disambiguating method directly. + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + AutoReactant(), DI_R.SecondOrder(AutoReactant(), AutoReactant()) + ) +end + +#= A `LogDensityProblemGradient` is a callable, so it clears the +`grad_backend === nothing` test that otherwise means "hand-written gradient", +but it is not Reactant-traceable. Checked on the type first, then through a real +LogDensityProblems model. Neither needs Reactant loaded: the check fires before +any gradient is resolved or anything compiled. =# +@testset "a LogDensityProblems gradient cannot pair with an AutoReactant hvp" begin + @test ParallelMCMC._check_reactant_hvp_source( + ParallelMCMC.LogDensityProblemGradient(nothing), AutoForwardDiff() + ) === nothing + @test_throws ArgumentError ParallelMCMC._check_reactant_hvp_source( + ParallelMCMC.LogDensityProblemGradient(nothing), AutoReactant() + ) + + struct _FakeLD end + LogDensityProblems.capabilities(::_FakeLD) = LogDensityProblems.LogDensityOrder{1}() + LogDensityProblems.dimension(::_FakeLD) = D_R + function LogDensityProblems.logdensity_and_gradient(::_FakeLD, x) + return logp_r(x), gradlogp_r(x) + end + + model = DensityModel(_FakeLD(); hvp=AutoReactant()) + @test_throws ArgumentError ParallelMCMC._prepare_model(model, zeros(D_R), 8, nothing) +end + +#= Outside the `reactant_ok` guard below. `_check_reactant_pair` now runs above +gradient resolution in `_prepare_model`, so a mismatched pair is caught before +anything is resolved or compiled and no Reactant install is needed. This used to +pay a full XLA compile per `@test_throws`, ~18s, to demonstrate a config error. =# +@testset "mixed pairs are refused at preparation" begin + x = zeros(D_R) + + reactant_grad = DensityModel(logp_r, AutoReactant(), D_R; hvp=AutoForwardDiff()) + @test_throws ArgumentError ParallelMCMC._prepare_model(reactant_grad, x, 8, nothing) + + reactant_hvp = DensityModel(logp_r, AutoForwardDiff(), D_R; hvp=AutoReactant()) + @test_throws ArgumentError ParallelMCMC._prepare_model(reactant_hvp, x, 8, nothing) +end + reactant_ok = try using Reactant: Reactant using Enzyme: Enzyme @@ -52,22 +127,54 @@ catch err false end +#= The `_REACTANT_LOAD_HINT` fallbacks (src/DEER/DEER.jl, and +`_reactant_resolve_gradient` / `_reactant_resolve_gradient_batch` in +src/interface.jl) can only be asserted in a session without Reactant: once +`ReactantExt` is loaded its more-specific methods shadow all of them and the +error can never fire. Hence the one testset here guarded on `!reactant_ok`. =# +if !reactant_ok + @testset "clear load-hint error without Reactant loaded" begin + #= Gradient slot: `_reactant_resolve_gradient`'s fallback. `backend` (the + 4th arg) is `AutoReactant()` too, as `_check_reactant_pair` requires, so + `_prepare_model` gets as far as gradient resolution rather than failing + first on the model having no HVP source — a correct failure, but not the + one being pinned here. =# + model_grad = DensityModel(logp_r, AutoReactant(), D_R) + @test_throws "AutoReactant requires Reactant.jl" ParallelMCMC._prepare_model( + model_grad, zeros(D_R), 8, AutoReactant() + ) + + # HVP slot over a hand-written gradient: `DEER._make_hvp_fn`'s + # `ReactantHVP` fallback, reached via `_hvp_strategy(::AutoReactant)`. + model_hvp = DensityModel(logp_r, gradlogp_r, D_R; hvp=AutoReactant()) + @test_throws "AutoReactant requires Reactant.jl" ParallelMCMC._prepare_model( + model_hvp, zeros(D_R), 8, nothing + ) + end +end + if reactant_ok @testset "extension is loaded" begin @test Base.get_extension(ParallelMCMC, :ReactantExt) !== nothing end - #= The same rule reached through `_prepare_model`, where the gradient slot - resolves first: a mixed pair must still surface as an ArgumentError and not - as whatever DI or Reactant would say downstream. =# - @testset "mixed pairs are refused at preparation" begin - x = zeros(D_R) + #= The extension does not honour `AutoReactant.mode` (the wrapped + `AutoEnzyme`): gradients always trace reverse, HVPs forward-over-that. A + non-default mode is rejected rather than ignored. =# + @testset "a non-default AutoReactant mode is rejected" begin + bad = AutoReactant(; mode=AutoEnzyme(; mode=Enzyme.Forward)) + model_grad = DensityModel(logp_r, bad, D_R) + @test_throws ArgumentError ParallelMCMC._prepare_model( + model_grad, zeros(D_R), 8, nothing + ) - reactant_grad = DensityModel(logp_r, AutoReactant(), D_R; hvp=AutoForwardDiff()) - @test_throws ArgumentError ParallelMCMC._prepare_model(reactant_grad, x, 8, nothing) + model_hvp = DensityModel(logp_r, gradlogp_r, D_R; hvp=bad) + @test_throws ArgumentError ParallelMCMC._prepare_model( + model_hvp, zeros(D_R), 8, nothing + ) - reactant_hvp = DensityModel(logp_r, AutoForwardDiff(), D_R; hvp=AutoReactant()) - @test_throws ArgumentError ParallelMCMC._prepare_model(reactant_hvp, x, 8, nothing) + # The default is, of course, fine. + @test DensityModel(logp_r, AutoReactant(), D_R) isa DensityModel end @testset "HVP matches analytic" begin @@ -107,14 +214,144 @@ if reactant_ok @test m_p.grad_logdensity_batch(X) ≈ gradlogp_batch_r(X) @test m_p.hvp_batch(X, V) ≈ Hv_cols end + + #= A hand-written batched gradient with `hvp_batch=AutoReactant()`, the + batched analogue of the "forward over user gradient" case above. Routes + through + `DEER._make_hvp_batch_fn(::ReactantHVP, grad_batch, ::AutoReactant, ...)`, + which nothing else here reaches. =# + @testset "forward over user batched gradient (hvp_batch=AutoReactant())" begin + T = 8 + X = randn(rng, D_R, T) + V = randn(rng, D_R, T) + Hv_cols = reduce(hcat, [hvp_r(X[:, t], V[:, t]) for t in 1:T]) + + model = DensityModel( + logp_r, + gradlogp_r, + D_R; + hvp=hvp_r, + logdensity_batch=logp_batch_r, + grad_logdensity_batch=gradlogp_batch_r, + hvp_batch=AutoReactant(), + ) + m_p = ParallelMCMC._prepare_model(model, X[:, 1], T, nothing) + @test m_p.hvp_batch(X, V) ≈ Hv_cols + end + + #= Edge shapes. Both slots AutoReactant, as in the + "forward-over-reverse from logp alone" case above, but at the smallest + sizes DEER ever prepares. =# + @testset "edge shapes" begin + @testset "D=1" begin + x1 = randn(rng, 1) + v1 = randn(rng, 1) + model = DensityModel(logp_r, AutoReactant(), 1; hvp=AutoReactant()) + m_p = ParallelMCMC._prepare_model(model, x1, 8, nothing) + @test m_p.grad_logdensity(x1) ≈ gradlogp_r(x1) + @test m_p.hvp(x1, v1) ≈ hvp_r(x1, v1) + end + + @testset "T=1" begin + T = 1 + X = reshape(x, D_R, T) + V = reshape(v, D_R, T) + model = DensityModel( + logp_r, + AutoReactant(), + D_R; + logdensity_batch=logp_batch_r, + grad_logdensity_batch=AutoReactant(), + hvp=AutoReactant(), + hvp_batch=AutoReactant(), + ) + m_p = ParallelMCMC._prepare_model(model, x, T, nothing) + @test m_p.grad_logdensity_batch(X) ≈ gradlogp_batch_r(X) + @test m_p.hvp_batch(X, V) ≈ reshape(hvp_r(x, v), D_R, T) + end + end + + #= Float32 on the CPU path. Previously only exercised inside the + CuArray-only block below, so it never ran without a functional CUDA. =# + @testset "Float32 on CPU" begin + x32 = Float32.(x) + v32 = Float32.(v) + model = DensityModel(logp_r32, AutoReactant(), D_R; hvp=AutoReactant()) + m_p = ParallelMCMC._prepare_model(model, x32, 8, nothing) + + g = m_p.grad_logdensity(x32) + @test eltype(g) === Float32 + @test g ≈ gradlogp_r(x32) + + Hv = m_p.hvp(x32, v32) + @test eltype(Hv) === Float32 + @test Hv ≈ hvp_r(x32, v32) + end + end + + #= Pins the "captured data is frozen at preparation time" limitation from + ext/ReactantExt.jl's module docstring and docs/src/15-gpu.md: `@compile` + bakes a captured plain `Array` in as a compile-time constant, so mutating it + afterwards has NO effect on later calls, with no error and no warning. Here + so that a future Reactant which does detect this gets noticed. Model code + still should not close over mutable data. + =# + @testset "documented caveat: captured data is frozen at preparation time" begin + data = [1.0, 1.0, 1.0, 1.0] + f(x) = -0.5 * sum(abs2, x .- data) # ∇f(x) = data - x + model = DensityModel(f, AutoReactant(), D_R) + # Two-argument `_prepare_model` only resolves `grad_logdensity`, which + # is all this test needs (no HVP involved). + m_p = ParallelMCMC._prepare_model(model, zeros(D_R)) + + x0 = zeros(D_R) + g_before = m_p.grad_logdensity(x0) + @test g_before ≈ [1.0, 1.0, 1.0, 1.0] + + data .= 5.0 # mutate the captured array *after* preparation + + g_after = m_p.grad_logdensity(x0) + #= Frozen at the pre-mutation value: NOT [5, 5, 5, 5], which is what a + correct re-evaluation against the mutated `data` would give. =# + @test g_after ≈ [1.0, 1.0, 1.0, 1.0] end + #= `size(chain) == (N,1)` and `all(isfinite, ...)` pass even for a badly + wrong HVP: DEER's Newton iteration just fails to converge and `DEER.solve` + returns the non-converged trajectory, no NaNs involved. Assert the posterior + mean against a target with a known mean instead (the convention in + test-GPU-AD-HVP.jl), and cross-check against the same model driven by the + analytic HVP on the same noise tape. + =# @testset "end-to-end sampling with AutoReactant" begin - model = DensityModel(logp_r, AutoReactant(), D_R) - s = ParallelMALASampler(0.02; T=16, backend=AutoReactant()) - chain = sample(MersenneTwister(72), model, s, 64; chain_type=CT_R, progress=false) - @test size(chain) == (64, 1) - @test all(x -> all(isfinite, x), chain[:x]) + @testset "posterior mean recovery (standard Gaussian, mean 0)" begin + model = DensityModel(logp_gauss, AutoReactant(), D_R) + s = ParallelMALASampler(0.3; T=16, backend=AutoReactant()) + n_samples, n_burn = 2000, 500 + chain = sample( + MersenneTwister(72), model, s, n_samples; chain_type=CT_R, progress=false + ) + @test size(chain) == (n_samples, 1) + + xs = chain[:x] + @test all(x -> all(isfinite, x), xs) + post_mean = vec(mean(reduce(hcat, xs[(n_burn + 1):end]); dims=2)) + @test maximum(abs, post_mean) < 0.3 + end + + @testset "matches analytic-HVP DEER on the same noise tape" begin + s = ParallelMALASampler(0.1; T=16, backend=AutoReactant()) + model_r = DensityModel(logp_gauss, AutoReactant(), D_R) + model_an = DensityModel(logp_gauss, gradlogp_gauss, D_R; hvp=hvp_gauss) + + c_r = sample( + MersenneTwister(99), model_r, s, 64; chain_type=CT_R, progress=false + ) + c_an = sample( + MersenneTwister(99), model_an, s, 64; chain_type=CT_R, progress=false + ) + @test c_r[:x] ≈ c_an[:x] + end end reactant_gpu_ok = try @@ -128,12 +365,13 @@ if reactant_ok @info "Reactant HVP test: CUDA not functional — skipping CuArray boundary" else #= - The compiled executable lives in Reactant's own (XLA) device memory; - what's checked here is the CuArray <-> Reactant marshalling boundary: - CuArray in, CuArray out, values matching the analytic HVP. + The compiled executable lives in Reactant's own (XLA) device memory, so + what this checks is the CuArray <-> Reactant marshalling boundary: + CuArray in, CuArray out, values matching the analytic HVP. It does NOT + establish that the HVP executes on the GPU, which depends on Reactant's + default XLA client (see docs/src/15-gpu.md) and is neither controlled + nor asserted here. =# - logp_r32(x) = -0.25f0 * sum(abs2, x)^2 - @testset "CuArray boundary" begin rng = MersenneTwister(73) x_h = randn(rng, Float32, D_R) From 3bea841a0daf72097f1ad7145da36fb068cfc592 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Wed, 12 Aug 2026 21:29:31 -0400 Subject: [PATCH 06/11] remove a few dead pieces --- CHANGELOG.md | 4 ++-- docs/src/15-gpu.md | 2 +- ext/ReactantExt.jl | 20 +++++--------------- src/interface.jl | 7 ++----- test/test-Reactant-HVP.jl | 2 +- 5 files changed, 11 insertions(+), 24 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bde0fe4..44c6194 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,7 +25,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 batched gradient derived from it when `grad_logdensity` is a backend, instead of leaving the batched DEER path switched off. `hvp_batch` can be a backend in that case too, and differentiates the derived gradient (#52). -- An HVP backend over an AD-derived gradient is now real second-order AD, +- An HVP backend over an AD-derived gradient is now true second-order AD, `DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` handed to `DI.hvp`, rather than an outer AD pass over the prepared DI gradient — which dropped out of its preparation as soon as tangents were pushed through it @@ -39,7 +39,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - New `ReactantExt`. `ADTypes.AutoReactant()` in a derivative slot, or as the sampler `backend`, traces the derivative with Enzyme-MLIR and compiles it to an XLA executable via Reactant.jl, off Enzyme's LLVM pipeline and off - DifferentiationInterface entirely, which yields a genuine second-order HVP for + DifferentiationInterface entirely, which yields a true second-order HVP for a log-density-only model (#37, #52). Requires `using Reactant` and a Reactant-traceable log-density. `AutoReactant` does not pair with a DifferentiationInterface backend, a `LogDensityProblems` gradient, or a diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index 318477a..1e40f4d 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -240,7 +240,7 @@ DEER needs a Hessian–vector product $H v$ at every Newton step. `DensityModel ### Reactant HVPs, off the DI path -`ADTypes.AutoReactant()` (requires `using Reactant`) goes another way. The derivative is traced with Enzyme-MLIR and compiled to an XLA executable by [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl), never touching Enzyme's LLVM pipeline, so neither the gc-transition abort nor the `pmcmc_*` wrappers above apply. It gives a genuine second-order HVP for a log-density-only model: +`ADTypes.AutoReactant()` (requires `using Reactant`) goes another way. The derivative is traced with Enzyme-MLIR and compiled to an XLA executable by [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl), never touching Enzyme's LLVM pipeline, so neither the gc-transition abort nor the `pmcmc_*` wrappers above apply. It gives a true second-order HVP for a log-density-only model: ```julia using Reactant, ADTypes diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl index 6c14c01..df82cd7 100644 --- a/ext/ReactantExt.jl +++ b/ext/ReactantExt.jl @@ -6,7 +6,7 @@ derivative slot of `DensityModel`, or as the sampler `backend`. Derivatives are traced with Enzyme-MLIR and compiled to XLA executables by `Reactant.@compile`, which keeps them off Enzyme's LLVM pipeline and so off the GPU `cuMemcpyDtoHAsync_v2` gc-transition abort, and off DifferentiationInterface -entirely. For a log-density-only model it yields a genuine second-order HVP. +entirely. For a log-density-only model it yields a true second-order HVP. Two silent failure modes, ahead of the ordinary limitations: @@ -144,26 +144,16 @@ end _rev_gradient(f, x) = Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(f), x)[1] #= -Gradient slots. They hold only the compiled callable; the HVP factories below -get `logdensity` from `_resolve_hvp`, which already has it. +Gradient slots. The HVP factories below get `logdensity` from `_resolve_hvp`, +which already has it. =# -struct _ReactantGradient{C} - compiled::C -end -(g::_ReactantGradient)(x) = g.compiled(x) - function ParallelMCMC._reactant_resolve_gradient( logdensity, backend::AutoReactant, x_template::AbstractVector ) _check_reactant_mode(backend) core = Base.Fix1(_rev_gradient, logdensity) - return _ReactantGradient(_compiled(core, x_template)) -end - -struct _ReactantGradientBatch{C} - compiled::C + return _compiled(core, x_template) end -(g::_ReactantGradientBatch)(X) = g.compiled(X) function ParallelMCMC._reactant_resolve_gradient_batch( logdensity_batch, backend::AutoReactant, X_template::AbstractMatrix @@ -173,7 +163,7 @@ function ParallelMCMC._reactant_resolve_gradient_batch( # gradient uses (src/interface.jl), so both batched paths differentiate the # same thing. core = Base.Fix1(_rev_gradient, ParallelMCMC._BatchLogdensitySum(logdensity_batch)) - return _ReactantGradientBatch(_compiled(core, X_template)) + return _compiled(core, X_template) end #= diff --git a/src/interface.jl b/src/interface.jl index b664b5b..cf77a17 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -254,12 +254,10 @@ the drift term the MALA step uses. below collapses an `AutoReactant` pair to a single `AutoReactant()` rather than a `DI.SecondOrder`, and `_hvp_strategy(::AutoReactant)` sends the hand-written-gradient case to `ReactantHVP`; `ReactantExt` supplies both matching -methods. `_check_reactant_pair` runs first, so a mismatched pairing is reported -before either path pays for a gradient resolution or an HVP compile. +methods. The pairing is already checked by `_prepare_model` before either path +here pays for a gradient resolution or an HVP compile. =# function _resolve_hvp(logdensity, grad, grad_backend, hvp_backend, x_template) - _check_reactant_pair(grad_backend, hvp_backend) - _check_reactant_hvp_source(grad, hvp_backend) if hvp_backend isa DI.SecondOrder return DEER._make_hvp_fn_second_order(logdensity, hvp_backend, x_template) elseif grad_backend !== nothing @@ -277,7 +275,6 @@ end function _resolve_hvp_batch( logdensity_batch, grad_batch, grad_batch_backend, hvp_backend, X_template ) - _check_reactant_pair(grad_batch_backend, hvp_backend) if hvp_backend isa DI.SecondOrder return DEER._make_hvp_batch_fn_second_order( _BatchLogdensitySum(logdensity_batch), hvp_backend, X_template diff --git a/test/test-Reactant-HVP.jl b/test/test-Reactant-HVP.jl index 562e33a..83be205 100644 --- a/test/test-Reactant-HVP.jl +++ b/test/test-Reactant-HVP.jl @@ -173,7 +173,7 @@ if reactant_ok model_hvp, zeros(D_R), 8, nothing ) - # The default is, of course, fine. + # The default still works. @test DensityModel(logp_r, AutoReactant(), D_R) isa DensityModel end From ff5241442b3f3b70a950b8f7a73353ea0de12983 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Sat, 8 Aug 2026 22:40:21 -0400 Subject: [PATCH 07/11] Initial reactant attempt. --- CHANGELOG.md | 8 ++ Project.toml | 5 +- docs/src/10-getting-started.md | 2 +- docs/src/15-gpu.md | 23 ++++- ext/ReactantExt.jl | 152 ++++++++++++++++++++++++++++++++ src/DEER/DEER.jl | 33 +++++++ src/ParallelMCMC.jl | 2 +- src/interface.jl | 77 +++++++++++++++- test/Project.toml | 1 + test/test-HVP-Strategy.jl | 8 ++ test/test-Reactant-HVP.jl | 156 +++++++++++++++++++++++++++++++++ 11 files changed, 461 insertions(+), 6 deletions(-) create mode 100644 ext/ReactantExt.jl create mode 100644 test/test-Reactant-HVP.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index d2d204a..851d2e1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 batched gradient derived for it when `grad_logdensity` is a backend, rather than leaving the batched DEER path switched off (#52). `hvp_batch` can be a backend in that case too, and differentiates the derived gradient. +- New `ReactantExt`: `ADTypes.AutoReactant()` in any derivative slot, or as the + sampler `backend`, traces the derivative with Enzyme-MLIR and compiles it to + an XLA executable via Reactant.jl. This bypasses Enzyme's LLVM pipeline and + is currently the only path that computes a genuine second-order HVP on GPU, + so `ParallelMALASampler` now works on CUDA for log-density-only models + (#37, #52). Requires `using Reactant` and a Reactant-traceable log-density; + `AutoReactant` cannot be paired with a DifferentiationInterface backend + across the two passes of an HVP, which raises an `ArgumentError`. - An HVP backend over an AD-derived gradient is now taken as true second-order AD, `DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` handed to `DI.hvp`, instead of an outer AD pass over the prepared DI gradient (#37). diff --git a/Project.toml b/Project.toml index 82d250a..7f7f77a 100644 --- a/Project.toml +++ b/Project.toml @@ -20,15 +20,17 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [extensions] DynamicPPLExt = ["DynamicPPL", "FlexiChains", "LogDensityProblems"] EnzymeExt = "Enzyme" LogDensityProblemsExt = "LogDensityProblems" +ReactantExt = ["Reactant", "Enzyme"] [compat] -ADTypes = "1.21.0" +ADTypes = "1.22.0" AbstractMCMC = "5.10.0" CUDA = "5.11.0" CUDA_Runtime_jll = "0.21" @@ -42,6 +44,7 @@ LogDensityProblems = "2" Mooncake = "0.5.26" OrderedCollections = "1" Random = "1" +Reactant = "0.2.278" Statistics = "1" Zygote = "0.7.10" julia = "1.10" diff --git a/docs/src/10-getting-started.md b/docs/src/10-getting-started.md index b24e81e..b257131 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -49,7 +49,7 @@ Backends become prepared [DifferentiationInterface](https://github.com/JuliaDiff Naming both passes yourself is the one route that ignores the gradient slot, hand-written or not. It is also the only AD route to an HVP for a Turing or LogDensityProblems model, whose gradient arrives already prepared and cannot be differentiated again. -Which pairs work is up to the backends. On CPU, ForwardDiff, ReverseDiff, Zygote and Enzyme all serve a log-density-only model. `AutoMooncake` serves neither direction: it has no reverse-over-reverse, and its gradient rejects an outer pass's tangents. Give Mooncake a hand-written `grad_logdensity` instead. No second-order pair works on GPU yet (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37) and the [GPU page](15-gpu.md)). +Which pairs work is up to the backends. On CPU, ForwardDiff, ReverseDiff, Zygote and Enzyme all serve a log-density-only model. `AutoMooncake` serves neither direction: it has no reverse-over-reverse, and its gradient rejects an outer pass's tangents. Give Mooncake a hand-written `grad_logdensity` instead. On GPU, no DI-driven second-order pair works yet (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)); `AutoReactant()` in both slots is the exception, and compiles the pair to a single XLA program. See the [GPU page](15-gpu.md). The batched pair works the same way, on `sum(logdensity_batch(X))`. That sum's gradient is the stacked per-column gradients only because the columns are independent, so `logdensity_batch` must not couple them. Omitting `grad_logdensity_batch` derives one when `grad_logdensity` is a backend; with a hand-written gradient the batched path stays off and the unbatched update covers it. Both batched derivative slots require `logdensity_batch`, which is also useful on its own for scoring a whole trajectory at once. diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index ea99f64..1f8c5a2 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -233,11 +233,32 @@ DEER needs a Hessian–vector product $H v$ at every Newton step. `DensityModel - **You only supply `gradlogp` / `grad_logdensity_batch`.** The sampler builds the HVP by differentiating your gradient — either a forward-mode pushforward of `gradlogp` ([`ForwardOnGrad`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/src/DEER/DEER.jl), the default for most backends) or a reverse-mode gradient of `x -> dot(gradlogp(x), v)` ([`ReverseOnGrad`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/src/DEER/DEER.jl), used for `AutoMooncake` and `AutoZygote`). This is the **AD-HVP fallback**, and it is what the logistic-regression example above uses. !!! warning "Log-density-only models on GPU" - `grad_logdensity` can itself be an AD backend (`DensityModel(logp, AutoEnzyme(), dim)`, see [Getting started](10-getting-started.md)), but don't do that with `ParallelMALASampler` on GPU. The HVP becomes `SecondOrder(hvp_backend, grad_backend)` on your log-density, which currently fails on GPU with both Enzyme and Mooncake (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)). The same goes for passing a `SecondOrder` explicitly. Write `gradlogp` out by hand for DEER, so the HVP is a single pass over it. The sequential samplers only need the gradient, so log-density-only models work there. + `grad_logdensity` can itself be an AD backend (`DensityModel(logp, AutoEnzyme(), dim)`, see [Getting started](10-getting-started.md)), but don't do that with `ParallelMALASampler` on GPU for the DI-driven backends. The HVP becomes `SecondOrder(hvp_backend, grad_backend)` on your log-density, which currently fails on GPU with both Enzyme and Mooncake (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)). The same goes for passing a `SecondOrder` explicitly. Write `gradlogp` out by hand for DEER, so the HVP is a single pass over it — or use `AutoReactant()` (below), the one backend where log-density-only DEER works on GPU. The sequential samplers only need the gradient, so log-density-only models work there. !!! note "A backend in `grad_logdensity` reaches `logdensity_batch` too" The batched path needs a batched gradient, and derives one from `logdensity_batch` when `grad_logdensity` is a backend. That puts `logdensity_batch` under the same restrictions as the rest of your AD-visible code. Supply `grad_logdensity_batch` to avoid it. +### Reactant: genuine second-order HVPs on GPU + +`ADTypes.AutoReactant()` (requires `using Reactant`) takes a different route entirely. The derivative is traced with Enzyme-MLIR and compiled to an XLA executable by [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl), bypassing Enzyme's LLVM pipeline — and with it the gc-transition abort and the `pmcmc_*` wrapper requirements above. It is currently the only path that computes a genuine second-order HVP on GPU, so log-density-only models work with DEER: + +```julia +using Reactant, ADTypes + +model = DensityModel(logp, AutoReactant(), D) # no hand-written gradient +sampler = ParallelMALASampler(0.005f0; T=16, backend=AutoReactant()) +``` + +When both the gradient slot and the HVP source are `AutoReactant()`, the HVP compiles as explicit forward-over-reverse from the raw log-density, in a single fused XLA program. With a hand-written `gradlogp`, it compiles a forward pushforward of your gradient instead. + +Three caveats: + +- The traced function must be **Reactant-traceable**: plain array operations. DynamicPPL-built log-densities do not trace as-is. +- **Reactant does not mix with the DI backends across the two passes of an HVP.** Both `grad_logdensity` and the HVP source take `AutoReactant()`, or neither does; a hand-written gradient pairs with either. Mixing raises an `ArgumentError` at preparation time, since Reactant cannot trace a DI-prepared gradient and DI cannot differentiate a compiled XLA executable. +- Calls cross a marshalling boundary (package arrays ↔ Reactant's XLA device memory) on every invocation, and executables are shape-specialized at preparation time. Correct everywhere, but it leaves fusion on the table relative to a future end-to-end Reactant pipeline. + +DifferentiationInterface cannot drive Reactant yet; when that support lands ([DI#918](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/918)), this path folds into the standard `hvp_mode` routing with no user-facing change. + ### When the fallback is the right call - **Complex or composed models.** Bayesian neural nets, hierarchical models with many transformations, mixtures, or anything where the Hessian has no convenient closed form. Deriving and maintaining `hvp` by hand for these is error-prone; AD removes a whole class of bugs. diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl new file mode 100644 index 0000000..21df47a --- /dev/null +++ b/ext/ReactantExt.jl @@ -0,0 +1,152 @@ +module ReactantExt + +#= +Reactant-compiled derivative paths, selected by `ADTypes.AutoReactant()` in +any derivative slot of `DensityModel` (or as the sampler `backend`). +Derivatives are traced with Enzyme-MLIR and compiled to XLA executables via +`Reactant.@compile`, bypassing Enzyme's LLVM pipeline — and with it the GPU +`cuMemcpyDtoHAsync_v2` gc-transition abort. This is the only path that +computes a genuine second-order HVP on GPU. + +Requirements / limitations: + - The traced function (`logdensity` / `gradlogp` / batched forms) must be + Reactant-traceable: plain array ops. DynamicPPL-built log-densities do + NOT trace as-is. + - Executables are shape-specialized to the preparation templates, and + arguments are marshalled to/from Reactant's own (XLA) device memory on + every call. A boundary copy, not a fused in-place path — a target for + later optimization. + - HVPs are explicit forward-over-(gradient) compositions; NEVER + `Enzyme.hvp`, which silently returns zeros under `@compile`. + - `AutoReactant.mode` (the wrapped `AutoEnzyme`) is currently ignored: + gradients always trace as Enzyme reverse, HVPs as forward-over-that. +=# + +using ParallelMCMC: ParallelMCMC +using ParallelMCMC.DEER: DEER +using ADTypes: AutoReactant +using Reactant: Reactant, @compile +using Enzyme: Enzyme + +# Host-materialize for the Reactant boundary: Reactant manages its own (XLA) +# device memory, so we round-trip through a plain host Array regardless of +# where the package's arrays live (Vector / CuArray). +_host(x::Array) = x +_host(x::AbstractArray) = Array(x) + +function _from_host(template::AbstractArray, out) + res = similar(template, size(out)) + copyto!(res, Array(out)) + return res +end + +#= +Compile `core` for the template shapes once and return a closure that +marshals package arrays <-> Reactant arrays. +=# +function _compiled(core, t1::AbstractArray) + r1 = Reactant.to_rarray(_host(t1)) + compiled = @compile core(r1) + return x -> _from_host(x, compiled(Reactant.to_rarray(_host(x)))) +end + +function _compiled(core, t1::AbstractArray, t2::AbstractArray) + r1 = Reactant.to_rarray(_host(t1)) + r2 = Reactant.to_rarray(_host(t2)) + compiled = @compile core(r1, r2) + return function (x, v) + out = compiled(Reactant.to_rarray(_host(x)), Reactant.to_rarray(_host(v))) + return _from_host(x, out) + end +end + +# Forward-mode JVP of `g` in direction `v`, i.e. J(g)·v. For g = gradlogp +# this is the HVP H·v. `Const(g)` so Enzyme doesn't treat captures as active. +function _jvp(g, x, v) + return only(Enzyme.autodiff(Enzyme.Forward, Enzyme.Const(g), Enzyme.Duplicated(x, v))) +end + +_rev_gradient(f, x) = Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(f), x)[1] + +# Columns of X are independent samples, so ∇_X sum(logp_batch(X)) stacks the +# per-column gradients. +_sumbatch(f, X) = sum(f(X)) + +#= +Gradient slots. The wrappers keep the raw log-density so the HVP factories +below can re-trace forward-over-reverse from it, instead of trying to trace +through an already-compiled executable. +=# +struct _ReactantGradient{F,C} + logdensity::F + compiled::C +end +(g::_ReactantGradient)(x) = g.compiled(x) + +function ParallelMCMC._reactant_resolve_gradient( + logdensity, backend::AutoReactant, x_template::AbstractVector +) + core = Base.Fix1(_rev_gradient, logdensity) + return _ReactantGradient(logdensity, _compiled(core, x_template)) +end + +struct _ReactantGradientBatch{F,C} + logdensity_batch::F + compiled::C +end +(g::_ReactantGradientBatch)(X) = g.compiled(X) + +function ParallelMCMC._reactant_resolve_gradient_batch( + logdensity_batch, backend::AutoReactant, X_template::AbstractMatrix +) + core = Base.Fix1(_rev_gradient, Base.Fix1(_sumbatch, logdensity_batch)) + return _ReactantGradientBatch(logdensity_batch, _compiled(core, X_template)) +end + +#= +HVP factories. Two tracings depending on where the gradient came from: + + - `_ReactantGradient` (the gradient slot was itself AutoReactant): + re-trace from the raw log-density as explicit forward-over-reverse — + genuine second-order AD fused into one XLA program. + - any other callable (hand-written gradient): forward JVP over it, + provided it is traceable. +=# +function DEER._make_hvp_fn( + ::DEER.ReactantHVP, + gradlogp::_ReactantGradient, + backend::AutoReactant, + x_template::AbstractVector, +) + inner = Base.Fix1(_rev_gradient, gradlogp.logdensity) + core(x, v) = _jvp(inner, x, v) + return _compiled(core, x_template, x_template) +end + +function DEER._make_hvp_fn( + ::DEER.ReactantHVP, gradlogp, backend::AutoReactant, x_template::AbstractVector +) + core(x, v) = _jvp(gradlogp, x, v) + return _compiled(core, x_template, x_template) +end + +function DEER._make_hvp_batch_fn( + ::DEER.ReactantHVP, + grad_batch::_ReactantGradientBatch, + backend::AutoReactant, + X_template::AbstractMatrix, +) + inner = Base.Fix1(_rev_gradient, Base.Fix1(_sumbatch, grad_batch.logdensity_batch)) + core(X, V) = _jvp(inner, X, V) + return _compiled(core, X_template, X_template) +end + +function DEER._make_hvp_batch_fn( + ::DEER.ReactantHVP, grad_batch, backend::AutoReactant, X_template::AbstractMatrix +) + # Column-independent batched gradient ⇒ forward JVP is the columnwise HVP. + core(X, V) = _jvp(grad_batch, X, V) + return _compiled(core, X_template, X_template) +end + +end # module diff --git a/src/DEER/DEER.jl b/src/DEER/DEER.jl index f0d677b..3abe628 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -203,10 +203,23 @@ abstract type HVPStrategy end struct ForwardOnGrad <: HVPStrategy end struct ReverseOnGrad <: HVPStrategy end +#= +ReactantHVP — trace the HVP with Enzyme-MLIR and compile it to an XLA +executable via Reactant.jl (see `ext/ReactantExt.jl`). Selected by +`ADTypes.AutoReactant()`, which DI cannot drive yet, so it short-circuits ahead +of the `hvp_mode` routing above; once DI gains Reactant support the +`AutoReactant` specializations can be deleted. Reactant bypasses Enzyme's LLVM +pipeline, avoiding the GPU gc-transition abort — the only path that computes a +genuine second-order HVP on GPU. The traced functions must be Reactant-traceable +(plain array ops). +=# +struct ReactantHVP <: HVPStrategy end + _strategy_from(::DI.ForwardOverAnything) = ForwardOnGrad() _strategy_from(::DI.HVPMode) = ReverseOnGrad() _hvp_strategy(backend::AbstractADType) = _strategy_from(DI.hvp_mode(backend)) +_hvp_strategy(::ADTypes.AutoReactant) = ReactantHVP() #= Hook for backend-specific normalization of the user's `backend`, applied on every @@ -300,6 +313,26 @@ function _make_hvp_batch_fn( return (X, V) -> _batch_hvp_via_grad_reverse_prepared(prep, X, V) end +#= +`ReactantHVP` fallbacks. `ReactantExt` adds methods with `backend` pinned to +`ADTypes.AutoReactant` (strictly more specific — no method overwriting, which +precompilation forbids); without Reactant loaded these give a clear error +instead of a `MethodError`. +=# +const _REACTANT_LOAD_HINT = "AutoReactant requires Reactant.jl: add `using Reactant` to load ParallelMCMC's ReactantExt." + +function _make_hvp_fn( + ::ReactantHVP, gradlogp, backend::AbstractADType, x_template::AbstractVector +) + return error(_REACTANT_LOAD_HINT) +end + +function _make_hvp_batch_fn( + ::ReactantHVP, grad_batch, backend::AbstractADType, X_template::AbstractMatrix +) + return error(_REACTANT_LOAD_HINT) +end + #= --------------------------------------------------------------------------- Second-order HVP, for a model whose gradient is itself AD-derived. `DI.hvp` diff --git a/src/ParallelMCMC.jl b/src/ParallelMCMC.jl index e98eebb..7041e90 100644 --- a/src/ParallelMCMC.jl +++ b/src/ParallelMCMC.jl @@ -1,7 +1,7 @@ module ParallelMCMC using AbstractMCMC -using ADTypes: AbstractADType +using ADTypes: ADTypes, AbstractADType using CUDA using DifferentiationInterface: DifferentiationInterface using FlexiChains diff --git a/src/interface.jl b/src/interface.jl index 69ede25..4e89cda 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -211,6 +211,29 @@ function _resolve_gradient_batch( ) end +#= +`AutoReactant` gradients bypass DI (which cannot drive Reactant yet) and go +through hook functions that `ReactantExt` fills in with strictly more specific +methods; the untyped fallbacks give a clear load-order error. +=# +function _resolve_gradient( + logdensity, backend::ADTypes.AutoReactant, x_template::AbstractVector +) + return _reactant_resolve_gradient(logdensity, backend, x_template) +end +function _reactant_resolve_gradient(logdensity, backend, x_template) + return error(DEER._REACTANT_LOAD_HINT) +end + +function _resolve_gradient_batch( + logdensity_batch, backend::ADTypes.AutoReactant, X_template::AbstractMatrix +) + return _reactant_resolve_gradient_batch(logdensity_batch, backend, X_template) +end +function _reactant_resolve_gradient_batch(logdensity_batch, backend, X_template) + return error(DEER._REACTANT_LOAD_HINT) +end + #= Resolve an HVP slot given as a backend. `grad_backend` is the backend that produced `grad`, or nothing when the gradient slot held a callable. Dispatch is @@ -220,9 +243,18 @@ statically known. A `SecondOrder` bypasses the gradient slot even when that slot is hand-written: naming both passes asks for two derivatives of `logdensity`. The slot is still the drift term the MALA step uses. + +`AutoReactant` short-circuits ahead of both: DI cannot drive Reactant, so it can +neither build the `SecondOrder` nor route through `hvp_mode`. `ReactantExt` takes +over the second-order case by dispatching on `grad` instead — a Reactant-resolved +gradient carries the raw `logdensity` with it and re-traces forward-over-reverse +from there. =# function _resolve_hvp(logdensity, grad, grad_backend, hvp_backend, x_template) - if hvp_backend isa DI.SecondOrder + _check_reactant_pair(grad_backend, hvp_backend) + if hvp_backend isa ADTypes.AutoReactant + return DEER._make_hvp_fn(DEER.ReactantHVP(), grad, hvp_backend, x_template) + elseif hvp_backend isa DI.SecondOrder return DEER._make_hvp_fn_second_order(logdensity, hvp_backend, x_template) elseif grad_backend !== nothing return DEER._make_hvp_fn_second_order( @@ -239,7 +271,12 @@ end function _resolve_hvp_batch( logdensity_batch, grad_batch, grad_batch_backend, hvp_backend, X_template ) - if hvp_backend isa DI.SecondOrder + _check_reactant_pair(grad_batch_backend, hvp_backend) + if hvp_backend isa ADTypes.AutoReactant + return DEER._make_hvp_batch_fn( + DEER.ReactantHVP(), grad_batch, hvp_backend, X_template + ) + elseif hvp_backend isa DI.SecondOrder return DEER._make_hvp_batch_fn_second_order( _BatchLogdensitySum(logdensity_batch), hvp_backend, X_template ) @@ -256,6 +293,42 @@ function _resolve_hvp_batch( end end +#= +Reactant does not pair with a DI backend across the two passes of an HVP: the +compiled gradient is an opaque XLA executable DI cannot differentiate, and a +DI-prepared gradient is not Reactant-traceable. Both slots take `AutoReactant` +or neither does; a hand-written gradient pairs with either. + +Dispatch rather than a runtime `isa` chain, so the check folds away with the rest +of `_resolve_hvp`'s branching. +=# +_check_reactant_pair(grad_backend, hvp_backend) = nothing +_check_reactant_pair(::ADTypes.AutoReactant, ::ADTypes.AutoReactant) = nothing +_check_reactant_pair(::Nothing, ::ADTypes.AutoReactant) = nothing + +function _check_reactant_pair(grad_backend::ADTypes.AutoReactant, hvp_backend) + return throw( + ArgumentError( + "an AutoReactant gradient needs an AutoReactant Hessian-vector product: " * + "got hvp backend $(hvp_backend). Reactant compiles the gradient to an XLA " * + "executable, which DifferentiationInterface cannot differentiate. Set the " * + "model's `hvp` (or the sampler's `backend`) to AutoReactant() as well.", + ), + ) +end + +function _check_reactant_pair(grad_backend, hvp_backend::ADTypes.AutoReactant) + return throw( + ArgumentError( + "an AutoReactant Hessian-vector product needs an AutoReactant or " * + "hand-written gradient: got gradient backend $(grad_backend). Reactant " * + "traces the HVP from the log-density (or from your gradient) and cannot " * + "trace a DifferentiationInterface-prepared gradient. Set " * + "`grad_logdensity` to AutoReactant() or supply a callable.", + ), + ) +end + """ _prepare_model(model, x_template) -> PreppedDensityModel _prepare_model(model, x_template, T::Int, backend) -> PreppedDensityModel diff --git a/test/Project.toml b/test/Project.toml index 7dbcce5..2892146 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -15,6 +15,7 @@ LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" ParallelMCMC = "1a970f40-4406-51c9-a967-cb3143c111e8" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" diff --git a/test/test-HVP-Strategy.jl b/test/test-HVP-Strategy.jl index e56997e..d28a02b 100644 --- a/test/test-HVP-Strategy.jl +++ b/test/test-HVP-Strategy.jl @@ -24,6 +24,14 @@ const DI_STRAT = ParallelMCMC.DEER.DI DEER_STRAT.ReverseOnGrad end + @testset "AutoReactant short-circuits hvp_mode" begin + #= DI cannot drive Reactant, so `AutoReactant` never reaches `DI.hvp_mode` + and the strategy is picked by dispatch instead. The method lives in DEER + rather than in ReactantExt, so this holds with Reactant unloaded. =# + @test DEER_STRAT._hvp_strategy(AutoReactant()) isa DEER_STRAT.ReactantHVP + @test @inferred(DEER_STRAT._hvp_strategy(AutoReactant())) isa DEER_STRAT.ReactantHVP + end + @testset "SecondOrder follows hvp_mode's composition" begin so_fwd_outer = DI_STRAT.SecondOrder(AutoForwardDiff(), AutoZygote()) @test DEER_STRAT._hvp_strategy(so_fwd_outer) isa DEER_STRAT.ForwardOnGrad diff --git a/test/test-Reactant-HVP.jl b/test/test-Reactant-HVP.jl new file mode 100644 index 0000000..c2b03eb --- /dev/null +++ b/test/test-Reactant-HVP.jl @@ -0,0 +1,156 @@ +using Test +using Random +using LinearAlgebra +using FlexiChains + +using ParallelMCMC +using ADTypes +# Stands in for "some backend that is not Reactant" in the pairing tests below. +using ForwardDiff: ForwardDiff + +#= +Reactant-compiled derivative paths (`AutoReactant`), see ext/ReactantExt.jl. + +The quartic target keeps second-order structure honest: logp = -0.25‖x‖⁴ has +H = -(‖x‖² I + 2 x xᵀ), so an HVP that silently drops the second-order term +(the failure mode of `Enzyme.hvp` under `@compile`) is caught, unlike a +Gaussian where H is constant. +=# +logp_r(x) = -0.25 * sum(abs2, x)^2 +gradlogp_r(x) = -sum(abs2, x) .* x +hvp_r(x, v) = -(sum(abs2, x) .* v .+ 2 .* dot(x, v) .* x) +logp_batch_r(X) = vec(-0.25 .* sum(abs2, X; dims=1) .^ 2) +gradlogp_batch_r(X) = -X .* sum(abs2, X; dims=1) + +const D_R = 4 +const CT_R = FlexiChains.FlexiChain{Symbol} + +#= The pairing rule lives in `_resolve_hvp`, not in the extension, so its +dispatch table is checked whether or not Reactant loads. =# +@testset "Reactant does not pair with a DI backend" begin + # Both slots Reactant, or a hand-written gradient (`nothing`), are accepted. + @test ParallelMCMC._check_reactant_pair(AutoReactant(), AutoReactant()) === nothing + @test ParallelMCMC._check_reactant_pair(nothing, AutoReactant()) === nothing + @test ParallelMCMC._check_reactant_pair(AutoForwardDiff(), AutoForwardDiff()) === + nothing + + # One of each is refused in both directions. + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + AutoReactant(), AutoForwardDiff() + ) + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + AutoForwardDiff(), AutoReactant() + ) +end + +reactant_ok = try + using Reactant: Reactant + using Enzyme: Enzyme + true +catch err + @warn "Reactant not available — skipping Reactant HVP tests" err + false +end + +if reactant_ok + @testset "extension is loaded" begin + @test Base.get_extension(ParallelMCMC, :ReactantExt) !== nothing + end + + #= The same rule reached through `_prepare_model`, where the gradient slot + resolves first: a mixed pair must still surface as an ArgumentError and not + as whatever DI or Reactant would say downstream. =# + @testset "mixed pairs are refused at preparation" begin + x = zeros(D_R) + + reactant_grad = DensityModel(logp_r, AutoReactant(), D_R; hvp=AutoForwardDiff()) + @test_throws ArgumentError ParallelMCMC._prepare_model(reactant_grad, x, 8, nothing) + + reactant_hvp = DensityModel(logp_r, AutoForwardDiff(), D_R; hvp=AutoReactant()) + @test_throws ArgumentError ParallelMCMC._prepare_model(reactant_hvp, x, 8, nothing) + end + + @testset "HVP matches analytic" begin + rng = MersenneTwister(71) + x = randn(rng, D_R) + v = randn(rng, D_R) + + @testset "forward over user gradient (sampler backend)" begin + model = DensityModel(logp_r, gradlogp_r, D_R) + m_p = ParallelMCMC._prepare_model(model, x, 8, AutoReactant()) + @test m_p.hvp(x, v) ≈ hvp_r(x, v) + end + + @testset "forward-over-reverse from logp alone (both slots AutoReactant)" begin + model = DensityModel(logp_r, AutoReactant(), D_R; hvp=AutoReactant()) + m_p = ParallelMCMC._prepare_model(model, x, 8, nothing) + @test m_p.grad_logdensity(x) ≈ gradlogp_r(x) + @test m_p.hvp(x, v) ≈ hvp_r(x, v) + end + + @testset "batched slots" begin + T = 8 + X = randn(rng, D_R, T) + V = randn(rng, D_R, T) + Hv_cols = reduce(hcat, [hvp_r(X[:, t], V[:, t]) for t in 1:T]) + + model = DensityModel( + logp_r, + AutoReactant(), + D_R; + logdensity_batch=logp_batch_r, + grad_logdensity_batch=AutoReactant(), + hvp=AutoReactant(), + hvp_batch=AutoReactant(), + ) + m_p = ParallelMCMC._prepare_model(model, X[:, 1], T, nothing) + @test m_p.grad_logdensity_batch(X) ≈ gradlogp_batch_r(X) + @test m_p.hvp_batch(X, V) ≈ Hv_cols + end + end + + @testset "end-to-end sampling with AutoReactant" begin + model = DensityModel(logp_r, AutoReactant(), D_R) + s = ParallelMALASampler(0.02; T=16, backend=AutoReactant()) + chain = sample(MersenneTwister(72), model, s, 64; chain_type=CT_R, progress=false) + @test size(chain) == (64, 1) + @test all(x -> all(isfinite, x), chain[:x]) + end + + reactant_gpu_ok = try + using CUDA: CUDA + CUDA.functional() && (CUDA.CuArray([1.0f0]); true) + catch + false + end + + if !reactant_gpu_ok + @info "Reactant HVP test: CUDA not functional — skipping CuArray boundary" + else + #= + The compiled executable lives in Reactant's own (XLA) device memory; + what's checked here is the CuArray <-> Reactant marshalling boundary: + CuArray in, CuArray out, values matching the analytic HVP. + =# + logp_r32(x) = -0.25f0 * sum(abs2, x)^2 + + @testset "CuArray boundary" begin + rng = MersenneTwister(73) + x_h = randn(rng, Float32, D_R) + v_h = randn(rng, Float32, D_R) + x_d = CUDA.CuArray(x_h) + v_d = CUDA.CuArray(v_h) + + model = DensityModel(logp_r32, AutoReactant(), D_R; hvp=AutoReactant()) + m_p = ParallelMCMC._prepare_model(model, x_d, 8, nothing) + + g = m_p.grad_logdensity(x_d) + @test g isa CUDA.CuArray + @test Array(g) ≈ gradlogp_r(x_h) + + Hv = m_p.hvp(x_d, v_d) + @test Hv isa CUDA.CuArray + @test Array(Hv) ≈ hvp_r(x_h, v_h) + end + end +end From fb05fe2eb8f6e706677552e0e23746771bc59881 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Wed, 12 Aug 2026 17:36:06 -0400 Subject: [PATCH 08/11] Some rewrites to text --- CHANGELOG.md | 184 +++++++++++---------- Project.toml | 2 +- docs/src/10-getting-started.md | 4 +- docs/src/15-gpu.md | 22 ++- docs/src/95-reference.md | 9 ++ ext/DynamicPPLExt.jl | 20 +-- ext/EnzymeExt.jl | 19 +-- ext/LogDensityProblemsExt.jl | 6 +- ext/ReactantExt.jl | 167 +++++++++++++------ src/DEER/DEER.jl | 120 ++++++++------ src/ParallelMCMC.jl | 4 + src/interface.jl | 259 ++++++++++++++++++++--------- test/Project.toml | 8 +- test/runtests.jl | 11 ++ test/test-HVP-Strategy.jl | 7 +- test/test-Reactant-HVP.jl | 288 ++++++++++++++++++++++++++++++--- 16 files changed, 796 insertions(+), 334 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 851d2e1..bde0fe4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,115 +10,123 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - The derivative slots of `DensityModel` (`grad_logdensity`, `hvp`, - `grad_logdensity_batch`, `hvp_batch`) now take an `ADTypes.AbstractADType` - in place of a callable, so `DensityModel(logp, AutoForwardDiff(), dim)` - builds a model from the log-density alone. Backends are turned into - prepared DifferentiationInterface callables when sampling starts, and that - preparation is reused across steps (#40, #52). The prepared model rides - along in the sampler state to get that reuse; a state handed to `step` for a - different model, which `initial_state` allows, is re-prepared from the model - passed rather than reused, so the model given to `sample` is the one sampled. -- `ParallelMALASampler`'s `backend` keyword is now optional. It is only the - fallback source of Hessian-vector products, so a `DensityModel` carrying - its own `hvp` / `hvp_batch` does not need it (#52). With no sampler - `backend`, a batched HVP is derived from the model's own `hvp` backend. + `grad_logdensity_batch`, `hvp_batch`) now accept an `ADTypes.AbstractADType` + in place of a callable, so `DensityModel(logp, AutoForwardDiff(), dim)` builds + a model from the log-density alone (#40, #52). Backends become prepared + DifferentiationInterface callables when sampling starts, and the prepared + model rides along in the sampler state so the preparation is reused. A state + handed to `step` for a different model — which `initial_state` allows — is + re-prepared, so the model given to `sample` is the one sampled. +- `ParallelMALASampler`'s `backend` keyword is now optional: it is only the + fallback source of Hessian-vector products, so a `DensityModel` carrying its + own `hvp` / `hvp_batch` does not need one. Without it, a batched HVP comes + from the model's own `hvp` backend (#52). - A `logdensity_batch` given without a `grad_logdensity_batch` now has the - batched gradient derived for it when `grad_logdensity` is a backend, rather - than leaving the batched DEER path switched off (#52). `hvp_batch` can be a - backend in that case too, and differentiates the derived gradient. -- New `ReactantExt`: `ADTypes.AutoReactant()` in any derivative slot, or as the + batched gradient derived from it when `grad_logdensity` is a backend, instead + of leaving the batched DEER path switched off. `hvp_batch` can be a backend in + that case too, and differentiates the derived gradient (#52). +- An HVP backend over an AD-derived gradient is now real second-order AD, + `DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` handed to + `DI.hvp`, rather than an outer AD pass over the prepared DI gradient — which + dropped out of its preparation as soon as tangents were pushed through it + (#37). Around 10x fewer allocations for a `logdensity`-only model. A backend + over a hand-written gradient still differentiates that gradient once. +- `hvp` / `hvp_batch` accept a `SecondOrder` with both halves honoured, so the + log-density is differentiated twice and the gradient slot is not the inner + pass. The inner half used to be discarded. This is the one AD route to an HVP + for a Turing or LogDensityProblems model, whose gradient arrives already + prepared and cannot be differentiated again. +- New `ReactantExt`. `ADTypes.AutoReactant()` in a derivative slot, or as the sampler `backend`, traces the derivative with Enzyme-MLIR and compiles it to - an XLA executable via Reactant.jl. This bypasses Enzyme's LLVM pipeline and - is currently the only path that computes a genuine second-order HVP on GPU, - so `ParallelMALASampler` now works on CUDA for log-density-only models - (#37, #52). Requires `using Reactant` and a Reactant-traceable log-density; - `AutoReactant` cannot be paired with a DifferentiationInterface backend - across the two passes of an HVP, which raises an `ArgumentError`. -- An HVP backend over an AD-derived gradient is now taken as true second-order - AD, `DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` handed - to `DI.hvp`, instead of an outer AD pass over the prepared DI gradient (#37). - That nesting dropped out of its preparation as soon as the outer pass pushed - tangents in, so the composed operator is both what was asked for and cheaper: - around 10x fewer allocations for a `logdensity`-only model. A backend over a - hand-written gradient still differentiates that gradient once, as before. -- `hvp` / `hvp_batch` accept a `SecondOrder` with both halves honoured, meaning - the log-density is differentiated twice and the gradient slot is not the inner - pass. Previously the inner half was silently discarded and only the outer used. - This is the one AD route to an HVP for a Turing or LogDensityProblems model, - whose gradient arrives already prepared and so cannot be differentiated again. + an XLA executable via Reactant.jl, off Enzyme's LLVM pipeline and off + DifferentiationInterface entirely, which yields a genuine second-order HVP for + a log-density-only model (#37, #52). Requires `using Reactant` and a + Reactant-traceable log-density. `AutoReactant` does not pair with a + DifferentiationInterface backend, a `LogDensityProblems` gradient, or a + `SecondOrder` across the two passes of an HVP, and doing so raises an + `ArgumentError` before any AD runs, as does a non-default + `AutoReactant(; mode=...)`. Two caveats, spelled out in `ext/ReactantExt.jl`'s + module docstring and the Reactant section of `docs/src/15-gpu.md`: a traced + log-density must be pure with respect to the data it captures, since a captured + array mutated after preparation stays frozen at its old value in every + derivative compiled from it, and the compiled derivative runs on whatever + device Reactant's XLA client targets, which need not be the GPU the model's + arrays live on. - The `DensityModel` constructors in `DynamicPPLExt` and `LogDensityProblemsExt` forward `logdensity_batch`, `grad_logdensity_batch` and `hvp_batch`, so a Turing or LogDensityProblems model can reach the batched DEER path. Neither provides a batched log-density, so `logdensity_batch` has to be written by hand. - Adds `JuliaFormatter` testing which was forgotten (#60). -### Fixed - -- The reverse-on-grad HVP path differentiated with `DI.inner(backend)` while - its strategy was routed on `DI.outer(backend)`, so a - `DifferentiationInterface.SecondOrder` ran the wrong half of the pair. A - `SecondOrder` now goes to the true second-order path instead of either - strategy, and the half-selecting helper it still uses agrees: normalization - applies to the outer pass. Unwrapping to that half happens before the - normalization hook is dispatched on, so a `SecondOrder(AutoEnzyme(), ...)` still - reaches `EnzymeExt` and gets its function annotation filled in rather than - running as a bare `AutoEnzyme()`. -- Backend normalization no longer picks a differentiation mode on the user's - behalf (#62). `EnzymeExt` pinned `mode=Enzyme.Forward` (with - `set_runtime_activity`) onto an `AutoEnzyme()` left mode-agnostic, on the - grounds that reverse mode hit a gc-transition abort on GPU and that composed - `pmcmc_matmul` calls needed runtime activity. The `pmcmc_*` Enzyme rules keep - Enzyme off both paths on their own now, so the pin bought nothing — and it cost - correctness, because it silently rewrote the direction of a `SecondOrder`'s - outer half. `SecondOrder(AutoEnzyme(), AutoForwardDiff())` is - reverse-over-forward to `hvp_mode`, its inner half being forward-only, and came - out forward-over-forward. Normalization now fills in only - `function_annotation=Enzyme.Const`, which is about this package's own read-only - HVP wrappers rather than about Enzyme's mode, and leaves `mode` exactly as given - — unset included, for DI to resolve from the operator it runs. `hvp_mode` is - therefore identical before and after normalization for every backend pair. - - A mode set explicitly was never overridden, so only mode-agnostic backends were - affected, and the HVP was a correct HVP either way; what changes is that the - composition asked for is the one that runs. The two normalization hooks - (`_hvp_forward_backend`, `_hvp_closure_backend`) collapse into a single - `_normalized_backend`, since without a mode to choose they no longer differ. - Users relying on a plain `AutoEnzyme()` being run forward should now pass - `AutoEnzyme(; mode=Enzyme.Forward)` explicitly. - ### Changed - The AD-HVP fallback strategy (forward-on-grad vs reverse-on-grad) now comes from DifferentiationInterface's `hvp_mode` trait rather than a hardcoded per-backend list, so `AutoEnzyme(mode=Enzyme.Reverse)` routes to the reverse-on-grad path (#38). -- Because a `logdensity_batch` without a `grad_logdensity_batch` now has the - batched gradient derived rather than switching the batched DEER path off, a - model whose `grad_logdensity` is a backend runs the batched update where it - used to run the unbatched one, and AD is applied to its `logdensity_batch`. On - GPU that subjects a function nothing was differentiating before to the - backend's restrictions (`pmcmc_*` wrappers for Enzyme). Supply - `grad_logdensity_batch` to keep AD out of the batched path. A model with a - hand-written `grad_logdensity` is unaffected: nothing derives a batched - gradient for it, so the batched path stays off as before. +- A model whose `grad_logdensity` is a backend now runs the batched update where + it used to run the unbatched one, since a `logdensity_batch` without a + `grad_logdensity_batch` has one derived for it. On GPU that puts the backend's + restrictions (`pmcmc_*` wrappers for Enzyme) on a `logdensity_batch` nothing + was differentiating before; supply `grad_logdensity_batch` to keep AD out of + it. A hand-written `grad_logdensity` is unaffected. - `ParallelMALASampler`'s `backend` no longer derives a batched gradient, only Hessian-vector products. It could previously switch the batched DEER path on - for a model with a hand-written gradient, which made a keyword that reads as - an HVP fallback decide which update path ran and put AD on a - `logdensity_batch` the user had not opted into differentiating. Models that - relied on that should pass `grad_logdensity_batch` explicitly, or a backend in - `grad_logdensity` for one to be derived from. -- Both batched derivative slots now require `logdensity_batch`, which the - batched update evaluates directly, and the constructor rejects them without - one. A callable `grad_logdensity_batch` or `hvp_batch` supplied on its own - used to be accepted and then silently ignored. `logdensity_batch` alone is - still valid and still used to score whole trajectories. + for a model with a hand-written gradient, which let a keyword that reads as an + HVP fallback decide which update path ran. Pass `grad_logdensity_batch` + explicitly, or a backend in `grad_logdensity` to derive one from. +- Both batched derivative slots now require `logdensity_batch`, which the batched + update evaluates directly, and the constructor rejects them without one. A + callable `grad_logdensity_batch` or `hvp_batch` on its own used to be accepted + and then ignored. `logdensity_batch` alone still scores whole trajectories. - An `hvp_batch` that reaches sampling with no batched gradient to pair it with - now raises rather than silently falling back to the unbatched update. + now raises rather than falling back to the unbatched update. + +### Fixed + +- The reverse-on-grad HVP path differentiated with `DI.inner(backend)` while its + strategy was routed on `DI.outer(backend)`, so a + `DifferentiationInterface.SecondOrder` ran the wrong half of the pair. A + `SecondOrder` now goes to the second-order path instead of either strategy, and + normalization applies to the outer half after unwrapping, so + `SecondOrder(AutoEnzyme(), ...)` still reaches `EnzymeExt` rather than running + as a bare `AutoEnzyme()`. +- Backend normalization no longer picks a differentiation mode on the user's + behalf (#62). `EnzymeExt` pinned `mode=Enzyme.Forward`, with + `set_runtime_activity`, onto a mode-agnostic `AutoEnzyme()`, against a + gc-transition abort on GPU and an `EnzymeRuntimeActivityError` on composed + `pmcmc_matmul` calls. The `pmcmc_*` Enzyme rules keep Enzyme off both paths on + their own, so the pin bought nothing and cost correctness: it rewrote the + direction of a `SecondOrder`'s outer half, turning + `SecondOrder(AutoEnzyme(), AutoForwardDiff())` — reverse-over-forward to + `hvp_mode`, its inner half being forward-only — into forward-over-forward. + Normalization now fills in `function_annotation=Enzyme.Const` and leaves `mode` + exactly as given, unset included, so `hvp_mode` reads the same before and after + it for every pair. Only mode-agnostic backends were affected and the HVP was + correct either way; what changes is that the composition asked for is the one + that runs. The two hooks `_hvp_forward_backend` and `_hvp_closure_backend` + collapse into one `_normalized_backend`. Pass + `AutoEnzyme(; mode=Enzyme.Forward)` to keep the old direction. +- `_check_reactant_pair` now runs before `_prepare_model` resolves a gradient, so + a mismatched `AutoReactant` pairing is reported immediately instead of after a + full XLA compile, and `AutoReactant` nested inside a `SecondOrder` is rejected + there too rather than reaching `DI.prepare_hvp`. An `AutoReactant` gradient + composes through the same second-order branch as any other AD-derived gradient, + and `_hvp_strategy(::AutoReactant)` is now live rather than a dead branch. +- `test/test-Reactant-HVP.jl` asserts a posterior mean against a target with a + known mean, and cross-checks against the analytic HVP on the same noise tape. + `size(chain)` and `all(isfinite, ...)` pass even for a badly wrong HVP, since + DEER's Newton iteration then just fails to converge rather than producing + `NaN`s. ### Removed - `DynamicPPLExt` no longer requires `ForwardDiff` as a triggering library to load. +- `Reactant` moved from `test/Project.toml`'s `[deps]` to `[extras]`, and + `test/test-Reactant-HVP.jl` is skipped unless `PARALLELMCMC_TEST_REACTANT` is + set. `Reactant_jll` ships a prebuilt XLA and is a large download most CI runs + and most local `Pkg.test()` calls should not have to pay for. Opt in by setting + the env var and adding `Reactant` to the test environment. ## [0.2.0] - 2026-06-29 diff --git a/Project.toml b/Project.toml index 7f7f77a..40b34e9 100644 --- a/Project.toml +++ b/Project.toml @@ -30,7 +30,7 @@ LogDensityProblemsExt = "LogDensityProblems" ReactantExt = ["Reactant", "Enzyme"] [compat] -ADTypes = "1.22.0" +ADTypes = "1.21.0" AbstractMCMC = "5.10.0" CUDA = "5.11.0" CUDA_Runtime_jll = "0.21" diff --git a/docs/src/10-getting-started.md b/docs/src/10-getting-started.md index b257131..50f4d4a 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -49,11 +49,11 @@ Backends become prepared [DifferentiationInterface](https://github.com/JuliaDiff Naming both passes yourself is the one route that ignores the gradient slot, hand-written or not. It is also the only AD route to an HVP for a Turing or LogDensityProblems model, whose gradient arrives already prepared and cannot be differentiated again. -Which pairs work is up to the backends. On CPU, ForwardDiff, ReverseDiff, Zygote and Enzyme all serve a log-density-only model. `AutoMooncake` serves neither direction: it has no reverse-over-reverse, and its gradient rejects an outer pass's tangents. Give Mooncake a hand-written `grad_logdensity` instead. On GPU, no DI-driven second-order pair works yet (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)); `AutoReactant()` in both slots is the exception, and compiles the pair to a single XLA program. See the [GPU page](15-gpu.md). +Which pairs work is up to the backends. On CPU, ForwardDiff, ReverseDiff, Zygote and Enzyme all serve a log-density-only model. `AutoMooncake` serves neither direction: it has no reverse-over-reverse, and its gradient rejects an outer pass's tangents, so give Mooncake a hand-written `grad_logdensity`. On GPU no DI-driven second-order pair works yet (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)). `AutoReactant()` in both slots is the exception: it skips DI and traces forward-over-reverse straight from the log-density as one compiled XLA program. See the [GPU page](15-gpu.md). The batched pair works the same way, on `sum(logdensity_batch(X))`. That sum's gradient is the stacked per-column gradients only because the columns are independent, so `logdensity_batch` must not couple them. Omitting `grad_logdensity_batch` derives one when `grad_logdensity` is a backend; with a hand-written gradient the batched path stays off and the unbatched update covers it. Both batched derivative slots require `logdensity_batch`, which is also useful on its own for scoring a whole trajectory at once. -`backend` on [`ParallelMALASampler`](@ref) supplies Hessian-vector products for a model that brings no `hvp` / `hvp_batch` of its own, and nothing else. A model carrying its own can leave it out. +`backend` on [`ParallelMALASampler`](@ref) supplies Hessian-vector products for a model that brings no `hvp` / `hvp_batch` of its own, and nothing else, so a model carrying its own can leave it out. --- diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index 1f8c5a2..318477a 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -233,14 +233,14 @@ DEER needs a Hessian–vector product $H v$ at every Newton step. `DensityModel - **You only supply `gradlogp` / `grad_logdensity_batch`.** The sampler builds the HVP by differentiating your gradient — either a forward-mode pushforward of `gradlogp` ([`ForwardOnGrad`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/src/DEER/DEER.jl), the default for most backends) or a reverse-mode gradient of `x -> dot(gradlogp(x), v)` ([`ReverseOnGrad`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/src/DEER/DEER.jl), used for `AutoMooncake` and `AutoZygote`). This is the **AD-HVP fallback**, and it is what the logistic-regression example above uses. !!! warning "Log-density-only models on GPU" - `grad_logdensity` can itself be an AD backend (`DensityModel(logp, AutoEnzyme(), dim)`, see [Getting started](10-getting-started.md)), but don't do that with `ParallelMALASampler` on GPU for the DI-driven backends. The HVP becomes `SecondOrder(hvp_backend, grad_backend)` on your log-density, which currently fails on GPU with both Enzyme and Mooncake (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)). The same goes for passing a `SecondOrder` explicitly. Write `gradlogp` out by hand for DEER, so the HVP is a single pass over it — or use `AutoReactant()` (below), the one backend where log-density-only DEER works on GPU. The sequential samplers only need the gradient, so log-density-only models work there. + `grad_logdensity` can itself be an AD backend (`DensityModel(logp, AutoEnzyme(), dim)`, see [Getting started](10-getting-started.md)), but not with `ParallelMALASampler` on GPU for the DI-driven backends. The HVP then becomes `SecondOrder(hvp_backend, grad_backend)` on your log-density, which currently fails on GPU with both Enzyme and Mooncake (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)); passing a `SecondOrder` explicitly hits the same wall. Write `gradlogp` out by hand so the HVP is a single pass over it, or use `AutoReactant()` below. The sequential samplers only need the gradient, so log-density-only models are fine there. !!! note "A backend in `grad_logdensity` reaches `logdensity_batch` too" The batched path needs a batched gradient, and derives one from `logdensity_batch` when `grad_logdensity` is a backend. That puts `logdensity_batch` under the same restrictions as the rest of your AD-visible code. Supply `grad_logdensity_batch` to avoid it. -### Reactant: genuine second-order HVPs on GPU +### Reactant HVPs, off the DI path -`ADTypes.AutoReactant()` (requires `using Reactant`) takes a different route entirely. The derivative is traced with Enzyme-MLIR and compiled to an XLA executable by [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl), bypassing Enzyme's LLVM pipeline — and with it the gc-transition abort and the `pmcmc_*` wrapper requirements above. It is currently the only path that computes a genuine second-order HVP on GPU, so log-density-only models work with DEER: +`ADTypes.AutoReactant()` (requires `using Reactant`) goes another way. The derivative is traced with Enzyme-MLIR and compiled to an XLA executable by [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl), never touching Enzyme's LLVM pipeline, so neither the gc-transition abort nor the `pmcmc_*` wrappers above apply. It gives a genuine second-order HVP for a log-density-only model: ```julia using Reactant, ADTypes @@ -249,15 +249,21 @@ model = DensityModel(logp, AutoReactant(), D) # no hand-written gradient sampler = ParallelMALASampler(0.005f0; T=16, backend=AutoReactant()) ``` -When both the gradient slot and the HVP source are `AutoReactant()`, the HVP compiles as explicit forward-over-reverse from the raw log-density, in a single fused XLA program. With a hand-written `gradlogp`, it compiles a forward pushforward of your gradient instead. +With both the gradient slot and the HVP source `AutoReactant()`, the HVP compiles as explicit forward-over-reverse from the raw log-density, as its own XLA program. That is a separate executable from the plain gradient's: what gets fused is the HVP's two AD passes, not the gradient and the HVP. With a hand-written `gradlogp` it compiles a forward pushforward of your gradient instead. -Three caveats: +!!! warning "Two silent failure modes" + **Captured data is frozen at compile time.** `Reactant.@compile` bakes any plain `Array` or `Ref` reached through the closure into the executable as a constant. A `logdensity` written as `x -> f(x, data)` whose `data` you later mutate (`data .= new_values`) keeps returning the pre-mutation derivative from every executable compiled before the mutation, with no error and no warning, which biases the samples. Traced functions must be pure with respect to what they capture; pass data that can change in as an argument. + + **"GPU" here describes your array type, not the device Reactant runs on.** Which XLA client Reactant compiles for is a Reactant/`Reactant_jll`-wide setting (`Reactant.set_default_backend`) that nothing in this package controls, and it is `"cpu"` unless a GPU client was selected explicitly. Preparing a model whose `x_template` is a `CuArray` against a CPU client still compiles and still gives correct answers, but every call round-trips `CuArray → host → XLA-CPU → host → CuArray`. `_prepare_model` warns on that combination and can do nothing else about it; point Reactant at a GPU client yourself. + +Other caveats: - The traced function must be **Reactant-traceable**: plain array operations. DynamicPPL-built log-densities do not trace as-is. -- **Reactant does not mix with the DI backends across the two passes of an HVP.** Both `grad_logdensity` and the HVP source take `AutoReactant()`, or neither does; a hand-written gradient pairs with either. Mixing raises an `ArgumentError` at preparation time, since Reactant cannot trace a DI-prepared gradient and DI cannot differentiate a compiled XLA executable. -- Calls cross a marshalling boundary (package arrays ↔ Reactant's XLA device memory) on every invocation, and executables are shape-specialized at preparation time. Correct everywhere, but it leaves fusion on the table relative to a future end-to-end Reactant pipeline. +- **Reactant does not mix with the DI backends across the two passes of an HVP.** Both `grad_logdensity` and the HVP source take `AutoReactant()`, or neither does; a hand-written gradient pairs with either. Mixing raises an `ArgumentError` at preparation time, since Reactant cannot trace a DI-prepared gradient and DI cannot differentiate a compiled XLA executable. Same for a `LogDensityProblems`/Turing gradient, and for nesting `AutoReactant()` inside a `DifferentiationInterface.SecondOrder`. +- **`AutoReactant(; mode=...)` raises** rather than being ignored. Derivatives always trace as Enzyme reverse-mode with the HVP as forward-over-that, whatever `mode` says. +- **Every call crosses a marshalling boundary** between package arrays and Reactant's XLA device memory: two host round-trips and a handful of allocations, not a fused in-place path. Executables are shape-specialized at preparation time, and `@compile` does not memoize across preparations, so every `sample()` call recompiles every `AutoReactant` slot rather than only the first one in the process. -DifferentiationInterface cannot drive Reactant yet; when that support lands ([DI#918](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/918)), this path folds into the standard `hvp_mode` routing with no user-facing change. +DifferentiationInterface cannot interoperate with Reactant. When that support lands ([DI#918](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/918)) this path folds into the ordinary `hvp_mode` routing with no user-facing change. ### When the fallback is the right call diff --git a/docs/src/95-reference.md b/docs/src/95-reference.md index b7f2fb9..aac80ed 100644 --- a/docs/src/95-reference.md +++ b/docs/src/95-reference.md @@ -6,6 +6,15 @@ CurrentModule = ParallelMCMC This page documents all public types and functions exported by ParallelMCMC.jl. +## Reactant + +Loading `Reactant` (`using Reactant`) enables `ADTypes.AutoReactant()` as a +derivative-slot backend on `DensityModel` and as `ParallelMALASampler`'s +`backend`. It is not a `DifferentiationInterface` backend like the others. +[GPU Execution](15-gpu.md) covers what it does, its pairing rules, and its +caveats, in particular that traced log-densities must be pure with respect +to any data they capture. + ## Extension constructors `DensityModel` also has extension constructors for common probabilistic-programming interfaces: diff --git a/ext/DynamicPPLExt.jl b/ext/DynamicPPLExt.jl index 064c6d6..f824f8c 100644 --- a/ext/DynamicPPLExt.jl +++ b/ext/DynamicPPLExt.jl @@ -18,17 +18,19 @@ computation via DynamicPPL's `adtype` interface. Requires `DynamicPPL` and `LogDensityProblems` to be loaded (these are the weak-dependency triggers for this extension), plus any AD backend that is used. -`ad_backend` is DynamicPPL's own `adtype`, not a `DensityModel` slot: it goes to -the `LogDensityFunction` that fills the log-density and gradient slots, which is -why it takes a backend only and never a callable. The rest are `DensityModel` -slots forwarded unchanged. +`ad_backend` is DynamicPPL's own `adtype` rather than a `DensityModel` slot: it +goes to the `LogDensityFunction` that fills the log-density and gradient slots, +which is why it takes a backend and never a callable. The remaining keywords are +`DensityModel` slots, forwarded unchanged. -`ParallelMALASampler` also needs an HVP. Give it a callable or a +`ParallelMALASampler` also wants an HVP. Give it a callable or a `DifferentiationInterface.SecondOrder`, which differentiates the log-density and -so bypasses DynamicPPL's gradient. A plain backend fails, since it would -differentiate the gradient `ad_backend` produced and that preparation rejects an -outer pass's tangents. DynamicPPL supplies no batched log-density either, so -reaching the batched DEER path means writing `logdensity_batch` by hand. +so bypasses DynamicPPL's gradient. A plain backend fails: it would differentiate +the gradient `ad_backend` produced, whose preparation rejects an outer pass's +tangents. `ADTypes.AutoReactant()` fails as well, bypassing DI or not, since +Reactant cannot trace DynamicPPL's model evaluation. DynamicPPL supplies no +batched log-density either, so reaching the batched DEER path means writing +`logdensity_batch` by hand. # Example ```julia diff --git a/ext/EnzymeExt.jl b/ext/EnzymeExt.jl index 82e63f8..33899fd 100644 --- a/ext/EnzymeExt.jl +++ b/ext/EnzymeExt.jl @@ -23,19 +23,14 @@ using Enzyme.EnzymeCore.EnzymeRules: Normalization of the user's `AutoEnzyme` for DEER's AD-HVP paths: fill in `function_annotation=Enzyme.Const` when they left it open, so Enzyme doesn't throw `EnzymeMutabilityException` on the read-only `_HvpReverseClosure` / -`_BatchHvpReverseClosure` wrappers, which capture `gradlogp`. Those wrapper types -belong to this package, so declaring them constant is this package's business. +`_BatchHvpReverseClosure` wrappers that capture `gradlogp`. -`mode` is passed through exactly as given, unset included. Choosing a direction -on the user's behalf is not our call: a mode they set is a decision, and an unset -one is DI's to resolve from the operator it runs. - -An earlier version pinned `mode=Enzyme.Forward` here (with -`set_runtime_activity`) against the gc-transition abort on GPU and -`EnzymeRuntimeActivityError` on composed `pmcmc_matmul` calls. The rules below -keep Enzyme off both paths on their own, so the pin bought nothing and cost -correctness: it silently rewrote the direction of a `SecondOrder`'s outer half -(see `DEER._normalized_backend`). +`mode` passes through exactly as given, unset included. An earlier version pinned +`mode=Enzyme.Forward` here, with `set_runtime_activity`, against the GPU +gc-transition abort and the `EnzymeRuntimeActivityError` on composed +`pmcmc_matmul` calls. The rules below keep Enzyme off both paths by themselves, +so the pin bought nothing and cost correctness: it rewrote the direction of a +`SecondOrder`'s outer half (see `DEER._normalized_backend`). =# function DEER._normalized_backend(backend::ADTypes.AutoEnzyme{M,A}) where {M,A} A === Nothing || return backend diff --git a/ext/LogDensityProblemsExt.jl b/ext/LogDensityProblemsExt.jl index 6f59578..b146b2b 100644 --- a/ext/LogDensityProblemsExt.jl +++ b/ext/LogDensityProblemsExt.jl @@ -24,8 +24,10 @@ parameter named `:x` will be chosen, unless you also pass `param_names` to `samp and keep their meaning there, with one caveat. `ld` fills the gradient slot, and a gradient `ld` computes by AD carries a preparation tied to its input type, so it rejects the tangents a plain `hvp` backend would push through it. Use a -callable, or a `DifferentiationInterface.SecondOrder` which differentiates the -log-density instead. Same for `hvp_batch`. +callable, or a `DifferentiationInterface.SecondOrder`, which differentiates the +log-density instead. Same for `hvp_batch`. `ADTypes.AutoReactant()` is out here +too, DI or no DI: Reactant cannot trace the DynamicPPL/LogDensityProblems +machinery `ld`'s gradient dispatches into. `ld` supplies no batched log-density, so reaching the batched DEER path means writing `logdensity_batch` by hand. diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl index 21df47a..6c14c01 100644 --- a/ext/ReactantExt.jl +++ b/ext/ReactantExt.jl @@ -1,12 +1,30 @@ module ReactantExt #= -Reactant-compiled derivative paths, selected by `ADTypes.AutoReactant()` in -any derivative slot of `DensityModel` (or as the sampler `backend`). -Derivatives are traced with Enzyme-MLIR and compiled to XLA executables via -`Reactant.@compile`, bypassing Enzyme's LLVM pipeline — and with it the GPU -`cuMemcpyDtoHAsync_v2` gc-transition abort. This is the only path that -computes a genuine second-order HVP on GPU. +Reactant-compiled derivative paths, selected by `ADTypes.AutoReactant()` in any +derivative slot of `DensityModel`, or as the sampler `backend`. Derivatives are +traced with Enzyme-MLIR and compiled to XLA executables by `Reactant.@compile`, +which keeps them off Enzyme's LLVM pipeline and so off the GPU +`cuMemcpyDtoHAsync_v2` gc-transition abort, and off DifferentiationInterface +entirely. For a log-density-only model it yields a genuine second-order HVP. + +Two silent failure modes, ahead of the ordinary limitations: + +Captured data is frozen at compile time. `@compile` bakes any plain `Array` / +`Ref` reached through the traced closure into the executable as a constant. A +`logdensity` written as `x -> f(x, data)` whose `data` is mutated afterwards +keeps handing back the pre-mutation derivative from every executable compiled +before the mutation, with no error and no warning. Traced functions must be pure +with respect to what they capture; pass data that can change in as an argument. + +The compiled program runs wherever Reactant's XLA client points, which need not +be a GPU. That client is a Reactant/`Reactant_jll`-wide setting +(`Reactant.set_default_backend`) and has nothing to do with where the package's +own `Vector`s / `CuArray`s live. Preparing on `CuArray` parameters while the +client targets `"cpu"` still compiles and still gives correct answers, but every +call round-trips `CuArray -> host -> XLA-CPU -> host -> CuArray`. +`_warn_reactant_host_roundtrip` below warns once per compiled slot on that +combination; it cannot fix it. Requirements / limitations: - The traced function (`logdensity` / `gradlogp` / batched forms) must be @@ -15,28 +33,83 @@ Requirements / limitations: - Executables are shape-specialized to the preparation templates, and arguments are marshalled to/from Reactant's own (XLA) device memory on every call. A boundary copy, not a fused in-place path — a target for - later optimization. + later optimization. `@compile` also does not memoize across preparations: + every `sample()` call recompiles every `AutoReactant` slot the model uses. - HVPs are explicit forward-over-(gradient) compositions; NEVER `Enzyme.hvp`, which silently returns zeros under `@compile`. - - `AutoReactant.mode` (the wrapped `AutoEnzyme`) is currently ignored: - gradients always trace as Enzyme reverse, HVPs as forward-over-that. + - `AutoReactant.mode` (the wrapped `AutoEnzyme`) is not honoured: gradients + always trace as Enzyme reverse, HVPs as forward-over-that. A non-default + `mode` is rejected outright (`_check_reactant_mode`) rather than ignored. =# +# Both `Reactant` and `Enzyme` trigger this extension (see Project.toml): the +# traced derivatives call `Enzyme.autodiff` / `Enzyme.gradient` themselves +# rather than reaching Enzyme-MLIR through Reactant. using ParallelMCMC: ParallelMCMC using ParallelMCMC.DEER: DEER -using ADTypes: AutoReactant +using ADTypes: ADTypes, AutoReactant, AutoEnzyme using Reactant: Reactant, @compile using Enzyme: Enzyme +#= +`AutoReactant()` defaults to `AutoReactant(; mode=AutoEnzyme())`, i.e. +`AutoReactant{AutoEnzyme{Nothing,Nothing}}` — the exact type matched below. +Anything else names an Enzyme mode or annotation, and this extension honours +neither. +=# +_check_reactant_mode(::AutoReactant{AutoEnzyme{Nothing,Nothing}}) = nothing +function _check_reactant_mode(backend::AutoReactant) + return throw( + ArgumentError( + "AutoReactant(; mode=$(backend.mode)) is not supported: gradients always " * + "trace as Enzyme reverse-mode and HVPs as forward-over-that, whatever " * + "`mode` says. Use the default AutoReactant().", + ), + ) +end + +#= +Warn once per compiled slot when the template is not a plain `Array` (so looks +like it lives on a GPU) while Reactant's default XLA client targets "cpu": every +call then pays a host round trip on top of the usual marshalling. Guarded, so a +Reactant version without `XLA.platform_name` / `XLA.default_backend` degrades to +no warning instead of erroring out of `_prepare_model`. +=# +function _reactant_client_platform() + return try + string(Reactant.XLA.platform_name(Reactant.XLA.default_backend())) + catch + nothing + end +end + +_warn_reactant_host_roundtrip(::Array) = nothing +function _warn_reactant_host_roundtrip(x::AbstractArray) + if _reactant_client_platform() == "cpu" + @warn "AutoReactant: preparing on a $(typeof(x)), but Reactant's default XLA " * + "client targets \"cpu\". Every call will round-trip to the host and back " * + "instead of running where the array lives, which is likely slower than not " * + "using Reactant at all. Point Reactant at a GPU client with " * + "`Reactant.set_default_backend(\"gpu\")` if one is available." maxlog = 1 + end + return nothing +end + # Host-materialize for the Reactant boundary: Reactant manages its own (XLA) # device memory, so we round-trip through a plain host Array regardless of -# where the package's arrays live (Vector / CuArray). +# where the package's arrays live (Vector / CuArray / SubArray). _host(x::Array) = x _host(x::AbstractArray) = Array(x) +#= +Eltype comes from `out`, what Reactant actually computed, not from `template`: a +traced computation that promotes internally would otherwise be narrowed back to +the template's eltype on the way out. +=# function _from_host(template::AbstractArray, out) - res = similar(template, size(out)) - copyto!(res, Array(out)) + out_h = Array(out) + res = similar(template, eltype(out_h), size(out_h)) + copyto!(res, out_h) return res end @@ -45,12 +118,14 @@ Compile `core` for the template shapes once and return a closure that marshals package arrays <-> Reactant arrays. =# function _compiled(core, t1::AbstractArray) + _warn_reactant_host_roundtrip(t1) r1 = Reactant.to_rarray(_host(t1)) compiled = @compile core(r1) return x -> _from_host(x, compiled(Reactant.to_rarray(_host(x)))) end function _compiled(core, t1::AbstractArray, t2::AbstractArray) + _warn_reactant_host_roundtrip(t1) r1 = Reactant.to_rarray(_host(t1)) r2 = Reactant.to_rarray(_host(t2)) compiled = @compile core(r1, r2) @@ -68,17 +143,11 @@ end _rev_gradient(f, x) = Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(f), x)[1] -# Columns of X are independent samples, so ∇_X sum(logp_batch(X)) stacks the -# per-column gradients. -_sumbatch(f, X) = sum(f(X)) - #= -Gradient slots. The wrappers keep the raw log-density so the HVP factories -below can re-trace forward-over-reverse from it, instead of trying to trace -through an already-compiled executable. +Gradient slots. They hold only the compiled callable; the HVP factories below +get `logdensity` from `_resolve_hvp`, which already has it. =# -struct _ReactantGradient{F,C} - logdensity::F +struct _ReactantGradient{C} compiled::C end (g::_ReactantGradient)(x) = g.compiled(x) @@ -86,12 +155,12 @@ end function ParallelMCMC._reactant_resolve_gradient( logdensity, backend::AutoReactant, x_template::AbstractVector ) + _check_reactant_mode(backend) core = Base.Fix1(_rev_gradient, logdensity) - return _ReactantGradient(logdensity, _compiled(core, x_template)) + return _ReactantGradient(_compiled(core, x_template)) end -struct _ReactantGradientBatch{F,C} - logdensity_batch::F +struct _ReactantGradientBatch{C} compiled::C end (g::_ReactantGradientBatch)(X) = g.compiled(X) @@ -99,26 +168,32 @@ end function ParallelMCMC._reactant_resolve_gradient_batch( logdensity_batch, backend::AutoReactant, X_template::AbstractMatrix ) - core = Base.Fix1(_rev_gradient, Base.Fix1(_sumbatch, logdensity_batch)) - return _ReactantGradientBatch(logdensity_batch, _compiled(core, X_template)) + _check_reactant_mode(backend) + # `_BatchLogdensitySum` is the same sum-over-columns the DI-driven batched + # gradient uses (src/interface.jl), so both batched paths differentiate the + # same thing. + core = Base.Fix1(_rev_gradient, ParallelMCMC._BatchLogdensitySum(logdensity_batch)) + return _ReactantGradientBatch(_compiled(core, X_template)) end #= -HVP factories. Two tracings depending on where the gradient came from: - - - `_ReactantGradient` (the gradient slot was itself AutoReactant): - re-trace from the raw log-density as explicit forward-over-reverse — - genuine second-order AD fused into one XLA program. - - any other callable (hand-written gradient): forward JVP over it, - provided it is traceable. +HVP factories. `_resolve_hvp` / `_resolve_hvp_batch` (src/interface.jl) route to +one of two shapes, the same two every other backend gets. + + - Both `grad_logdensity` and the HVP source `AutoReactant`: the AD-derived + gradient case, with `_second_order` collapsing the pair to one + `AutoReactant()` since DI cannot form a `SecondOrder` from it. + `_make_hvp_fn_second_order` here traces forward-over-reverse from + `logdensity` as a single XLA program. + - A hand-written `gradlogp` with an `AutoReactant` HVP source: routed by + `_hvp_strategy(::AutoReactant) = ReactantHVP()` in `DEER.jl` to + `_make_hvp_fn` below, a forward JVP over that callable. =# -function DEER._make_hvp_fn( - ::DEER.ReactantHVP, - gradlogp::_ReactantGradient, - backend::AutoReactant, - x_template::AbstractVector, +function DEER._make_hvp_fn_second_order( + logdensity, backend::AutoReactant, x_template::AbstractVector ) - inner = Base.Fix1(_rev_gradient, gradlogp.logdensity) + _check_reactant_mode(backend) + inner = Base.Fix1(_rev_gradient, logdensity) core(x, v) = _jvp(inner, x, v) return _compiled(core, x_template, x_template) end @@ -126,17 +201,16 @@ end function DEER._make_hvp_fn( ::DEER.ReactantHVP, gradlogp, backend::AutoReactant, x_template::AbstractVector ) + _check_reactant_mode(backend) core(x, v) = _jvp(gradlogp, x, v) return _compiled(core, x_template, x_template) end -function DEER._make_hvp_batch_fn( - ::DEER.ReactantHVP, - grad_batch::_ReactantGradientBatch, - backend::AutoReactant, - X_template::AbstractMatrix, +function DEER._make_hvp_batch_fn_second_order( + logdensity_batch_sum, backend::AutoReactant, X_template::AbstractMatrix ) - inner = Base.Fix1(_rev_gradient, Base.Fix1(_sumbatch, grad_batch.logdensity_batch)) + _check_reactant_mode(backend) + inner = Base.Fix1(_rev_gradient, logdensity_batch_sum) core(X, V) = _jvp(inner, X, V) return _compiled(core, X_template, X_template) end @@ -145,6 +219,7 @@ function DEER._make_hvp_batch_fn( ::DEER.ReactantHVP, grad_batch, backend::AutoReactant, X_template::AbstractMatrix ) # Column-independent batched gradient ⇒ forward JVP is the columnwise HVP. + _check_reactant_mode(backend) core(X, V) = _jvp(grad_batch, X, V) return _compiled(core, X_template, X_template) end diff --git a/src/DEER/DEER.jl b/src/DEER/DEER.jl index 3abe628..07919ea 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -164,7 +164,7 @@ We bundle the closure with the prep so `prepare_gradient` and `gradient` see the same function instance (DI keys preparations on function identity). --------------------------------------------------------------------------- =# -import ..ParallelMCMC: pmcmc_dot, pmcmc_dotsum +import ..ParallelMCMC: pmcmc_dot, pmcmc_dotsum, _REACTANT_LOAD_HINT struct _HvpReverseClosure{F} grad::F @@ -177,17 +177,15 @@ end (c::_BatchHvpReverseClosure)(X, V) = pmcmc_dotsum(c.grad_batch(X), V) #= -Pick the AD-HVP fallback strategy from the user's backend. These two apply when -the HVP is one AD pass over a gradient we already have, i.e., a hand-written -`gradlogp`, which neither of them differentiates twice: +Pick the AD-HVP fallback strategy from the user's backend. Both are one AD pass +over a hand-written `gradlogp`; an AD-derived gradient goes to +`_make_hvp_fn_second_order` instead. ForwardOnGrad() — `pushforward(gradlogp, x, v)`. Routes through the `pmcmc_matmul` frule. ReverseOnGrad() — `gradient(x -> pmcmc_dot(gradlogp(x), v))`. Routes through the matmul and dot/sum rrules. -An AD-derived gradient takes neither and goes to `_make_hvp_fn_second_order`. - These are singleton types rather than symbols so the choice dispatches statically — `_make_hvp_fn(_hvp_strategy(backend), ...)` resolves to one concrete method (and one concrete return type) at compile time, without @@ -204,14 +202,17 @@ struct ForwardOnGrad <: HVPStrategy end struct ReverseOnGrad <: HVPStrategy end #= -ReactantHVP — trace the HVP with Enzyme-MLIR and compile it to an XLA -executable via Reactant.jl (see `ext/ReactantExt.jl`). Selected by -`ADTypes.AutoReactant()`, which DI cannot drive yet, so it short-circuits ahead -of the `hvp_mode` routing above; once DI gains Reactant support the -`AutoReactant` specializations can be deleted. Reactant bypasses Enzyme's LLVM -pipeline, avoiding the GPU gc-transition abort — the only path that computes a -genuine second-order HVP on GPU. The traced functions must be Reactant-traceable -(plain array ops). +ReactantHVP traces the HVP with Enzyme-MLIR and compiles it to an XLA +executable, off Enzyme's LLVM pipeline and so off the GPU gc-transition abort. +DI cannot drive Reactant, so `AutoReactant` short-circuits the `hvp_mode` +routing above; drop these specializations once it can. What the traced function +has to look like, and which device the compiled program actually runs on, are in +`ext/ReactantExt.jl`'s module docstring. + +Only reached over a hand-written `gradlogp`. An `AutoReactant` gradient slot is +an AD-derived gradient like any other and goes to `_make_hvp_fn_second_order` / +`_make_hvp_batch_fn_second_order`, which dispatch on the backend rather than on +the resolved gradient's type (see `_resolve_hvp`). =# struct ReactantHVP <: HVPStrategy end @@ -222,26 +223,21 @@ _hvp_strategy(backend::AbstractADType) = _strategy_from(DI.hvp_mode(backend)) _hvp_strategy(::ADTypes.AutoReactant) = ReactantHVP() #= -Hook for backend-specific normalization of the user's `backend`, applied on every -AD-HVP path before the backend reaches DI. - -It supplies what the wrappers DEER differentiates need, and nothing else. Those -wrapper types are ours, so annotating them is ours to do: EnzymeExt specializes -this to fill `function_annotation=Enzyme.Const`, without which Enzyme throws +Hook for backend-specific normalization, applied on every AD-HVP path before the +backend reaches DI. It fills in what the wrappers DEER differentiates need and +nothing else: those wrapper types are ours, so EnzymeExt sets +`function_annotation=Enzyme.Const` on them, without which Enzyme throws `EnzymeMutabilityException` on the read-only `_HvpReverseClosure` / -`_BatchHvpReverseClosure`, which capture `gradlogp`. - -It deliberately does not choose a differentiation mode. Which direction a pass -runs is the user's call when they state one and DI's to resolve from the operator -when they don't; this package is not an AD package and has no business overriding -either. Picking one here also used to corrupt a `SecondOrder`, whose halves carry -directions of their own (see `_normalized_second_order`). - -Callers hand this a single pass, never a `SecondOrder`: the strategy paths below -run one AD pass over a hand-written `gradlogp`, and `_resolve_hvp` sends every -`SecondOrder` to `_make_hvp_fn_second_order` before they are reached. -`_normalized_second_order` is the one caller that starts from a pair, and it -selects the outer half itself. +`_BatchHvpReverseClosure` that capture `gradlogp`. + +It does not choose a differentiation mode. A mode the user set is a decision, an +unset one is DI's to resolve from the operator it runs, and substituting one here +used to rewrite the outer half of a `SecondOrder` out from under `hvp_mode` (see +`_normalized_second_order`). + +Only ever handed a single pass. `_resolve_hvp` sends every `SecondOrder` to +`_make_hvp_fn_second_order` before the strategy paths below are reached, and +`_normalized_second_order` picks the outer half itself. =# _normalized_backend(backend::AbstractADType) = backend @@ -314,13 +310,11 @@ function _make_hvp_batch_fn( end #= -`ReactantHVP` fallbacks. `ReactantExt` adds methods with `backend` pinned to -`ADTypes.AutoReactant` (strictly more specific — no method overwriting, which -precompilation forbids); without Reactant loaded these give a clear error -instead of a `MethodError`. +`ReactantHVP` fallbacks, so a missing `using Reactant` gives the load hint +rather than a `MethodError`. `ReactantExt` pins `backend` to +`ADTypes.AutoReactant`, which is strictly more specific, so nothing is +overwritten — precompilation forbids that. =# -const _REACTANT_LOAD_HINT = "AutoReactant requires Reactant.jl: add `using Reactant` to load ParallelMCMC's ReactantExt." - function _make_hvp_fn( ::ReactantHVP, gradlogp, backend::AbstractADType, x_template::AbstractVector ) @@ -333,23 +327,43 @@ function _make_hvp_batch_fn( return error(_REACTANT_LOAD_HINT) end +#= +Fallback for the "both slots `AutoReactant`" second-order path, which +`_second_order` in `interface.jl` routes here with `backend::AutoReactant` +rather than a `DI.SecondOrder`. Signature is `AbstractADType` and not +`AutoReactant` because `ReactantExt`'s method is `AutoReactant` exactly, and +precompilation refuses an identical signature; less specific still loses to it. +JET needs a method here too, since it cannot see a conditionally-loaded +extension and would otherwise flag `_resolve_hvp`'s `AutoReactant` branch as +having none. +=# +function _make_hvp_fn_second_order( + logdensity, backend::AbstractADType, x_template::AbstractVector +) + return error(_REACTANT_LOAD_HINT) +end + +function _make_hvp_batch_fn_second_order( + logdensity_batch_sum, backend::AbstractADType, X_template::AbstractMatrix +) + return error(_REACTANT_LOAD_HINT) +end + #= --------------------------------------------------------------------------- Second-order HVP, for a model whose gradient is itself AD-derived. `DI.hvp` -takes both passes over the log-density, so these never touch the gradient slot. -Preferred over pushing tangents through a prepared DI gradient, which drops out -of its preparation once the outer pass hands it an unexpected tangent type. - -Both halves are passed to `DI.hvp` as the user composed them, so the direction -each one runs in is theirs and DI's, not ours. Normalization touches only the -outer half, and only to fill in annotations for the wrappers being -differentiated; because it never substitutes a mode, `DI.hvp_mode` of the pair is -the same before and after. The inner half is a plain first-order gradient over -the user's own `logdensity` and is passed straight through. - -The batched form differentiates `sum(logdensity_batch(X))`, whose Hessian is -block-diagonal by column independence, so its HVP along `V` is the columnwise -HVP. Same argument the batched gradient rests on. +takes both passes over the log-density, so the gradient slot is never touched. +The alternative — pushing tangents through the prepared DI gradient — drops out +of that preparation the moment the outer pass hands it an unexpected tangent +type. + +Both halves reach `DI.hvp` as the user composed them. Normalization touches the +outer one, and only to fill in annotations, so `DI.hvp_mode` of the pair reads +the same before and after; the inner half is a first-order gradient over the +user's own `logdensity` and goes through untouched. + +The batched form differentiates `sum(logdensity_batch(X))`. Column independence +makes that Hessian block-diagonal, so its HVP along `V` is the columnwise HVP. --------------------------------------------------------------------------- =# function _normalized_second_order(backend::DI.SecondOrder) diff --git a/src/ParallelMCMC.jl b/src/ParallelMCMC.jl index 7041e90..126e64f 100644 --- a/src/ParallelMCMC.jl +++ b/src/ParallelMCMC.jl @@ -28,6 +28,10 @@ pmcmc_matmul(A::AbstractVecOrMat, B::AbstractVecOrMat) = A * B pmcmc_dot(a::AbstractVector, b::AbstractVector) = dot(a, b) pmcmc_dotsum(A::AbstractVecOrMat, B::AbstractVecOrMat) = sum(A .* B) +#= Lives here rather than in `DEER` because both DEER's `ReactantHVP` fallbacks +and `interface.jl`'s gradient hooks report it. =# +const _REACTANT_LOAD_HINT = "AutoReactant requires Reactant.jl: add `using Reactant` to load ParallelMCMC's ReactantExt." + include("MALA/MALA.jl") include("DEER/DEERScan.jl") include("DEER/DEER.jl") diff --git a/src/interface.jl b/src/interface.jl index 4e89cda..b664b5b 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -11,20 +11,28 @@ Defines model/sampler/state/transition types and implements Wraps a log-density function, its gradient, and optional Hessian-vector product helpers for use with ParallelMCMC samplers. -The derivative slots (`grad_logdensity`, `hvp`, `grad_logdensity_batch`, -`hvp_batch`) take either a callable or an `ADTypes.AbstractADType`. Backends -are turned into prepared DifferentiationInterface callables when sampling -starts, and any AD failure surfaces there. So a model needs nothing beyond -the log-density: +Each derivative slot (`grad_logdensity`, `hvp`, `grad_logdensity_batch`, +`hvp_batch`) takes a callable or an `ADTypes.AbstractADType`, so a model can be +built from the log-density alone: DensityModel(logp, AutoForwardDiff(), dim) -How a backend in `hvp` / `hvp_batch` gets its second derivative depends on the -gradient slot. Over a hand-written gradient it is a single AD pass across your -own code. Over an AD-derived one it is +Backends become prepared DifferentiationInterface callables when sampling starts, +and an AD failure surfaces there. `ADTypes.AutoReactant()` is the one backend DI +cannot drive; it is traced with Enzyme-MLIR and compiled to an XLA executable by +Reactant.jl instead, which needs `using Reactant` and brings requirements of its +own — see the GPU guide, `docs/src/15-gpu.md`, and `ext/ReactantExt.jl`'s module +docstring. + +A backend in `hvp` / `hvp_batch` over a hand-written gradient is a single AD pass +across your own code. Over an AD-derived one it is `DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)`, taken through DI's second-order operator. Passing a `SecondOrder` yourself always means the latter, and bypasses the gradient slot even when you wrote it by hand. +`AutoReactant` sits outside both: it cannot go inside a `SecondOrder`, and +`grad_logdensity` and `hvp`/`hvp_batch` must either both be `AutoReactant()` or +neither, since mixing it with a DifferentiationInterface backend across the two +passes of an HVP raises an `ArgumentError` at preparation time. - `logdensity(x::AbstractVector) -> Real` - `grad_logdensity` — callable `x -> AbstractVector`, or a backend to @@ -49,12 +57,11 @@ latter, and bypasses the gradient slot even when you wrote it by hand. information. Both batched derivative slots require `logdensity_batch`, which the batched -update evaluates directly. `logdensity_batch` alone is allowed and scores whole -trajectories at once without switching the batched update on. - -`ParallelMALASampler` runs the batched DEER update once it has a -`logdensity_batch` and a batched gradient: either `grad_logdensity_batch`, or one -derived from `logdensity_batch` when `grad_logdensity` is a backend. +update evaluates directly. `ParallelMALASampler` runs that update once it has a +`logdensity_batch` and a batched gradient — `grad_logdensity_batch`, or one +derived from `logdensity_batch` when `grad_logdensity` is a backend. A +`logdensity_batch` on its own is fine too: it scores whole trajectories at once +and leaves the batched update off. """ struct DensityModel{F,G,H,FB,GB,HB,PN} <: AbstractMCMC.AbstractModel logdensity::F @@ -95,15 +102,14 @@ function DensityModel( "grad_logdensity must be a callable or an ADTypes.AbstractADType backend" ), ) - #= The batched DEER update evaluates `logdensity_batch` itself, so neither - batched derivative is usable without one: a backend would have nothing to - differentiate, and a callable would never be reached. Rejected here rather - than silently ignored. A `logdensity_batch` on its own is allowed, and - `_prepare_model` decides from the gradient slot whether the path can run. =# + #= The batched DEER update evaluates `logdensity_batch` itself, so without one + a batched derivative slot is unusable either way: a backend has nothing to + differentiate and a callable is never reached. Rejected rather than ignored. + A `logdensity_batch` on its own is fine; `_prepare_model` decides from the + gradient slot whether the batched path can run. =# _batch_needs_logp(name) = throw( ArgumentError( - "$name requires logdensity_batch: the batched DEER path evaluates the " * - "batched log-density, so it cannot run without one", + "$name requires logdensity_batch, which the batched DEER update evaluates" ), ) if logdensity_batch === nothing @@ -123,12 +129,13 @@ function DensityModel( end """ -A [`DensityModel`](@ref) with its backend slots resolved to prepared -DifferentiationInterface callables. `_prepare_model` builds these, and the -sampler internals take them rather than a `DensityModel`, so no slot of one -of these ever holds an `AbstractADType`. Which slots are filled depends on -the sampler: the sequential samplers only need `grad_logdensity` and get -`nothing` for the DEER-only slots, DEER fills the rest. +A [`DensityModel`](@ref) with its backend slots resolved to prepared callables — +DifferentiationInterface ones, or a compiled Reactant/XLA executable for +`AutoReactant()`. `_prepare_model` builds these, and the sampler internals take +them rather than a `DensityModel`, so no slot here ever holds an +`AbstractADType`. Which slots are filled depends on the sampler: the sequential +samplers only need `grad_logdensity` and get `nothing` for the DEER-only slots, +DEER fills the rest. `source` is the `DensityModel` this was prepared from. A sampler state carries a prepped model so the preparation is reused across steps, and `initial_state` @@ -158,8 +165,7 @@ _prepped_for(prepped::PreppedDensityModel, model::DensityModel) = prepped.source #= Resolved gradient wrappers. Structs rather than anonymous closures since DI keys preparations on function identity. `TX` is the input type the prep was made for; -anything else falls back to unprepared `DI.gradient` rather than failing. In a -normal run the prepared branch is the one that fires. +anything else falls back to an unprepared `DI.gradient` rather than failing. =# struct _ADGradient{F,B<:AbstractADType,P,TX} logdensity::F @@ -222,7 +228,7 @@ function _resolve_gradient( return _reactant_resolve_gradient(logdensity, backend, x_template) end function _reactant_resolve_gradient(logdensity, backend, x_template) - return error(DEER._REACTANT_LOAD_HINT) + return error(_REACTANT_LOAD_HINT) end function _resolve_gradient_batch( @@ -231,7 +237,7 @@ function _resolve_gradient_batch( return _reactant_resolve_gradient_batch(logdensity_batch, backend, X_template) end function _reactant_resolve_gradient_batch(logdensity_batch, backend, X_template) - return error(DEER._REACTANT_LOAD_HINT) + return error(_REACTANT_LOAD_HINT) end #= @@ -244,21 +250,21 @@ A `SecondOrder` bypasses the gradient slot even when that slot is hand-written: naming both passes asks for two derivatives of `logdensity`. The slot is still the drift term the MALA step uses. -`AutoReactant` short-circuits ahead of both: DI cannot drive Reactant, so it can -neither build the `SecondOrder` nor route through `hvp_mode`. `ReactantExt` takes -over the second-order case by dispatching on `grad` instead — a Reactant-resolved -gradient carries the raw `logdensity` with it and re-traces forward-over-reverse -from there. +`AutoReactant` takes the same three branches as anything else. `_second_order` +below collapses an `AutoReactant` pair to a single `AutoReactant()` rather than a +`DI.SecondOrder`, and `_hvp_strategy(::AutoReactant)` sends the +hand-written-gradient case to `ReactantHVP`; `ReactantExt` supplies both matching +methods. `_check_reactant_pair` runs first, so a mismatched pairing is reported +before either path pays for a gradient resolution or an HVP compile. =# function _resolve_hvp(logdensity, grad, grad_backend, hvp_backend, x_template) _check_reactant_pair(grad_backend, hvp_backend) - if hvp_backend isa ADTypes.AutoReactant - return DEER._make_hvp_fn(DEER.ReactantHVP(), grad, hvp_backend, x_template) - elseif hvp_backend isa DI.SecondOrder + _check_reactant_hvp_source(grad, hvp_backend) + if hvp_backend isa DI.SecondOrder return DEER._make_hvp_fn_second_order(logdensity, hvp_backend, x_template) elseif grad_backend !== nothing return DEER._make_hvp_fn_second_order( - logdensity, DI.SecondOrder(hvp_backend, grad_backend), x_template + logdensity, _second_order(hvp_backend, grad_backend), x_template ) else return DEER._make_hvp_fn( @@ -272,18 +278,14 @@ function _resolve_hvp_batch( logdensity_batch, grad_batch, grad_batch_backend, hvp_backend, X_template ) _check_reactant_pair(grad_batch_backend, hvp_backend) - if hvp_backend isa ADTypes.AutoReactant - return DEER._make_hvp_batch_fn( - DEER.ReactantHVP(), grad_batch, hvp_backend, X_template - ) - elseif hvp_backend isa DI.SecondOrder + if hvp_backend isa DI.SecondOrder return DEER._make_hvp_batch_fn_second_order( _BatchLogdensitySum(logdensity_batch), hvp_backend, X_template ) elseif grad_batch_backend !== nothing return DEER._make_hvp_batch_fn_second_order( _BatchLogdensitySum(logdensity_batch), - DI.SecondOrder(hvp_backend, grad_batch_backend), + _second_order(hvp_backend, grad_batch_backend), X_template, ) else @@ -293,6 +295,16 @@ function _resolve_hvp_batch( end end +#= The HVP backend composed with the backend that produced the gradient under it. +For a DI-driven pair that is literally `DI.SecondOrder(hvp_backend, +grad_backend)`, taken through `DI.hvp`. Two `AutoReactant`s are not a pair DI +could run at all, so they collapse to the one backend that traces +forward-over-reverse from `logdensity` itself (`ReactantExt`'s +`_make_hvp_fn_second_order`). A mixed pair is already out by the time this runs, +via `_check_reactant_pair`. =# +_second_order(hvp_backend, grad_backend) = DI.SecondOrder(hvp_backend, grad_backend) +_second_order(::ADTypes.AutoReactant, ::ADTypes.AutoReactant) = ADTypes.AutoReactant() + #= Reactant does not pair with a DI backend across the two passes of an HVP: the compiled gradient is an opaque XLA executable DI cannot differentiate, and a @@ -329,6 +341,50 @@ function _check_reactant_pair(grad_backend, hvp_backend::ADTypes.AutoReactant) ) end +#= Would otherwise be ambiguous between the two methods above, and wants its own +message anyway: "an AutoReactant Hessian-vector product" does not describe a +`SecondOrder`. =# +function _check_reactant_pair(::ADTypes.AutoReactant, hvp_backend::DI.SecondOrder) + return throw( + ArgumentError( + "an AutoReactant gradient needs a bare AutoReactant Hessian-vector " * + "product: got hvp backend $(hvp_backend). An AutoReactant gradient " * + "compiles to an XLA executable, which DifferentiationInterface's " * + "SecondOrder cannot drive. Set the model's `hvp` (or the sampler's " * + "`backend`) to AutoReactant(), unwrapped.", + ), + ) +end + +#= `SecondOrder(AutoReactant(), AutoReactant())` is a natural thing to try, given +the pairing table in `10-getting-started.md`, and would otherwise land in +`DI.prepare_hvp` several frames deep with no Reactant support. Checked whatever +the gradient slot holds, since a `SecondOrder` bypasses it anyway (see +`_resolve_hvp`). =# +function _check_reactant_pair(grad_backend, hvp_backend::DI.SecondOrder) + _second_order_has_reactant(hvp_backend) || return nothing + return throw( + ArgumentError( + "AutoReactant cannot go inside a DifferentiationInterface SecondOrder: " * + "got hvp backend $(hvp_backend). DifferentiationInterface has no Reactant " * + "support at all. Set `hvp` (or the sampler's `backend`) to a bare " * + "AutoReactant() instead.", + ), + ) +end + +function _second_order_has_reactant(so::DI.SecondOrder) + return DI.outer(so) isa ADTypes.AutoReactant || DI.inner(so) isa ADTypes.AutoReactant +end + +#= A `LogDensityProblemGradient` (defined below) is a callable, so it clears the +`grad_backend === nothing` test that otherwise means "hand-written gradient" — +but it dispatches into DynamicPPL/LogDensityProblems and is not +Reactant-traceable. `_check_reactant_pair` only sees the backend, which is +`nothing` for both, so this checks `grad`'s type instead. The specialization has +to wait for `LogDensityProblemGradient` to exist and sits further down. =# +_check_reactant_hvp_source(grad, hvp_backend) = nothing + """ _prepare_model(model, x_template) -> PreppedDensityModel _prepare_model(model, x_template, T::Int, backend) -> PreppedDensityModel @@ -354,10 +410,9 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector) else model.grad_logdensity end - #= The derivative slots DEER alone reads are dropped rather than passed - along: a sequential sampler never looks at them, and carrying an - unresolved backend would break the invariant that no slot here holds an - `AbstractADType`. =# + #= The DEER-only slots are dropped rather than passed along: a sequential + sampler never reads them, and an unresolved backend sitting in one would + break the invariant that no slot here holds an `AbstractADType`. =# return PreppedDensityModel( model.logdensity, grad, @@ -374,31 +429,48 @@ end function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, backend) grad_backend = model.grad_logdensity isa AbstractADType ? model.grad_logdensity : nothing - grad = if grad_backend !== nothing - _resolve_gradient(model.logdensity, grad_backend, x_template) - else - model.grad_logdensity - end - hvp = if model.hvp === nothing || model.hvp isa AbstractADType - hvp_backend = model.hvp === nothing ? backend : model.hvp - hvp_backend === nothing && throw( + #= Settle the HVP backend and check its pairing with `grad_backend` before + resolving the gradient. Otherwise a mismatched `AutoReactant` pair, or the + LogDensityProblems-gradient case, surfaces only once `_resolve_gradient` has + paid for an XLA compile: 18+ seconds to report a config error. Neither check + needs the resolved gradient. `grad_backend` is known already, and + `_check_reactant_hvp_source` only looks at `model.grad_logdensity`'s type, + which is `grad` unchanged whenever `grad_backend` is `nothing`. =# + needs_hvp = model.hvp === nothing || model.hvp isa AbstractADType + hvp_backend = if needs_hvp + hb = model.hvp === nothing ? backend : model.hvp + hb === nothing && throw( ArgumentError( "ParallelMALASampler needs a Hessian-vector product: supply `hvp` " * "on the DensityModel (callable or AD backend), or pass `backend=` " * "to ParallelMALASampler", ), ) + _check_reactant_pair(grad_backend, hb) + _check_reactant_hvp_source(model.grad_logdensity, hb) + hb + else + nothing + end + + grad = if grad_backend !== nothing + _resolve_gradient(model.logdensity, grad_backend, x_template) + else + model.grad_logdensity + end + + hvp = if needs_hvp _resolve_hvp(model.logdensity, grad, grad_backend, hvp_backend, x_template) else model.hvp end #= A batched log-density with no batched gradient gets one from the model's - own gradient backend, never the sampler's: a model with a hand-written - gradient has not opted into AD, and deriving one anyway would let `backend=` - decide which update path runs. Failing to derive leaves the path off rather - than raising, since `_trajectory_logps` uses `logdensity_batch` regardless. =# + own gradient backend, never the sampler's: a hand-written gradient has not + opted into AD, and deriving one anyway would let `backend=` decide which + update path runs. Not deriving one leaves the batched path off rather than + raising, since `_trajectory_logps` uses `logdensity_batch` either way. =# grad_batch = model.grad_logdensity_batch if grad_batch === nothing && model.logdensity_batch !== nothing grad_batch = grad_backend @@ -413,26 +485,35 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, X_template = similar(x_template, length(x_template), T) X_template .= x_template - if grad_batch_backend !== nothing - grad_batch = _resolve_gradient_batch( - model.logdensity_batch, grad_batch_backend, X_template - ) - end - - if hvp_batch === nothing || hvp_batch isa AbstractADType + # Same reasoning as the unbatched case above: check before compiling. + needs_hvp_batch = hvp_batch === nothing || hvp_batch isa AbstractADType + hvp_batch_backend = if needs_hvp_batch # The model's own HVP backend if it has one, else the sampler's. - hvp_batch_backend = if hvp_batch === nothing + hbb = if hvp_batch === nothing model.hvp isa AbstractADType ? model.hvp : backend else hvp_batch end - hvp_batch_backend === nothing && throw( + hbb === nothing && throw( ArgumentError( "the batched DEER path needs a batched Hessian-vector product: " * "supply `hvp_batch` on the DensityModel (callable or AD backend), " * "or pass `backend=` to ParallelMALASampler", ), ) + _check_reactant_pair(grad_batch_backend, hbb) + hbb + else + nothing + end + + if grad_batch_backend !== nothing + grad_batch = _resolve_gradient_batch( + model.logdensity_batch, grad_batch_backend, X_template + ) + end + + if needs_hvp_batch hvp_batch = _resolve_hvp_batch( model.logdensity_batch, grad_batch, @@ -467,10 +548,10 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, ) end -#= Callable structs that allow us to dispatch on the type of the LogDensityProblems object in -the postprocessing stage. Ideally these would be defined in the LogDensityProblemsExt. -However, structs defined in extensions are hard to get hold of so we define them here. -The callable behaviour itself is implemented in LogDensityProblemsExt =# +# Callable structs that allow us to dispatch on the type of the LogDensityProblems object in +# the postprocessing stage. Ideally these would be defined in the LogDensityProblemsExt. +# However, structs defined in extensions are hard to get hold of so we define them here. +# The callable behaviour itself is implemented in LogDensityProblemsExt. struct LogDensityProblemPrimal{L} ld::L end @@ -478,6 +559,20 @@ struct LogDensityProblemGradient{L} ld::L end +# The `_check_reactant_hvp_source` specialization promised further up. +function _check_reactant_hvp_source( + ::LogDensityProblemGradient, hvp_backend::ADTypes.AutoReactant +) + return throw( + ArgumentError( + "an AutoReactant Hessian-vector product needs an AutoReactant or " * + "hand-written gradient: got a LogDensityProblems-derived gradient. Reactant " * + "cannot trace DynamicPPL/LogDensityProblems machinery. Set `grad_logdensity` " * + "to AutoReactant(), or supply a Reactant-traceable callable.", + ), + ) +end + """ MALASampler(epsilon; cholM=nothing) @@ -654,11 +749,15 @@ DEER-parallelized MALA sampler. Supported Jacobian modes are `:stoch_diag` (the default Hutchinson diagonal estimator) and `:diag` (exact diagonal via `D` JVPs). -`backend` is the fallback source of Hessian-vector products, used when the -`DensityModel` brings no `hvp` / `hvp_batch` of its own. That is all it does: it -never supplies a gradient, so it cannot change which update path runs or put AD -on a function the model did not already have a backend for. A model carrying its -own HVPs does not need it. +`backend` supplies Hessian-vector products when the `DensityModel` brings no +`hvp` / `hvp_batch` of its own, and does nothing else. It never supplies a +gradient, so it cannot decide which update path runs, nor put AD on a function +the model had no backend for. Leave it out for a model that carries its own HVPs. + +`backend = ADTypes.AutoReactant()` constrains the model too, since +`AutoReactant` does not pair with a DifferentiationInterface backend across the +two passes of an HVP: `grad_logdensity` must then be `AutoReactant()` as well, or +a hand-written callable. Mixing raises an `ArgumentError` at preparation time. """ struct ParallelMALASampler{FP<:AbstractFloat,CM,AD} <: AbstractMCMC.AbstractSampler epsilon::FP diff --git a/test/Project.toml b/test/Project.toml index 2892146..a2370e2 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -15,7 +15,6 @@ LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" ParallelMCMC = "1a970f40-4406-51c9-a967-cb3143c111e8" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" -Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" @@ -23,8 +22,9 @@ TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" +[compat] +JuliaFormatter = "2.12" + [extras] CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" - -[compat] -JuliaFormatter = "2.12" \ No newline at end of file +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" diff --git a/test/runtests.jl b/test/runtests.jl index 45cecaf..4d5b1d8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,13 @@ using ParallelMCMC using Test +#= `Reactant_jll` ships a prebuilt XLA and is a large download, so `Reactant` sits +in test/Project.toml's `[extras]` and `test-Reactant-HVP.jl` runs only when this +is set. Add Reactant to the test environment as well; the file's own +`try ... using Reactant ... catch` skips its testsets if it still won't load. =# +const _RUN_REACTANT_TESTS = + lowercase(get(ENV, "PARALLELMCMC_TEST_REACTANT", "false")) in ("1", "true", "yes") + @testset verbose=true "ParallelMCMC" begin #= Don't add your tests to runtests.jl. Instead, create files named @@ -14,6 +21,10 @@ using Test if isnothing(match(r"^test-.*\.jl$", file)) continue end + if file == "test-Reactant-HVP.jl" && !_RUN_REACTANT_TESTS + @info "Skipping $file (set PARALLELMCMC_TEST_REACTANT=true to opt in)" + continue + end title = titlecase(replace(splitext(file[6:end])[1], "-" => " ")) @testset verbose=true "$title" begin include(joinpath(root, file)) # robust if walkdir recurses diff --git a/test/test-HVP-Strategy.jl b/test/test-HVP-Strategy.jl index d28a02b..44305c1 100644 --- a/test/test-HVP-Strategy.jl +++ b/test/test-HVP-Strategy.jl @@ -46,8 +46,7 @@ const DI_STRAT = ParallelMCMC.DEER.DI @testset "normalization supplies Const but never a mode" begin #= The wrappers DEER differentiates are its own types, so annotating them - `Const` is its business. The mode is not: one the user set is a decision, - and an unset one is DI's to resolve from the operator it runs. =# + `Const` is its business. The mode is not. =# bare = DEER_STRAT._normalized_backend(AutoEnzyme()) @test bare isa AutoEnzyme{<:Any,Enzyme.Const} @test bare.mode === nothing @@ -70,9 +69,9 @@ const DI_STRAT = ParallelMCMC.DEER.DI @testset "normalizing a SecondOrder keeps the composition DI resolved" begin #= Regression for #62. Normalization used to route the outer half through a forward-only hook, which pinned `Enzyme.Forward` onto it. For a pair - `hvp_mode` resolves to reverse -- `SecondOrder(AutoEnzyme(), + `hvp_mode` resolves to reverse — `SecondOrder(AutoEnzyme(), AutoForwardDiff())` is reverse-over-forward, its inner half being - forward-only -- that silently made it forward-over-forward. =# + forward-only — that made it forward-over-forward. =# for so in ( DI_STRAT.SecondOrder(AutoEnzyme(), AutoForwardDiff()), DI_STRAT.SecondOrder(AutoEnzyme(), AutoZygote()), diff --git a/test/test-Reactant-HVP.jl b/test/test-Reactant-HVP.jl index c2b03eb..562e33a 100644 --- a/test/test-Reactant-HVP.jl +++ b/test/test-Reactant-HVP.jl @@ -1,32 +1,46 @@ using Test using Random using LinearAlgebra +using Statistics using FlexiChains using ParallelMCMC using ADTypes # Stands in for "some backend that is not Reactant" in the pairing tests below. using ForwardDiff: ForwardDiff +using LogDensityProblems: LogDensityProblems + +const DI_R = ParallelMCMC.DEER.DI #= Reactant-compiled derivative paths (`AutoReactant`), see ext/ReactantExt.jl. -The quartic target keeps second-order structure honest: logp = -0.25‖x‖⁴ has -H = -(‖x‖² I + 2 x xᵀ), so an HVP that silently drops the second-order term -(the failure mode of `Enzyme.hvp` under `@compile`) is caught, unlike a -Gaussian where H is constant. +The quartic target keeps the second-order structure honest: logp = -0.25‖x‖⁴ has +H = -(‖x‖² I + 2 x xᵀ), so an HVP that drops the second-order term — the failure +mode of `Enzyme.hvp` under `@compile` — is caught, where a Gaussian's constant H +would hide it. Every derivative-accuracy testset below uses it; do not swap in a +Gaussian. =# logp_r(x) = -0.25 * sum(abs2, x)^2 gradlogp_r(x) = -sum(abs2, x) .* x hvp_r(x, v) = -(sum(abs2, x) .* v .+ 2 .* dot(x, v) .* x) logp_batch_r(X) = vec(-0.25 .* sum(abs2, X; dims=1) .^ 2) gradlogp_batch_r(X) = -X .* sum(abs2, X; dims=1) +logp_r32(x) = -0.25f0 * sum(abs2, x)^2 + +#= Standard Gaussian, constant Hessian (H = -I). Only for the end-to-end sampling +tests, which ask whether the sampler converged to the right posterior; whether +the HVP is second-order-correct is the quartic target's job. =# +logp_gauss(x) = -0.5 * sum(abs2, x) +gradlogp_gauss(x) = -x +hvp_gauss(x, v) = -v const D_R = 4 const CT_R = FlexiChains.FlexiChain{Symbol} -#= The pairing rule lives in `_resolve_hvp`, not in the extension, so its -dispatch table is checked whether or not Reactant loads. =# +#= The pairing rule lives in `_resolve_hvp` / `_prepare_model`, not in the +extension, so its dispatch table is checked whether or not Reactant loads — +none of these calls resolve a gradient or compile anything. =# @testset "Reactant does not pair with a DI backend" begin # Both slots Reactant, or a hand-written gradient (`nothing`), are accepted. @test ParallelMCMC._check_reactant_pair(AutoReactant(), AutoReactant()) === nothing @@ -43,6 +57,67 @@ dispatch table is checked whether or not Reactant loads. =# ) end +#= `SecondOrder(AutoReactant(), AutoReactant())`, or `AutoReactant()` paired with +a `SecondOrder` at all, is a natural thing to try given the pairing table in +10-getting-started.md, and would otherwise land in `DI.prepare_hvp` several +frames deep with no Reactant support. A dispatch-table property, so this runs +whether or not Reactant is loaded. =# +@testset "AutoReactant cannot appear inside a SecondOrder" begin + # Sanity: an ordinary SecondOrder is unaffected. + @test ParallelMCMC._check_reactant_pair( + nothing, DI_R.SecondOrder(AutoForwardDiff(), AutoForwardDiff()) + ) === nothing + + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + nothing, DI_R.SecondOrder(AutoReactant(), AutoReactant()) + ) + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + AutoForwardDiff(), DI_R.SecondOrder(AutoReactant(), AutoForwardDiff()) + ) + # (AutoReactant, SecondOrder) exercises the disambiguating method directly. + @test_throws ArgumentError ParallelMCMC._check_reactant_pair( + AutoReactant(), DI_R.SecondOrder(AutoReactant(), AutoReactant()) + ) +end + +#= A `LogDensityProblemGradient` is a callable, so it clears the +`grad_backend === nothing` test that otherwise means "hand-written gradient", +but it is not Reactant-traceable. Checked on the type first, then through a real +LogDensityProblems model. Neither needs Reactant loaded: the check fires before +any gradient is resolved or anything compiled. =# +@testset "a LogDensityProblems gradient cannot pair with an AutoReactant hvp" begin + @test ParallelMCMC._check_reactant_hvp_source( + ParallelMCMC.LogDensityProblemGradient(nothing), AutoForwardDiff() + ) === nothing + @test_throws ArgumentError ParallelMCMC._check_reactant_hvp_source( + ParallelMCMC.LogDensityProblemGradient(nothing), AutoReactant() + ) + + struct _FakeLD end + LogDensityProblems.capabilities(::_FakeLD) = LogDensityProblems.LogDensityOrder{1}() + LogDensityProblems.dimension(::_FakeLD) = D_R + function LogDensityProblems.logdensity_and_gradient(::_FakeLD, x) + return logp_r(x), gradlogp_r(x) + end + + model = DensityModel(_FakeLD(); hvp=AutoReactant()) + @test_throws ArgumentError ParallelMCMC._prepare_model(model, zeros(D_R), 8, nothing) +end + +#= Outside the `reactant_ok` guard below. `_check_reactant_pair` now runs above +gradient resolution in `_prepare_model`, so a mismatched pair is caught before +anything is resolved or compiled and no Reactant install is needed. This used to +pay a full XLA compile per `@test_throws`, ~18s, to demonstrate a config error. =# +@testset "mixed pairs are refused at preparation" begin + x = zeros(D_R) + + reactant_grad = DensityModel(logp_r, AutoReactant(), D_R; hvp=AutoForwardDiff()) + @test_throws ArgumentError ParallelMCMC._prepare_model(reactant_grad, x, 8, nothing) + + reactant_hvp = DensityModel(logp_r, AutoForwardDiff(), D_R; hvp=AutoReactant()) + @test_throws ArgumentError ParallelMCMC._prepare_model(reactant_hvp, x, 8, nothing) +end + reactant_ok = try using Reactant: Reactant using Enzyme: Enzyme @@ -52,22 +127,54 @@ catch err false end +#= The `_REACTANT_LOAD_HINT` fallbacks (src/DEER/DEER.jl, and +`_reactant_resolve_gradient` / `_reactant_resolve_gradient_batch` in +src/interface.jl) can only be asserted in a session without Reactant: once +`ReactantExt` is loaded its more-specific methods shadow all of them and the +error can never fire. Hence the one testset here guarded on `!reactant_ok`. =# +if !reactant_ok + @testset "clear load-hint error without Reactant loaded" begin + #= Gradient slot: `_reactant_resolve_gradient`'s fallback. `backend` (the + 4th arg) is `AutoReactant()` too, as `_check_reactant_pair` requires, so + `_prepare_model` gets as far as gradient resolution rather than failing + first on the model having no HVP source — a correct failure, but not the + one being pinned here. =# + model_grad = DensityModel(logp_r, AutoReactant(), D_R) + @test_throws "AutoReactant requires Reactant.jl" ParallelMCMC._prepare_model( + model_grad, zeros(D_R), 8, AutoReactant() + ) + + # HVP slot over a hand-written gradient: `DEER._make_hvp_fn`'s + # `ReactantHVP` fallback, reached via `_hvp_strategy(::AutoReactant)`. + model_hvp = DensityModel(logp_r, gradlogp_r, D_R; hvp=AutoReactant()) + @test_throws "AutoReactant requires Reactant.jl" ParallelMCMC._prepare_model( + model_hvp, zeros(D_R), 8, nothing + ) + end +end + if reactant_ok @testset "extension is loaded" begin @test Base.get_extension(ParallelMCMC, :ReactantExt) !== nothing end - #= The same rule reached through `_prepare_model`, where the gradient slot - resolves first: a mixed pair must still surface as an ArgumentError and not - as whatever DI or Reactant would say downstream. =# - @testset "mixed pairs are refused at preparation" begin - x = zeros(D_R) + #= The extension does not honour `AutoReactant.mode` (the wrapped + `AutoEnzyme`): gradients always trace reverse, HVPs forward-over-that. A + non-default mode is rejected rather than ignored. =# + @testset "a non-default AutoReactant mode is rejected" begin + bad = AutoReactant(; mode=AutoEnzyme(; mode=Enzyme.Forward)) + model_grad = DensityModel(logp_r, bad, D_R) + @test_throws ArgumentError ParallelMCMC._prepare_model( + model_grad, zeros(D_R), 8, nothing + ) - reactant_grad = DensityModel(logp_r, AutoReactant(), D_R; hvp=AutoForwardDiff()) - @test_throws ArgumentError ParallelMCMC._prepare_model(reactant_grad, x, 8, nothing) + model_hvp = DensityModel(logp_r, gradlogp_r, D_R; hvp=bad) + @test_throws ArgumentError ParallelMCMC._prepare_model( + model_hvp, zeros(D_R), 8, nothing + ) - reactant_hvp = DensityModel(logp_r, AutoForwardDiff(), D_R; hvp=AutoReactant()) - @test_throws ArgumentError ParallelMCMC._prepare_model(reactant_hvp, x, 8, nothing) + # The default is, of course, fine. + @test DensityModel(logp_r, AutoReactant(), D_R) isa DensityModel end @testset "HVP matches analytic" begin @@ -107,14 +214,144 @@ if reactant_ok @test m_p.grad_logdensity_batch(X) ≈ gradlogp_batch_r(X) @test m_p.hvp_batch(X, V) ≈ Hv_cols end + + #= A hand-written batched gradient with `hvp_batch=AutoReactant()`, the + batched analogue of the "forward over user gradient" case above. Routes + through + `DEER._make_hvp_batch_fn(::ReactantHVP, grad_batch, ::AutoReactant, ...)`, + which nothing else here reaches. =# + @testset "forward over user batched gradient (hvp_batch=AutoReactant())" begin + T = 8 + X = randn(rng, D_R, T) + V = randn(rng, D_R, T) + Hv_cols = reduce(hcat, [hvp_r(X[:, t], V[:, t]) for t in 1:T]) + + model = DensityModel( + logp_r, + gradlogp_r, + D_R; + hvp=hvp_r, + logdensity_batch=logp_batch_r, + grad_logdensity_batch=gradlogp_batch_r, + hvp_batch=AutoReactant(), + ) + m_p = ParallelMCMC._prepare_model(model, X[:, 1], T, nothing) + @test m_p.hvp_batch(X, V) ≈ Hv_cols + end + + #= Edge shapes. Both slots AutoReactant, as in the + "forward-over-reverse from logp alone" case above, but at the smallest + sizes DEER ever prepares. =# + @testset "edge shapes" begin + @testset "D=1" begin + x1 = randn(rng, 1) + v1 = randn(rng, 1) + model = DensityModel(logp_r, AutoReactant(), 1; hvp=AutoReactant()) + m_p = ParallelMCMC._prepare_model(model, x1, 8, nothing) + @test m_p.grad_logdensity(x1) ≈ gradlogp_r(x1) + @test m_p.hvp(x1, v1) ≈ hvp_r(x1, v1) + end + + @testset "T=1" begin + T = 1 + X = reshape(x, D_R, T) + V = reshape(v, D_R, T) + model = DensityModel( + logp_r, + AutoReactant(), + D_R; + logdensity_batch=logp_batch_r, + grad_logdensity_batch=AutoReactant(), + hvp=AutoReactant(), + hvp_batch=AutoReactant(), + ) + m_p = ParallelMCMC._prepare_model(model, x, T, nothing) + @test m_p.grad_logdensity_batch(X) ≈ gradlogp_batch_r(X) + @test m_p.hvp_batch(X, V) ≈ reshape(hvp_r(x, v), D_R, T) + end + end + + #= Float32 on the CPU path. Previously only exercised inside the + CuArray-only block below, so it never ran without a functional CUDA. =# + @testset "Float32 on CPU" begin + x32 = Float32.(x) + v32 = Float32.(v) + model = DensityModel(logp_r32, AutoReactant(), D_R; hvp=AutoReactant()) + m_p = ParallelMCMC._prepare_model(model, x32, 8, nothing) + + g = m_p.grad_logdensity(x32) + @test eltype(g) === Float32 + @test g ≈ gradlogp_r(x32) + + Hv = m_p.hvp(x32, v32) + @test eltype(Hv) === Float32 + @test Hv ≈ hvp_r(x32, v32) + end + end + + #= Pins the "captured data is frozen at preparation time" limitation from + ext/ReactantExt.jl's module docstring and docs/src/15-gpu.md: `@compile` + bakes a captured plain `Array` in as a compile-time constant, so mutating it + afterwards has NO effect on later calls, with no error and no warning. Here + so that a future Reactant which does detect this gets noticed. Model code + still should not close over mutable data. + =# + @testset "documented caveat: captured data is frozen at preparation time" begin + data = [1.0, 1.0, 1.0, 1.0] + f(x) = -0.5 * sum(abs2, x .- data) # ∇f(x) = data - x + model = DensityModel(f, AutoReactant(), D_R) + # Two-argument `_prepare_model` only resolves `grad_logdensity`, which + # is all this test needs (no HVP involved). + m_p = ParallelMCMC._prepare_model(model, zeros(D_R)) + + x0 = zeros(D_R) + g_before = m_p.grad_logdensity(x0) + @test g_before ≈ [1.0, 1.0, 1.0, 1.0] + + data .= 5.0 # mutate the captured array *after* preparation + + g_after = m_p.grad_logdensity(x0) + #= Frozen at the pre-mutation value: NOT [5, 5, 5, 5], which is what a + correct re-evaluation against the mutated `data` would give. =# + @test g_after ≈ [1.0, 1.0, 1.0, 1.0] end + #= `size(chain) == (N,1)` and `all(isfinite, ...)` pass even for a badly + wrong HVP: DEER's Newton iteration just fails to converge and `DEER.solve` + returns the non-converged trajectory, no NaNs involved. Assert the posterior + mean against a target with a known mean instead (the convention in + test-GPU-AD-HVP.jl), and cross-check against the same model driven by the + analytic HVP on the same noise tape. + =# @testset "end-to-end sampling with AutoReactant" begin - model = DensityModel(logp_r, AutoReactant(), D_R) - s = ParallelMALASampler(0.02; T=16, backend=AutoReactant()) - chain = sample(MersenneTwister(72), model, s, 64; chain_type=CT_R, progress=false) - @test size(chain) == (64, 1) - @test all(x -> all(isfinite, x), chain[:x]) + @testset "posterior mean recovery (standard Gaussian, mean 0)" begin + model = DensityModel(logp_gauss, AutoReactant(), D_R) + s = ParallelMALASampler(0.3; T=16, backend=AutoReactant()) + n_samples, n_burn = 2000, 500 + chain = sample( + MersenneTwister(72), model, s, n_samples; chain_type=CT_R, progress=false + ) + @test size(chain) == (n_samples, 1) + + xs = chain[:x] + @test all(x -> all(isfinite, x), xs) + post_mean = vec(mean(reduce(hcat, xs[(n_burn + 1):end]); dims=2)) + @test maximum(abs, post_mean) < 0.3 + end + + @testset "matches analytic-HVP DEER on the same noise tape" begin + s = ParallelMALASampler(0.1; T=16, backend=AutoReactant()) + model_r = DensityModel(logp_gauss, AutoReactant(), D_R) + model_an = DensityModel(logp_gauss, gradlogp_gauss, D_R; hvp=hvp_gauss) + + c_r = sample( + MersenneTwister(99), model_r, s, 64; chain_type=CT_R, progress=false + ) + c_an = sample( + MersenneTwister(99), model_an, s, 64; chain_type=CT_R, progress=false + ) + @test c_r[:x] ≈ c_an[:x] + end end reactant_gpu_ok = try @@ -128,12 +365,13 @@ if reactant_ok @info "Reactant HVP test: CUDA not functional — skipping CuArray boundary" else #= - The compiled executable lives in Reactant's own (XLA) device memory; - what's checked here is the CuArray <-> Reactant marshalling boundary: - CuArray in, CuArray out, values matching the analytic HVP. + The compiled executable lives in Reactant's own (XLA) device memory, so + what this checks is the CuArray <-> Reactant marshalling boundary: + CuArray in, CuArray out, values matching the analytic HVP. It does NOT + establish that the HVP executes on the GPU, which depends on Reactant's + default XLA client (see docs/src/15-gpu.md) and is neither controlled + nor asserted here. =# - logp_r32(x) = -0.25f0 * sum(abs2, x)^2 - @testset "CuArray boundary" begin rng = MersenneTwister(73) x_h = randn(rng, Float32, D_R) From 3e4447b936ada35ddc614c498e3507746e03c51d Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Wed, 12 Aug 2026 21:29:31 -0400 Subject: [PATCH 09/11] remove a few dead pieces --- CHANGELOG.md | 4 ++-- docs/src/15-gpu.md | 2 +- ext/ReactantExt.jl | 20 +++++--------------- src/interface.jl | 7 ++----- test/test-Reactant-HVP.jl | 2 +- 5 files changed, 11 insertions(+), 24 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bde0fe4..44c6194 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,7 +25,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 batched gradient derived from it when `grad_logdensity` is a backend, instead of leaving the batched DEER path switched off. `hvp_batch` can be a backend in that case too, and differentiates the derived gradient (#52). -- An HVP backend over an AD-derived gradient is now real second-order AD, +- An HVP backend over an AD-derived gradient is now true second-order AD, `DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` handed to `DI.hvp`, rather than an outer AD pass over the prepared DI gradient — which dropped out of its preparation as soon as tangents were pushed through it @@ -39,7 +39,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - New `ReactantExt`. `ADTypes.AutoReactant()` in a derivative slot, or as the sampler `backend`, traces the derivative with Enzyme-MLIR and compiles it to an XLA executable via Reactant.jl, off Enzyme's LLVM pipeline and off - DifferentiationInterface entirely, which yields a genuine second-order HVP for + DifferentiationInterface entirely, which yields a true second-order HVP for a log-density-only model (#37, #52). Requires `using Reactant` and a Reactant-traceable log-density. `AutoReactant` does not pair with a DifferentiationInterface backend, a `LogDensityProblems` gradient, or a diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index 318477a..1e40f4d 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -240,7 +240,7 @@ DEER needs a Hessian–vector product $H v$ at every Newton step. `DensityModel ### Reactant HVPs, off the DI path -`ADTypes.AutoReactant()` (requires `using Reactant`) goes another way. The derivative is traced with Enzyme-MLIR and compiled to an XLA executable by [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl), never touching Enzyme's LLVM pipeline, so neither the gc-transition abort nor the `pmcmc_*` wrappers above apply. It gives a genuine second-order HVP for a log-density-only model: +`ADTypes.AutoReactant()` (requires `using Reactant`) goes another way. The derivative is traced with Enzyme-MLIR and compiled to an XLA executable by [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl), never touching Enzyme's LLVM pipeline, so neither the gc-transition abort nor the `pmcmc_*` wrappers above apply. It gives a true second-order HVP for a log-density-only model: ```julia using Reactant, ADTypes diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl index 6c14c01..df82cd7 100644 --- a/ext/ReactantExt.jl +++ b/ext/ReactantExt.jl @@ -6,7 +6,7 @@ derivative slot of `DensityModel`, or as the sampler `backend`. Derivatives are traced with Enzyme-MLIR and compiled to XLA executables by `Reactant.@compile`, which keeps them off Enzyme's LLVM pipeline and so off the GPU `cuMemcpyDtoHAsync_v2` gc-transition abort, and off DifferentiationInterface -entirely. For a log-density-only model it yields a genuine second-order HVP. +entirely. For a log-density-only model it yields a true second-order HVP. Two silent failure modes, ahead of the ordinary limitations: @@ -144,26 +144,16 @@ end _rev_gradient(f, x) = Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(f), x)[1] #= -Gradient slots. They hold only the compiled callable; the HVP factories below -get `logdensity` from `_resolve_hvp`, which already has it. +Gradient slots. The HVP factories below get `logdensity` from `_resolve_hvp`, +which already has it. =# -struct _ReactantGradient{C} - compiled::C -end -(g::_ReactantGradient)(x) = g.compiled(x) - function ParallelMCMC._reactant_resolve_gradient( logdensity, backend::AutoReactant, x_template::AbstractVector ) _check_reactant_mode(backend) core = Base.Fix1(_rev_gradient, logdensity) - return _ReactantGradient(_compiled(core, x_template)) -end - -struct _ReactantGradientBatch{C} - compiled::C + return _compiled(core, x_template) end -(g::_ReactantGradientBatch)(X) = g.compiled(X) function ParallelMCMC._reactant_resolve_gradient_batch( logdensity_batch, backend::AutoReactant, X_template::AbstractMatrix @@ -173,7 +163,7 @@ function ParallelMCMC._reactant_resolve_gradient_batch( # gradient uses (src/interface.jl), so both batched paths differentiate the # same thing. core = Base.Fix1(_rev_gradient, ParallelMCMC._BatchLogdensitySum(logdensity_batch)) - return _ReactantGradientBatch(_compiled(core, X_template)) + return _compiled(core, X_template) end #= diff --git a/src/interface.jl b/src/interface.jl index b664b5b..cf77a17 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -254,12 +254,10 @@ the drift term the MALA step uses. below collapses an `AutoReactant` pair to a single `AutoReactant()` rather than a `DI.SecondOrder`, and `_hvp_strategy(::AutoReactant)` sends the hand-written-gradient case to `ReactantHVP`; `ReactantExt` supplies both matching -methods. `_check_reactant_pair` runs first, so a mismatched pairing is reported -before either path pays for a gradient resolution or an HVP compile. +methods. The pairing is already checked by `_prepare_model` before either path +here pays for a gradient resolution or an HVP compile. =# function _resolve_hvp(logdensity, grad, grad_backend, hvp_backend, x_template) - _check_reactant_pair(grad_backend, hvp_backend) - _check_reactant_hvp_source(grad, hvp_backend) if hvp_backend isa DI.SecondOrder return DEER._make_hvp_fn_second_order(logdensity, hvp_backend, x_template) elseif grad_backend !== nothing @@ -277,7 +275,6 @@ end function _resolve_hvp_batch( logdensity_batch, grad_batch, grad_batch_backend, hvp_backend, X_template ) - _check_reactant_pair(grad_batch_backend, hvp_backend) if hvp_backend isa DI.SecondOrder return DEER._make_hvp_batch_fn_second_order( _BatchLogdensitySum(logdensity_batch), hvp_backend, X_template diff --git a/test/test-Reactant-HVP.jl b/test/test-Reactant-HVP.jl index 562e33a..83be205 100644 --- a/test/test-Reactant-HVP.jl +++ b/test/test-Reactant-HVP.jl @@ -173,7 +173,7 @@ if reactant_ok model_hvp, zeros(D_R), 8, nothing ) - # The default is, of course, fine. + # The default still works. @test DensityModel(logp_r, AutoReactant(), D_R) isa DensityModel end From d084cc8380cd9cd60bbe97433325e38e0b48bf26 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Wed, 12 Aug 2026 23:48:04 -0400 Subject: [PATCH 10/11] review --- docs/src/15-gpu.md | 3 --- ext/ReactantExt.jl | 21 +++++++------------ src/DEER/DEER.jl | 31 ++------------------------- src/interface.jl | 44 +++++++++++++++------------------------ test/test-Reactant-HVP.jl | 13 ++++++------ 5 files changed, 32 insertions(+), 80 deletions(-) diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index 1e40f4d..29aab89 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -232,9 +232,6 @@ DEER needs a Hessian–vector product $H v$ at every Newton step. `DensityModel - **You supply `hvp` / `hvp_batch`.** These run as plain kernels. The AD backend is never invoked for HVPs. - **You only supply `gradlogp` / `grad_logdensity_batch`.** The sampler builds the HVP by differentiating your gradient — either a forward-mode pushforward of `gradlogp` ([`ForwardOnGrad`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/src/DEER/DEER.jl), the default for most backends) or a reverse-mode gradient of `x -> dot(gradlogp(x), v)` ([`ReverseOnGrad`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/src/DEER/DEER.jl), used for `AutoMooncake` and `AutoZygote`). This is the **AD-HVP fallback**, and it is what the logistic-regression example above uses. -!!! warning "Log-density-only models on GPU" - `grad_logdensity` can itself be an AD backend (`DensityModel(logp, AutoEnzyme(), dim)`, see [Getting started](10-getting-started.md)), but not with `ParallelMALASampler` on GPU for the DI-driven backends. The HVP then becomes `SecondOrder(hvp_backend, grad_backend)` on your log-density, which currently fails on GPU with both Enzyme and Mooncake (see [#37](https://github.com/rsenne/ParallelMCMC.jl/issues/37)); passing a `SecondOrder` explicitly hits the same wall. Write `gradlogp` out by hand so the HVP is a single pass over it, or use `AutoReactant()` below. The sequential samplers only need the gradient, so log-density-only models are fine there. - !!! note "A backend in `grad_logdensity` reaches `logdensity_batch` too" The batched path needs a batched gradient, and derives one from `logdensity_batch` when `grad_logdensity` is a backend. That puts `logdensity_batch` under the same restrictions as the rest of your AD-visible code. Supply `grad_logdensity_batch` to avoid it. diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl index df82cd7..aab4500 100644 --- a/ext/ReactantExt.jl +++ b/ext/ReactantExt.jl @@ -2,11 +2,7 @@ module ReactantExt #= Reactant-compiled derivative paths, selected by `ADTypes.AutoReactant()` in any -derivative slot of `DensityModel`, or as the sampler `backend`. Derivatives are -traced with Enzyme-MLIR and compiled to XLA executables by `Reactant.@compile`, -which keeps them off Enzyme's LLVM pipeline and so off the GPU -`cuMemcpyDtoHAsync_v2` gc-transition abort, and off DifferentiationInterface -entirely. For a log-density-only model it yields a true second-order HVP. +derivative slot of `DensityModel`, or as the sampler `backend`. Two silent failure modes, ahead of the ordinary limitations: @@ -42,14 +38,12 @@ Requirements / limitations: `mode` is rejected outright (`_check_reactant_mode`) rather than ignored. =# -# Both `Reactant` and `Enzyme` trigger this extension (see Project.toml): the -# traced derivatives call `Enzyme.autodiff` / `Enzyme.gradient` themselves -# rather than reaching Enzyme-MLIR through Reactant. using ParallelMCMC: ParallelMCMC using ParallelMCMC.DEER: DEER using ADTypes: ADTypes, AutoReactant, AutoEnzyme using Reactant: Reactant, @compile using Enzyme: Enzyme +using CUDA: CUDA #= `AutoReactant()` defaults to `AutoReactant(; mode=AutoEnzyme())`, i.e. @@ -70,10 +64,7 @@ end #= Warn once per compiled slot when the template is not a plain `Array` (so looks -like it lives on a GPU) while Reactant's default XLA client targets "cpu": every -call then pays a host round trip on top of the usual marshalling. Guarded, so a -Reactant version without `XLA.platform_name` / `XLA.default_backend` degrades to -no warning instead of erroring out of `_prepare_model`. +like it lives on a GPU) while Reactant's default XLA client targets "cpu". =# function _reactant_client_platform() return try @@ -83,8 +74,10 @@ function _reactant_client_platform() end end -_warn_reactant_host_roundtrip(::Array) = nothing -function _warn_reactant_host_roundtrip(x::AbstractArray) +# Only `CuArray` implies a device round-trip; other non-`Array` templates (e.g. +# `SubArray` views) already live on the host and would be false positives. +_warn_reactant_host_roundtrip(::AbstractArray) = nothing +function _warn_reactant_host_roundtrip(x::CUDA.CuArray) if _reactant_client_platform() == "cpu" @warn "AutoReactant: preparing on a $(typeof(x)), but Reactant's default XLA " * "client targets \"cpu\". Every call will round-trip to the host and back " * diff --git a/src/DEER/DEER.jl b/src/DEER/DEER.jl index 07919ea..12690e4 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -201,19 +201,7 @@ abstract type HVPStrategy end struct ForwardOnGrad <: HVPStrategy end struct ReverseOnGrad <: HVPStrategy end -#= -ReactantHVP traces the HVP with Enzyme-MLIR and compiles it to an XLA -executable, off Enzyme's LLVM pipeline and so off the GPU gc-transition abort. -DI cannot drive Reactant, so `AutoReactant` short-circuits the `hvp_mode` -routing above; drop these specializations once it can. What the traced function -has to look like, and which device the compiled program actually runs on, are in -`ext/ReactantExt.jl`'s module docstring. - -Only reached over a hand-written `gradlogp`. An `AutoReactant` gradient slot is -an AD-derived gradient like any other and goes to `_make_hvp_fn_second_order` / -`_make_hvp_batch_fn_second_order`, which dispatch on the backend rather than on -the resolved gradient's type (see `_resolve_hvp`). -=# +# Need separate HVPStrategy for Reactant; DI does not support and so needs separate logic struct ReactantHVP <: HVPStrategy end _strategy_from(::DI.ForwardOverAnything) = ForwardOnGrad() @@ -309,12 +297,7 @@ function _make_hvp_batch_fn( return (X, V) -> _batch_hvp_via_grad_reverse_prepared(prep, X, V) end -#= -`ReactantHVP` fallbacks, so a missing `using Reactant` gives the load hint -rather than a `MethodError`. `ReactantExt` pins `backend` to -`ADTypes.AutoReactant`, which is strictly more specific, so nothing is -overwritten — precompilation forbids that. -=# +# Fallbacks in case a user forgets `using Reactant` function _make_hvp_fn( ::ReactantHVP, gradlogp, backend::AbstractADType, x_template::AbstractVector ) @@ -327,16 +310,6 @@ function _make_hvp_batch_fn( return error(_REACTANT_LOAD_HINT) end -#= -Fallback for the "both slots `AutoReactant`" second-order path, which -`_second_order` in `interface.jl` routes here with `backend::AutoReactant` -rather than a `DI.SecondOrder`. Signature is `AbstractADType` and not -`AutoReactant` because `ReactantExt`'s method is `AutoReactant` exactly, and -precompilation refuses an identical signature; less specific still loses to it. -JET needs a method here too, since it cannot see a conditionally-loaded -extension and would otherwise flag `_resolve_hvp`'s `AutoReactant` branch as -having none. -=# function _make_hvp_fn_second_order( logdensity, backend::AbstractADType, x_template::AbstractVector ) diff --git a/src/interface.jl b/src/interface.jl index cf77a17..a515ff4 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -217,11 +217,7 @@ function _resolve_gradient_batch( ) end -#= -`AutoReactant` gradients bypass DI (which cannot drive Reactant yet) and go -through hook functions that `ReactantExt` fills in with strictly more specific -methods; the untyped fallbacks give a clear load-order error. -=# +# `AutoReactant` gradients bypass DI function _resolve_gradient( logdensity, backend::ADTypes.AutoReactant, x_template::AbstractVector ) @@ -275,6 +271,7 @@ end function _resolve_hvp_batch( logdensity_batch, grad_batch, grad_batch_backend, hvp_backend, X_template ) + _check_reactant_pair(grad_batch_backend, hvp_backend) if hvp_backend isa DI.SecondOrder return DEER._make_hvp_batch_fn_second_order( _BatchLogdensitySum(logdensity_batch), hvp_backend, X_template @@ -295,22 +292,18 @@ end #= The HVP backend composed with the backend that produced the gradient under it. For a DI-driven pair that is literally `DI.SecondOrder(hvp_backend, grad_backend)`, taken through `DI.hvp`. Two `AutoReactant`s are not a pair DI -could run at all, so they collapse to the one backend that traces -forward-over-reverse from `logdensity` itself (`ReactantExt`'s -`_make_hvp_fn_second_order`). A mixed pair is already out by the time this runs, -via `_check_reactant_pair`. =# +could run at all, so they collapse to the `hvp_backend` of the two: the one +backend that traces forward-over-reverse from `logdensity` itself (`ReactantExt`'s +`_make_hvp_fn_second_order`). Returning `hvp_backend` rather than a fresh +`AutoReactant()` keeps a non-default `mode` on it reachable by +`_check_reactant_mode`. A mixed pair is already out by the time this runs, via +`_check_reactant_pair`. =# _second_order(hvp_backend, grad_backend) = DI.SecondOrder(hvp_backend, grad_backend) -_second_order(::ADTypes.AutoReactant, ::ADTypes.AutoReactant) = ADTypes.AutoReactant() - -#= -Reactant does not pair with a DI backend across the two passes of an HVP: the -compiled gradient is an opaque XLA executable DI cannot differentiate, and a -DI-prepared gradient is not Reactant-traceable. Both slots take `AutoReactant` -or neither does; a hand-written gradient pairs with either. +function _second_order(hvp_backend::ADTypes.AutoReactant, ::ADTypes.AutoReactant) + return hvp_backend +end -Dispatch rather than a runtime `isa` chain, so the check folds away with the rest -of `_resolve_hvp`'s branching. -=# +# Check Reactant isn't passed with DI _check_reactant_pair(grad_backend, hvp_backend) = nothing _check_reactant_pair(::ADTypes.AutoReactant, ::ADTypes.AutoReactant) = nothing _check_reactant_pair(::Nothing, ::ADTypes.AutoReactant) = nothing @@ -430,10 +423,7 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, #= Settle the HVP backend and check its pairing with `grad_backend` before resolving the gradient. Otherwise a mismatched `AutoReactant` pair, or the LogDensityProblems-gradient case, surfaces only once `_resolve_gradient` has - paid for an XLA compile: 18+ seconds to report a config error. Neither check - needs the resolved gradient. `grad_backend` is known already, and - `_check_reactant_hvp_source` only looks at `model.grad_logdensity`'s type, - which is `grad` unchanged whenever `grad_backend` is `nothing`. =# + paid for an XLA compile: 18+ seconds to report a config error. =# needs_hvp = model.hvp === nothing || model.hvp isa AbstractADType hvp_backend = if needs_hvp hb = model.hvp === nothing ? backend : model.hvp @@ -545,10 +535,10 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, ) end -# Callable structs that allow us to dispatch on the type of the LogDensityProblems object in -# the postprocessing stage. Ideally these would be defined in the LogDensityProblemsExt. -# However, structs defined in extensions are hard to get hold of so we define them here. -# The callable behaviour itself is implemented in LogDensityProblemsExt. +#= Callable structs that allow us to dispatch on the type of the LogDensityProblems object in +the postprocessing stage. Ideally these would be defined in the LogDensityProblemsExt. +However, structs defined in extensions are hard to get hold of so we define them here. +The callable behaviour itself is implemented in LogDensityProblemsExt =# struct LogDensityProblemPrimal{L} ld::L end diff --git a/test/test-Reactant-HVP.jl b/test/test-Reactant-HVP.jl index 83be205..e85a3e2 100644 --- a/test/test-Reactant-HVP.jl +++ b/test/test-Reactant-HVP.jl @@ -15,11 +15,10 @@ const DI_R = ParallelMCMC.DEER.DI #= Reactant-compiled derivative paths (`AutoReactant`), see ext/ReactantExt.jl. -The quartic target keeps the second-order structure honest: logp = -0.25‖x‖⁴ has -H = -(‖x‖² I + 2 x xᵀ), so an HVP that drops the second-order term — the failure -mode of `Enzyme.hvp` under `@compile` — is caught, where a Gaussian's constant H -would hide it. Every derivative-accuracy testset below uses it; do not swap in a -Gaussian. +The quartic target is needed to test second order structure: logp = -0.25‖x‖⁴ has +H = -(‖x‖² I + 2 x xᵀ), so an HVP that drops the second-order term is caught, +where a Gaussian's constant H would hide it. Every derivative-accuracy testset +below uses it; do not swap in a Gaussian. =# logp_r(x) = -0.25 * sum(abs2, x)^2 gradlogp_r(x) = -sum(abs2, x) .* x @@ -39,8 +38,8 @@ const D_R = 4 const CT_R = FlexiChains.FlexiChain{Symbol} #= The pairing rule lives in `_resolve_hvp` / `_prepare_model`, not in the -extension, so its dispatch table is checked whether or not Reactant loads — -none of these calls resolve a gradient or compile anything. =# +extension, so its dispatch table is checked whether or not Reactant loads. +None of these calls resolve a gradient or compile anything. =# @testset "Reactant does not pair with a DI backend" begin # Both slots Reactant, or a hand-written gradient (`nothing`), are accepted. @test ParallelMCMC._check_reactant_pair(AutoReactant(), AutoReactant()) === nothing From 4dc9a626cf49c136a9280e20367ac0d713b1106b Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Thu, 13 Aug 2026 09:51:43 -0400 Subject: [PATCH 11/11] cut down chaff --- ext/DynamicPPLExt.jl | 14 +++-------- ext/ReactantExt.jl | 58 +++++++++++++++++--------------------------- src/DEER/DEER.jl | 7 ++---- src/interface.jl | 55 ++++++++++++++--------------------------- 4 files changed, 46 insertions(+), 88 deletions(-) diff --git a/ext/DynamicPPLExt.jl b/ext/DynamicPPLExt.jl index f824f8c..1df483b 100644 --- a/ext/DynamicPPLExt.jl +++ b/ext/DynamicPPLExt.jl @@ -21,16 +21,10 @@ triggers for this extension), plus any AD backend that is used. `ad_backend` is DynamicPPL's own `adtype` rather than a `DensityModel` slot: it goes to the `LogDensityFunction` that fills the log-density and gradient slots, which is why it takes a backend and never a callable. The remaining keywords are -`DensityModel` slots, forwarded unchanged. - -`ParallelMALASampler` also wants an HVP. Give it a callable or a -`DifferentiationInterface.SecondOrder`, which differentiates the log-density and -so bypasses DynamicPPL's gradient. A plain backend fails: it would differentiate -the gradient `ad_backend` produced, whose preparation rejects an outer pass's -tangents. `ADTypes.AutoReactant()` fails as well, bypassing DI or not, since -Reactant cannot trace DynamicPPL's model evaluation. DynamicPPL supplies no -batched log-density either, so reaching the batched DEER path means writing -`logdensity_batch` by hand. +`DensityModel` slots, forwarded unchanged — see the `LogDensityProblemsExt` +constructor's docstring for the `hvp`/`hvp_batch` caveats that come with an +AD-derived gradient (no plain backend, no `AutoReactant()`, no batched +log-density without writing one by hand). # Example ```julia diff --git a/ext/ReactantExt.jl b/ext/ReactantExt.jl index aab4500..116b0dd 100644 --- a/ext/ReactantExt.jl +++ b/ext/ReactantExt.jl @@ -4,38 +4,26 @@ module ReactantExt Reactant-compiled derivative paths, selected by `ADTypes.AutoReactant()` in any derivative slot of `DensityModel`, or as the sampler `backend`. -Two silent failure modes, ahead of the ordinary limitations: - -Captured data is frozen at compile time. `@compile` bakes any plain `Array` / -`Ref` reached through the traced closure into the executable as a constant. A -`logdensity` written as `x -> f(x, data)` whose `data` is mutated afterwards -keeps handing back the pre-mutation derivative from every executable compiled -before the mutation, with no error and no warning. Traced functions must be pure -with respect to what they capture; pass data that can change in as an argument. - -The compiled program runs wherever Reactant's XLA client points, which need not -be a GPU. That client is a Reactant/`Reactant_jll`-wide setting -(`Reactant.set_default_backend`) and has nothing to do with where the package's -own `Vector`s / `CuArray`s live. Preparing on `CuArray` parameters while the -client targets `"cpu"` still compiles and still gives correct answers, but every -call round-trips `CuArray -> host -> XLA-CPU -> host -> CuArray`. -`_warn_reactant_host_roundtrip` below warns once per compiled slot on that -combination; it cannot fix it. - -Requirements / limitations: - - The traced function (`logdensity` / `gradlogp` / batched forms) must be - Reactant-traceable: plain array ops. DynamicPPL-built log-densities do - NOT trace as-is. - - Executables are shape-specialized to the preparation templates, and - arguments are marshalled to/from Reactant's own (XLA) device memory on - every call. A boundary copy, not a fused in-place path — a target for - later optimization. `@compile` also does not memoize across preparations: - every `sample()` call recompiles every `AutoReactant` slot the model uses. - - HVPs are explicit forward-over-(gradient) compositions; NEVER - `Enzyme.hvp`, which silently returns zeros under `@compile`. - - `AutoReactant.mode` (the wrapped `AutoEnzyme`) is not honoured: gradients - always trace as Enzyme reverse, HVPs as forward-over-that. A non-default - `mode` is rejected outright (`_check_reactant_mode`) rather than ignored. +Two silent failure modes: + - Captured data is frozen at compile time. `@compile` bakes any plain + `Array`/`Ref` reached through the closure in as a constant, so mutating it + later keeps returning the pre-mutation derivative with no error or + warning. Traced functions must be pure w.r.t. their captures; pass mutable + data in as an argument instead. + - HVPs are explicit forward-over-(gradient) compositions, NEVER `Enzyme.hvp`, + which silently returns zeros under `@compile`. + +Other limitations: + - The traced function must be Reactant-traceable (plain array ops). + DynamicPPL-built log-densities do NOT trace as-is. + - Reactant's XLA client (`Reactant.set_default_backend`) is process-wide and + independent of where package arrays live. Preparing on a `CuArray` while it + targets `"cpu"` still compiles and is still correct, but every call + round-trips through the host; `_warn_reactant_host_roundtrip` warns once + and can't fix it. + - Executables are shape-specialized to the preparation templates and are not + memoized across preparations, so every `sample()` call recompiles every + `AutoReactant` slot the model uses. =# using ParallelMCMC: ParallelMCMC @@ -164,10 +152,8 @@ HVP factories. `_resolve_hvp` / `_resolve_hvp_batch` (src/interface.jl) route to one of two shapes, the same two every other backend gets. - Both `grad_logdensity` and the HVP source `AutoReactant`: the AD-derived - gradient case, with `_second_order` collapsing the pair to one - `AutoReactant()` since DI cannot form a `SecondOrder` from it. - `_make_hvp_fn_second_order` here traces forward-over-reverse from - `logdensity` as a single XLA program. + gradient case (see `_second_order` in `interface.jl`). Traced here as + forward-over-reverse from `logdensity`, a single XLA program. - A hand-written `gradlogp` with an `AutoReactant` HVP source: routed by `_hvp_strategy(::AutoReactant) = ReactantHVP()` in `DEER.jl` to `_make_hvp_fn` below, a forward JVP over that callable. diff --git a/src/DEER/DEER.jl b/src/DEER/DEER.jl index 12690e4..098300c 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -212,11 +212,8 @@ _hvp_strategy(::ADTypes.AutoReactant) = ReactantHVP() #= Hook for backend-specific normalization, applied on every AD-HVP path before the -backend reaches DI. It fills in what the wrappers DEER differentiates need and -nothing else: those wrapper types are ours, so EnzymeExt sets -`function_annotation=Enzyme.Const` on them, without which Enzyme throws -`EnzymeMutabilityException` on the read-only `_HvpReverseClosure` / -`_BatchHvpReverseClosure` that capture `gradlogp`. +backend reaches DI. It fills in only what DEER's own wrapper types need; see +`ext/EnzymeExt.jl` for the one specialization that exists. It does not choose a differentiation mode. A mode the user set is a decision, an unset one is DI's to resolve from the operator it runs, and substituting one here diff --git a/src/interface.jl b/src/interface.jl index a515ff4..97aaec3 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -102,11 +102,7 @@ function DensityModel( "grad_logdensity must be a callable or an ADTypes.AbstractADType backend" ), ) - #= The batched DEER update evaluates `logdensity_batch` itself, so without one - a batched derivative slot is unusable either way: a backend has nothing to - differentiate and a callable is never reached. Rejected rather than ignored. - A `logdensity_batch` on its own is fine; `_prepare_model` decides from the - gradient slot whether the batched path can run. =# + # `_prepare_model` decides from the gradient slot whether the batched path runs. _batch_needs_logp(name) = throw( ArgumentError( "$name requires logdensity_batch, which the batched DEER update evaluates" @@ -154,12 +150,10 @@ struct PreppedDensityModel{F,G,H,FB,GB,HB,PN,SM<:DensityModel} source::SM end -#= -Whether a prepped model carried in a state was built from the model a `step` +#= Whether a prepped model carried in a state was built from the model a `step` was handed. Identity, not equality: an equal-but-distinct `DensityModel` just -costs one re-preparation, whereas treating a different model as a match would -silently sample the wrong target. -=# +costs one re-preparation, whereas a false match would silently sample the wrong +target. =# _prepped_for(prepped::PreppedDensityModel, model::DensityModel) = prepped.source === model #= @@ -245,13 +239,6 @@ statically known. A `SecondOrder` bypasses the gradient slot even when that slot is hand-written: naming both passes asks for two derivatives of `logdensity`. The slot is still the drift term the MALA step uses. - -`AutoReactant` takes the same three branches as anything else. `_second_order` -below collapses an `AutoReactant` pair to a single `AutoReactant()` rather than a -`DI.SecondOrder`, and `_hvp_strategy(::AutoReactant)` sends the -hand-written-gradient case to `ReactantHVP`; `ReactantExt` supplies both matching -methods. The pairing is already checked by `_prepare_model` before either path -here pays for a gradient resolution or an HVP compile. =# function _resolve_hvp(logdensity, grad, grad_backend, hvp_backend, x_template) if hvp_backend isa DI.SecondOrder @@ -289,15 +276,13 @@ function _resolve_hvp_batch( end end -#= The HVP backend composed with the backend that produced the gradient under it. -For a DI-driven pair that is literally `DI.SecondOrder(hvp_backend, -grad_backend)`, taken through `DI.hvp`. Two `AutoReactant`s are not a pair DI -could run at all, so they collapse to the `hvp_backend` of the two: the one -backend that traces forward-over-reverse from `logdensity` itself (`ReactantExt`'s -`_make_hvp_fn_second_order`). Returning `hvp_backend` rather than a fresh -`AutoReactant()` keeps a non-default `mode` on it reachable by -`_check_reactant_mode`. A mixed pair is already out by the time this runs, via -`_check_reactant_pair`. =# +#= The HVP backend composed with the backend that produced the gradient under it, +via `DI.SecondOrder(hvp_backend, grad_backend)` and `DI.hvp`. Two `AutoReactant`s +are not a pair DI could run, so they collapse to `hvp_backend`: the backend +`ReactantExt` traces forward-over-reverse from `logdensity` in +`_make_hvp_fn_second_order`. Returning it, not a fresh `AutoReactant()`, keeps a +non-default `mode` reachable by `_check_reactant_mode`. A mixed pair is already +out by the time this runs, via `_check_reactant_pair`. =# _second_order(hvp_backend, grad_backend) = DI.SecondOrder(hvp_backend, grad_backend) function _second_order(hvp_backend::ADTypes.AutoReactant, ::ADTypes.AutoReactant) return hvp_backend @@ -400,9 +385,8 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector) else model.grad_logdensity end - #= The DEER-only slots are dropped rather than passed along: a sequential - sampler never reads them, and an unresolved backend sitting in one would - break the invariant that no slot here holds an `AbstractADType`. =# + # Dropped, not passed along: a sequential sampler never reads them, and an + # unresolved backend there would break `PreppedDensityModel`'s invariant. return PreppedDensityModel( model.logdensity, grad, @@ -453,11 +437,8 @@ function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, model.hvp end - #= A batched log-density with no batched gradient gets one from the model's - own gradient backend, never the sampler's: a hand-written gradient has not - opted into AD, and deriving one anyway would let `backend=` decide which - update path runs. Not deriving one leaves the batched path off rather than - raising, since `_trajectory_logps` uses `logdensity_batch` either way. =# + # Leaves the batched path off rather than raising: `_trajectory_logps` uses + # `logdensity_batch` either way. grad_batch = model.grad_logdensity_batch if grad_batch === nothing && model.logdensity_batch !== nothing grad_batch = grad_backend @@ -1093,9 +1074,9 @@ function _construct_flexichain( param_names::Any, model::DensityModel, ) where {TKey} - #= Wrap user-supplied names in `Parameter`. This allows people to specify, e.g., - `param_names=(:x, :y, :z=>(2,))` without faffing with `Parameter` themselves. Also - 'upgrade' symbol parameter names to VarNames if the user requested a VNChain. =# + # Wrap user-supplied names in `Parameter`. This allows people to specify, e.g., + # `param_names=(:x, :y, :z=>(2,))` without faffing with `Parameter` themselves. Also + # 'upgrade' symbol parameter names to VarNames if the user requested a VNChain. to_parameter(vn::VarName) = FlexiChains.Parameter(vn) to_parameter(s::Symbol) = FlexiChains.Parameter(TKey <: VarName ? VarName{s}() : s)