From 5f3d5cc08116c00769c8af3d1d0c90d25763f6d5 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Tue, 15 Sep 2026 12:51:40 -0400 Subject: [PATCH 1/5] Initial text changes. --- docs/src/15-gpu.md | 14 +++---- docs/src/95-reference.md | 6 +-- src/DEER/DEER.jl | 82 +++++++++++----------------------------- src/DEER/DEERScan.jl | 7 +--- src/ParallelMCMC.jl | 10 ++--- src/interface.jl | 6 +-- 6 files changed, 41 insertions(+), 84 deletions(-) diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index baf920d..d52f23d 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -1,6 +1,6 @@ # GPU Execution -DEER's prefix-scan and per-step batch evaluations are array-type-agnostic, so the same `ParallelMALASampler` runs on either CPU or GPU. This page covers when GPU pays off, the current limitations, and a worked example. +`ParallelMALASampler` can use CPU or GPU arrays. This page explains how to write a GPU model and which backend restrictions to check. ParallelMCMC does not depend on CUDA.jl. `using CUDA` loads the `CUDAExt` extension, which is required for `CuArray` parameters. @@ -32,7 +32,7 @@ pmcmc_dot(a, b) = dot(a, b) pmcmc_dotsum(A, B) = sum(A .* B) ``` -These exist purely so the package can register Enzyme rules against them without committing type piracy on `Base.*`, `Base.dot`, `Base.sum`. The rules compute primal and cotangents with plain operators *outside* Enzyme's IR rewriter, which keeps them off the reverse-mode path that aborts on GPU. +These wrappers have Enzyme differentiation rules defined by ParallelMCMC. Use them in GPU code differentiated by Enzyme to avoid this compilation failure: ``` unsupported tag gc-transition for @@ -50,11 +50,11 @@ When you do **not** need them: - GPU code with `AutoMooncake(; config=nothing)`: Mooncake's own CUDA extension handles these operations natively. - GPU code with `AutoZygote()`: Zygote uses ChainRules adjoints, which have cuBLAS-backed rules for `*` / `dot` / `sum` on `CuArray`. -These wrappers are a workaround pending upstream Enzyme support for GPU matmul rules. Once that lands, plain `A * B` on `CuArray`s will work and the wrappers will be deprecated. +The wrappers work around this Enzyme compilation failure. -### 2. Some Enzyme GPU gradients hit a `gc-transition` abort. Staging them as single ops is a workaround +### 2. Enzyme may need broadcasts split into separate expressions -Certain `CuArray` gradient expressions abort during Enzyme compilation on GPU with the same `gc-transition` failure as limitation 1. In practice the fused form below aborts while the staged form compiles, so **writing the gradient as single-op stages is a reliable workaround**: +Some `CuArray` gradients fail during Enzyme compilation with the same `gc-transition` error. If you encounter it, try splitting broadcasts into separate expressions, as in this example: ```julia # ABORTS during Enzyme compile on GPU @@ -148,7 +148,7 @@ chain = sample(model, sampler, 1_600; Posterior mean recovery error `‖β_post − β_true‖ / ‖β_true‖` should land in the 0.1–0.2 range after a few hundred post-warmup samples. -### Enzyme backend (requires `pmcmc_matmul) +### Enzyme backend (requires `pmcmc_matmul`) Same model, with the GPU-Enzyme restrictions applied: every `*` becomes `pmcmc_matmul`, every `dot` becomes `pmcmc_dot`, and every gradient broadcast is expanded into single-op stages: @@ -217,7 +217,7 @@ chain = sample(model, sampler, 1_600; chain_type=VNChain) ``` -The two snippets sample the same posterior; the difference is purely in what the AD backend can chew on. +Both examples describe the same posterior. The Enzyme version separates operations to avoid the compilation failures above. --- diff --git a/docs/src/95-reference.md b/docs/src/95-reference.md index 167a0fc..996d007 100644 --- a/docs/src/95-reference.md +++ b/docs/src/95-reference.md @@ -9,7 +9,7 @@ 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 for derivatives in `DensityModel` and as `ParallelMALASampler`'s `backend`. See [GPU Execution](15-gpu.md) for its constraints. ## CUDA and other device arrays @@ -54,8 +54,8 @@ ParallelMALATransition ## Low-level namespaces -These lower-level building blocks power the public samplers and are useful if you -want to work with taped recursions or the diagonal affine scan directly. +Use these functions and workspaces to solve taped recursions or run the +diagonal affine scan directly. ```@docs DEER.TapedRecursion diff --git a/src/DEER/DEER.jl b/src/DEER/DEER.jl index ed93e64..b54d208 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -39,11 +39,11 @@ function TapedRecursion( end """ -Reusable buffers for allocation-light DEER updates and solves. +Reusable buffers for DEER updates and solves. -All buffers are created with the same array type / device placement as `S_template` -(or `s0_template` for vector buffers), so the workspace is GPU-compatible when -constructed from device-array templates. +Matrix buffers use the array type and device of `S_template`; vector buffers +use those of `s0_template`. Device arrays that require host staging also get +CPU buffers for generating random probes. """ struct DEERWorkspace{M,V,SW,HZ,H} A::M @@ -147,24 +147,11 @@ function _batch_hvp_from_grad_prepared( return res isa Tuple ? first(res) : res end -#= ---------------------------------------------------------------------------- -Reverse-on-grad HVP. Computes Hv as the gradient of -`x -> pmcmc_dot(gradlogp(x), v)`, with `v` carried as a DI Constant context -so a single `prepare_gradient` covers all subsequent Newton steps. The -batched variant uses `B -> pmcmc_dotsum(grad_batch(B), V)` and exploits the -column-independence of the user's batched gradient — its gradient w.r.t. B -is the columnwise HVP. - -The reductions go through `pmcmc_dot`/`pmcmc_dotsum` (rather than `dot`/`sum`) -so that on GPU the EnzymeExt reverse rules intercept them — otherwise Enzyme -reverse-mode hits the `cuMemcpyDtoHAsync_v2` gc-transition abort (see the -Enzyme rules in `ext/EnzymeExt.jl`). - -We bundle the closure with the prep so `prepare_gradient` and `gradient` -see the same function instance (DI keys preparations on function identity). ---------------------------------------------------------------------------- -=# +# Compute H*v by differentiating dot(gradlogp(x), v). DI.Constant lets us +# change v without preparing the gradient again. Batched gradients must treat +# columns independently for this calculation to give columnwise HVPs. +# Use pmcmc reductions so Enzyme's GPU rules avoid the gc-transition abort. +# Keep the closure with its preparation: DI requires the same function instance. import ..ParallelMCMC: pmcmc_dot, pmcmc_dotsum, _REACTANT_LOAD_HINT struct _HvpReverseClosure{F} @@ -177,27 +164,10 @@ struct _BatchHvpReverseClosure{F} end (c::_BatchHvpReverseClosure)(X, V) = pmcmc_dotsum(c.grad_batch(X), V) -#= -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. - -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 -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`. The mode is whatever the user's backend carries; -nothing here substitutes one (see `_normalized_backend`). -=# +# Differentiate a supplied gradient in the direction chosen by DI.hvp_mode: +# forward mode computes J*v; reverse mode differentiates dot(gradlogp(x), v). +# Strategy types let dispatch select a concrete closure type. AD-derived +# gradients use _make_hvp_fn_second_order instead. abstract type HVPStrategy end struct ForwardOnGrad <: HVPStrategy end struct ReverseOnGrad <: HVPStrategy end @@ -248,12 +218,7 @@ function _batch_hvp_via_grad_reverse_prepared( return DI.gradient(f, prep, eff_backend, X, DI.Constant(V)) end -#= -Strategy-dispatched factories. Each method returns a closure with one -concrete type, so the call site `_make_hvp_fn(_hvp_strategy(backend), ...)` -is type-stable: dispatch on the singleton `HVPStrategy` picks the method -at compile time, and the returned closure type is statically known. -=# +# Each strategy returns its own concrete closure type. function _make_hvp_fn( ::ForwardOnGrad, gradlogp, backend::AbstractADType, x_template::AbstractVector ) @@ -412,8 +377,8 @@ end """ In-place DEER update. -Writes the updated trajectory into `S_out` using `S_in` as the current iterate and -workspace buffers for all internal temporaries that can be controlled here. +Write one update of `S_in` into `S_out`, reusing buffers in `ws`. +Both trajectories must have size `(length(s0_in), length(rec.tape))`. """ function deer_update!( ws::DEERWorkspace, @@ -535,16 +500,15 @@ function deer_update( end """ -Run DEER iterations until convergence. +Run DEER until the change between iterates meets the tolerances or `maxiter` +is reached. Set `return_info=true` to return `(trajectory, info)` and check +`info.converged`; reaching `maxiter` does not throw an error. -When no workspace is supplied, one is created automatically. Supplying a -pre-allocated `DEERWorkspace` is the intended path for repeated GPU solves. +Pass a `DEERWorkspace` to reuse buffers across solves. Otherwise, one is created. -If `workspace` is supplied, `copy_result` defaults to `true` so the returned -trajectory is owned by the caller and will not be overwritten by a later solve -using the same workspace. Set `copy_result=false` for allocation-sensitive -internal loops; in that mode the returned trajectory may alias workspace-owned -buffers such as `workspace.S_tmp`. +With a supplied workspace, the result is copied by default so later solves +cannot overwrite it. Set `copy_result=false` to skip that copy, but use or copy +the result before reusing the workspace. """ function solve( rec::TapedRecursion, diff --git a/src/DEER/DEERScan.jl b/src/DEER/DEERScan.jl index 57f251e..3db4f72 100644 --- a/src/DEER/DEERScan.jl +++ b/src/DEER/DEERScan.jl @@ -11,9 +11,7 @@ export AffineScanWorkspace, """ Reusable workspace for the diagonal affine scan. -Buffers are allocated with the same array type / device placement as the template -matrix used to construct the workspace, so this works on GPU arrays and CPU -arrays alike. +Buffers use the same array type and device as the template matrix. """ struct AffineScanWorkspace{M} alpha::M @@ -45,8 +43,7 @@ Inputs: - `s0 :: length-D` - `S :: D×T` output buffer -This is the ground-truth implementation to compare against the parallel scan path. -It is intentionally simple and should be used in tests. +Use this sequential implementation to check the parallel scan in tests. """ function solve_affine_seq!( S::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix, s0::AbstractVector diff --git a/src/ParallelMCMC.jl b/src/ParallelMCMC.jl index 69a5873..d8aafce 100644 --- a/src/ParallelMCMC.jl +++ b/src/ParallelMCMC.jl @@ -11,13 +11,9 @@ using Statistics const DI = DifferentiationInterface -#= -Owned wrappers: identical semantics to their Base counterparts, but provide -stable function identities for backend-specific AD rules in `ext/EnzymeExt.jl` -without committing type piracy on `Base.*` / `Base.dot` / `Base.sum`. User -model code that wants those rules to fire (notably on GPU) should call these -instead. See `ext/EnzymeExt.jl` for the gc-transition abort they work around. -=# +# These wrappers let EnzymeExt define GPU differentiation rules without +# changing rules for Julia's standard operators. GPU models using Enzyme +# should call them to avoid the compilation failure described in the GPU guide. 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) diff --git a/src/interface.jl b/src/interface.jl index 8935f84..68aaaa8 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -10,7 +10,7 @@ Defines model/sampler/state/transition types and implements Wrap a log-density and its derivatives for ParallelMCMC samplers. -Each derivative slot (`grad_logdensity`, `hvp`, `grad_logdensity_batch`, +Each derivative argument (`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: @@ -19,7 +19,7 @@ built from the log-density alone: Backends are prepared when sampling starts. An HVP backend differentiates a callable gradient once. With an AD-derived gradient it forms `DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` on the -log-density. An explicit `SecondOrder` also bypasses the gradient slot. +log-density. An explicit `SecondOrder` also differentiates the log-density directly. - `logdensity(x::AbstractVector) -> Real` - `grad_logdensity` — callable `x -> AbstractVector`, or an AD backend. @@ -38,7 +38,7 @@ log-density. An explicit `SecondOrder` also bypasses the gradient slot. See the [`Parameter names`](@ref parameter-names) section of the docs for more information. -Both batched derivative slots require `logdensity_batch`. `AutoReactant()` has +Both batched derivative arguments require `logdensity_batch`. `AutoReactant()` has additional pairing and tracing constraints; see the GPU guide. """ struct DensityModel{F,G,H,FB,GB,HB,PN} <: AbstractMCMC.AbstractModel From 76e0c7169b3ea9bf7a24c289fb400cc19401009d Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Sun, 27 Sep 2026 15:04:20 -0400 Subject: [PATCH 2/5] Text Changes --- .JuliaFormatter.toml | 2 +- .github/workflows/Test.yml | 4 +- .gitignore | 2 +- README.md | 64 ++++--- docs/make.jl | 68 ++------ docs/src/10-getting-started.md | 309 ++++++++++++++++----------------- docs/src/12-models.md | 205 ++++++++++++++++++++++ docs/src/15-gpu.md | 59 +++++-- docs/src/20-algorithms.md | 219 ++++++++++++++--------- docs/src/95-reference.md | 64 ++++--- docs/src/index.md | 156 ++++++++--------- src/interface.jl | 137 +++++++++------ 12 files changed, 785 insertions(+), 504 deletions(-) create mode 100644 docs/src/12-models.md diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml index 98a6700..61598dd 100644 --- a/.JuliaFormatter.toml +++ b/.JuliaFormatter.toml @@ -4,4 +4,4 @@ normalize_line_endings = "unix" v2_stable_multiline_strings = true conditional_to_if = false pipe_to_function_call = false -style="blue" \ No newline at end of file +style="blue" diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index 5d00841..ce484a3 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -25,8 +25,8 @@ jobs: - "1" os: - ubuntu-latest - - + + arch: - x64 allow_failure: [false] diff --git a/.gitignore b/.gitignore index 27204ce..015020e 100644 --- a/.gitignore +++ b/.gitignore @@ -29,4 +29,4 @@ CODEX.md # random scripts for debugging /scripts -/dev \ No newline at end of file +/dev diff --git a/README.md b/README.md index b443e61..87a4205 100644 --- a/README.md +++ b/README.md @@ -23,28 +23,34 @@ ## What this package does -**ParallelMCMC.jl** implements *parallel-across-the-sequence* MCMC in Julia: instead of generating samples one at a time, an entire trajectory of $T$ correlated steps is solved *simultaneously*. This makes wall-clock time per sample sublinear in chain length on multi-core CPUs and GPUs, where conventional sequential MCMC scales linearly. +ParallelMCMC.jl implements MCMC methods that parallelize *within* a chain. Its +main sampler uses DEER to solve a block of MALA transitions together instead of +waiting for each transition before starting the next one. -The flagship algorithm is **DEER** (Lim et al. 2024; Gonzalez et al. 2024), which reformulates a chain of $T$ MALA steps as a fixed-point problem and solves it with Newton iterations. Each iteration linearizes the per-step transition around the current trajectory guess and resolves the resulting linear recursion in $O(\log T)$ parallel work via an associative prefix scan. With shared input randomness, DEER converges to the exact sequential MALA trace up to a numerical tolerance — typically in tens of iterations even for chains of tens of thousands of samples. +This is parallel-across-the-sequence MCMC, not just several independent chains +running at once. It is most useful when density and derivative evaluations are +expensive enough to keep parallel hardware busy. Small CPU targets may still be +faster with ordinary sequential MALA. -The approach and its scaling tricks (stochastic Hutchinson Jacobian estimators, damping, sliding windows) are described in: +The method is described in: > Zoltowski, D. M., Wu, S., Gonzalez, X., Kozachkov, L., & Linderman, S. W. (2025). > **Parallelizing MCMC Across the Sequence Length.** *NeurIPS 2025.* > [arXiv:2508.18413](https://arxiv.org/abs/2508.18413) -### Samplers +## Samplers | Sampler | Role | |---|---| -| [`ParallelMALASampler`](src/interface.jl) | **Primary** — parallel-across-sequence MALA via DEER; $O(\log T)$ per solve | -| [`MALASampler`](src/interface.jl) | Baseline — sequential MALA with a fixed step size | -| [`AdaptiveMALASampler`](src/interface.jl) | Baseline — sequential MALA with dual-averaging step-size adaptation | +| [`ParallelMALASampler`](docs/src/95-reference.md) | Parallel-across-sequence MALA via DEER | +| [`AdaptiveMALASampler`](docs/src/95-reference.md) | Sequential MALA with step-size adaptation | +| [`MALASampler`](docs/src/95-reference.md) | Sequential MALA with a fixed step size | -All samplers implement the [AbstractMCMC](https://github.com/TuringLang/AbstractMCMC.jl) interface and return [`FlexiChains`](https://pysm.dev/FlexiChains.jl) objects, so they slot into existing AbstractMCMC workflows. +All three implement the +[AbstractMCMC](https://github.com/TuringLang/AbstractMCMC.jl) interface and +return [FlexiChains](https://pysm.dev/FlexiChains.jl) objects. - -### Quick start +## Quick start ParallelMCMC can be installed via Julia's package manager. In the Julia REPL, press `]` to enter pkg mode, then run: @@ -53,22 +59,34 @@ pkg> add ParallelMCMC ``` ```julia -using ParallelMCMC, FlexiChains -using ADTypes, Enzyme - -logp(x) = -0.5 * sum(abs2, x) # 2-D standard normal -grad_logp(x) = -x - -model = DensityModel(logp, grad_logp, 2; param_names=[:x1, :x2]) -sampler = ParallelMALASampler(0.1; T=64, jacobian=:stoch_diag, - backend=AutoEnzyme()) - -chain = sample(model, sampler, 500; chain_type=VNChain) +using ParallelMCMC, FlexiChains, Random + +logdensity(x) = -sum(abs2, x) / 2 +gradient(x) = -x +hvp(x, v) = -v + +model = DensityModel( + logdensity, + gradient, + 2; + hvp=hvp, + param_names=[:x1, :x2], +) + +sampler = ParallelMALASampler(0.1; T=64) +chain = sample( + MersenneTwister(42), model, sampler, 1_000; + initial_params=zeros(2), + chain_type=VNChain, +) ``` -See the [Getting Started guide](docs/src/10-getting-started.md) for worked examples, GPU usage, and step-size tuning. +Start with the [documentation](https://rsenne.github.io/ParallelMCMC.jl/dev/) +or the repository's [getting-started guide](docs/src/10-getting-started.md). +The docs cover step-size tuning, model and AD setup, Turing integration, GPU +execution, troubleshooting, and the DEER algorithm. -## How to Cite +## How to cite If you use ParallelMCMC.jl in your work, please cite using the reference given in [CITATION.cff](https://github.com/rsenne/ParallelMCMC.jl/blob/main/CITATION.cff). diff --git a/docs/make.jl b/docs/make.jl index 79c38fb..39ac842 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -4,60 +4,18 @@ using Documenter.Remotes: GitHub DocMeta.setdocmeta!(ParallelMCMC, :DocTestSetup, :(using ParallelMCMC); recursive=true) -# Add titles of sections and overrides page titles -const titles = Dict( - # "10-tutorials" => "Tutorials", # example folder title - "91-developer.md" => "Developer docs", -) - -function recursively_list_pages(folder; path_prefix="") - pages_list = Any[] - for file in readdir(folder) - if file == "index.md" - # We add index.md separately to make sure it is the first in the list - continue - end - # this is the relative path according to our prefix, not @__DIR__, i.e., relative to `src` - relpath = joinpath(path_prefix, file) - # full path of the file - fullpath = joinpath(folder, relpath) - - if isdir(fullpath) - # If this is a folder, enter the recursion case - subsection = recursively_list_pages(fullpath; path_prefix=relpath) - - # Ignore empty folders - if length(subsection) > 0 - title = if haskey(titles, relpath) - titles[relpath] - else - @error "Bad usage: '$relpath' does not have a title set. Fix in 'docs/make.jl'" - relpath - end - push!(pages_list, title => subsection) - end - - continue - end - - if splitext(file)[2] != ".md" # non .md files are ignored - continue - elseif haskey(titles, relpath) # case 'title => path' - push!(pages_list, titles[relpath] => relpath) - else # case 'title' - push!(pages_list, relpath) - end - end - - return pages_list -end - -function list_pages() - root_dir = joinpath(@__DIR__, "src") - pages_list = recursively_list_pages(root_dir) - - return ["index.md"; pages_list] -end +const pages = [ + "Home" => "index.md", + "Getting started" => "10-getting-started.md", + "Defining models" => "12-models.md", + "GPU execution" => "15-gpu.md", + "How the algorithm works" => "20-algorithms.md", + "API reference" => "95-reference.md", + "Contributing" => [ + "Contributor guide" => "90-contributing.md", + "Developer guide" => "91-developer.md", + ], +] makedocs(; modules=[ParallelMCMC], @@ -70,7 +28,7 @@ makedocs(; repolink="https://github.com/rsenne/ParallelMCMC.jl", edit_link="main", ), - pages=list_pages(), + pages=pages, ) deploydocs(; repo="github.com/rsenne/ParallelMCMC.jl", devbranch="main") diff --git a/docs/src/10-getting-started.md b/docs/src/10-getting-started.md index 7475b75..4ab0deb 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -1,211 +1,202 @@ -# Getting Started +# Getting started -ParallelMCMC.jl is built around **parallel-across-the-sequence** MCMC. The core algorithm is [`ParallelMALASampler`](@ref), which solves an entire trajectory of correlated steps simultaneously. [`MALASampler`](@ref) and [`AdaptiveMALASampler`](@ref) are sequential MALA baselines, included mainly for correctness testing and step-size exploration before running DEER. +This guide takes a model from a log-density to a tuned parallel MALA run. The +example uses a banana-shaped target because it is nonlinear enough to exercise +the derivative path while still being easy to inspect. -## Defining a model +## Install an AD backend -All samplers take a [`DensityModel`](@ref) as their first argument. A `DensityModel` wraps: +ParallelMCMC treats automatic differentiation packages as optional +dependencies. This guide uses ForwardDiff to build the Hessian-vector products +that parallel MALA needs. -- `logdensity(x) -> Real` — log-density of the target (up to a constant) -- `grad_logdensity(x) -> AbstractVector` — its gradient -- optional `hvp`, `logdensity_batch`, `grad_logdensity_batch`, and `hvp_batch` helpers for faster DEER Jacobian updates -- `dim::Int` — dimension of the parameter space - -```julia -using ParallelMCMC, FlexiChains -using ADTypes, Enzyme - -# Banana-shaped target in 2-D -function logp(x) - -0.5 * (x[1]^2 / 100 + (x[2] + 0.1*x[1]^2 - 1)^2) -end - -function grad_logp(x) - dx1 = -x[1]/100 - 0.2*x[1]*(x[2] + 0.1*x[1]^2 - 1) - dx2 = -(x[2] + 0.1*x[1]^2 - 1) - [dx1, dx2] -end - -model = DensityModel(logp, grad_logp, 2; param_names=[:x1, :x2]) +```julia-repl +pkg> add ParallelMCMC ForwardDiff ``` -### AD backends in derivative slots - -Derivative slots (`grad_logdensity`, `hvp`, `grad_logdensity_batch`, and `hvp_batch`) accept callables or `ADTypes.AbstractADType` backends. This is enough to build a model from only its log-density: - ```julia -model = DensityModel(logp, AutoEnzyme(), 2; param_names=[:x1, :x2]) +using AbstractMCMC: MCMCThreads +using ADTypes, FlexiChains, ForwardDiff, ParallelMCMC, Random ``` -Hand-written and AD-derived slots can be mixed. +## Define the target -### What a backend in `hvp` differentiates +A [`DensityModel`](@ref) needs a log-density, a gradient, and the number of +unconstrained parameters. The log-density may omit its normalizing constant. -| `grad_logdensity` | `hvp` | what runs | -|---|---|---| -| callable | backend | one AD pass over your gradient | -| backend | backend | `SecondOrder(hvp, grad_logdensity)` on `logdensity` | -| either | `SecondOrder(...)` | that pair on `logdensity`, gradient slot unused | - -An explicit `SecondOrder` bypasses the gradient slot. Use it when the gradient can't be differentiated. - -See [GPU Execution](15-gpu.md) for backend compatibility. +```julia +function banana_logdensity(x) + a = x[2] + 0.1 * x[1]^2 - 1 + return -(x[1]^2 / 100 + a^2) / 2 +end -The batched pair differentiates `sum(logdensity_batch(X))`, so each output must depend only on the corresponding column of `X`. When `grad_logdensity` is a backend, omitting `grad_logdensity_batch` derives it automatically. Both batched derivative slots require `logdensity_batch`. +function banana_gradient(x) + a = x[2] + 0.1 * x[1]^2 - 1 + return [-x[1] / 100 - 0.2 * x[1] * a, -a] +end -[`ParallelMALASampler`](@ref)'s `backend` supplies missing HVPs. Omit it when the model supplies them. +model = DensityModel( + banana_logdensity, + banana_gradient, + 2; + param_names=[:x1, :x2], +) +``` ---- +The gradient is hand-written, but it does not have to be. See +[Defining models](12-models.md) for AD-derived gradients, analytical +Hessian-vector products, batched functions, and Turing models. -## ParallelMALASampler +## Tune the step size -[`ParallelMALASampler`](@ref) reformulates a trajectory of `T` MALA steps as a fixed-point problem and solves it via Newton iterations, each of which costs $O(\log T)$ parallel work via an associative prefix scan. Wall-clock time per sample is therefore sublinear in chain length on multi-core CPUs and GPUs. +MALA is sensitive to its step size. A value that is too large causes frequent +rejections; one that is too small produces highly correlated samples. Start +with [`AdaptiveMALASampler`](@ref), which tunes the step size during warmup. ```julia -sampler = ParallelMALASampler(0.1; T=64, jacobian=:stoch_diag, damping=0.5, - backend=AutoEnzyme()) - -chain = sample(model, sampler, 500; - chain_type=VNChain, progress=true) +rng = MersenneTwister(42) +tuning = sample( + rng, + model, + AdaptiveMALASampler(0.1; n_warmup=500), + 750; + initial_params=zeros(2), + chain_type=VNChain, + discard_warmup=true, +) + +epsilon = last(vec(tuning[:step_size])) ``` -`sample` requests 500 total samples. Internally, DEER solves trajectories of length `T=64` and returns each column of the solved trajectory as a separate sample. When the trajectory is exhausted a new noise tape is drawn and DEER re-solves from the last state. - -### Trajectory length `T` - -A larger `T` means more samples per DEER solve and better parallelism. On GPU or with many CPU threads the cost per sample decreases roughly as $1/T$ (up to the $O(\log T)$ scan overhead). A typical starting point is `T=64`; increase to `T=128` or `T=256` for long sampling runs on parallel hardware. +The requested 750 transitions include warmup; `discard_warmup=true` removes +every transition marked as warmup from the returned chain. `:step_size`, +`:accepted`, and `:logp` are stored as extra fields alongside the parameters. -### Jacobian modes +!!! tip "Treat tuning as a diagnostic" + Check that the tuned step size settles to a finite, positive value and that + the chain moves through the target. Automatic tuning cannot rescue an + incorrect gradient or a badly chosen parameterization. -The `jacobian` keyword controls how the per-step Jacobian is approximated during each Newton iteration: +## Run parallel MALA -| Mode | Cost per step | Notes | -|---|---|---| -| `:stoch_diag` (default) | `probes` JVPs | Hutchinson estimator; scalable default for high dimensions | -| `:diag` | `D` JVPs | Exact diagonal; useful for low-dimensional checks | - -For high-dimensional targets, `:stoch_diag` with a small number of `probes` is a good trade-off: +Pass the tuned value to [`ParallelMALASampler`](@ref). Because the model does +not provide an analytical Hessian-vector product, `backend=AutoForwardDiff()` +tells the sampler how to construct one from `banana_gradient`. ```julia -sampler = ParallelMALASampler(0.1; T=128, jacobian=:stoch_diag, probes=2, - backend=AutoEnzyme()) +sampler = ParallelMALASampler( + epsilon; + T=64, + backend=AutoForwardDiff(), +) + +chain = sample( + MersenneTwister(43), + model, + sampler, + 2_000; + initial_params=zeros(2), + chain_type=VNChain, +) ``` -### Damping - -Setting `damping < 1` blends the Newton update with the previous iterate, which improves convergence on poorly conditioned targets. The default is `0.5`; values in `[0.3, 0.7]` are typical. - -### GPU execution - -The prefix-scan kernel runs as pure array broadcasts and is array-type-agnostic. To run DEER on GPU: - -1. Implement `logdensity` and `grad_logdensity` using GPU-compatible operations. -2. Pass for example, `backend = ADTypes.AutoEnzyme()` to `ParallelMALASampler`. -3. Pass a `CuVector` as `initial_params` to `sample`. +The important arguments are: -GPU execution has its own caveats (`AutoEnzyme()` currently needs the `pmcmc_matmul` / `pmcmc_dot` / `pmcmc_dotsum` wrappers, gradients should avoid scalar indexing). See [GPU Execution](15-gpu.md) for the worked example and the full set of limitations. +- `epsilon`: the MALA step size. +- `T`: the number of transitions in each DEER solve. It does not change the + number of samples returned by `sample`. +- `backend`: the AD backend used for missing Hessian-vector products. Leave it + out when the model provides `hvp` and, when applicable, `hvp_batch`. +- `initial_params`: the first state. Its array type also selects CPU or GPU + storage. +- `chain_type`: use `VNChain` for `VarName` keys or `SymChain` for `Symbol` + keys. Turing models require `VNChain`. ---- +## Choose `T` and the DEER controls -## LogDensityProblems models +`T=64` is a reasonable first experiment. Larger blocks expose more work to +parallel hardware but use more memory and can make the nonlinear solve harder. +Increase `T` only after the model works at a smaller value. -With `LogDensityProblems` loaded, `DensityModel` accepts any object that implements the [LogDensityProblems](https://github.com/tpapp/LogDensityProblems.jl) interface with gradients: +The defaults are intended to be a usable starting point: ```julia -using LogDensityProblems, ParallelMCMC, FlexiChains - -# μ ~ Normal(0, 1), y ~ Normal(μ, 0.5) -struct NormalPosterior - y::Float64 -end -LogDensityProblems.capabilities(::Type{NormalPosterior}) = - LogDensityProblems.LogDensityOrder{1}() -LogDensityProblems.dimension(::NormalPosterior) = 1 -LogDensityProblems.logdensity(ld::NormalPosterior, x) = -0.5 * x[1]^2 - 2 * (ld.y - x[1])^2 -function LogDensityProblems.logdensity_and_gradient(ld::NormalPosterior, x) - return LogDensityProblems.logdensity(ld, x), [-x[1] + 4 * (ld.y - x[1])] -end - -model = DensityModel(NormalPosterior(1.5); param_names=[:μ]) -chain = sample(model, AdaptiveMALASampler(0.3; n_warmup=500), 2_000; - chain_type=VNChain, discard_warmup=true) +ParallelMALASampler( + epsilon; + T=64, + jacobian=:stoch_diag, + probes=1, + damping=0.5, + maxiter=200, + tol_abs=1e-6, + tol_rel=1e-5, + backend=AutoForwardDiff(), +) ``` -The object must implement `capabilities` (at least `LogDensityOrder{1}`), `dimension`, `logdensity`, and `logdensity_and_gradient`. `hvp` and the batched slots are forwarded to `DensityModel`. - -A backend in `hvp` differentiates `logdensity_and_gradient`, so that function must be differentiable by the backend. If it is not (for example, it calls AD internally), pass an explicit `SecondOrder`, which differentiates `logdensity` twice instead. `AutoReactant()` is not supported for these models. - ---- +- `jacobian=:stoch_diag` uses a Hutchinson estimate of the Jacobian diagonal. + This is the scalable default. +- `probes` controls the number of random vectors in that estimate. More probes + reduce estimator noise but require more derivative work. +- `damping` blends each DEER update with the previous trajectory. Lower it when + the solve is unstable; raise it cautiously when convergence is already + reliable. +- `maxiter`, `tol_abs`, and `tol_rel` control the nonlinear solve. Tightening + tolerances costs more iterations. -## Parallel chains +For a low-dimensional correctness check, `jacobian=:diag` computes the exact +diagonal with one JVP per dimension. It is usually too expensive for large +models. -All samplers support `MCMCThreads()`. Start Julia with multiple threads (e.g. `julia -t 4`): +## Use the returned chain -```julia -chain = sample(model, ParallelMALASampler(0.1; T=64, backend=AutoEnzyme()), - MCMCThreads(), 500, 4; - chain_type=VNChain) -``` +ParallelMCMC returns a [FlexiChains](https://pysm.dev/FlexiChains.jl/) chain. +With the parameter names above, values are available as `chain[:x1]` and +`chain[:x2]`. Sampler diagnostics live in extra columns such as +`chain[:accepted]` and `chain[:logp]`. -`FlexiChains` computes R-hat and ESS across chains automatically: +Use several chains for convergence diagnostics. Start Julia with multiple +threads, for example `julia -t 4`, then run: ```julia -ess(chain) +chains = sample( + model, + sampler, + MCMCThreads(), + 2_000, + 4; + initial_params=zeros(2), + chain_type=VNChain, +) + +ess(chains) ``` -To calculate intra-chain metrics, you can pass `dims=:iter` [as described in the FlexiChains docs](https://pysm.dev/FlexiChains.jl/stable/summarising/#Individual-statistics): - -```julia -ess(chain; dims=:iter) -``` +`MCMCThreads()` runs independent chains concurrently. That is separate from +the within-chain parallelism controlled by `T`. ---- +## When something goes wrong -## [Specifying parameter names](@id parameter-names) +### The DEER solve takes many iterations -For manually constructed `DensityModel`s, you can optionally specify parameter names with the `param_names` keyword. -The resulting `FlexiChain` object will then have named entries for each parameter. +First reduce `epsilon` or `T`. If the model otherwise behaves well, lower +`damping` from `0.5` toward `0.3`. Do not loosen the convergence tolerances +until you have ruled out a derivative error. -If you do not specify `param_names`, the chain will store a single vector-valued parameter called `x` of length `D`. +### Parallel MALA is slower than sequential MALA -You can specify parameter names as a collection of either: +That is expected for small, cheap targets on a CPU. Parallel MALA has setup, +derivative, and scan overhead. It is designed for targets with enough batched +work to benefit from parallel hardware. Benchmark the complete sampling call, +including compilation separately from steady-state execution. -- `Symbol`s (e.g. `[:x1, :x2]`); -- `VarName`s (e.g. `[@varname(x[1]), @varname(x[2])]`) if `chain_type=VNChain`; or -- A tuple of the above, *plus* a size. In this case, a total of `prod(size)` entries will be allocated to the named parameter, and the results in the chain will be reshaped to that size. +### The chain does not move or has many rejections -The above can be mixed and matched as desired, as long as the total number of parameters matches the dimension of the model. -For example: +Run [`MALASampler`](@ref) or [`AdaptiveMALASampler`](@ref) first. Check the +gradient against finite differences or an independent AD backend, then try a +smaller step size. -```julia -# Three scalar parameters, called `x1` through `x3` -model = DensityModel(...; param_names=[:x1, :x2, :x3]) - -# One scalar parameter called `x`, and a 1x2 matrix parameter called `y` -model = DensityModel(...; param_names=[:x, (:y, (1, 2))]) -``` - ---- - -## Sequential MALA baselines - -[`MALASampler`](@ref) and [`AdaptiveMALASampler`](@ref) are standard sequential MALA implementations. They are useful for: - -- Verifying that a model and step size are set up correctly before running DEER -- Obtaining a reference chain to compare against -- Step-size exploration (run `AdaptiveMALASampler` first, read off the frozen `ε̄`, pass it to `ParallelMALASampler`) - -```julia -# Step 1: find a good step size with adaptive MALA -baseline = sample(model, AdaptiveMALASampler(0.1; n_warmup=500), 600; - chain_type=VNChain, discard_warmup=true) - -# Read off the frozen step size from the last internal value -eps_tuned = baseline[end, :step_size, 1] - -# Step 2: run DEER with the tuned step size -chain = sample(model, ParallelMALASampler(eps_tuned; T=64, backend=AutoEnzyme()), 2_000; - chain_type=VNChain) -``` +### GPU code reports scalar indexing or AD compilation errors -See the [`MALASampler`](@ref) and [`AdaptiveMALASampler`](@ref) reference pages for the full keyword listing. +The model must be written in GPU-compatible array operations, and backend +support differs. Continue with [GPU execution](15-gpu.md). diff --git a/docs/src/12-models.md b/docs/src/12-models.md new file mode 100644 index 0000000..403a5ee --- /dev/null +++ b/docs/src/12-models.md @@ -0,0 +1,205 @@ +# Defining models + +Every sampler in ParallelMCMC works with a [`DensityModel`](@ref). The model is +the boundary between your statistical problem and the sampling algorithm: it +describes the target density, its derivatives, its dimension, and optionally +how parameters should be named in the result. + +## The basic contract + +The main constructor is: + +```julia +DensityModel(logdensity, grad_logdensity, dim; kwargs...) +``` + +For a parameter vector `x` of length `dim`: + +- `logdensity(x)` returns one real number. The density may be unnormalized. +- `grad_logdensity(x)` returns a vector with the same shape as `x`. +- `dim` is the number of unconstrained parameters. + +The derivative slot may contain either a callable or an +`ADTypes.AbstractADType` backend. + +## Option 1: write the derivatives + +Hand-written derivatives keep the execution path predictable and are often the +best choice for GPU models. + +```julia +using ParallelMCMC + +logdensity(x) = -sum(abs2, x) / 2 +gradient(x) = -x +hvp(x, v) = -v + +model = DensityModel( + logdensity, + gradient, + 4; + hvp=hvp, + param_names=[:a, :b, :c, :d], +) +``` + +Sequential MALA needs the gradient. Parallel MALA also needs +Hessian-vector products (HVPs). You can supply `hvp(x, v)` as above, or let the +sampler construct it with an AD backend. + +## Option 2: derive gradients with AD + +Pass an AD backend in the gradient slot to build a model from the log-density +alone: + +```julia +using ADTypes, ForwardDiff, ParallelMCMC + +logdensity(x) = -sum(abs2, x) / 2 +model = DensityModel(logdensity, AutoForwardDiff(), 4) +``` + +Backends are prepared when sampling begins, once the package knows the element +type and storage of `initial_params`. + +For parallel MALA, also provide a backend to the sampler unless the model has +an `hvp`: + +```julia +sampler = ParallelMALASampler(0.1; backend=AutoForwardDiff()) +``` + +You can mix approaches. For example, keep a hand-written gradient and use AD +only for the HVP: + +```julia +model = DensityModel(logdensity, gradient, 4) +sampler = ParallelMALASampler(0.1; backend=AutoForwardDiff()) +``` + +### Which function does an HVP backend differentiate? + +| Gradient slot | HVP slot | What ParallelMCMC differentiates | +|---|---|---| +| callable | backend | the callable gradient | +| backend | backend | the log-density with a second-order backend pair | +| either | `SecondOrder(...)` | the log-density with that explicit pair | + +An explicit `DifferentiationInterface.SecondOrder` takes precedence over the +gradient slot. Backend combinations are not equally useful on every device; +see [GPU execution](15-gpu.md) before selecting one for GPU work. + +## Batched functions + +DEER evaluates corresponding operations across all `T` columns of a trajectory. +If your model has an efficient matrix implementation, provide: + +```julia +logdensity_batch(X) # one log-density per column +grad_logdensity_batch(X) # one gradient per column +hvp_batch(X, V) # H(X[:, j]) * V[:, j] per column +``` + +Pass them as keywords: + +```julia +model = DensityModel( + logdensity, + gradient, + dim; + hvp=hvp, + logdensity_batch=logdensity_batch, + grad_logdensity_batch=grad_logdensity_batch, + hvp_batch=hvp_batch, +) +``` + +Both batched derivative slots require `logdensity_batch`. When the gradient +slot contains an AD backend, ParallelMCMC can derive a batched gradient by +differentiating `sum(logdensity_batch(X))`. For that to be correct, each output +of `logdensity_batch` must depend only on the corresponding column of `X`. + +The scalar functions remain part of the contract even when all batched +functions are present. + +## Parameter names + +Without `param_names`, the chain contains one vector-valued parameter named +`:x`. Pass names when separate scalar or array-valued entries would be easier to +work with. + +```julia +# Three scalar parameters +DensityModel(logdensity, gradient, 3; param_names=[:a, :b, :c]) + +# One scalar and one 1 x 2 matrix +DensityModel(logdensity, gradient, 3; param_names=[:a, (:B, (1, 2))]) +``` + +Names may be `Symbol`s or `VarName`s. The declared shapes must account for +exactly `dim` scalar values. `VNChain` accepts `VarName` keys, while `SymChain` +uses `Symbol` keys. + +## Turing models + +Loading Turing activates a convenience constructor that handles unconstraining, +parameter names, and conversion back to the model's original parameter space. +The AD backend is required. + +```julia +using ADTypes, FlexiChains, ForwardDiff, ParallelMCMC, Turing + +@model function normal_location(y) + mu ~ Normal(0, 1) + y ~ Normal(mu, 0.5) +end + +model = DensityModel( + normal_location(1.5); + ad_backend=AutoForwardDiff(), +) + +chain = sample( + model, + AdaptiveMALASampler(0.3; n_warmup=500), + 2_000; + chain_type=VNChain, + discard_warmup=true, +) +``` + +Turing-backed models require `VNChain`. Their output includes `:logjoint`, +`:logprior`, and `:loglikelihood`; a manually constructed model has a single +`:logp` field. + +To use parallel MALA, supply an HVP or an appropriate sampler backend: + +```julia +using Enzyme + +sampler = ParallelMALASampler(0.1; T=64, backend=AutoEnzyme()) +chain = sample(model, sampler, 1_000; chain_type=VNChain) +``` + +Do not assume that the model's gradient backend is also a valid HVP backend. +In particular, applying ForwardDiff around a DynamicPPL gradient that already +uses ForwardDiff creates unsupported nested dual numbers. Enzyme is the +supported default for differentiating that gradient; an explicit second-order +backend pair or analytical HVP can also be used. + +Turing models currently run on the CPU. For a GPU target, write the density and +derivatives directly with GPU-compatible array operations. + +## LogDensityProblems models + +Any object implementing the gradient-capable +[LogDensityProblems](https://github.com/tpapp/LogDensityProblems.jl) interface +can be wrapped directly: + +```julia +model = DensityModel(logdensity_problem; param_names=[:a, :b]) +``` + +This constructor gets the dimension, log-density, and gradient from the +interface. The same HVP and batching keywords accepted by the main constructor +are available when you need them. diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index d52f23d..a9f6545 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -1,28 +1,46 @@ -# GPU Execution +# GPU execution -`ParallelMALASampler` can use CPU or GPU arrays. This page explains how to write a GPU model and which backend restrictions to check. +[`ParallelMALASampler`](@ref) follows the storage of `initial_params`, so the +same sampler can work with CPU vectors or supported device vectors. Moving the +initial state is not enough, though: the target density, its derivatives, and +the selected AD backend must all support the device. -ParallelMCMC does not depend on CUDA.jl. `using CUDA` loads the `CUDAExt` extension, which is required for `CuArray` parameters. +CUDA is an optional dependency. `using CUDA` activates ParallelMCMC's extension +for `CuArray` storage without making CUDA part of a CPU-only installation. --- -## When to use GPU +## Decide whether a GPU is worthwhile -GPU pays off when the per-Newton-step work i.e., evaluating `logdensity_batch` and `grad_logdensity_batch` across the `T`-wide trajectory and the Hutchinson `probes` is large enough to saturate the device. +A GPU helps only when each DEER iteration provides enough work to amortize +kernel launches and data movement. The important workload is the evaluation of +batched densities, gradients, and Hessian-vector products across a `T`-column +trajectory. | Regime | Reach for | |---|---| -| `D` ≲ a few hundred, lightweight `logp` | **CPU.** Per-step cost is dominated by overhead; GPU launch latency wins out. | -| Large `D`, or `logp` is itself a batched linear-algebra / kernel workload | **GPU.** `T` × `probes` Jacobian work fills the device. | -| Many independent chains via `MCMCThreads()` | **CPU**, scaled with threads. GPU chains do not currently share a device pool. | +| Small `D` and a cheap target | CPU; launch overhead is likely to dominate. | +| Large `D` or a target built from substantial batched linear algebra | GPU; the model is more likely to keep the device occupied. | +| Many small independent chains | CPU threads; GPU chains do not currently share a device pool. | -A useful heuristic: if your CPU `grad_logdensity_batch` already saturates BLAS at the batch sizes DEER passes it (`T` columns, plus `probes` HVP probes), GPU is likely to help. If not, profile before switching. +A useful rule is to benchmark before and after moving the model. Include one +run for compilation, then compare warmed-up sampling calls at the same `T`, +step size, and tolerances. --- -## Limitations +## Current limitations -### 1. Enzyme on GPU currently needs `pmcmc_matmul` / `pmcmc_dot` / `pmcmc_dotsum` +### Turing models are CPU-only + +The Turing constructor and its DynamicPPL/LogDensityProblems path currently +execute with CPU vectors. They cannot run a `@model` directly against a +`CuArray`. + +For GPU sampling, write the density and derivatives with GPU-compatible array +operations and pass them to [`DensityModel`](@ref), as in the example below. + +### Enzyme needs ParallelMCMC's linear-algebra wrappers `ParallelMCMC` exports three thin wrappers: @@ -32,27 +50,30 @@ pmcmc_dot(a, b) = dot(a, b) pmcmc_dotsum(A, B) = sum(A .* B) ``` -These wrappers have Enzyme differentiation rules defined by ParallelMCMC. Use them in GPU code differentiated by Enzyme to avoid this compilation failure: +ParallelMCMC defines Enzyme rules for these wrappers without changing the rules +for Julia's base operators. Use them in GPU code differentiated by Enzyme to +avoid failures such as: -``` +```text unsupported tag gc-transition for call i32 @cuMemcpyDtoHAsync_v2(...) [ "jl_roots"(...), "gc-transition"() ] UNREACHABLE executed at .../Enzyme/GradientUtils.cpp:309 (signal 6) ``` -When you should use them: +Use the wrappers when: - You are on GPU **and** using `AutoEnzyme()` as the AD backend. -When you do **not** need them: +Keep the ordinary operators when: - CPU code: plain `*`, `dot`, `sum` are faster and clearer. - GPU code with `AutoMooncake(; config=nothing)`: Mooncake's own CUDA extension handles these operations natively. - GPU code with `AutoZygote()`: Zygote uses ChainRules adjoints, which have cuBLAS-backed rules for `*` / `dot` / `sum` on `CuArray`. -The wrappers work around this Enzyme compilation failure. +The wrappers are an Enzyme-specific compatibility layer; they do not make an +otherwise CPU-only function GPU compatible. -### 2. Enzyme may need broadcasts split into separate expressions +### Enzyme may need broadcasts split into separate expressions Some `CuArray` gradients fail during Enzyme compilation with the same `gc-transition` error. If you encounter it, try splitting broadcasts into separate expressions, as in this example: @@ -247,7 +268,7 @@ sampler = ParallelMALASampler(0.005f0; T=16, backend=AutoReactant()) **Reactant chooses the execution device independently of the input array.** With a CUDA client, `CuArray` inputs are staged through pinned host memory and results are copied back on the device. With the CPU client, both directions go through the host and the sampler warns when it prepares the model. Select a GPU client with `Reactant.set_default_backend` when available. -- The log-density must be Reactant-traceable. +- The log-density must be Reactant-traceable. DynamicPPL-built log-densities are not. - Across an HVP's two AD passes, use `AutoReactant()` for both or neither. A hand-written gradient may pair with it. `SecondOrder` cannot contain `AutoReactant()`. - Only the default `AutoReactant()` mode is supported. @@ -262,7 +283,7 @@ sampler = ParallelMALASampler(0.005f0; T=16, backend=AutoReactant()) - **The HVP has a clean closed form.** Quadratic priors, Gaussian likelihoods, GLMs (logistic, Poisson, probit) — the second derivative is a known function of intermediate quantities you already compute in `gradlogp`. A few extra lines and you skip the AD pipeline entirely. - **Performance matters and the AD compile is heavy.** Enzyme and Mooncake both pay a one-shot compilation cost on the user's gradient. For long-running chains this amortizes, but for many short runs the analytical HVP wins. -- **You're hitting AD-backend-specific GPU restrictions.** The [Enzyme limitations](#1-enzyme-on-gpu-currently-needs-pmcmc_matmul-pmcmc_dot-pmcmc_dotsum) above (`pmcmc_*` wrappers, staged broadcasts) only matter when the AD backend is invoked. Supplying analytical HVP sidesteps them — your `gradlogp` and `hvp` can use plain `*`, `dot`, `sum`, and the sampler's `backend` can be omitted entirely because no AD is ever invoked. +- **You're hitting AD-backend-specific GPU restrictions.** The [Enzyme limitations](#enzyme-needs-parallelmcmcs-linear-algebra-wrappers) above (`pmcmc_*` wrappers, staged broadcasts) only matter when the AD backend is invoked. Supplying an analytical HVP sidesteps them: your `gradlogp` and `hvp` can use plain `*`, `dot`, and `sum`, and the sampler's `backend` can be omitted because no AD is invoked. - **You can reuse intermediates between gradient and HVP.** When `hvp` shares $X\beta$, $\sigma(X\beta)$, or similar with the gradient computation, an analytical version can be both faster *and* shorter than what AD produces. ### Same example with analytical HVP diff --git a/docs/src/20-algorithms.md b/docs/src/20-algorithms.md index df9ab9a..f4b6650 100644 --- a/docs/src/20-algorithms.md +++ b/docs/src/20-algorithms.md @@ -1,147 +1,198 @@ -# Algorithm Details +# How the algorithm works -This page describes the mathematics behind the DEER algorithm, the stochastic diagonal Jacobian estimator, and how both are applied to MALA. The primary reference is Zoltowski et al. (2025) [^1]. +This page explains the implementation behind [`ParallelMALASampler`](@ref). +The main idea comes from Zoltowski et al. (2025) [^1]: turn a sequential Markov +recursion into a nonlinear system, then solve that system with parallel +operations. -## The core idea +You do not need this material to use the sampler. It is useful when choosing +`T`, diagnosing convergence, or working with the low-level DEER API. -A Markov chain can be written as a nonlinear recursion +## From a chain to a nonlinear system + +Once its random draws have been fixed, a Markov chain is a deterministic +recursion: ```math -s_t = f_t(s_{t-1}), \quad t = 1, \ldots, T, +s_t = f_t(s_{t-1}), \qquad t = 1, \ldots, T. ``` -where $f_t$ is the transition kernel (e.g., one MALA step) and the subscript $t$ indicates that the kernel may depend on pre-generated randomness (noise vectors, uniform draws) that is fixed before solving. Generating this randomness ahead of time is called **taping** the chain. +The subscript on $f_t$ accounts for the random values used by transition $t$. +ParallelMCMC draws those values in advance and stores them in a *tape*. -Treating the full trajectory $s_{1:T}$ as the unknown and defining the residual +Instead of applying the transitions one after another, collect the unknown +states into $S = (s_1, \ldots, s_T)$ and define ```math -r_t(s_{1:T}) = s_t - f_t(s_{t-1}), +r_t(S) = s_t - f_t(s_{t-1}). ``` -finding the trajectory is equivalent to solving $r(s_{1:T}) = 0$. Sequentially running the chain is one way to solve this; DEER is another — and a parallelisable one. - ---- +The taped sequential trajectory is the root of $r(S)=0$. DEER searches for the +same root while exposing work across all $T$ transitions. -## DEER: Newton's method for nonlinear recursions +## One DEER iteration -Apply one step of Newton's method to the residual $r = 0$ starting from a trajectory guess $S^{(i)}$: +Let $S^{(i)}$ be the current guess. Linearize each transition around the +previous state in that guess: ```math -s^{(i+1)}_t = J_t \, s^{(i+1)}_{t-1} + \underbrace{\bigl[f_t(s^{(i)}_{t-1}) - J_t \, s^{(i)}_{t-1}\bigr]}_{u_t}, +J_t = \nabla f_t\!\left(s^{(i)}_{t-1}\right), ``` -where $J_t = \nabla_{s} f_t(s^{(i)}_{t-1})$ is the Jacobian of the transition at the current guess. - -The key observation is that once $J_t$ and $u_t$ are computed from the previous iterate $S^{(i)}$ — which can be done **in parallel over $t$** — the update is a **linear** time-varying recursion - ```math -s^{(i+1)}_t = J_t \, s^{(i+1)}_{t-1} + u_t, \quad s^{(i+1)}_0 = s_0, +u_t = f_t\!\left(s^{(i)}_{t-1}\right) + - J_t s^{(i)}_{t-1}. ``` -which can itself be solved in $O(\log T)$ steps via a parallel prefix (associative) scan. - -The full DEER loop repeats until the change between iterates falls below a tolerance: +The next trajectory solves the affine recursion ```math -\max_t \| s^{(i+1)}_t - s^{(i)}_t \|_\infty \;\leq\; \delta_\text{abs} + \delta_\text{rel} \cdot \max_t \| s^{(i+1)}_t \|_\infty. +s^{(i+1)}_t = J_t s^{(i+1)}_{t-1} + u_t, +\qquad s^{(i+1)}_0 = s_0. ``` -Newton's method converges **quadratically** once near the fixed point, so the number of iterations scales only logarithmically with $T$ in practice. - ---- +All $J_t$ and $u_t$ values depend only on the old guess, so they can be +evaluated together. The remaining affine recursion can be solved by an +associative scan rather than a serial loop. -## The parallel prefix scan - -When $J_t$ is diagonal (quasi-DEER), the linear recursion decouples dimension-by-dimension into $D$ independent scalar affine recursions +With damping, the package replaces the raw update with ```math -s^{(i+1)}_{d,t} = a_{d,t} \, s^{(i+1)}_{d,t-1} + b_{d,t}, +S^{(i+1)} \leftarrow +(1 - \lambda) S^{(i)} + \lambda S^{(i+1)}, +\qquad 0 < \lambda \leq 1. ``` -where $a_{d,t} = [J_t]_{dd}$ and $b_{d,t} = u_{d,t}$. Written in matrix form: given $A, B \in \mathbb{R}^{D \times T}$, solve +This is the `damping` keyword. Smaller values often make difficult solves more +stable, at the cost of more iterations. -```math -S[:,t] = A[:,t] \odot S[:,t-1] + B[:,t], \quad S[:,0] = s_0. -``` +## Why the diagonal approximation matters -The associative operator for combining two adjacent segments $(\alpha_1, \beta_1)$ and $(\alpha_2, \beta_2)$ is +A dense $D \times D$ Jacobian at every transition is expensive to form, store, +and combine. The public sampler therefore uses a diagonal approximation: ```math -(\alpha_2, \beta_2) \circ (\alpha_1, \beta_1) = (\alpha_2 \odot \alpha_1,\; \alpha_2 \odot \beta_1 + \beta_2). +s^{(i+1)}_{d,t} += a_{d,t} s^{(i+1)}_{d,t-1} + b_{d,t}. ``` -This associativity means the recurrence can be solved by an inclusive parallel-prefix scan in $O(\log T)$ levels, each level consisting of a single broadcast over all $T$ columns and no per-timestep loops. The implementation in `ParallelMCMC.DEERScan.solve_affine_scan_diag!` is array-type-agnostic and runs identically on CPU `Matrix` and GPU `CuMatrix`. The CPU path is mainly useful for correctness checks; without a GPU or substantial parallel hardware, sequential MCMC is usually faster wall-clock. - ---- - -## Jacobian variants +Each parameter dimension is now a scalar affine recursion. This is often +called quasi-DEER. -### Full DEER +`ParallelMALASampler` offers two ways to get the diagonal: -Use the full $D \times D$ Jacobian $J_t$ at each timestep. The linear recursion becomes a sequence of dense matrix multiplications, solved sequentially (no scan shortcut for the general case). Cost per iteration: $O(TD^3)$. Memory: $O(TD^2)$. Accurate but impractical for large $D$. +- `jacobian=:diag` computes it exactly with $D$ Jacobian-vector products. This + is useful for small models and reference checks. +- `jacobian=:stoch_diag` uses a Hutchinson estimator. This is the default and + usually the only practical choice in high dimensions. -This package currently exposes the diagonal scan variants below as `jacobian` modes; full-matrix DEER is discussed here for context but is not available through the `jacobian` keyword. - -### Quasi-DEER (`:diag`) - -Replace $J_t$ with $\mathrm{diag}(J_t)$, retaining only the diagonal. The recursion reduces to the scalar affine scan described above, solved in $O(TD \log T)$ total work. The exact diagonal is computed with `D` Jacobian-vector products. This mode is useful for low-dimensional checks and reference runs. - -### Stochastic quasi-DEER (`:stoch_diag`) - -Computing the exact diagonal of $J_t$ requires a full Jacobian (or $D$ JVPs), which is expensive in high dimensions. The **Hutchinson–Rademacher estimator** provides an unbiased approximation using only JVPs: +For independent Rademacher vectors $z^{(k)}$, whose entries are equally likely +to be $-1$ or $1$, ```math -\mathrm{diag}(J_t) \approx \frac{1}{K} \sum_{k=1}^K z^{(k)} \odot J_t z^{(k)}, -\quad z^{(k)}_i \overset{\text{iid}}{\sim} \mathrm{Rademacher}(\pm 1). +\operatorname{diag}(J_t) +\approx \frac{1}{K}\sum_{k=1}^{K} +z^{(k)} \odot J_t z^{(k)}. ``` -Each probe $z^{(k)}$ costs a single Jacobian-vector product (one forward-mode or reverse-mode pass), so the total cost is $O(KTD)$ — linear in $D$ and $T$. In the limit $K \to \infty$ the estimate converges to the true diagonal. This is the default mode. See Zoltowski et al. (2025) [^1] for convergence analysis. - -In practice $K=1$ or $K=2$ probes works well; controlled by the `probes` argument of [`ParallelMALASampler`](@ref). +One probe needs one Jacobian-vector product instead of one per dimension. The +`probes` keyword sets $K$. More probes reduce the estimator's noise but add +derivative work to every DEER iteration. ---- +The diagonal is an approximation to the Newton direction, not a change to the +target transition. If the nonlinear solve converges, its fixed point still +satisfies the original taped recursion up to the requested numerical tolerance. -## Applying DEER to MALA +## The parallel affine scan -One step of MALA with pre-drawn noise $(\xi_t, u_t)$ is +Represent one scalar affine map as a pair $(a, b)$ acting on $x$ by +$a x + b$. Composing two maps gives ```math -\tilde{x}_t = x_{t-1} + \varepsilon \nabla \log p(x_{t-1}) + \sqrt{2\varepsilon}\, \xi_t, +(a_2, b_2) \circ (a_1, b_1) += (a_2 a_1,\; a_2 b_1 + b_2). ``` -```math -x_t = g_t \, \tilde{x}_t + (1 - g_t)\, x_{t-1}, \quad g_t = \mathbf{1}[\log u_t < \log \alpha_t], -``` - -where $\log \alpha_t = \log p(\tilde{x}_t) + \log q(x_{t-1} \mid \tilde{x}_t) - \log p(x_{t-1}) - \log q(\tilde{x}_t \mid x_{t-1})$ is the Metropolis log-acceptance ratio. +Composition is associative. A parallel-prefix scan can therefore combine the +maps for all prefixes in $O(\log T)$ dependent levels. Each level still does +work across the trajectory, so $O(\log T)$ describes the critical path of the +scan, not zero total work. -### The stop-gradient trick +The implementation in +`ParallelMCMC.DEERScan.solve_affine_scan_diag!` uses array operations over +$D \times T$ matrices. The same code works with ordinary matrices and supported +device matrices. -The indicator $g_t \in \{0, 1\}$ is not differentiable, so $\nabla_{x_{t-1}} x_t$ is undefined at accept/reject boundaries. DEER needs this Jacobian to form the Newton update. +## Applying the method to MALA -The solution is a **surrogate step** used only during Jacobian computation (the *forward* step always uses the exact indicator): +Using the convention in this package, a MALA proposal is ```math -x_t^\text{surrogate} = \hat{g}_t \, \tilde{x}_t + (1 - \hat{g}_t)\, x_{t-1}, -\quad \hat{g}_t = \sigma(\log \alpha_t - \log u_t) + \underbrace{\mathrm{sg}\!\left(\mathbf{1}[\cdots] - \sigma(\cdots)\right)}_{\text{stop-gradient}}, +\widetilde{x}_t = x_{t-1} ++ \varepsilon \nabla \log p(x_{t-1}) ++ \sqrt{2\varepsilon}\,\xi_t, ``` -where $\sigma$ is the logistic function. The stop-gradient term makes $\hat{g}_t$ equal to the exact indicator in the forward pass while routing gradients through $\sigma$ during the backward pass. This gives a well-defined, smooth Jacobian whose value at the operating point equals that of the relaxed step. - -In the implementation, the accept indicator $g_t$ is pre-computed from the previous-iterate state and passed as a **frozen constant** to the surrogate ([`MALA.mala_step_surrogate_sigmoid`](@ref MALA.mala_step_surrogate_sigmoid)), so differentiation never touches the discontinuity. +followed by the usual Metropolis-Hastings accept/reject decision. The tape holds +the Gaussian noise $\xi_t$ and the uniform draw used by that decision. -### Summary of the DEER–MALA loop +The hard accept/reject indicator is not differentiable. DEER still needs a +useful local derivative, so the implementation uses a smooth sigmoid surrogate +for differentiation while keeping the actual decision fixed in the forward +evaluation. In other words, the sampler does not replace Metropolis-Hastings +with a soft acceptance rule; the relaxation only supplies a Jacobian for the +nonlinear solver. -1. **Tape generation.** Draw $T$ noise pairs $(\xi_t, u_t)$ and store them. -2. **Jacobian and offset computation (parallel over $t$).** For each $t$, evaluate $f_t$ at the current guess $s^{(i)}_{t-1}$ to get $u_t$, and differentiate the surrogate to get $\mathrm{diag}(J_t)$ (or the Hutchinson approximation thereof). -3. **Parallel scan.** Solve the diagonal linear recursion $S^{(i+1)}[:,t] = A[:,t] \odot S^{(i+1)}[:,t-1] + B[:,t]$ in $O(\log T)$ levels. -4. **Convergence check.** If the change is below tolerance, return $S^{(i+1)}$; otherwise go to step 2. -5. **Sample delivery.** Return the $T$ columns of the converged trajectory as individual MCMC samples. +The relevant implementation is +[`MALA.mala_step_surrogate_sigmoid`](@ref MALA.mala_step_surrogate_sigmoid). -For low-level callers using `DEER.solve(...; workspace=ws)`, the returned trajectory is copied by default so it remains valid after later solves reuse `ws`. Set `copy_result=false` only when you are intentionally accepting workspace-owned output that may be overwritten by a later call. +## Convergence and sample delivery ---- +After each update, DEER compares the largest elementwise change with a mixed +absolute and relative tolerance: -## References +```math +\max |S^{(i+1)} - S^{(i)}| +\leq \texttt{tol_abs} ++ \texttt{tol_rel}\,\max |S^{(i+1)}|. +``` -[^1]: Zoltowski, D. M., Wu, S., Gonzalez, X., Kozachkov, L., & Linderman, S. W. (2025). *Parallelizing MCMC Across the Sequence Length*. NeurIPS 2025. [arXiv:2508.18413](https://arxiv.org/abs/2508.18413) +The solve stops when this condition is met or after `maxiter` iterations. Once +a block has been solved, its `T` columns are delivered as MCMC samples. If the +caller asks for more than `T` samples, the sampler starts another block from the +last state and draws a new tape. A final partial block is trimmed so that the +returned chain has the requested length. + +For low-level callers, `DEER.solve(...; return_info=true)` returns convergence +information with the trajectory. When reusing a `DEERWorkspace`, the result is +copied by default so a later solve cannot overwrite it. Set `copy_result=false` +only when accepting a workspace-owned result is intentional. + +## What controls performance? + +The scan has a logarithmic dependency depth, but the whole sampler also pays +for model evaluations, derivative products, nonlinear iterations, memory +traffic, and compilation. The most important practical factors are: + +- how efficiently the model evaluates a `T`-column batch; +- the number of DEER iterations needed for convergence; +- the Jacobian mode and number of probes; +- the trajectory length and parameter dimension; +- whether the hardware has enough parallel work to offset launch and setup + costs. + +This is why parallel MALA is not automatically faster than sequential MALA. +Measure both on the target and hardware that matter for your application. + +## One block at a glance + +1. Draw a tape of `T` MALA noise and uniform values. +2. Evaluate each taped transition at the current trajectory guess. +3. Estimate or compute each transition's Jacobian diagonal. +4. Solve the resulting affine recursion with a parallel scan. +5. Apply damping and test convergence; repeat if needed. +6. Return the solved states as samples and continue from the last one. + +[^1]: Zoltowski, D. M., Wu, S., Gonzalez, X., Kozachkov, L., & Linderman, + S. W. (2025). *Parallelizing MCMC Across the Sequence Length*. NeurIPS + 2025. [arXiv:2508.18413](https://arxiv.org/abs/2508.18413) diff --git a/docs/src/95-reference.md b/docs/src/95-reference.md index 996d007..4623ca4 100644 --- a/docs/src/95-reference.md +++ b/docs/src/95-reference.md @@ -1,34 +1,28 @@ -# [API Reference](@id reference) +# API reference ```@meta CurrentModule = ParallelMCMC ``` -This page documents all public types and functions exported by ParallelMCMC.jl. +This page collects the public API and the low-level DEER types that are useful +for custom solvers. For an end-to-end example, begin with +[Getting started](10-getting-started.md). -## Reactant - -Loading `Reactant` (`using Reactant`) enables `ADTypes.AutoReactant()` as a -backend for derivatives in `DensityModel` and as `ParallelMALASampler`'s -`backend`. See [GPU Execution](15-gpu.md) for its constraints. - -## CUDA and other device arrays - -`using CUDA` loads the `CUDAExt` extension; see [GPU Execution](15-gpu.md). +## Model ```@docs -ParallelMCMC.needs_host_staging +DensityModel ``` -## Extension constructor - -`DensityModel(ld; param_names=nothing)` wraps any `LogDensityProblems` object with gradients when `LogDensityProblems` is loaded. See [LogDensityProblems models](@ref) for an example. +Additional constructors become available when their packages are loaded: -## Model +- `DensityModel(ld; param_names=nothing, kwargs...)` wraps a gradient-capable + `LogDensityProblems` object. +- `DensityModel(turing_model; ad_backend, kwargs...)` wraps a DynamicPPL/Turing + model, extracts its parameter names, and converts output back to the original + parameter space. -```@docs -DensityModel -``` +See [Defining models](12-models.md) for examples and the derivative-slot rules. ## Samplers @@ -38,24 +32,40 @@ AdaptiveMALASampler ParallelMALASampler ``` -## Internal types +All three implement the AbstractMCMC interface. The usual entry point is +`sample(model, sampler, n; kwargs...)`, optionally with an RNG or an +AbstractMCMC ensemble such as `MCMCThreads()`. + +## Transition and state types -These types appear in the `AbstractMCMC` state/transition protocol. You generally do not need to construct them directly. +These types are part of the AbstractMCMC protocol. Most users receive them from +`AbstractMCMC.step` and do not construct them directly. ```@docs -MALATapeElement MALAState MALATransition AdaptiveMALAState AdaptiveMALATransition ParallelMALAState ParallelMALATransition +MALATapeElement +``` + +## Device support + +Loading CUDA activates the package extension for `CuArray` storage. Loading +Reactant activates `ADTypes.AutoReactant()` support for derivative slots and +the parallel sampler's `backend`. The restrictions and backend comparison are +documented in [GPU execution](15-gpu.md). + +```@docs +ParallelMCMC.needs_host_staging ``` -## Low-level namespaces +## Low-level DEER API -Use these functions and workspaces to solve taped recursions or run the -diagonal affine scan directly. +[`ParallelMALASampler`](@ref) builds on the types below. Use them directly when +solving a custom taped recursion or managing workspaces across repeated solves. ```@docs DEER.TapedRecursion @@ -65,7 +75,9 @@ ParallelMCMC.DEERScan.AffineScanWorkspace MALA.mala_step_surrogate_sigmoid ``` -The diagonal scan implementation itself lives in `ParallelMCMC.DEERScan.solve_affine_scan_diag!`. +The in-place diagonal scan is available as +`ParallelMCMC.DEERScan.solve_affine_scan_diag!`. See +[How the algorithm works](20-algorithms.md) before using the low-level API. ## Index diff --git a/docs/src/index.md b/docs/src/index.md index 0bd8784..652c8ac 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -10,112 +10,98 @@ CurrentModule = ParallelMCMC

