Skip to content

Fix BatchNorm crash under Mooncake in eval mode#740

Open
AstitvaAggarwal wants to merge 2 commits into
FluxML:masterfrom
AstitvaAggarwal:fix/batchnorm-cache-eval-mode
Open

Fix BatchNorm crash under Mooncake in eval mode#740
AstitvaAggarwal wants to merge 2 commits into
FluxML:masterfrom
AstitvaAggarwal:fix/batchnorm-cache-eval-mode

Conversation

@AstitvaAggarwal

@AstitvaAggarwal AstitvaAggarwal commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

While running BatchNorm through Mooncake in eval mode, the backward pass crashes with:

MethodError: no method matching unsafe_convert(::Type{CuPtr{Nothing}}, ::Nothing)

The cause: cudnnBNForward! only fills in the cache's mean/ivar when training=true — in eval mode it just leaves them unset. But cudnnBNBackward! doesn't check for that; it assumes any non-nothing cache is already populated and hands whatever's in it straight to cuDNN.

The Mooncake rule for batchnorm always builds a live cache, whether it's training or not, so eval-mode calls hit this every time. This is easy to trigger via Flux.BatchNorm, since its training/eval auto-detection (NNlib.within_gradient) never returns true under Mooncake — that trick only works for Zygote.

This PR makes two small fixes:

  • Only build the cache when training is actually happening (matches what the "no cache" path already does safely).
  • If the caller passes their own cache object, use it instead of quietly replacing it with our own.

I tested this directly on GPU: reproduced the crash on the released NNlib, confirmed it's gone with the patch, checked the gradient still matches Zygote, and ran through all four combinations of training/track_stats to make sure nothing else was missed.

@AstitvaAggarwal
AstitvaAggarwal marked this pull request as draft July 22, 2026 00:40
cudnnBNForward! only fills in a passed-in cache's mean/ivar when
training=true; in eval mode it leaves them unset, but cudnnBNBackward!
reads the cache anyway, crashing on the stale nothing.

The Mooncake batchnorm rule always built a live cache regardless of
mode, so any call that turned out to be eval mode crashed with
MethodError converting nothing to a CuPtr. This is easy to hit via
Flux.BatchNorm, whose auto mode-detection (NNlib.within_gradient)
never resolves true under Mooncake, since it's implemented purely as
a ChainRulesCore.rrule trick that only Zygote-style AD triggers.

Fix: only build the cache when training is actually happening.
@AstitvaAggarwal
AstitvaAggarwal force-pushed the fix/batchnorm-cache-eval-mode branch from a1af08e to bf330b9 Compare July 22, 2026 00:46
@AstitvaAggarwal AstitvaAggarwal changed the title Fix BatchNorm cache crash when training=false under Mooncake Fix BatchNorm cache crash when training=false under Mooncake Jul 22, 2026
The rule always replaced whatever `cache` value the caller passed
(e.g. via Flux.BatchNorm's own `cache` argument) with its own
_BNFwdCache. If a caller ever supplied their own cache object, it was
silently discarded rather than being written to, breaking any state
they expected to read back afterward.

Now only fall back to constructing a fresh cache when the caller
didn't supply one; otherwise use theirs.
@AstitvaAggarwal AstitvaAggarwal changed the title Fix BatchNorm cache crash when training=false under Mooncake Fix BatchNorm cache handling under Mooncake (crash on training=false, ignored caller-supplied cache) Jul 22, 2026
@AstitvaAggarwal AstitvaAggarwal changed the title Fix BatchNorm cache handling under Mooncake (crash on training=false, ignored caller-supplied cache) Fix BatchNorm crash under Mooncake in eval mode Jul 22, 2026
@AstitvaAggarwal
AstitvaAggarwal marked this pull request as ready for review July 22, 2026 10:53
@AstitvaAggarwal

Copy link
Copy Markdown
Contributor Author

@CarloLucibello I think this is good from my side, please let me know if this needs anything more. Thank you

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