Mooncake on CUDA issue tracker

オープン
  • 何か問題がありますか？
#702 コメント 1 件 リアクション 0 件 担当者 0 名 GitHub で見る

まだ誰も着手していません。

評価

難易度
4/5
見積もり時間
3〜5日
初心者へのやさしさ
45/100
issue の種類
機能追加
明瞭さ
おおむね明確
活発さ
静か
技術スタック
julia
領域
machine-learning, testing-qa

調査の方向性

まず、transform.jl にある既存の非微分可能マーカーと、提案されている GNNGraphs/ext/GNNGraphsMooncakeExt.jl の統合を読みます。次に、COO グラフと密グラフについて説明されている Mooncake の GPU 失敗を再現し、その後、勾配を Zygote と比較しながら、列挙されているレイヤーに対して extension を検証します。完了の条件は、リポジトリ内の Mooncake extension の範囲が定まり、サポート対象のケースが通ることです。別個の upstream の不足は、引き続き別のものとして扱います。

索引モデルが issue の本文から書いたものです。

説明

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.

主要言語
Julia
スター
310
フォーク
74
平均マージ
3日 23時間
マージ済み PR（30日）
2

環境構築

このプロジェクトには開発コンテナ、Dockerfile、コントリビューションガイドがありません。まず README を読み、一般的な手順ははじめてのコントリビューションガイドを参照してください。

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

JuliaGraphs/GraphNeuralNetworks.jl のほかの issue

JuliaGraphs/GraphNeuralNetworks.jl の issue をすべて見る

似ている issue

Julia の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。