Replacing Zygote with Mooncake

[yebai]

Mooncake is intended as a drop-in replacement for Zygote, and the Mooncake team is happy to address GPU-related issues.

[CarloLucibello]

Is mooncake ready for Flux's needs? In particular, what is the status of

• differentiating with respect to generic structs
• gpu (cuda/amdgpu/metal) support
• generic float types

[yebai]

It may be missing some GPU rules; otherwise, it should just work. GPU rules can often be easily adapted from Zygote or ChainRules.

[CarloLucibello]

Some preliminary experiments

using Mooncake
import DifferentiationInterface as DI
using CUDA
using Flux

### ON CPU
x = rand(Float32, 2)
model = Dense(2, 2)
@time g = DI.gradient(m -> sum(m(x)), DI.AutoMooncake(), model)
# 1st time: 104.655418 seconds (97.43 M allocations: 4.912 GiB, 1.46% gc time, 95.14% compilation time: <1% of which was recompilation)
# 2nd time:   1.691145 seconds (7.41 M allocations: 351.098 MiB, 7.55% gc time, 44.78% compilation time: <1% of which was recompilation)
g.fields # gradient is in g.fields

@time Flux.gradient(m -> sum(m(x)), model)[1] # Zygote
# 1st time:  11.746135 seconds (17.81 M allocations: 845.103 MiB, 1.39% gc time, 99.99% compilation time)
# 2nd time:    0.109257 seconds (146.23 k allocations: 6.656 MiB, 99.63% compilation time)

loss(m) = sum(m(x))
@time g = DI.gradient(loss, DI.AutoMooncake(), model)
# 1st time:   4.004522 seconds (6.83 M allocations: 323.982 MiB, 1.60% gc time, 10.12% compilation time: <1% of which was recompilation)
# 2nd time:     0.000358 seconds (328 allocations: 13.656 KiB)

prep = DI.prepare_gradient(loss, DI.AutoMooncake(), model)
@time g = DI.gradient(loss, prep, DI.AutoMooncake(), model)
# every time:   0.000075 seconds (68 allocations: 2.375 KiB)

### ON CUDA GPU
x_gpu = x |> gpu
model_gpu = model |> gpu
g_gpu = DI.gradient(m -> sum(m(x_gpu)), DI.AutoMooncake(), model_gpu) # had to interrupt, was taking several minutes

I used a closure as a loss since it is a very common pattern when using Flux. Can things be improved there?

[MasonProtter]

@CarloLucibello surely you don't typically run Flux over closures containing non-constant global variables though, right? That's gotta be a performance nightmare.

I think this is probably a more representative look

julia> using Mooncake
       using DifferentiationInterface: DifferentiationInterface as DI
       using Flux
       using BenchmarkTools

julia> let
       x = rand(Float32, 2)
       model = Dense(2, 2)

       cl = m -> sum(m(x))
       println("First Mooncake closure gradient")
       @time DI.gradient(cl, DI.AutoMooncake(), model)
       println("Warmed-up Mooncake closure gradient")
       @btime DI.gradient($cl, DI.AutoMooncake(), $model)
       println()

       prep = DI.prepare_gradient(cl, DI.AutoMooncake(), model)
       println("Mooncake prepared closure gradient")
       @btime DI.gradient($cl, $prep, DI.AutoMooncake(), $model)

       println()
       println("First Zygote closure gradient")
       @time Flux.gradient(cl, model)[1]
       println("Warmed-up Zygote closure gradient")
       @btime Flux.gradient($cl, $model)[1]

       end;

First Mooncake closure gradient
 51.280655 seconds (109.97 M allocations: 5.220 GiB, 1.00% gc time, 95.70% compilation time: <1% of which was recompilation)
Warmed-up Mooncake closure gradient
  6.969 μs (197 allocations: 8.61 KiB)

Mooncake prepared closure gradient
  311.485 ns (15 allocations: 528 bytes)

First Zygote closure gradient
  6.210035 seconds (19.73 M allocations: 961.157 MiB, 1.00% gc time, 99.99% compilation time)
Warmed-up Zygote closure gradient
  4.197 μs (63 allocations: 4.53 KiB)

Once we have chalk-lab/Mooncake.jl#900 working with ReverseMode, preparation won't be so important for Mooncake, and it'll be more like Zygote in that respect.

[CarloLucibello]

