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..02b60b3 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, 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..c80cf2b 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -4,60 +4,19 @@ 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", + "Troubleshooting" => "16-troubleshooting.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 +29,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..5c3da71 100644 --- a/docs/src/10-getting-started.md +++ b/docs/src/10-getting-started.md @@ -1,211 +1,175 @@ -# 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: +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. -- `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. - -### What a backend in `hvp` differentiates - -| `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. - -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`. +## Define the target -[`ParallelMALASampler`](@ref)'s `backend` supplies missing HVPs. Omit it when the model supplies them. - ---- - -## ParallelMALASampler - -[`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. +A [`DensityModel`](@ref) needs a log-density, a gradient, and the number of +unconstrained parameters. The log-density may omit its normalizing constant. ```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) -``` - -`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. - -### Jacobian modes - -The `jacobian` keyword controls how the per-step Jacobian is approximated during each Newton iteration: - -| 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: - -```julia -sampler = ParallelMALASampler(0.1; T=128, jacobian=:stoch_diag, probes=2, - backend=AutoEnzyme()) -``` - -### 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`. - -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. - ---- - -## LogDensityProblems models - -With `LogDensityProblems` loaded, `DensityModel` accepts any object that implements the [LogDensityProblems](https://github.com/tpapp/LogDensityProblems.jl) interface with gradients: - -```julia -using LogDensityProblems, ParallelMCMC, FlexiChains - -# μ ~ Normal(0, 1), y ~ Normal(μ, 0.5) -struct NormalPosterior - y::Float64 +function banana_logdensity(x) + a = x[2] + 0.1 * x[1]^2 - 1 + return -(x[1]^2 / 100 + a^2) / 2 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])] + +function banana_gradient(x) + a = x[2] + 0.1 * x[1]^2 - 1 + return [-x[1] / 100 - 0.2 * x[1] * a, -a] 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) +model = DensityModel( + banana_logdensity, + banana_gradient, + 2; + param_names=[:x1, :x2], +) ``` -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. - ---- +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. -## Parallel chains +## Tune the step size -All samplers support `MCMCThreads()`. Start Julia with multiple threads (e.g. `julia -t 4`): +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 -chain = sample(model, ParallelMALASampler(0.1; T=64, backend=AutoEnzyme()), - MCMCThreads(), 500, 4; - chain_type=VNChain) +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])) ``` -`FlexiChains` computes R-hat and ESS across chains automatically: +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. -```julia -ess(chain) -``` +!!! 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. + +## Run parallel MALA -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): +Pass the tuned value to [`ParallelMALASampler`](@ref). ```julia -ess(chain; dims=:iter) +sampler = ParallelMALASampler( + epsilon; + T=64, + backend=AutoForwardDiff(), +) + +chain = sample( + MersenneTwister(43), + model, + sampler, + 2_000; + initial_params=zeros(2), + chain_type=VNChain, +) ``` ---- +The important arguments are: -## [Specifying parameter names](@id parameter-names) +- `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. +- `chain_type`: use `VNChain` for `VarName` keys or `SymChain` for `Symbol` + keys. -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. +## Choose `T` and the DEER controls -If you do not specify `param_names`, the chain will store a single vector-valued parameter called `x` of length `D`. +`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. -You can specify parameter names as a collection of either: - -- `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 above can be mixed and matched as desired, as long as the total number of parameters matches the dimension of the model. -For example: +The defaults are intended to be a usable starting point: ```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))]) +ParallelMALASampler( + epsilon; + T=64, + jacobian=:stoch_diag, + probes=1, + damping=0.5, + maxiter=200, + tol_abs=1e-6, + tol_rel=1e-5, + backend=nothing, +) ``` ---- +- `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. -## Sequential MALA baselines +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. -[`MALASampler`](@ref) and [`AdaptiveMALASampler`](@ref) are standard sequential MALA implementations. They are useful for: +## Use the returned chain -- 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`) +ParallelMCMC returns a [FlexiChains](https://pysm.dev/FlexiChains.jl/) chain. +With the parameter names above, values are available as `chain[:x1]` and +`chain[:x2]`. The log-density of each sample is stored in the extra column +`chain[:logp]`. -```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] +Use several chains for convergence diagnostics. Start Julia with multiple +threads, for example `julia -t 4`, then run: -# Step 2: run DEER with the tuned step size -chain = sample(model, ParallelMALASampler(eps_tuned; T=64, backend=AutoEnzyme()), 2_000; - chain_type=VNChain) +```julia +chains = sample( + model, + sampler, + MCMCThreads(), + 2_000, + 4; + initial_params=fill(zeros(2), 4), + chain_type=VNChain, +) + +ess(chains) ``` -See the [`MALASampler`](@ref) and [`AdaptiveMALASampler`](@ref) reference pages for the full keyword listing. +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 new file mode 100644 index 0000000..5b46e3f --- /dev/null +++ b/docs/src/12-models.md @@ -0,0 +1,184 @@ +# 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 gradient argument may be either a callable or an `ADTypes.AbstractADType` +backend. + +## Option 1: write the gradient + +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 + +model = DensityModel( + logdensity, + gradient, + 4; + param_names=[:a, :b, :c, :d], +) +``` + +## Option 2: derive the gradient with AD + +Pass an AD backend as the gradient argument 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`. + +## 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 +hvp(x, v) = -v +model = DensityModel(logdensity, gradient, 4; hvp=hvp) +``` + +To use AD instead, provide a backend to the sampler: + +```julia +model = DensityModel(logdensity, gradient, 4) +sampler = ParallelMALASampler(0.1; backend=AutoForwardDiff()) +``` + +### Which function does an HVP backend differentiate? + +| `gradient = ...` | `hvp = ...` | HVP calculated by... | +|---|---|---| +| callable | backend | differentiating `gradient` with `backend` | +| backend1 | backend2 | differentiating log-density with `SecondOrder(backend2, backend1)` | +| either | `SecondOrder(...)` | differentiating log-density with `SecondOrder(...)` | + +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 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 must still be provided even when all batched +functions are present. + +## Parameter names + +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 +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. 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 + +Any object implementing the gradient-capable +[LogDensityProblems](https://github.com/tpapp/LogDensityProblems.jl) interface +can be wrapped directly: + +```julia +using FlexiChains, LogDensityProblems, ParallelMCMC + +# mu ~ Normal(0, 1), y ~ Normal(mu, 0.5) +struct NormalPosterior + y::Float64 +end + +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), + 2_000; + chain_type=VNChain, + discard_warmup=true, +) +``` + +This constructor gets the dimension, log-density, and gradient from the +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 baf920d..c81dbed 100644 --- a/docs/src/15-gpu.md +++ b/docs/src/15-gpu.md @@ -1,28 +1,38 @@ -# GPU Execution +# 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`](@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`. -ParallelMCMC does not depend on CUDA.jl. `using CUDA` loads the `CUDAExt` extension, which is required for `CuArray` parameters. +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. --- -## 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 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. | 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` +### Enzyme needs ParallelMCMC's linear-algebra wrappers `ParallelMCMC` exports three thin wrappers: @@ -32,29 +42,32 @@ 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. +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`. -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 are an Enzyme-specific compatibility layer; they do not make an +otherwise CPU-only function GPU compatible. -### 2. Some Enzyme GPU gradients hit a `gc-transition` abort. Staging them as single ops is a workaround +### 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 this error, try splitting broadcasts into separate expressions, as in this example: ```julia # ABORTS during Enzyme compile on GPU @@ -111,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)) @@ -148,7 +161,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 +230,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 only differs because it separates operations to avoid the compilation failures above. --- @@ -262,7 +275,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/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/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/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/docs/src/95-reference.md b/docs/src/95-reference.md index 167a0fc..99a2839 100644 --- a/docs/src/95-reference.md +++ b/docs/src/95-reference.md @@ -1,28 +1,12 @@ -# [API Reference](@id reference) +# API reference ```@meta CurrentModule = ParallelMCMC ``` -This page documents all public types and functions exported by ParallelMCMC.jl. - -## Reactant - -Loading `Reactant` (`using Reactant`) enables `ADTypes.AutoReactant()` as a -derivative-slot backend on `DensityModel` and as `ParallelMALASampler`'s -`backend`. 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). - -```@docs -ParallelMCMC.needs_host_staging -``` - -## Extension constructor - -`DensityModel(ld; param_names=nothing)` wraps any `LogDensityProblems` object with gradients when `LogDensityProblems` is loaded. See [LogDensityProblems models](@ref) for an example. +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). ## Model @@ -30,6 +14,11 @@ ParallelMCMC.needs_host_staging DensityModel ``` +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. + ## Samplers ```@docs @@ -38,24 +27,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 arguments 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 -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. +[`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 +70,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..3f2aacf 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -10,112 +10,99 @@ 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) | +| 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, and reading the result. +- [Defining models](12-models.md) covers hand-written derivatives, automatic + 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. +- [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/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/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/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!( 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..6f32d4b 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 slot (`grad_logdensity`, `hvp`, `grad_logdensity_batch`, -`hvp_batch`) takes a callable or an `ADTypes.AbstractADType`, so a model can be -built from the log-density alone: - - DensityModel(logp, AutoForwardDiff(), dim) - -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. - -- `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 slots 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, 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