``` -**ParallelMCMC.jl** is a Julia package for **parallel-across-the-sequence** MCMC — algorithms that solve an entire trajectory of $T$ correlated steps *simultaneously* rather than one at a time. +ParallelMCMC.jl implements Markov chain Monte Carlo methods that parallelize +*across the sequence of samples*. Its main sampler, [`ParallelMALASampler`](@ref), +uses DEER to solve a block of MALA transitions together instead of waiting for +each transition before starting the next one. -```@raw html -

- DEER trajectory estimates improving on a Julia-logo-shaped posterior -

- -

- DEER iterates on a synthetic Julia-logo-shaped posterior: orange trajectory estimates move toward the taped MALA path over repeated trajectory solves. -

-``` - -The flagship algorithm is **DEER** (Deterministic Equivalent-Expectation Recursion), which reformulates a chain of $T$ MALA steps as a fixed-point problem and solves it via Newton iterations, each costing $O(\log T)$ parallel work via an associative prefix scan. The result is that wall-clock time per sample is *sublinear* in chain length on multi-core CPUs and GPUs. - -The algorithm is described in: +This is different from running several independent chains at once. You can do +that too, but the package's main goal is to shorten the critical path *within* +each chain. Whether that produces a speedup depends on the target density, +trajectory length, AD backend, and hardware. -> Zoltowski, D. M., Wu, S., Gonzalez, X., Kozachkov, L., & Linderman, S. W. (2025). -> **Parallelizing MCMC Across the Sequence Length.** -> *NeurIPS 2025.* [arXiv:2508.18413](https://arxiv.org/abs/2508.18413) +## Is this package a good fit? -The included [`MALASampler`](@ref) and [`AdaptiveMALASampler`](@ref) are sequential MALA baselines — useful for correctness checks and step-size tuning, but not the primary focus of the package. - -## Samplers - -| Sampler | Role | +| If you want to... | Start with... | |---|---| -| [`ParallelMALASampler`](@ref) | **Primary** — parallel-across-sequence MALA via DEER; O(log T) per solve | -| [`MALASampler`](@ref) | Baseline — sequential MALA with a fixed step size | -| [`AdaptiveMALASampler`](@ref) | Baseline — sequential MALA with dual-averaging step-size adaptation | +| Try the parallel-across-sequence algorithm | [`ParallelMALASampler`](@ref) | +| Tune a step size automatically | [`AdaptiveMALASampler`](@ref) | +| Check a model against ordinary sequential MALA | [`MALASampler`](@ref) | +| Sample a Turing model | [`DensityModel`](@ref) with the Turing extension | +| Run a large, batch-friendly target on a GPU | [GPU execution](15-gpu.md) | -All samplers implement the [AbstractMCMC](https://github.com/TuringLang/AbstractMCMC.jl) interface and return [`FlexiChains`](https://pysm.dev/FlexiChains.jl) objects. +Parallel MALA is most compelling when evaluating the target and its derivatives +is expensive enough to keep parallel hardware busy. For a small target on a +CPU, sequential MALA will often finish sooner. -## Installation +## Install -To install the package into your current environment: +ParallelMCMC supports Julia 1.10 and later. ```julia-repl pkg> add ParallelMCMC ``` -## Quick start +AD backends and GPU packages are optional. Add one only when your model needs +it; the first example below supplies its derivatives directly. -### Define a model +## A complete first run -The simplest entry point is [`DensityModel`](@ref), which wraps a log-density and its gradient: +The target here is a two-dimensional standard normal. Its gradient and +Hessian-vector product are simple enough to write down, so the example has no +optional dependencies. ```julia -using ParallelMCMC, FlexiChains -using ADTypes, Enzyme - -# Example: 2-D standard normal -logp(x) = -0.5 * sum(abs2, x) -grad_logp(x) = -x - -model = DensityModel(logp, grad_logp, 2; - param_names=[:x1, :x2]) +using ParallelMCMC, FlexiChains, Random + +logdensity(x) = -sum(abs2, x) / 2 +gradient(x) = -x +hvp(x, v) = -v + +model = DensityModel( + logdensity, + gradient, + 2; + hvp=hvp, + param_names=[:x1, :x2], +) + +sampler = ParallelMALASampler(0.1; T=64) + +chain = sample( + MersenneTwister(42), + model, + sampler, + 1_000; + initial_params=zeros(2), + chain_type=VNChain, +) ``` -### DEER — parallel-across-sequence (primary algorithm) - -```julia -sampler = ParallelMALASampler(0.1; T=64, jacobian=:stoch_diag, - backend=AutoEnzyme()) +`T=64` means that DEER solves blocks of 64 transitions. The call still returns +exactly 1,000 samples; `T` controls the internal block size, not the requested +chain length. -chain = sample(model, sampler, 500; - chain_type=VNChain) -``` +## Where to go next -Each call to `sample` draws 500 samples by solving DEER trajectories of length `T=64` in parallel, re-solving from the last state when each trajectory is exhausted. +- [Getting started](10-getting-started.md) walks through step-size tuning, + running parallel MALA, reading the result, and common failure modes. +- [Defining models](12-models.md) covers hand-written derivatives, automatic + differentiation, batched functions, parameter names, and Turing models. +- [GPU execution](15-gpu.md) explains when a GPU is worthwhile and shows a + complete logistic-regression example. +- [How the algorithm works](20-algorithms.md) develops the DEER update and the + diagonal parallel scan. +- [API reference](95-reference.md) lists constructors and low-level building + blocks. -Specifying `chain_type=VNChain` returns a `FlexiChain{VarName}`, which has a parameter type of `VarName`. -This is intended for maximum ease of use; however, if you prefer parameter type of `Symbol` you can use `SymChain` instead. -See [the FlexiChains.jl docs](https://pysm.dev/FlexiChains.jl/) for more information about how to analyze and visualize chains. +## Citation -### Sequential MALA baseline +The parallel-across-sequence method is described in: -```julia -sampler = AdaptiveMALASampler(0.1; n_warmup=500) +> Zoltowski, D. M., Wu, S., Gonzalez, X., Kozachkov, L., & Linderman, S. W. +> (2025). *Parallelizing MCMC Across the Sequence Length*. NeurIPS 2025. +> [arXiv:2508.18413](https://arxiv.org/abs/2508.18413) -chain = sample(model, sampler, 2_000; - chain_type=VNChain, - discard_warmup=true, - progress=true) -``` - -See [Getting Started](10-getting-started.md) for worked examples and guidance on choosing samplers, and [Algorithm Details](20-algorithms.md) for the mathematics behind DEER. - -## Contributors - -```@raw html - - - - - - - - - - - - -
Ryan Senne
Ryan Senne
💻 🚧 ⚠️ 🤔 👀 📖		Penelope Yong
Penelope Yong
💻 ⚠️ 🤔 👀 📖		Guillaume Dalle
Guillaume Dalle
👀 🤔		William Moses
William Moses
👀
- - - - - -``` +If you use ParallelMCMC.jl in published work, please use the package citation +in [`CITATION.cff`](https://github.com/rsenne/ParallelMCMC.jl/blob/main/CITATION.cff). diff --git a/src/interface.jl b/src/interface.jl index 68aaaa8..b70e339 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -6,40 +6,42 @@ Defines model/sampler/state/transition types and implements =# """ - DensityModel(logdensity, grad_logdensity, dim; param_names, logdensity_batch, grad_logdensity_batch, hvp, hvp_batch) - -Wrap a log-density and its derivatives for ParallelMCMC samplers. - -Each derivative argument (`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) - -Backends are prepared when sampling starts. An HVP backend differentiates a -callable gradient once. With an AD-derived gradient it forms -`DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` on the -log-density. An explicit `SecondOrder` also differentiates the log-density directly. - -- `logdensity(x::AbstractVector) -> Real` -- `grad_logdensity` — callable `x -> AbstractVector`, or an AD backend. -- `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 - log-density over independent columns. -- `grad_logdensity_batch` — optional callable `X -> AbstractMatrix`, or a - backend. It is derived from `logdensity_batch` when `grad_logdensity` is a - backend. -- `hvp_batch` — optional callable `(X, V) -> AbstractMatrix`, or a backend, - resolved like `hvp`. -- `dim::Int` — dimensionality of the parameter space -- `param_names` — optional collection of parameter names used in `FlexiChains` output. If - `nothing` (the default), uses a single vector-valued parameter `:x` with shape `(dim,)`. - See the [`Parameter names`](@ref parameter-names) section of the docs for more - information. - -Both batched derivative arguments require `logdensity_batch`. `AutoReactant()` has -additional pairing and tracing constraints; see the GPU guide. + DensityModel(logdensity, grad_logdensity, dim; + param_names=nothing, hvp=nothing, + logdensity_batch=nothing, + grad_logdensity_batch=nothing, hvp_batch=nothing) + +Describe a target density and the derivatives used by ParallelMCMC's samplers. + +For a vector `x` of length `dim`, `logdensity(x)` must return a real scalar and +`grad_logdensity(x)` must return a vector like `x`. The log-density may omit its +normalizing constant. Instead of a gradient function, `grad_logdensity` may be +an `ADTypes.AbstractADType` backend: + +```julia +model = DensityModel(logdensity, AutoForwardDiff(), dim) +``` + +Sequential MALA needs only the gradient. Parallel MALA also needs +Hessian-vector products. Supply `hvp(x, v)`, put an AD backend in `hvp`, or pass +a `backend` to [`ParallelMALASampler`](@ref). AD backends are prepared when +sampling starts, after the element type and storage of the initial state are +known. + +The optional batched functions operate on independent columns: + +- `logdensity_batch(X)` returns one log-density per column of `X`. +- `grad_logdensity_batch(X)` returns a matrix shaped like `X`. +- `hvp_batch(X, V)` returns `H(X[:, j]) * V[:, j]` in column `j`. + +Both batched derivative slots require `logdensity_batch`. A derivative slot may +hold a callable or an AD backend. When a batched gradient is derived with AD, +each output of `logdensity_batch` must depend only on the matching input column. + +`param_names` controls the parameter keys in the returned `FlexiChain`. With +the default `nothing`, the chain contains one vector-valued parameter named +`:x`. See [Defining models](@ref) for naming, batching, Turing integration, and +the rules for combining derivative backends. """ struct DensityModel{F,G,H,FB,GB,HB,PN} <: AbstractMCMC.AbstractModel logdensity::F @@ -416,11 +418,14 @@ end """ MALASampler(epsilon; cholM=nothing) -Metropolis-Adjusted Langevin Algorithm sampler with step size `epsilon`. +Create a sequential Metropolis-adjusted Langevin (MALA) sampler with positive +step size `epsilon`. -Optionally pass `cholM = cholesky(M)` to use a mass matrix `M` as a -preconditioner. The proposal becomes `y = x + ε M ∇logp(x) + √(2ε) L ξ` -where `L` is the Cholesky factor of `M`. +This sampler is useful as a baseline for [`ParallelMALASampler`](@ref) and for +checking a model with an ordinary sequential transition. Pass +`cholM=cholesky(M)` to precondition proposals with a positive-definite mass +matrix `M`. With `M = L * L'`, the proposal is +`y = x + epsilon * M * gradient(x) + sqrt(2epsilon) * L * noise`. """ struct MALASampler{FP<:AbstractFloat,CM} <: AbstractMCMC.AbstractSampler epsilon::FP @@ -583,21 +588,37 @@ struct MALATapeElement{FP<:AbstractFloat,V<:AbstractVector{FP}} end """ - ParallelMALASampler(epsilon; T, maxiter, tol_abs, tol_rel, jacobian, damping, probes, cholM, backend) + ParallelMALASampler(epsilon; T=64, maxiter=200, + tol_abs=1e-6, tol_rel=1e-5, + jacobian=:stoch_diag, damping=0.5, + probes=1, cholM=nothing, backend=nothing) -DEER-parallelized MALA sampler. +Create a MALA sampler that solves blocks of transitions with DEER. -Supported Jacobian modes are `:stoch_diag` (the default Hutchinson diagonal -estimator) and `:diag` (exact diagonal via `D` JVPs). +`epsilon` is the positive MALA step size. `T` is the number of transitions in +each DEER block; it does not determine how many samples `sample` returns. -`backend` supplies Hessian-vector products when the `DensityModel` brings no -`hvp` / `hvp_batch` of its own. It does not supply gradients. Leave it out when -the model carries its own HVPs. +The nonlinear solve stops after `maxiter` updates or when its largest +elementwise change is no greater than +`tol_abs + tol_rel * maximum(abs, trajectory)`. `damping` controls how much of +each new update is applied and must lie in `(0, 1]`. -With `backend=ADTypes.AutoReactant()`, `grad_logdensity` must be `AutoReactant()` -or a callable. +`jacobian=:stoch_diag` estimates each transition Jacobian's diagonal with +`probes` Hutchinson vectors. `jacobian=:diag` computes the exact diagonal with +one Jacobian-vector product per parameter dimension. -`callback` and `initial_state` route through the generic step-wise loop instead. +`backend` constructs missing Hessian-vector products; it does not supply the +model's gradient. Leave it as `nothing` when the [`DensityModel`](@ref) provides +the required `hvp` and `hvp_batch` functions. `cholM` has the same +preconditioning meaning as in [`MALASampler`](@ref). With +`backend=ADTypes.AutoReactant()`, the model's `grad_logdensity` must be +`AutoReactant()` or a callable. + +`sample` normally runs whole DEER blocks. Passing `callback` or `initial_state` +switches it to the generic step-wise AbstractMCMC loop. + +See [Getting started](@ref) for tuning advice and [GPU execution](@ref) for +backend-specific device constraints. """ struct ParallelMALASampler{FP<:AbstractFloat,CM,AD} <: AbstractMCMC.AbstractSampler epsilon::FP @@ -1294,7 +1315,25 @@ for TKey in (Symbol, VarName) end """ - AdaptiveMALASampler(epsilon_init; n_warmup, target_accept, gamma, t0, kappa, cholM) + AdaptiveMALASampler(epsilon_init; n_warmup=1000, + target_accept=0.574, gamma=0.05, + t0=10.0, kappa=0.75, cholM=nothing) + +Create a sequential MALA sampler with dual-averaging step-size adaptation. + +Adaptation begins at the positive step size `epsilon_init`, targets the +acceptance probability `target_accept`, and stops after `n_warmup` warmup +transitions. The averaged step size is then held fixed. Set +`discard_warmup=true` in `sample` to omit warmup-marked states from the returned +chain. + +`gamma`, `t0`, and `kappa` control the dual-averaging schedule. Their defaults +are usually a better starting point than manual adjustment. `cholM` has the +same preconditioning meaning as in [`MALASampler`](@ref). + +The chain stores `:step_size`, `:accepted`, and `:is_warmup` as extra fields. +This sampler is a convenient way to choose an `epsilon` for +[`ParallelMALASampler`](@ref). """ struct AdaptiveMALASampler{FP<:AbstractFloat,CM} <: AbstractMCMC.AbstractSampler epsilon_init::FP From cf6c02b73a13200b01c5c5f479cd9fdb64277406 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Mon, 5 Oct 2026 10:08:25 -0400 Subject: [PATCH 3/5] Some reviews --- docs/make.jl | 1 + docs/src/10-getting-started.md | 56 ++++++--------------- docs/src/12-models.md | 90 +++++++++++++++++----------------- docs/src/15-gpu.md | 19 +++---- docs/src/16-troubleshooting.md | 29 +++++++++++ docs/src/95-reference.md | 4 +- docs/src/index.md | 3 +- 7 files changed, 103 insertions(+), 99 deletions(-) create mode 100644 docs/src/16-troubleshooting.md diff --git a/docs/make.jl b/docs/make.jl index 39ac842..c80cf2b 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -9,6 +9,7 @@ const pages = [ "Getting started" => "10-getting-started.md", "Defining models" => "12-models.md", "GPU execution" => "15-gpu.md", + "Troubleshooting" => "16-troubleshooting.md", "How the algorithm works" => "20-algorithms.md", "API reference" => "95-reference.md", "Contributing" => [ diff --git a/docs/src/10-getting-started.md b/docs/src/10-getting-started.md index 4ab0deb..c3a6f3c 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -6,9 +6,9 @@ the derivative path while still being easy to inspect. ## Install an AD backend -ParallelMCMC treats automatic differentiation packages as optional -dependencies. This guide uses ForwardDiff to build the Hessian-vector products -that parallel MALA needs. +Parallel MALA is the main algorithm in ParallelMCMC. It requires an automatic +differentiation backend to compute Hessian–vector products, so you will have to +load one yourself. In this guide we will use ForwardDiff. ```julia-repl pkg> add ParallelMCMC ForwardDiff @@ -43,15 +43,16 @@ model = DensityModel( ) ``` -The gradient is hand-written, but it does not have to be. See -[Defining models](12-models.md) for AD-derived gradients, analytical -Hessian-vector products, batched functions, and Turing models. +Here, the gradient is hand-written, but it does not have to be. See +[Defining models](12-models.md) for AD-derived gradients and analytical +Hessian-vector products. ## Tune the step size -MALA is sensitive to its step size. A value that is too large causes frequent -rejections; one that is too small produces highly correlated samples. Start -with [`AdaptiveMALASampler`](@ref), which tunes the step size during warmup. +MALA is sensitive to its step size, such that a value that is too large causes +frequent rejections, while one that is too small produces highly correlated +samples. Start with [`AdaptiveMALASampler`](@ref), which tunes the step size +during warmup. ```julia rng = MersenneTwister(42) @@ -79,9 +80,7 @@ every transition marked as warmup from the returned chain. `:step_size`, ## Run parallel MALA -Pass the tuned value to [`ParallelMALASampler`](@ref). Because the model does -not provide an analytical Hessian-vector product, `backend=AutoForwardDiff()` -tells the sampler how to construct one from `banana_gradient`. +Pass the tuned value to [`ParallelMALASampler`](@ref). ```julia sampler = ParallelMALASampler( @@ -107,8 +106,7 @@ The important arguments are: number of samples returned by `sample`. - `backend`: the AD backend used for missing Hessian-vector products. Leave it out when the model provides `hvp` and, when applicable, `hvp_batch`. -- `initial_params`: the first state. Its array type also selects CPU or GPU - storage. +- `initial_params`: the first state. - `chain_type`: use `VNChain` for `VarName` keys or `SymChain` for `Symbol` keys. Turing models require `VNChain`. @@ -152,7 +150,7 @@ models. ParallelMCMC returns a [FlexiChains](https://pysm.dev/FlexiChains.jl/) chain. With the parameter names above, values are available as `chain[:x1]` and -`chain[:x2]`. Sampler diagnostics live in extra columns such as +`chain[:x2]`. Sampler diagnostics are stored in extra columns such as `chain[:accepted]` and `chain[:logp]`. Use several chains for convergence diagnostics. Start Julia with multiple @@ -172,31 +170,5 @@ chains = sample( ess(chains) ``` -`MCMCThreads()` runs independent chains concurrently. That is separate from +`MCMCThreads()` runs independent chains concurrently, which is separate from the within-chain parallelism controlled by `T`. - -## When something goes wrong - -### The DEER solve takes many iterations - -First reduce `epsilon` or `T`. If the model otherwise behaves well, lower -`damping` from `0.5` toward `0.3`. Do not loosen the convergence tolerances -until you have ruled out a derivative error. - -### Parallel MALA is slower than sequential MALA - -That is expected for small, cheap targets on a CPU. Parallel MALA has setup, -derivative, and scan overhead. It is designed for targets with enough batched -work to benefit from parallel hardware. Benchmark the complete sampling call, -including compilation separately from steady-state execution. - -### The chain does not move or has many rejections - -Run [`MALASampler`](@ref) or [`AdaptiveMALASampler`](@ref) first. Check the -gradient against finite differences or an independent AD backend, then try a -smaller step size. - -### GPU code reports scalar indexing or AD compilation errors - -The model must be written in GPU-compatible array operations, and backend -support differs. Continue with [GPU execution](15-gpu.md). diff --git a/docs/src/12-models.md b/docs/src/12-models.md index 403a5ee..9f5cd97 100644 --- a/docs/src/12-models.md +++ b/docs/src/12-models.md @@ -19,10 +19,10 @@ For a parameter vector `x` of length `dim`: - `grad_logdensity(x)` returns a vector with the same shape as `x`. - `dim` is the number of unconstrained parameters. -The derivative slot may contain either a callable or an -`ADTypes.AbstractADType` backend. +The gradient argument may be either a callable or an `ADTypes.AbstractADType` +backend. -## Option 1: write the derivatives +## Option 1: write the gradient Hand-written derivatives keep the execution path predictable and are often the best choice for GPU models. @@ -32,25 +32,19 @@ using ParallelMCMC logdensity(x) = -sum(abs2, x) / 2 gradient(x) = -x -hvp(x, v) = -v model = DensityModel( logdensity, gradient, 4; - hvp=hvp, param_names=[:a, :b, :c, :d], ) ``` -Sequential MALA needs the gradient. Parallel MALA also needs -Hessian-vector products (HVPs). You can supply `hvp(x, v)` as above, or let the -sampler construct it with an AD backend. +## Option 2: derive the gradient with AD -## Option 2: derive gradients with AD - -Pass an AD backend in the gradient slot to build a model from the log-density -alone: +Pass an AD backend as the gradient argument to build a model from the +log-density alone: ```julia using ADTypes, ForwardDiff, ParallelMCMC @@ -62,15 +56,19 @@ model = DensityModel(logdensity, AutoForwardDiff(), 4) Backends are prepared when sampling begins, once the package knows the element type and storage of `initial_params`. -For parallel MALA, also provide a backend to the sampler unless the model has -an `hvp`: +## Hessian-vector products for parallel MALA + +Parallel MALA additionally requires Hessian-vector products (HVPs). Just like +for gradients, you can either hand-write them or use AD. + +A hand-written HVP is passed with the `hvp` keyword: ```julia -sampler = ParallelMALASampler(0.1; backend=AutoForwardDiff()) +hvp(x, v) = -v +model = DensityModel(logdensity, gradient, 4; hvp=hvp) ``` -You can mix approaches. For example, keep a hand-written gradient and use AD -only for the HVP: +To use AD instead, provide a backend to the sampler: ```julia model = DensityModel(logdensity, gradient, 4) @@ -79,14 +77,13 @@ sampler = ParallelMALASampler(0.1; backend=AutoForwardDiff()) ### Which function does an HVP backend differentiate? -| Gradient slot | HVP slot | What ParallelMCMC differentiates | +| `gradient = ...` | `hvp = ...` | HVP calculated by... | |---|---|---| -| callable | backend | the callable gradient | -| backend | backend | the log-density with a second-order backend pair | -| either | `SecondOrder(...)` | the log-density with that explicit pair | +| callable | backend | differentiating `gradient` with `backend` | +| backend1 | backend2 | differentiating log-density with `SecondOrder(backend2, backend1)` | +| either | `SecondOrder(...)` | differentiating log-density with `SecondOrder(...)` | -An explicit `DifferentiationInterface.SecondOrder` takes precedence over the -gradient slot. Backend combinations are not equally useful on every device; +Backend combinations are not equally useful on every device; see [GPU execution](15-gpu.md) before selecting one for GPU work. ## Batched functions @@ -114,19 +111,19 @@ model = DensityModel( ) ``` -Both batched derivative slots require `logdensity_batch`. When the gradient -slot contains an AD backend, ParallelMCMC can derive a batched gradient by +Both batched derivative arguments require `logdensity_batch`. When the +gradient argument is an AD backend, ParallelMCMC can derive a batched gradient by differentiating `sum(logdensity_batch(X))`. For that to be correct, each output of `logdensity_batch` must depend only on the corresponding column of `X`. -The scalar functions remain part of the contract even when all batched +The scalar functions must still be provided even when all batched functions are present. ## Parameter names -Without `param_names`, the chain contains one vector-valued parameter named -`:x`. Pass names when separate scalar or array-valued entries would be easier to -work with. +If `param_names` is not passed, the chain will contain one vector-valued +parameter named `:x`. Pass names when separate scalar or array-valued entries +would be easier to work with. ```julia # Three scalar parameters @@ -140,14 +137,16 @@ Names may be `Symbol`s or `VarName`s. The declared shapes must account for exactly `dim` scalar values. `VNChain` accepts `VarName` keys, while `SymChain` uses `Symbol` keys. -## Turing models +## DynamicPPL models -Loading Turing activates a convenience constructor that handles unconstraining, -parameter names, and conversion back to the model's original parameter space. -The AD backend is required. +DynamicPPL is the package that defines Turing's `@model` syntax. Loading it +activates a convenience constructor that handles unconstraining, parameter +names, and conversion back to the model's original parameter space. Unlike when +sampling with Turing, the AD backend must be explicitly passed, and is used to +calculate the gradient of the log-density. ```julia -using ADTypes, FlexiChains, ForwardDiff, ParallelMCMC, Turing +using ADTypes, Distributions, DynamicPPL, FlexiChains, ForwardDiff, ParallelMCMC @model function normal_location(y) mu ~ Normal(0, 1) @@ -168,11 +167,11 @@ chain = sample( ) ``` -Turing-backed models require `VNChain`. Their output includes `:logjoint`, -`:logprior`, and `:loglikelihood`; a manually constructed model has a single -`:logp` field. +DynamicPPL-backed models require `VNChain`. Their output includes `:logjoint`, +`:logprior`, and `:loglikelihood`; this is more informative than a manually +constructed model, which has a single `:logp` field. -To use parallel MALA, supply an HVP or an appropriate sampler backend: +To use parallel MALA, supply an analytical HVP or an appropriate AD backend: ```julia using Enzyme @@ -181,14 +180,15 @@ sampler = ParallelMALASampler(0.1; T=64, backend=AutoEnzyme()) chain = sample(model, sampler, 1_000; chain_type=VNChain) ``` -Do not assume that the model's gradient backend is also a valid HVP backend. -In particular, applying ForwardDiff around a DynamicPPL gradient that already -uses ForwardDiff creates unsupported nested dual numbers. Enzyme is the -supported default for differentiating that gradient; an explicit second-order -backend pair or analytical HVP can also be used. +You should not assume that a model's gradient backend is also a valid HVP +backend. In particular, applying ForwardDiff around a DynamicPPL gradient that +already uses ForwardDiff creates nested dual numbers, which are not supported. +For such a gradient, Enzyme is the supported default, but you may also specify +an explicit second-order backend pair or analytical HVP. -Turing models currently run on the CPU. For a GPU target, write the density and -derivatives directly with GPU-compatible array operations. +DynamicPPL models currently run on the CPU, because DynamicPPL evaluates the +model with CPU vectors. For a GPU target, write the density and derivatives +directly with GPU-compatible array operations. ## LogDensityProblems models diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index a9f6545..279416a 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -1,18 +1,19 @@ # GPU execution -[`ParallelMALASampler`](@ref) follows the storage of `initial_params`, so the -same sampler can work with CPU vectors or supported device vectors. Moving the -initial state is not enough, though: the target density, its derivatives, and -the selected AD backend must all support the device. +[`ParallelMALASampler`](@ref) automatically infers the backend depending on the +array type of `initial_params`, so the same sampler can work with CPU vectors or +supported device vectors. For example, `CuArray`s can be used after +ParallelMCMC's CUDA extension is loaded with `using CUDA`. -CUDA is an optional dependency. `using CUDA` activates ParallelMCMC's extension -for `CuArray` storage without making CUDA part of a CPU-only installation. +Changing the type of the initial state is not enough, though: the target +density, its derivatives, and the selected AD backend must all be written in a +device-compatible manner. --- ## Decide whether a GPU is worthwhile -A GPU helps only when each DEER iteration provides enough work to amortize +A GPU helps only when each DEER iteration is expensive enough to amortize kernel launches and data movement. The important workload is the evaluation of batched densities, gradients, and Hessian-vector products across a `T`-column trajectory. @@ -75,7 +76,7 @@ otherwise CPU-only function GPU compatible. ### Enzyme may need broadcasts split into separate expressions -Some `CuArray` gradients fail during Enzyme compilation with the same `gc-transition` error. If you encounter it, try splitting broadcasts into separate expressions, as in this example: +Some `CuArray` gradients fail during Enzyme compilation with the same `gc-transition` error. If you encounter this error, try splitting broadcasts into separate expressions, as in this example: ```julia # ABORTS during Enzyme compile on GPU @@ -238,7 +239,7 @@ chain = sample(model, sampler, 1_600; chain_type=VNChain) ``` -Both examples describe the same posterior. The Enzyme version separates operations to avoid the compilation failures above. +Both examples describe the same posterior; the Enzyme version only differs because it separates operations to avoid the compilation failures above. --- diff --git a/docs/src/16-troubleshooting.md b/docs/src/16-troubleshooting.md new file mode 100644 index 0000000..4dce5fa --- /dev/null +++ b/docs/src/16-troubleshooting.md @@ -0,0 +1,29 @@ +# Troubleshooting + +## The DEER solve takes many iterations + +First reduce `epsilon` or `T`. If the model otherwise behaves well, lower +`damping` from `0.5` toward `0.3`. Do not loosen the convergence tolerances +until you have ruled out a derivative error. + +## Parallel MALA is slower than sequential MALA + +That is expected for small, cheap targets on a CPU. Parallel MALA has setup, +derivative, and scan overhead, and is designed for targets with enough batched +work to benefit from parallel hardware. + +When comparing the two, benchmark the complete sampling call, and measure +compilation separately from steady-state execution. + +## The chain does not move or has many rejections + +Run [`MALASampler`](@ref) or [`AdaptiveMALASampler`](@ref) first. Check the +gradient against finite differences or an independent AD backend, then try a +smaller step size. + +## GPU code reports scalar indexing or AD compilation errors + +The model must be written in GPU-compatible array operations. + +AD backend support also differs on the GPU. See [GPU execution](15-gpu.md) for +the current limitations. diff --git a/docs/src/95-reference.md b/docs/src/95-reference.md index 4623ca4..4768cfe 100644 --- a/docs/src/95-reference.md +++ b/docs/src/95-reference.md @@ -22,7 +22,7 @@ Additional constructors become available when their packages are loaded: model, extracts its parameter names, and converts output back to the original parameter space. -See [Defining models](12-models.md) for examples and the derivative-slot rules. +See [Defining models](12-models.md) for examples and the meanings of the derivative arguments. ## Samplers @@ -54,7 +54,7 @@ MALATapeElement ## Device support Loading CUDA activates the package extension for `CuArray` storage. Loading -Reactant activates `ADTypes.AutoReactant()` support for derivative slots and +Reactant activates `ADTypes.AutoReactant()` support for derivative arguments and the parallel sampler's `backend`. The restrictions and backend comparison are documented in [GPU execution](15-gpu.md). diff --git a/docs/src/index.md b/docs/src/index.md index 652c8ac..7b274e7 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -85,11 +85,12 @@ chain length. ## Where to go next - [Getting started](10-getting-started.md) walks through step-size tuning, - running parallel MALA, reading the result, and common failure modes. + running parallel MALA, and reading the result. - [Defining models](12-models.md) covers hand-written derivatives, automatic differentiation, batched functions, parameter names, and Turing models. - [GPU execution](15-gpu.md) explains when a GPU is worthwhile and shows a complete logistic-regression example. +- [Troubleshooting](16-troubleshooting.md) covers common failure modes. - [How the algorithm works](20-algorithms.md) develops the DEER update and the diagonal parallel scan. - [API reference](95-reference.md) lists constructors and low-level building From cc566b38c4246b996f2336fae00d99fd989bebb5 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Tue, 13 Oct 2026 10:57:47 -0400 Subject: [PATCH 4/5] Drop Turing material from the refreshed docs Turing support was removed in #75. Move the runnable LogDensityProblems example and HVP note into the models page. --- README.md | 4 +- docs/src/10-getting-started.md | 2 +- docs/src/12-models.md | 76 ++++++++++++---------------------- docs/src/15-gpu.md | 11 +---- docs/src/95-reference.md | 9 +--- docs/src/index.md | 4 +- src/interface.jl | 4 +- 7 files changed, 37 insertions(+), 73 deletions(-) diff --git a/README.md b/README.md index 87a4205..02b60b3 100644 --- a/README.md +++ b/README.md @@ -83,8 +83,8 @@ chain = sample( Start with the [documentation](https://rsenne.github.io/ParallelMCMC.jl/dev/) or the repository's [getting-started guide](docs/src/10-getting-started.md). -The docs cover step-size tuning, model and AD setup, Turing integration, GPU -execution, troubleshooting, and the DEER algorithm. +The docs cover step-size tuning, model and AD setup, GPU execution, +troubleshooting, and the DEER algorithm. ## How to cite diff --git a/docs/src/10-getting-started.md b/docs/src/10-getting-started.md index c3a6f3c..f2dd90a 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -108,7 +108,7 @@ The important arguments are: out when the model provides `hvp` and, when applicable, `hvp_batch`. - `initial_params`: the first state. - `chain_type`: use `VNChain` for `VarName` keys or `SymChain` for `Symbol` - keys. Turing models require `VNChain`. + keys. ## Choose `T` and the DEER controls diff --git a/docs/src/12-models.md b/docs/src/12-models.md index 9f5cd97..84df5f5 100644 --- a/docs/src/12-models.md +++ b/docs/src/12-models.md @@ -137,27 +137,31 @@ Names may be `Symbol`s or `VarName`s. The declared shapes must account for exactly `dim` scalar values. `VNChain` accepts `VarName` keys, while `SymChain` uses `Symbol` keys. -## DynamicPPL models +## LogDensityProblems models -DynamicPPL is the package that defines Turing's `@model` syntax. Loading it -activates a convenience constructor that handles unconstraining, parameter -names, and conversion back to the model's original parameter space. Unlike when -sampling with Turing, the AD backend must be explicitly passed, and is used to -calculate the gradient of the log-density. +Any object implementing the gradient-capable +[LogDensityProblems](https://github.com/tpapp/LogDensityProblems.jl) interface +can be wrapped directly: ```julia -using ADTypes, Distributions, DynamicPPL, FlexiChains, ForwardDiff, ParallelMCMC +using FlexiChains, LogDensityProblems, ParallelMCMC -@model function normal_location(y) - mu ~ Normal(0, 1) - y ~ Normal(mu, 0.5) +# mu ~ Normal(0, 1), y ~ Normal(mu, 0.5) +struct NormalPosterior + y::Float64 end -model = DensityModel( - normal_location(1.5); - ad_backend=AutoForwardDiff(), -) +LogDensityProblems.capabilities(::Type{NormalPosterior}) = + LogDensityProblems.LogDensityOrder{1}() +LogDensityProblems.dimension(::NormalPosterior) = 1 +LogDensityProblems.logdensity(ld::NormalPosterior, x) = + -x[1]^2 / 2 - 2 * (ld.y - x[1])^2 +function LogDensityProblems.logdensity_and_gradient(ld::NormalPosterior, x) + return LogDensityProblems.logdensity(ld, x), [-x[1] + 4 * (ld.y - x[1])] +end + +model = DensityModel(NormalPosterior(1.5); param_names=[:mu]) chain = sample( model, AdaptiveMALASampler(0.3; n_warmup=500), @@ -167,39 +171,13 @@ chain = sample( ) ``` -DynamicPPL-backed models require `VNChain`. Their output includes `:logjoint`, -`:logprior`, and `:loglikelihood`; this is more informative than a manually -constructed model, which has a single `:logp` field. - -To use parallel MALA, supply an analytical HVP or an appropriate AD backend: - -```julia -using Enzyme - -sampler = ParallelMALASampler(0.1; T=64, backend=AutoEnzyme()) -chain = sample(model, sampler, 1_000; chain_type=VNChain) -``` - -You should not assume that a model's gradient backend is also a valid HVP -backend. In particular, applying ForwardDiff around a DynamicPPL gradient that -already uses ForwardDiff creates nested dual numbers, which are not supported. -For such a gradient, Enzyme is the supported default, but you may also specify -an explicit second-order backend pair or analytical HVP. - -DynamicPPL models currently run on the CPU, because DynamicPPL evaluates the -model with CPU vectors. For a GPU target, write the density and derivatives -directly with GPU-compatible array operations. - -## LogDensityProblems models - -Any object implementing the gradient-capable -[LogDensityProblems](https://github.com/tpapp/LogDensityProblems.jl) interface -can be wrapped directly: - -```julia -model = DensityModel(logdensity_problem; param_names=[:a, :b]) -``` - This constructor gets the dimension, log-density, and gradient from the -interface. The same HVP and batching keywords accepted by the main constructor -are available when you need them. +interface, so the object must implement `capabilities` (at least +`LogDensityOrder{1}`), `dimension`, `logdensity`, and +`logdensity_and_gradient`. The same HVP and batching keywords accepted by the +main constructor are available when you need them. + +An HVP backend differentiates `logdensity_and_gradient`, so that function must +be differentiable by the backend. If it is not, for example because it calls AD +internally, pass an explicit `SecondOrder`, which differentiates `logdensity` +twice instead. `AutoReactant()` is not supported for these models. diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index 279416a..7106393 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -32,15 +32,6 @@ step size, and tolerances. ## Current limitations -### Turing models are CPU-only - -The Turing constructor and its DynamicPPL/LogDensityProblems path currently -execute with CPU vectors. They cannot run a `@model` directly against a -`CuArray`. - -For GPU sampling, write the density and derivatives with GPU-compatible array -operations and pass them to [`DensityModel`](@ref), as in the example below. - ### Enzyme needs ParallelMCMC's linear-algebra wrappers `ParallelMCMC` exports three thin wrappers: @@ -269,7 +260,7 @@ sampler = ParallelMALASampler(0.005f0; T=16, backend=AutoReactant()) **Reactant chooses the execution device independently of the input array.** With a CUDA client, `CuArray` inputs are staged through pinned host memory and results are copied back on the device. With the CPU client, both directions go through the host and the sampler warns when it prepares the model. Select a GPU client with `Reactant.set_default_backend` when available. -- The log-density must be Reactant-traceable. DynamicPPL-built log-densities are not. +- The log-density must be Reactant-traceable. - Across an HVP's two AD passes, use `AutoReactant()` for both or neither. A hand-written gradient may pair with it. `SecondOrder` cannot contain `AutoReactant()`. - Only the default `AutoReactant()` mode is supported. diff --git a/docs/src/95-reference.md b/docs/src/95-reference.md index 4768cfe..99a2839 100644 --- a/docs/src/95-reference.md +++ b/docs/src/95-reference.md @@ -14,13 +14,8 @@ for custom solvers. For an end-to-end example, begin with DensityModel ``` -Additional constructors become available when their packages are loaded: - -- `DensityModel(ld; param_names=nothing, kwargs...)` wraps a gradient-capable - `LogDensityProblems` object. -- `DensityModel(turing_model; ad_backend, kwargs...)` wraps a DynamicPPL/Turing - model, extracts its parameter names, and converts output back to the original - parameter space. +Loading `LogDensityProblems` adds `DensityModel(ld; param_names=nothing, +kwargs...)`, which wraps a gradient-capable `LogDensityProblems` object. See [Defining models](12-models.md) for examples and the meanings of the derivative arguments. diff --git a/docs/src/index.md b/docs/src/index.md index 7b274e7..3f2aacf 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -27,7 +27,6 @@ trajectory length, AD backend, and hardware. | Try the parallel-across-sequence algorithm | [`ParallelMALASampler`](@ref) | | Tune a step size automatically | [`AdaptiveMALASampler`](@ref) | | Check a model against ordinary sequential MALA | [`MALASampler`](@ref) | -| Sample a Turing model | [`DensityModel`](@ref) with the Turing extension | | Run a large, batch-friendly target on a GPU | [GPU execution](15-gpu.md) | Parallel MALA is most compelling when evaluating the target and its derivatives @@ -87,7 +86,8 @@ chain length. - [Getting started](10-getting-started.md) walks through step-size tuning, running parallel MALA, and reading the result. - [Defining models](12-models.md) covers hand-written derivatives, automatic - differentiation, batched functions, parameter names, and Turing models. + differentiation, batched functions, parameter names, and LogDensityProblems + models. - [GPU execution](15-gpu.md) explains when a GPU is worthwhile and shows a complete logistic-regression example. - [Troubleshooting](16-troubleshooting.md) covers common failure modes. diff --git a/src/interface.jl b/src/interface.jl index b70e339..6f32d4b 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -40,8 +40,8 @@ each output of `logdensity_batch` must depend only on the matching input column. `param_names` controls the parameter keys in the returned `FlexiChain`. With the default `nothing`, the chain contains one vector-valued parameter named -`:x`. See [Defining models](@ref) for naming, batching, Turing integration, and -the rules for combining derivative backends. +`:x`. See [Defining models](@ref) for naming, batching, and the rules for +combining derivative backends. """ struct DensityModel{F,G,H,FB,GB,HB,PN} <: AbstractMCMC.AbstractModel logdensity::F From 9227e16ff11c08690c0098e497e9a05f6f1e318d Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:47:29 -0400 Subject: [PATCH 5/5] fix incorrect statements --- docs/src/10-getting-started.md | 15 ++++++++------- docs/src/12-models.md | 9 +++++---- docs/src/15-gpu.md | 2 +- docs/src/91-developer.md | 2 +- ext/EnzymeExt.jl | 3 ++- src/MALA/MALA.jl | 17 ++++++++++------- 6 files changed, 27 insertions(+), 21 deletions(-) diff --git a/docs/src/10-getting-started.md b/docs/src/10-getting-started.md index f2dd90a..5c3da71 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -6,9 +6,9 @@ the derivative path while still being easy to inspect. ## Install an AD backend -Parallel MALA is the main algorithm in ParallelMCMC. It requires an automatic -differentiation backend to compute Hessian–vector products, so you will have to -load one yourself. In this guide we will use ForwardDiff. +Parallel MALA is the main algorithm in ParallelMCMC. It needs Hessian-vector +products, which you can either write yourself or get from an automatic +differentiation backend. This guide uses ForwardDiff. ```julia-repl pkg> add ParallelMCMC ForwardDiff @@ -128,7 +128,7 @@ ParallelMALASampler( maxiter=200, tol_abs=1e-6, tol_rel=1e-5, - backend=AutoForwardDiff(), + backend=nothing, ) ``` @@ -150,8 +150,8 @@ models. ParallelMCMC returns a [FlexiChains](https://pysm.dev/FlexiChains.jl/) chain. With the parameter names above, values are available as `chain[:x1]` and -`chain[:x2]`. Sampler diagnostics are stored in extra columns such as -`chain[:accepted]` and `chain[:logp]`. +`chain[:x2]`. The log-density of each sample is stored in the extra column +`chain[:logp]`. Use several chains for convergence diagnostics. Start Julia with multiple threads, for example `julia -t 4`, then run: @@ -163,12 +163,13 @@ chains = sample( MCMCThreads(), 2_000, 4; - initial_params=zeros(2), + initial_params=fill(zeros(2), 4), chain_type=VNChain, ) ess(chains) ``` +With several chains, `initial_params` needs one starting point per chain. `MCMCThreads()` runs independent chains concurrently, which is separate from the within-chain parallelism controlled by `T`. diff --git a/docs/src/12-models.md b/docs/src/12-models.md index 84df5f5..5b46e3f 100644 --- a/docs/src/12-models.md +++ b/docs/src/12-models.md @@ -130,12 +130,13 @@ would be easier to work with. DensityModel(logdensity, gradient, 3; param_names=[:a, :b, :c]) # One scalar and one 1 x 2 matrix -DensityModel(logdensity, gradient, 3; param_names=[:a, (:B, (1, 2))]) +DensityModel(logdensity, gradient, 3; param_names=[:a, :B => (1, 2)]) ``` -Names may be `Symbol`s or `VarName`s. The declared shapes must account for -exactly `dim` scalar values. `VNChain` accepts `VarName` keys, while `SymChain` -uses `Symbol` keys. +Names may be `Symbol`s or `VarName`s. Give an array-valued entry as a pair of +its name and shape, as with `:B => (1, 2)` above. The declared shapes must +account for exactly `dim` scalar values. `VNChain` accepts `VarName` keys, while +`SymChain` uses `Symbol` keys. ## LogDensityProblems models diff --git a/docs/src/15-gpu.md b/docs/src/15-gpu.md index 7106393..c81dbed 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -124,7 +124,7 @@ y_gpu = CUDA.CuVector(y_cpu) ### Mooncake backend (plain operators) ```julia -using ParallelMCMC, FlexiChains +using ParallelMCMC, FlexiChains, LinearAlgebra using ADTypes, Mooncake softplus(z) = log1p(exp(-abs(z))) + max(z, zero(z)) diff --git a/docs/src/91-developer.md b/docs/src/91-developer.md index d0abf71..72211f5 100644 --- a/docs/src/91-developer.md +++ b/docs/src/91-developer.md @@ -63,7 +63,7 @@ pre-commit run -a We use `lychee` for link checking in CI. You can run it locally to avoid waiting for CI. First, [install lychee](https://github.com/lycheeverse/lychee?tab=readme-ov-file#installation), then run against the repository root using the project config: ```bash -lychee --no-progress --config lychee.toml . +lychee --no-progress --config .lychee.toml . ``` ## Testing diff --git a/ext/EnzymeExt.jl b/ext/EnzymeExt.jl index f19a461..dbeb746 100644 --- a/ext/EnzymeExt.jl +++ b/ext/EnzymeExt.jl @@ -17,7 +17,8 @@ using Enzyme.EnzymeCore.EnzymeRules: overwritten, width -# TODO: Implement matmul overloads upstream in Enzyme. See: https://github.com/EnzymeAD/Enzyme.jl/issues/3122 +# TODO: Enzyme has GPU matmul support since EnzymeAD/Enzyme.jl#3148, but forward mode +# on CuArray still fails (EnzymeAD/Enzyme.jl#3841). Revisit these rules once that is fixed. # Add the default function annotation without changing the user's mode. function DEER._normalized_backend(backend::ADTypes.AutoEnzyme{M,A}) where {M,A} diff --git a/src/MALA/MALA.jl b/src/MALA/MALA.jl index 465434f..dc6fdff 100644 --- a/src/MALA/MALA.jl +++ b/src/MALA/MALA.jl @@ -49,8 +49,7 @@ _logdet_M(cholM::Cholesky) = logdet(cholM) """ Reusable scratch buffers for allocation-light MALA primitives. -All vectors are created with the same array type / device placement as `x_template`, -so this workspace is GPU-compatible when constructed from a `CuVector`. +All vectors have the same array type and device as `x`. """ struct MALAWorkspace{V} g_x::V @@ -85,8 +84,8 @@ end """ Reusable scratch buffers for batched MALA primitives. -Matrix buffers have the same size, eltype, and device placement as `X_template`. -Vector buffers have length `size(X_template, 2)`. +Matrix buffers have the same size, eltype, and device as `X`. +Vector buffers have length `size(X, 2)`. """ struct MALABatchedWorkspace{M,V,RM,BV,BM} G_X::M @@ -411,7 +410,10 @@ function mala_step_with_logα!( end """ -Stop-gradient surrogate step used for Jacobians. +MALA step with the accept test replaced by a sigmoid, used for Jacobians. + +Returns `g * y + (1 - g) * x`, where `y` is the proposal and +`g = sigmoid(logα - log(u))`. """ function mala_step_surrogate_sigmoid( logp, gradlogp, x::AbstractVector, ϵ::Real, ξ::AbstractVector, u::Real; cholM=nothing @@ -548,9 +550,10 @@ function mala_step_batched( end """ -Batched fused surrogate forward step and JVP over columns. +Batched MALA forward step and JVP of the sigmoid surrogate over columns. -`X`, `Ξ`, and `Z` are all D×T, and `u` has length T. +`FT` is the exact step with a hard accept test. Only the JVP `Jt` uses the +sigmoid surrogate. `X`, `Ξ`, and `Z` are all D×T, and `u` has length T. Returns `(FT, Jt)` where both outputs are D×T. """ function mala_step_batched_fwd_and_jvp!(