Skip to content

Skip attention dropout in GATConv/GATv2Conv when dropout == 0 - #709

Open
Parvm1102 wants to merge 1 commit into
JuliaGraphs:masterfrom
Parvm1102:mooncake-cuda-status
Open

Skip attention dropout in GATConv/GATv2Conv when dropout == 0#709
Parvm1102 wants to merge 1 commit into
JuliaGraphs:masterfrom
Parvm1102:mooncake-cuda-status

Conversation

@Parvm1102

Copy link
Copy Markdown
Contributor

NNlib.dropout(A, p) fetches the RNG (CUDA.default_rng() for CuArray) before it checks p, so dropout = 0 still pulls the CUDA RNG into the traced program, which Mooncake cannot differentiate. Every GATConv/GATv2Conv gradient on the GPU failed under Mooncake (#702).

MWE:

using CUDA, Flux, Mooncake, NNlib
x = CUDA.rand(Float32, 3, 4)
Flux.gradient(x -> sum(NNlib.dropout(x, 0.0)), Flux.AutoMooncake(), x)
TypeError: in typeassert, expected CoDual{IdDict{Any,Any}, IdDict{Any,Any}},
           got CoDual{IdDict{Any,Any}, NoFData}

Fix: skip the dropout call when p == 0 in gat_conv/gatv2_conv (also saves an RNG launch at inference). With it, GATConv and GATv2Conv pass all 96 GPU gradient sites under Mooncake (≤ 2.5e-7 vs Zygote); CPU tests unchanged. dropout > 0 still needs an upstream Mooncake rule for CUDA.default_rng(), tracked in #702.

Signed-off-by: Parvm1102 <parvmittal31757@gmail.com>
@Parvm1102

Parvm1102 commented Aug 29, 2026

Copy link
Copy Markdown
Contributor Author

@CarloLucibello please review this.
Also my GSoC dashboard shows the mentor evaluation is still pending. Please submit it, the deadline is soon.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant