More CUDA compatibility

HISTORY.md

@@ -1,3 +1,19 @@
+# 0.4.1
+
+## Other changes
+
+`elbo_batch` now works when the samples live on a GPU.
+It used to assemble the ELBO from `Distributions.logpdf`, which evaluates the base distribution column by column and returns a host array.
+The base distribution's log-density is now taken with a whole-array form that stays on the device holding the samples, for any GPU array backend.
+`elbo` is unchanged and still goes through `Distributions.logpdf`, so it stays CPU only.
+
+Coupling flows (`realnvp`, `nsf`) still do not run on the GPU.
+`Bijectors.PartitionMask` holds host sparse matrices and `partition`/`combine` multiply against them, so the split is done on the host.
+`example/gpu/demo_gpu.jl` trains a planar flow instead.
+
+The base distribution is treated as a constant when its log-density is taken this way.
+Differentiating more than one use of a full covariance leaves a cotangent per use, and summing those indexes a device array element by element.
+
 # 0.4.0
 
 ## Breaking changes

Project.toml

@@ -1,18 +1,21 @@
 name = "NormalizingFlows"
 uuid = "50e4474d-9f12-44b7-af7a-91ab30ff6256"
-version = "0.4.0"
+version = "0.4.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
@@ -30,12 +33,15 @@ ADTypes = "1.5"
 AbstractPPL = "0.15.4"
 Bijectors = "0.14.2, 0.15, 0.16"
 CUDA = "5, 6.2"
+ChainRulesCore = "1"
 Distributions = "0.25"
 DocStringExtensions = "0.9"
 Flux = "0.16"
 Functors = "0.5.2"
+GPUArraysCore = "0.1, 0.2"
 LogExpFunctions = "0.3.3"
 Optimisers = "0.2.16, 0.3, 0.4"
+PDMats = "0.11"
 ProgressMeter = "1.0.0"
 ReverseDiff = "1"
 StatsBase = "0.33, 0.34"

example/README.md

@@ -16,3 +16,5 @@ Currently, all examples share the same [Julia project](https://pkgdocs.julialang
 using Pkg; Pkg.activate("."); Pkg.instantiate()
 ```
 This will install all needed packages, at the exact versions when the model was last updated. Then you can run the model code with `include("<example-to-run>.jl")`, or by running the example script line-by-line.
+
+`gpu/` holds the GPU demo and has its own project, so the examples above do not depend on CUDA. It needs a GPU; activate and instantiate `gpu/` to run it.

example/gpu/Project.toml

@@ -0,0 +1,22 @@
+[deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+NormalizingFlows = "50e4474d-9f12-44b7-af7a-91ab30ff6256"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[compat]
+ADTypes = "1.22.2"
+Bijectors = "0.16.2"
+CUDA = "6.2.1"
+DifferentiationInterface = "0.7.20"
+Distributions = "0.25.129"
+Functors = "0.5.2"
+NormalizingFlows = "0.4.1"
+Optimisers = "0.4.7"
+Zygote = "0.7.11"

example/gpu/demo_gpu.jl

@@ -0,0 +1,86 @@
+# Training a planar flow on the GPU. This demo has its own project so that the CPU examples
+# do not pull in CUDA. Run it from `example/gpu` after
+#
+#     using Pkg; Pkg.activate("."); Pkg.develop(; path="../.."); Pkg.instantiate()
+#
+# The `develop` is needed until 0.4.1 is registered, because the device log-density this
+# demo relies on landed in that version.
+#
+# Coupling flows (RealNVP, NSF) do not run on the GPU yet: `Bijectors.PartitionMask` holds
+# host sparse matrices and `partition`/`combine` multiply against them.
+
+using Distributions, LinearAlgebra
+using Bijectors
+using Functors
+using Optimisers, ADTypes, Zygote
+using CUDA
+using NormalizingFlows
+# loads the AbstractPPL extension that routes `AutoZygote` through DifferentiationInterface
+using DifferentiationInterface
+
+# Bijectors' planar layer broadcasts in a way that CUDA cannot fuse, and reads `flow.b` back
+# from the device.
+# https://github.com/TuringLang/Bijectors.jl/blob/93cb25563043c527905519d81d6dee7917af4dbe/src/bijectors/planar_layer.jl#L65-L110
+function Bijectors.get_u_hat(u::CuVector{T}, w::CuVector{T}) where {T<:Real}
+    wT_u = dot(w, u)
+    scale = (Bijectors.LogExpFunctions.log1pexp(-wT_u) - 1) / sum(abs2, w)
+    û = CUDA.broadcast(+, u, CUDA.broadcast(*, scale, w))
+    wT_û = Bijectors.LogExpFunctions.log1pexp(wT_u) - 1
+    return û, wT_û
+end
+function Bijectors._transform(flow::Bijectors.PlanarLayer, z::CuArray{T}) where {T<:Real}
+    w = CuArray(flow.w)
+    û, wT_û = Bijectors.get_u_hat(CuArray(flow.u), w)
+    wT_z = Bijectors.aT_b(w, z)
+    # `flow.b` holds one element, so broadcasting it is the same as Bijectors' `first(flow.b)`
+    # without reading back from the device.
+    tanh_term = CUDA.tanh.(CUDA.broadcast(+, wT_z, flow.b))
+    transformed = CUDA.broadcast(+, z, CUDA.broadcast(*, û, tanh_term))
+    return (transformed=transformed, wT_û=wT_û, wT_z=wT_z)
+end
+function Bijectors.with_logabsdet_jacobian(
+    flow::Bijectors.PlanarLayer, z::CuMatrix{T}
+) where {T<:Real}
+    transformed, wT_û, wT_z = Bijectors._transform(flow, z)
+    logjac = log1p.(wT_û .* abs2.(sech.(vec(wT_z) .+ flow.b)))
+    return (result=transformed, logabsdetjac=logjac)
+end
+
+rng = CUDA.default_rng()
+T = Float32
+d = 2
+
+@leaf MvNormal
+q0 = MvNormal(CUDA.zeros(T, d), Diagonal(CUDA.ones(T, d)))
+
+# `logp` takes the whole `(d, n)` batch and returns one value per column. Writing it with
+# array operations keeps it on the device, where `logpdf` would gather the columns onto the
+# host. The normaliser is kept so the reported ELBO is the true one.
+const μ_target = cu(T[2, -1])
+const logZ = T(d * log(2 * π))
+logp(z) = vec(-(logZ .+ sum(abs2, z .- μ_target; dims=1)) ./ 2)
+
+layers = [
+    Bijectors.PlanarLayer(CUDA.rand(T, d), CUDA.rand(T, d), CUDA.rand(T, 1)) for _ in 1:4
+]
+flow = create_flow(layers, q0)
+
+sample_per_iter = 64
+flow_trained, stats, _ = train_flow(
+    rng,
+    elbo_batch,
+    flow,
+    logp,
+    sample_per_iter;
+    max_iters=2_000,
+    optimiser=Optimisers.Adam(one(T) / 100),
+    # Zygote rather than Mooncake: Mooncake 0.5.48 fails on this flow inside its own CUDA
+    # kernel launch, a `CoDual` type assertion in `Adapt.adapt_storage`.
+    ADbackend=ADTypes.AutoZygote(),
+)
+
+losses = map(x -> x.loss, stats)
+@info "ELBO" start = -losses[1] final = -losses[end]
+
+ys = NormalizingFlows._device_specific_rand(rng, flow_trained, 1_000)
+@info "posterior mean (target is $(Array(μ_target)))" mean(Array(ys); dims=2)

ext/NormalizingFlowsCUDAExt.jl

@@ -3,6 +3,7 @@ module NormalizingFlowsCUDAExt
 using CUDA
 using NormalizingFlows
 using NormalizingFlows: Bijectors, Distributions, Random
+using ChainRulesCore: @non_differentiable
 
 function NormalizingFlows._device_specific_rand(
     rng::CUDA.RNG,
@@ -23,21 +24,28 @@ function _cuda_rand(
     rng::CUDA.RNG,
     s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
 )
-    return @inbounds Distributions.rand!(
-        rng, Distributions.sampler(s), CuArray{float(eltype(s))}(undef, size(s))
-    )
+    return _cuda_draw(rng, s, size(s))
 end
 
 function _cuda_rand(
     rng::CUDA.RNG,
     s::Distributions.Sampleable{<:Distributions.ArrayLikeVariate,Distributions.Continuous},
     n::Int,
 )
+    return _cuda_draw(rng, s, (size(s)..., n))
+end
+
+function _cuda_draw(rng::CUDA.RNG, s, dims::Tuple)
     return @inbounds Distributions.rand!(
-        rng, Distributions.sampler(s), CuArray{float(eltype(s))}(undef, size(s)..., n)
+        rng, Distributions.sampler(s), CuArray{float(eltype(s))}(undef, dims)
     )
 end
 
+# Zygote cannot trace a `CuArray` allocation: it descends into the CUDA allocator and fails
+# to compile. Nothing is lost by hiding the draw, since `rand!` carries no gradient on the
+# host path either. Only the draw is opaque; a flow's transform stays differentiable.
+@non_differentiable _cuda_draw(::Any, ::Any, ::Any)
+
 # ! this is type piracy
 # replacing original function with scalar indexing
 function Distributions._rand!(rng::CUDA.RNG, d::Distributions.MvNormal, x::CuVecOrMat)

src/NormalizingFlows.jl

@@ -12,6 +12,9 @@ using Bijectors: PartitionMask, Inverse, combine, partition
 using Functors
 using AbstractPPL: AbstractPPL
 using LogExpFunctions: LogExpFunctions
+using ChainRulesCore: ignore_derivatives
+using GPUArraysCore: AbstractGPUMatrix
+using PDMats: PDMat, whiten
 
 using DocStringExtensions
 
@@ -135,6 +138,43 @@ function _device_specific_rand(
     return Random.rand(rng, td, n)
 end
 
+"""
+    _device_specific_logpdf(d, xs)
+
+Log-density of `d` at each column of `xs`, left on the device holding `xs`.
+`Distributions.logpdf` maps over the columns and materialises a host array, so the ELBO
+cannot be assembled from it when the samples live on a GPU.
+"""
+_device_specific_logpdf(d, xs::AbstractMatrix) = logpdf(d, xs)
+
+function _device_specific_logpdf(d::Distributions.MvNormal, xs::AbstractGPUMatrix)
+    return _batched_mvnormal_logpdf(d, xs)
+end
+
+# `logdet(::Cholesky)` accumulates `factors[i, i]` in a host loop, which a GPU array rejects.
+# Gathering the diagonal keeps it to one kernel. The other covariance types reduce over a
+# scalar or a vector in PDMats, so they need no help.
+_cov_logdet(Σ) = logdet(Σ)
+_cov_logdet(Σ::PDMat) = 2 * sum(log, diag(cholesky(Σ).factors))
+
+# Whole-array form of the multivariate normal log-density, so it runs wherever `xs` lives.
+# `whiten` stays on the device and does not mutate, unlike the `sqmahal` behind `logpdf`; a
+# solve against `d.Σ` would leave a `PDMats` tangent that AD cannot accumulate.
+#
+# `d` is held constant. Differentiating more than one use of a full covariance leaves a
+# cotangent per use, a `Diagonal` from the log-determinant and an `UpperTriangular` from the
+# whitening, and summing those two indexes a device array element by element. Base
+# distributions are leaves and targets are fixed, so no gradient is owed for `d` here, and
+# returning none beats returning a wrong one.
+function _batched_mvnormal_logpdf(d::Distributions.MvNormal, xs::AbstractMatrix)
+    T = eltype(xs)
+    μ = ignore_derivatives(d.μ)
+    Σ = ignore_derivatives(d.Σ)
+    c = ignore_derivatives(T(length(d) * log(2 * π)) + _cov_logdet(Σ))
+    q = sum(abs2, whiten(Σ, xs .- μ); dims=1)
+    return vec(-(c .+ q) ./ 2)
+end
+
 # interface of contructing common flow layers
 include("flows/utils.jl")
 include("flows/planar_radial.jl")

src/flows/neuralspline.jl

@@ -226,9 +226,11 @@ Returns
 - `Bijectors.TransformedDistribution` representing the NSF flow.
 
 !!! note
-    The rational quadratic spline is written with whole-array operations, so the flow runs
-    on the CPU and the GPU and is differentiable by `Zygote`, `ForwardDiff`, `ReverseDiff`,
-    `Mooncake`, and, on Julia 1.11 and newer, `Enzyme`.
+    The rational quadratic spline is written with whole-array operations, so it runs on the
+    CPU and the GPU and is differentiable by `Zygote`, `ForwardDiff`, `ReverseDiff`,
+    `Mooncake`, and, on Julia 1.11 and newer, `Enzyme`. The assembled flow is CPU only: the
+    coupling layer splits its input through `Bijectors.PartitionMask`, which holds host
+    sparse matrices.
 
 !!! note
     When training the flow, mark the base distribution as a leaf first (`Functors.@leaf

src/objectives/elbo.jl

@@ -65,7 +65,7 @@ Returns
 function _batched_elbos(flow::Bijectors.MultivariateTransformed, logp, xs::AbstractMatrix)
     # requires the flow transformation to be able to handle batched inputs
     ys, logabsdetjac = with_logabsdet_jacobian(flow.transform, xs) 
-    elbos = logp(ys) .- logpdf(flow.dist, xs) .+ logabsdetjac
+    elbos = logp(ys) .- _device_specific_logpdf(flow.dist, xs) .+ logabsdetjac
     return elbos
 end
 

test/Project.toml

@@ -7,10 +7,13 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NormalizingFlows = "50e4474d-9f12-44b7-af7a-91ab30ff6256"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
@@ -25,8 +28,11 @@ Enzyme = "0.13.186 - 0.13.188"
 Flux = "0.16.10"
 ForwardDiff = "1.4.1"
 Functors = "0.5.2"
+GPUArraysCore = "0.2"
+JLArrays = "0.3.3"
 Mooncake = "0.5"
 NormalizingFlows = "0.4.0"
 Optimisers = "0.4.7"
+PDMats = "0.11"
 ReverseDiff = "1.17.0"
 Zygote = "0.7.11"

test/device.jl

@@ -0,0 +1,57 @@
+# JLArrays is the GPUArrays reference backend: it runs on the CPU but rejects scalar
+# indexing, so these tests reach device-only failures that would otherwise need a GPU.
+
+# Mirror what cuSOLVER returns: factors on the device with `uplo = 'U'`, which is what makes
+# PDMats wrap them in an `Adjoint` when whitening.
+function device_pdmat(A::AbstractMatrix)
+    c = cholesky(A)
+    return PDMats.PDMat(jl(A), Cholesky(jl(Matrix(c.factors)), c.uplo, c.info))
+end
+
+@testset "batched MvNormal log-density on device arrays" begin
+    GPUArraysCore.allowscalar(false)
+
+    @testset "$T" for T in (Float32, Float64)
+        rtol = T == Float32 ? 1.0f-4 : 1.0e-8
+        xs = randn(T, 2, 5)
+        covariances = (
+            PDMats.ScalMat(2, one(T)),
+            PDMats.PDiagMat(T[2, 0.5]),
+            PDMats.PDMat(T[2 0.3; 0.3 1]),
+        )
+
+        @testset "$(nameof(typeof(Σ)))" for Σ in covariances
+            host = MvNormal(zeros(T, 2), Σ)
+            Σ_dev = if Σ isa PDMats.PDMat
+                device_pdmat(Matrix(Σ))
+            elseif Σ isa PDMats.PDiagMat
+                PDMats.PDiagMat(jl(Σ.diag))
+            else
+                Σ
+            end
+            dev = MvNormal(jl(zeros(T, 2)), Σ_dev)
+            xs_dev = jl(xs)
+
+            batched = NormalizingFlows._batched_mvnormal_logpdf(dev, xs_dev)
+            @test batched isa JLArray
+            @test Array(batched) ≈ logpdf(host, xs) rtol = rtol
+
+            # the dispatch is on any GPU array, not just CUDA, so a second backend routes here
+            @test NormalizingFlows._device_specific_logpdf(dev, xs_dev) ≈ batched rtol =
+                rtol
+            # and the host path is still Distributions
+            @test NormalizingFlows._device_specific_logpdf(host, xs) == logpdf(host, xs)
+
+            # The gradient is the part that used to fail: a full covariance leaves one
+            # cotangent per differentiated use and adding them indexes the device array.
+            g = only(
+                Zygote.gradient(
+                    x -> sum(NormalizingFlows._batched_mvnormal_logpdf(dev, x)), xs_dev
+                ),
+            )
+            @test g isa JLArray
+            g_ref = ForwardDiff.gradient(x -> sum(logpdf(host, reshape(x, 2, 5))), vec(xs))
+            @test Array(g) ≈ reshape(g_ref, 2, 5) rtol = rtol
+        end
+    end
+end