feat(differentiation): GPU default ad_backend AutoZygote -> AutoMooncake#501
Merged
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Swaps the GPU default
:ad_backendfromAutoZygote()toAutoMooncake(), based on an H200 measurement of every reverse-mode DI candidate end to end: CTFlows PR #353.Why
AutoMooncakecleared the entire matrix, including the realFlow(ocp)Hamiltonian call on aCuArraywith a numerically correct result.AutoZygoteturned out to fail unexpectedly on a non-mutatinghamiltonian_gradientdevice 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.AutoZygoteremains 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_backendoption'sdescriptionis now P-dependent instead of one static string shared by CPU/GPU.describe(:di, registry)now shows a GPU-specific message — verified manually:ADTypes.AutoMooncake()is a marker type fromADTypes(already a hard dep), same patternAutoZygote()used before it.0.28.1-beta→0.28.2-beta;CHANGELOG.md/BREAKING.mdentries (non-breaking: same option name, override mechanism,computed=trueflag — 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