More CUDA compatibility

[zuhengxu]

This PR aims to get a complete normalizing flow example/test on GPU (see #49).

After #11, we now support batch-wise training/inference for some flows (e.g., RealNVP), which is crucial for getting performance gain on GPU. PR #25 then got the cuda sampling for the reference dist working. The only missing piece is to ensure logpdf computation for both the reference and flow to work, and to ensure AD working.

I've identified some issue here:

using CUDA
const NF = NormalizingFlows
rng_g = CUDA.default_rng() # use GPU RNG if available

CUDA.allowscalar(true)
n_samples = 100
q0 = MvNormal(CUDA.zeros(2), cu([1f0 0f0; 0f0 1f0]))
# gpu sample from the reference
xs = NF._device_specific_rand(rng_g, q0, n_samples) # working properly

# Although xs is a cuda array, logpdf(q0, xs) returns a cpu array
logpdf(q0, xs) # returns a CPU array

I think a quick solution is to define a NF._device_specific_logpdf(...).

[yebai]

@shravanngoswamii, can you help to push this across the finish line?

[shravanngoswamii]

@shravanngoswamii, can you help to push this across the finish line?

I am on it.

[github-actions]

NormalizingFlows.jl documentation for PR #52 is available at:
https://TuringLang.github.io/NormalizingFlows.jl/previews/PR52/

[shravanngoswamii]

The GPU training test uses AutoZygote rather than AutoMooncake.

Mooncake 0.5.48 fails on the planar flow with TypeError: in typeassert, expected Mooncake.CoDual{CuPtr{Float32}, CuPtr{Float32}}, got Mooncake.CoDual{CuPtr{Float32}, Mooncake.NoFData} inside Adapt.adapt_storage during the kernel launch, on both Julia 1.10 and 1.12.

Reduced to a CuArray returned by a @non_differentiable function and then used in a differentiated broadcast:

using CUDA, Mooncake
using ChainRulesCore: @non_differentiable

draw() = CUDA.rand(Float32, 2, 8)
@non_differentiable draw()

x = CUDA.rand(Float32, 2, 8)
t = CUDA.rand(Float32, 5)

f(t) = sum(t[1:2] .* x)        # ok
g(t) = sum(t[1:2] .* draw())   # TypeError

Mooncake.value_and_gradient!!(Mooncake.build_rrule(g, t), g, t)

Captured arrays differentiate fine, drawn ones do not: TypeError: in typeassert, expected Mooncake.CoDual{IdDict{Any, Any}, IdDict{Any, Any}}, got Mooncake.CoDual{IdDict{Any, Any}, Mooncake.NoFData}. So a @non_differentiable CuArray reaches the CUDA rules as NoFData where a tangent is expected.

This package hides the GPU draw that way because Zygote otherwise traces into the CUDA allocator, so the two backends want opposite things here.

0.5.48 added an explicit refusal for broadcasts that carry captured state: MooncakeCUDAExt.jl#L2095-L2098. That refusal does not fire here, so this looks like a separate bug.

Is Zygote fine for the GPU path, and is it worth opening a Mooncake issue with this case? cc @yebai

[yebai]

I think you need Mooncake.@zero_derivative instead of ChainRulesCore.@non_differentiable.