Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

Mooncake on CUDA issue tracker

Aperta
  • Qualcosa non va?
#702 1 commento 0 reazioni 0 assegnatari Vedi su GitHub

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
45/100
Tipo di issue
Funzionalità
Chiarezza
Abbastanza chiara
Stato di attività
Tranquilla
Stack tecnologico
julia
Ambito
machine-learning, testing-qa

Direzione di ricerca

Inizia leggendo il marker non differenziabile esistente in transform.jl e l’integrazione proposta di GNNGraphs/ext/GNNGraphsMooncakeExt.jl. Riproduci i fallimenti di Mooncake su GPU descritti per i grafi COO e densi, quindi valida l’estensione rispetto ai layer elencati, confrontando i gradienti con Zygote. Il lavoro è completo quando l’estensione Mooncake inclusa nel repository è circoscritta e i casi supportati superano i test; le lacune separate di upstream restano distinte.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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.

Lingua principale
Julia
Stelle
310
Fork
74
Merge medio
3g 23h
PR unite (30g)
2

Preparare l'ambiente

Questo progetto non fornisce container di sviluppo, Dockerfile né guida per i contributori, quindi l'ambiente è a tuo carico: parti dal suo README e consulta la nostra guida al primo contributo per i passaggi generali.

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di JuliaGraphs/GraphNeuralNetworks.jl

Tutte le issue di JuliaGraphs/GraphNeuralNetworks.jl

Issue simili

Altre issue su Julia

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.