surely you don't typically run Flux over closures containing non-constant global variables though, right?

right, your benchmark is more representative.

[wsmoses]

To offer a different viewpoint, perhaps alternatively this is an opportunity to devote resources to Reactant. The performance gains for both primal (aka inference) and derivative (aka training) for Lux have already been show to be quite substantial: https://github.com/LuxDL/Lux.jl/tree/main/perf, and it natively supports single and distributed CPU, GPU (CUDA + ROCM), TPU, and is onboarding TensorTorrent and Trainium.

[wsmoses]

Reactant's performance is at least comparable on this small of a model (and will improve on larger sizes since it has a larger constant time launch overhead), though not quite as good as pure Enzyme (which seems to beat all of the above).

(base) wmoses@hydra:~/git/Reactant.jl$ julia +1.12 --project=flux
The latest version of Julia in the `1.12` channel is 1.12.3+0.x64.linux.gnu. You currently have `1.12.1+0.x64.linux.gnu` installed. Run:

  juliaup update

in your terminal shell to install Julia 1.12.3+0.x64.linux.gnu and update the `1.12` channel to that version.
               _
   _       _ _(_)_     |  Documentation: https://docs.julialang.org
  (_)     | (_) (_)    |
   _ _   _| |_  __ _   |  Type "?" for help, "]?" for Pkg help.
  | | | | | | |/ _` |  |
  | | |_| | | | (_| |  |  Version 1.12.1 (2025-10-17)
 _/ |\__'_|_|_|\__'_|  |  Official https://julialang.org release
|__/                   |

julia> using Mooncake


julia> using DifferentiationInterface: DifferentiationInterface as DI

julia> using Flux

julia> using BenchmarkTools

julia> using Enzyme, Reactant

julia> let
           x = rand(Float32, 2)
           model = Dense(2, 2)

          cl = m -> sum(m(x))

           println("First Mooncake closure gradient")
           @time DI.gradient(cl, DI.AutoMooncake(), model)

           println("First Enzyme closure gradient")
           @time Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(cl), model)

           println("Warmed-up Mooncake closure gradient")
           @btime DI.gradient($cl, DI.AutoMooncake(), $model)

           println("Warmed-up Enzyme closure gradient")
           @btime Enzyme.gradient(Enzyme.Reverse, $(Enzyme.Const(cl)), $model)
           println()

           # Prep allocates the shadow
           prep = DI.prepare_gradient(cl, DI.AutoMooncake(), model)
           println("Mooncake prepared closure gradient")
           @btime DI.gradient($cl, $prep, DI.AutoMooncake(), $model)

           dmodel = Enzyme.Duplicated(model)
           println("First Enzyme \"prepared\" closure gradient")
           @time Enzyme.autodiff(Enzyme.Reverse, Enzyme.Const(cl), dmodel)
           println("Warmed-up Enzyme \"prepared\" closure gradient")
           @btime Enzyme.autodiff(Enzyme.Reverse, $(Enzyme.Const(cl)), $dmodel)

           println()

           Reactant.set_default_backend("cpu")
           cl_r = Reactant.to_rarray(cl)
           r_model = Reactant.to_rarray(dmodel)

           compiled_grad = Reactant.@compile Enzyme.autodiff(Enzyme.Reverse, Enzyme.Const(cl_r), r_model)

           println("First Reactant closure gradient")
           @time compiled_grad(Enzyme.Reverse, Enzyme.Const(cl_r), r_model)

           println("Warmed-up Reactant closure gradient")
           @btime $compiled_grad($(Enzyme.Reverse), $(Enzyme.Const(cl_r)), $r_model)

           println()
           println("First Zygote closure gradient")
           @time Flux.gradient(cl, model)[1]
           println("Warmed-up Zygote closure gradient")
           @btime Flux.gradient($cl, $model)[1]

       end
First Mooncake closure gradient
 58.545394 seconds (109.92 M allocations: 5.216 GiB, 1.84% gc time, 96.55% compilation time: <1% of which was recompilation)
First Enzyme closure gradient
 50.714129 seconds (124.72 M allocations: 6.031 GiB, 3.28% gc time, 100.00% compilation time)
Warmed-up Mooncake closure gradient
  7.569 μs (197 allocations: 8.61 KiB)
Warmed-up Enzyme closure gradient
  636.110 ns (21 allocations: 976 bytes)

Mooncake prepared closure gradient
  367.302 ns (15 allocations: 528 bytes)
First Enzyme "prepared" closure gradient
  0.007966 seconds (20.58 k allocations: 1.039 MiB, 99.58% compilation time)
Warmed-up Enzyme "prepared" closure gradient
  93.810 ns (4 allocations: 128 bytes)

WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1766538762.173659 2624266 service.cc:153] XLA service 0xcf496d0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
I0000 00:00:1766538762.173702 2624266 service.cc:161]   StreamExecutor device (0): NVIDIA GeForce RTX 5090, Compute Capability 12.0a
I0000 00:00:1766538762.173731 2624266 service.cc:161]   StreamExecutor device (1): NVIDIA GeForce RTX 5090, Compute Capability 12.0a
I0000 00:00:1766538762.174991 2624266 se_gpu_pjrt_client.cc:1440] Using BFC allocator.
I0000 00:00:1766538762.175017 2624266 gpu_helpers.cc:141] XLA backend allocating 25249628160 bytes on device 0 for BFCAllocator.
I0000 00:00:1766538762.175056 2624266 gpu_helpers.cc:141] XLA backend allocating 25251987456 bytes on device 1 for BFCAllocator.
I0000 00:00:1766538762.175069 2624266 gpu_helpers.cc:182] XLA backend will use up to 8416542720 bytes on device 0 for CollectiveBFCAllocator.
I0000 00:00:1766538762.175080 2624266 gpu_helpers.cc:182] XLA backend will use up to 8417329152 bytes on device 1 for CollectiveBFCAllocator.
I0000 00:00:1766538762.200852 2624266 cuda_dnn.cc:461] Loaded cuDNN version 91400
First Reactant closure gradient
  0.279245 seconds (745.97 k allocations: 37.078 MiB, 99.94% compilation time)
Warmed-up Reactant closure gradient
  3.257 μs (26 allocations: 1.02 KiB)

First Zygote closure gradient
  7.109280 seconds (21.26 M allocations: 1.009 GiB, 1.39% gc time, 100.00% compilation time)
Warmed-up Zygote closure gradient
  4.890 μs (63 allocations: 4.53 KiB)
(weight = Float32[0.6111169 0.7883759; 0.6111169 0.7883759], bias = Float32[1.0, 1.0], σ = nothing)

[CarloLucibello]

enzyme and reactant are super promising, but so far for flux we have seen intermittent beakage too often (e.g. right now enzyme gpu tests are broken). Moreover, each new julia release takes months and a lot of work to be supported, so I'm a bit concerned by the long-term sustainability of the project.

In any case, we can support both ad backends and reactant, we don't really have to choose. In fact, we already provide some support for Enzyme.

A good step forward is to extend Flux.gradient to take in an ADType.

[wsmoses]

Hm I'm unaware of any breakages. Reactant explicitly tests flux (which has remained passing), and I don't think the non-Reactant Flux GPU tests ever passed (though for accelerators Reactant should almost always be used)? In any case if you want better guarantees/stability, I pinged you a while back on adding Flux as reverse integration to Enzyme (which would ensure things don't regress).

[CarloLucibello]

It is nice to see that these tests for Mooncake are passing

using Functors, ADTypes, Mooncake, Zygote
import DifferentiationInterface as DI
using Test
using Flux
using Random, Statistics, LinearAlgebra

function mooncake_withgradient(loss, x...)
    y, gs = DI.value_and_gradient(x -> loss(x...), AutoMooncake(), x)
    return y, _unwrap.(gs)
end

_unwrap(x) = fmap(_unwrap, x, exclude= y -> y isa Mooncake.Tangent)
_unwrap(x::Mooncake.Tangent) = _unwrap(x.fields)
_unwrap(x::Mooncake.MutableTangent) = _unwrap(x.fields)

function check_equal_leaves(a, b; rtol=1e-4, atol=1e-4)
    fmapstructure_with_path(a, b) do kp, x, y
        # @show kp
        if x isa AbstractArray
            @test x ≈ y rtol=rtol atol=atol
        elseif x isa Number
            @test x ≈ y rtol=rtol atol=atol
        end
    end
end

function test_mooncake_gradients(
            f, 
            xs...;
            rtol=1e-4, atol=1e-4,
            test_gpu = false,
            test_grad_f = true,
            test_grad_x = true,
            loss = (f, xs...) -> mean(f(xs...)),
            )


    Flux.trainmode!(f) # for layers like BatchNorm

    ## Let's make sure first that the forward pass works.
    l = loss(f, xs...)
    @test l isa Number
    if test_gpu
        gpu_dev = gpu_device(force=true)
        cpu_dev = cpu_device()
        xs_gpu = xs |> gpu_dev
        f_gpu = f |> gpu_dev
        l_gpu = loss(f_gpu, xs_gpu...)
        @test l_gpu isa Number
    end

    if test_grad_x
        y, g = Zygote.withgradient((xs...) -> loss(f, xs...), xs...)
        
        y_ez, g_ez = mooncake_withgradient((xs...) -> loss(f, xs...), xs...)
        @test y ≈ y_ez rtol=rtol atol=atol
        check_equal_leaves(g, g_ez; rtol, atol)
        
        if test_gpu
            y_gpu, g_gpu = mooncake_withgradient((xs...) -> loss(f_gpu, xs...), xs_gpu...)
            @test get_device(g_gpu) == get_device(xs_gpu)
            @test y_gpu ≈ y rtol=rtol atol=atol
            check_equal_leaves(g_gpu |> cpu_dev, g; rtol, atol)
        end
    end

    if test_grad_f
        y, g = Zygote.withgradient(f -> loss(f, xs...), f)

        y_ez, g_ez = mooncake_withgradient(f -> loss(f, xs...), f)
        @test y ≈ y_ez rtol=rtol atol=atol
        check_equal_leaves(g, g_ez; rtol, atol)

        if test_gpu
            y_gpu, g_gpu = mooncake_withgradient(f -> loss(f, xs_gpu...), f_gpu)
            # @test get_device(g_gpu) == get_device(xs_gpu)
            @test y_gpu ≈ y rtol=rtol atol=atol
            check_equal_leaves(g_gpu |> cpu_dev, g; rtol, atol)
        end
    end
    return true
end

loss(model, x) = mean(model(x))

const models_xs = [
    (Dense(2=>4), randn(Float32, 2), "Dense"),
    (Chain(Dense(2=>4, tanh), Dense(4=>3)), randn(Float32, 2), "Chain(Dense, Dense)"),
    (f64(Chain(Dense(2=>4), Dense(4=>2))), randn(Float64, 2, 1), "f64(Chain(Dense, Dense))"),
    (Flux.Scale([1.0f0 2.0f0 3.0f0 4.0f0], true, abs2), randn(Float32, 2), "Flux.Scale"),
    (Conv((3, 3), 2 => 3), randn(Float32, 3, 3, 2, 1), "Conv"),
    (Chain(Conv((3, 3), 2 => 3, ), Conv((3, 3), 3 => 1, tanh)), rand(Float32, 5, 5, 2, 1), "Chain(Conv, Conv)"),
    (Chain(Conv((4, 4), 2 => 2, pad=SamePad()), MeanPool((5, 5), pad=SamePad())), rand(Float32, 5, 5, 2, 2), "Chain(Conv, MeanPool)"),
    (Maxout(() -> Dense(5 => 4, tanh), 3), randn(Float32, 5, 1), "Maxout"),
    (SkipConnection(Dense(2 => 2), vcat), randn(Float32, 2, 3), "SkipConnection"),
    (Flux.Bilinear((2, 2) => 3), randn(Float32, 2, 1), "Bilinear"),
    (ConvTranspose((3, 3), 3 => 2, stride=2), rand(Float32, 5, 5, 3, 1), "ConvTranspose"),
    (first ∘ LayerNorm(2), randn(Float32, 2, 10), "LayerNorm"),
    (BatchNorm(2), randn(Float32, 2, 10), "BatchNorm"),
    (first ∘ MultiHeadAttention(16), randn32(16, 20, 2), "MultiHeadAttention"),
]

# model, x = models_xs[1][1:2]
# y, gs = mooncake_withgradient(loss, model, x)

for (model, x, name) in models_xs
    @testset "Mooncake grad check $name" begin
        @test test_mooncake_gradients(model, x; loss, test_gpu=false)
    end
end

[CarloLucibello]

Zygote is still the default AD. At some point, it should be replaced by either Mooncake or Enzyme. At the moment though, both frameworks have outstanding problems with CUDA.