Skip to content

Master CI red on Julia 1.12 (Core / DeepBSDE): SIGSEGV in Zygote ensemble adjoint codegen (SciMLBaseZygoteExt.∇responsible_map) #153

Description

@ChrisRackauckas-Claude

Summary

Master CI is red on the julia 1 (= 1.12.6) and julia pre channels (Core test set), across ubuntu/macos/windows. The julia lts (1.10) channel is green. The failing test is test/DeepBSDE.jl and it does not fail with an assertion — it segfaults (Package HighDimPDE errored during testing (received signal: 11), process exit 139).

I reproduced this locally on Julia 1.12.6 (the exact channel CI runs as "julia 1"):

$ julia +1.12.6 --project=. -e 'include("test/DeepBSDE.jl")'
DeepBSDE_tests
DeepBSDE
[----] signal 11 (1): Segmentation fault
in expression starting at .../test/DeepBSDE.jl:19
emit_unboxed_coercion at src/intrinsics.cpp:394 [inlined]
emit_unbox          at src/intrinsics.cpp:458
emit_condition      at src/codegen.cpp:6089
emit_function       at src/codegen.cpp:9294
jl_emit_code        ...
...
∇responsible_map     at SciMLBase/ext/SciMLBaseZygoteExt.jl:332 [inlined]
adjoint              at SciMLBase/ext/SciMLBaseZygoteExt.jl:373 [inlined]
solve_batch          at SciMLBase/src/ensemble/basic_ensemble_solve.jl:496
...
neural_sde           at src/DeepBSDE.jl:157 [inlined]
loss_n_sde           at src/DeepBSDE.jl:185 [inlined]
gradient             at Zygote/.../interface.jl:156
solve                at src/DeepBSDE.jl:193
EXIT=139

Root cause (verified)

The crash is inside Julia 1.12.6's LLVM codegen (emit_unboxed_coercionemit_unboxemit_condition), while compiling the Zygote reverse-pass closure for SciMLBase's ensemble adjoint (SciMLBaseZygoteExt.∇responsible_map). DeepBSDE differentiates an EnsembleProblem solve via Flux.gradient / TrackerAdjoint, and that pullback is what 1.12.6 fails to compile.

This is not a HighDimPDE bug. I reduced it to a ~50-line MWE with no HighDimPDE code (only Flux + Zygote + StochasticDiffEq + SciMLBase/SciMLSensitivity + an EnsembleProblem) and it segfaults identically, in the same ∇responsible_map codegen.

Julia core has already triaged this exact emit_unboxed_coercion/emit_condition signature as Zygote emitting invalid LLVM IR, not a Julia bug:

Versions (all latest released)

Julia 1.12.6 · SciMLBase 3.21.0 · SciMLSensitivity 7.112.0 · Zygote 0.7.11 · Flux 0.16.10 · StochasticDiffEq 7.0.0

MWE (segfaults on Julia 1.12.6, no HighDimPDE code)

click to expand
using Flux, Zygote
import StochasticDiffEq
using LinearAlgebra
using SciMLBase, SciMLSensitivity
import Tracker

d = 1; tspan = (0.0f0, 5.0f0); dt = 0.5f0; x0 = Float32[11.0]
net = Flux.Chain(Dense(d + 1, 11, relu), Dense(11, d))
p2, re2 = Flux.destructure(net); p3 = copy(p2); ps = Flux.params(p3)
μ(X, p, t) = X * 0.0f0
σ(X, p, t) = Matrix(Diagonal(ones(eltype(X), d)))
function F(h::Tracker.TrackedArray, p, t)
    X = h[1:(end-1)].data; t_ = eltype(X)(t); re2(p)([X; t_])' |> collect
    return Tracker.TrackedArray(vcat(μ(X, p, t), [eltype(X)(0.0)]))
end
function G(h::Tracker.TrackedArray, p, t)
    X = h[1:(end-1)].data; t_ = eltype(X)(t); s = re2(p)([X; t_])' |> collect
    return Tracker.TrackedArray(vcat(σ(X, p, t), s))
end
noise = zeros(Float32, d + 1, d)
sde_prob = SDEProblem{false}(F, G, [x0; 0.0f0], tspan, p3, noise_rate_prototype = noise)
function neural_sde(ic)
    sol = solve(EnsembleProblem(remake(sde_prob, u0 = ic)), StochasticDiffEq.EM(), EnsembleSerial();
        u0 = ic, trajectories = 5, dt = dt, p = p3,
        sensealg = SciMLSensitivity.TrackerAdjoint(), save_everystep = false)
    arr = Array(sol); return [arr[:, end, i] for i in axes(arr, 3)]
end
loss() = sum(sum(abs2, pr) for pr in neural_sde([x0; 0.0f0]))
gs = Flux.gradient(ps) do; loss(); end   # SIGSEGV in ∇responsible_map codegen

Status / options

This needs an upstream fix — it cannot be correctly fixed inside HighDimPDE without a fragile, unverified workaround around a compiler/Zygote-IR bug. Realistic paths:

  1. Zygote stops emitting invalid IR for the ∇responsible_map ensemble pullback on 1.12 (Assertion errors on IR verification on upcoming 1.12 FluxML/Zygote.jl#1580 follow-up), or
  2. SciMLBase's responsible_map ensemble adjoint avoids the closure pattern Zygote miscompiles, or
  3. DeepBSDE's AD path moves off Zygote (e.g. Mooncake), which is the ecosystem direction Julia/Zygote maintainers are recommending.

Filing as the tracking issue for the master CI red. I can test candidate upstream fixes locally on 1.12.6.

cc @ChrisRackauckas

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions