Mooncake on CUDA issue tracker

Open
  • Is something wrong?
#702 1 comment 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 2 days

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
45/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Quiet
Tech stack
julia
Domain
machine-learning, testing-qa

Research direction

Start by reading the existing non-differentiable marker in transform.jl and the proposed GNNGraphs/ext/GNNGraphsMooncakeExt.jl integration. Reproduce the Mooncake GPU failures described for COO and dense graphs, then validate the extension against the listed layers with gradients compared to Zygote. Done means the in-repository Mooncake extension is scoped and its supported cases pass; the separate upstream gaps remain distinct.

Written by the indexing model from the issue text.

Description

flux gnngraphs gnnlib

I ran Mooncake on GPU where the test sites for CPU were enabled (Mooncake 0.5.41, CUDA.jl 6.2.1, Flux 0.16.10, Julia 1.12.5).
CPU is fully green, but many fail on GPU. Here are my findings:

Already works on GPU (COO storage): GraphConv, SAGEConv, GINConv, EdgeConv, NNConv, ResGatedGraphConv, CGConv, MEGNetConv, GlobalPool, GlobalAttentionPool, GCNConv with add_self_loops=false, and all the propagate fast paths (copy_xj with +/mean, e_mul_xj, w_mul_xj). The rules in GNNlibMooncakeExt.

Zygote uses ChainRules for non-differentiable functions, Mooncake doesn't, so it the hits GPU code it can't differentiate:

  • add_self_loops for COO does nodes = convert(typeof(s), [1:n;]), i.e. builds a CPU vector and copies it to the GPU. That lands on a CPU→GPU unsafe_copyto! containing try/catch, which Mooncake can't trace (its cross-device copy rule only covers float arrays, and here the destination is an integer index array). It affects all layers with self loops: GCNConv (default), GATConv, GATv2Conv, TransformerConv, AGNNConv, SGConv, TAGConv, TGCN.
  • On :dense graphs, edge_index/to_coo call _findnz_idx, i.e. findall on a CuArray{Bool} (no Mooncake rule), and right after that v = A[nz] indexes with a vector of CartesianIndex (also no rule). This kills every layer on dense storage.

We can add GNNGraphs/ext/GNNGraphsMooncakeExt.jl (Mooncake as weakdep) giving Mooncake the same semantics the ChainRules markers give Zygote:

Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(add_self_loops), GNNGraph}
Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(GNNGraphs._findnz_idx), Any}
Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(to_coo), GNNGraphs.ADJMAT_T}
Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(Core.kwcall), NamedTuple, typeof(to_coo), GNNGraphs.ADJMAT_T}

The Core.kwcall rule is required because all in-repo call sites call to_coo with keyword arguments. I validated all of this by defining the rules via type piracy in a test script: with them in place, GCNConv (coo and dense), EdgeConv/dense, ResGatedGraphConv/dense, SGConv, AGNNConv, GMMConv and GatedGraphConv pass on GPU with gradients matching Zygote to ~1e-7.

Two caveats worth recording in the ext:

  • add_self_loops: same caveat as the existing @non_differentiable in transform.jl (the graph carries feature arrays; fine in practice since layers take x explicitly).
  • to_coo: Zygote only skips _findnz_idx and still differentiates A[nz], so it delivers dA for weighted dense adjacencies. A zero-derivative to_coo drops that under Mooncake. If Mooncake gains a CartesianIndex getindex rule upstream (item 4 below), we can drop the to_coo rules and keep full parity.

Upstream Mooncake.jl gaps found (I think I can open some issues on Mooncake regarding this):

  1. Correctness bug: broadcasts over non-contiguous SubArrays of CuArrays silently drop gradients. This is why GatedGraphConv returns wrong gradients with no error on GPU. Flux's GRUCell slices its gates with chunk(Wi*x, 3, dims=1).
  2. No kwcall rule for sum(x::CuArray; dims). It breaks GMMConv and other attention score computation.
  3. Cross-device unsafe_copyto! rule covers float destinations only, not integer.
  4. No getindex(::CuArray, ::Vector{CartesianIndex}) rule.
  5. CuMatrix + CuMatrix → cuBLAS.geam! has no rule (CuRef tangent type missing).
  6. No repeat(::CuArray, ...) rule (hits TGCNCell's repeat(h, 1, g.num_nodes)).
  7. task_local_storage() IdDict fdata mismatch, reached through CUDA's default RNG by dropout.

I will work on these once it gets green flag.

Dominant language
Julia
Stars
310
Forks
75
Avg merge
2d 15h
Merged PRs (30d)
10

Getting set up

This project ships no dev container, Dockerfile or contributing guide, so setting up is up to you: start from its README, and see our first-contribution guide for the general steps.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from JuliaGraphs/GraphNeuralNetworks.jl

All issues in JuliaGraphs/GraphNeuralNetworks.jl

Similar issues

More Julia issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.