World age issues in nested autograd

[yebai]

The following misty closure heuristic appears to introduce genuine problems, likely limited to Julia 1.11:

Mooncake.jl/src/rules/misty_closures.jl

Lines 20 to 38 in f937521

	# Build a forward-mode rule for a MistyClosure using its original world age.
	#
	# We cannot use the current world age because the MistyClosure's IR (p.ir[]) has a
	# valid_worlds range set at creation time. On Julia 1.12+, generate_dual_ir calls
	# set_valid_world!(ir, interp.world), which throws if the world is outside this range.
	# If methods were defined after the MistyClosure was created, the current world would
	# fall outside valid_worlds and cause an error.
	#
	# Using the original world age is safe because lookup_ir for MistyClosure returns mc.ir[]
	# directly, bypassing method table lookups. Nested non-primitive calls use LazyFRule or
	# DynamicFRule, which obtain a current-world interpreter via get_interpreter() at runtime.
	# We pass skip_world_age_check=true since build_frule's safety check would incorrectly
	# reject our intentionally-older interpreter.
	#
	function _dual_mc(p::MistyClosure)
	mc_world = UInt(p.oc.world)
	interp = MooncakeInterpreter(DefaultCtx, ForwardMode; world=mc_world)
	return build_frule(interp, p; skip_world_age_check=true)
	end

Some examples:

• DI second-order tests fail (CI run). The failure pattern is inconsistent with an out-of-memory condition. Instead, Mooncake likely encounters an opaque-closure bug in Julia 1.11. The issue is resolved when transitioning to Julia 1.12.
  • The DI crash is reproducible after removing the misty closure heuristic (d857774), adding weight that this is an opaque-closure bug in Julia 1.11. This is potentially the same issue as discussed in compat fix for Julia 1.12 MistyClosures.jl#12 (comment)
• In DI Hessian computations, repeated preparation (e.g. DontPrepareInner()) triggers a world-age error (issue discussion). This is fixed by Fix inner preparation behavior for Mooncake JuliaDiff/DifferentiationInterface.jl#948.
  • One can modify MistyClosure to carry a back reference of the original Julia method so we can re-derive rules in the current world age, instead of the heuristic. Re-deriving rules would, in theory, be the right thing to do.

[yebai]

In DI Hessian computations, repeated preparation (e.g. DontPrepareInner()) triggers a world-age error (#632 (comment)). This is fixed by JuliaDiff/DifferentiationInterface.jl#948.

@sunxd3, is it clear why repeated DI preparation still causes world age issues even with the _dual_mc trick? It would be great to get to the bottom of it.

[sunxd3]

Not to me immediately, but digging deeper is definitely needed