Skip to content

feat(differentiation): GPU default ad_backend AutoZygote -> AutoMooncake#501

Merged
ocots merged 1 commit into
mainfrom
feat/gpu-ad-default-mooncake
Jul 22, 2026
Merged

feat(differentiation): GPU default ad_backend AutoZygote -> AutoMooncake#501
ocots merged 1 commit into
mainfrom
feat/gpu-ad-default-mooncake

Conversation

@ocots

@ocots ocots commented Jul 22, 2026

Copy link
Copy Markdown
Member

Swaps the GPU default :ad_backend from AutoZygote() to AutoMooncake(), based on an H200 measurement of every reverse-mode DI candidate end to end: CTFlows PR #353.

Why

AutoMooncake cleared the entire matrix, including the real Flow(ocp) Hamiltonian call on a CuArray with a numerically correct result. AutoZygote turned out to fail unexpectedly on a non-mutating hamiltonian_gradient device call (KernelException) — using the exact same Hamiltonian formula that succeeds through a different call shape elsewhere. Its device behaviour is call-context dependent, not a stable pass/fail — a materially weaker guarantee than what a shipped default should offer. AutoZygote remains selectable explicitly; it's simply no longer the default.

What changed

  • src/Differentiation/default.jl: __ad_backend(::Type{GPU}) = AutoMooncake().
  • src/Differentiation/differentiation_interface.jl: the :ad_backend option's description is now P-dependent instead of one static string shared by CPU/GPU. describe(:di, registry) now shows a GPU-specific message — verified manually:
    ad_backend (backend, ad)::ADTypes.AbstractADType (default: ADTypes.AutoMooncake() [computed])
       description: DifferentiationInterface.jl backend for GPU execution. Default: AutoMooncake(),
       validated end-to-end on CuArray including through a mutating in-place RHS. AutoForwardDiff()
       does not work on GPU (scalar-indexes a CuArray). AutoZygote() can be selected explicitly; it
       was found unreliable in some device call contexts (an unexpected KernelException was
       observed on a non-mutating call).
    
  • No new dependencyADTypes.AutoMooncake() is a marker type from ADTypes (already a hard dep), same pattern AutoZygote() used before it.
  • Tests: GPU default assertions updated; new assertions pin the GPU-specific description text and confirm CPU's is unaffected.
  • Version 0.28.1-beta0.28.2-beta; CHANGELOG.md/BREAKING.md entries (non-breaking: same option name, override mechanism, computed=true flag — only the computed default value changes, on the opt-in GPU path).

Verification

Targeted suite/differentiation: 110/110. Full suite: 4758/4758.

🤖 Generated with Claude Code

Roadmap-v4 §5 phase 4e (CTFlows PR #353) measured every reverse-mode DI
candidate end to end on an H200. AutoMooncake cleared the entire matrix,
including the real Flow(ocp) Hamiltonian call on a CuArray with a
numerically correct result. AutoZygote turned out to fail unexpectedly on a
non-mutating hamiltonian_gradient device call (KernelException) using the
exact same Hamiltonian formula that succeeds through a different call
shape elsewhere -- its device behaviour is call-context dependent, not a
stable pass/fail. AutoZygote remains selectable explicitly; it is simply no
longer the default.

- src/Differentiation/default.jl: __ad_backend(::Type{GPU}) = AutoMooncake().
- src/Differentiation/differentiation_interface.jl: the :ad_backend
  description is now P-dependent (new _ad_backend_description(P)) instead of
  one static string shared by CPU and GPU, so describe(:di, registry) shows a
  GPU-specific message explaining that AutoForwardDiff doesn't work on GPU
  and that AutoZygote was found unreliable in some device call contexts.
  Manually verified the rendered describe() output.
- No new dependency: ADTypes.AutoMooncake() is a marker type from ADTypes
  (already a hard dep), same pattern as AutoZygote() before it.
- test_di_parameter.jl: GPU default assertions updated; new assertions pin
  the GPU-specific description text and confirm CPU's is unaffected.
- Version 0.28.1-beta -> 0.28.2-beta; CHANGELOG.md and BREAKING.md entries
  (non-breaking: same option name, override mechanism, and computed=true
  flag -- only the computed default value changes, on the opt-in GPU path).

Targeted suite/differentiation: 110/110. Full suite: 4758/4758.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@ocots
ocots merged commit bae6ebd into main Jul 22, 2026
6 checks passed
@ocots
ocots deleted the feat/gpu-ad-default-mooncake branch July 22, 2026 17:22
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