Skip to content

Add opt-in PreserveDiffModel to keep the diff model across Parameter-only re-solves#371

Open
yeabbratz wants to merge 3 commits into
jump-dev:masterfrom
yeabbratz:preserve-diff-model
Open

Add opt-in PreserveDiffModel to keep the diff model across Parameter-only re-solves#371
yeabbratz wants to merge 3 commits into
jump-dev:masterfrom
yeabbratz:preserve-diff-model

Conversation

@yeabbratz

Copy link
Copy Markdown

MOI.optimize!(::DiffOpt.Optimizer) unconditionally sets model.diff = nothing, so the first
differentiation after every re-solve re-instantiates the differentiation model and re-runs a
full MOI.copy_to of the problem. That is the right default — a re-solve invalidates the
point-dependent state of diff — but when the only change since the last instantiation is
MOI.Parameter values (the standard parametric training / sensitivity loop:
set_parameter_valueoptimize! → differentiate, repeatedly), the structure inside diff
is still valid, and the re-instantiate + copy_to is pure per-solve overhead.

This PR adds an opt-in optimizer attribute, default false, stock behavior unchanged:

MOI.set(model, DiffOpt.PreserveDiffModel(), true)

When enabled, optimize! keeps diff and refreshes exactly the point-dependent state instead
of discarding everything:

  • the input sensitivities of diff are cleared and the new primal/dual solution is re-copied
    after each solve (_copy_dual — the same call the rebuild path runs after copy_to, fed by
    the same source);
  • parameter values are re-written into diff at the next differentiation call
    (_refresh_parameters): both their MOI.ConstraintSet and their slot in the primal point
    (_copy_dual copies MOI.VariablePrimal of every variable, and for a parameter that is its
    value).

Why a surviving diff is structurally in sync by construction: every structural edit on
DiffOpt.Optimizer (add/delete variable or constraint, function/set modification, objective
change, copy_to, …) already resets diff = nothing — this PR only narrows the
MOI.ConstraintSet setter for MOI.Parameter to not discard it when the attribute is enabled.
So "only parameter values changed" needs no separate dirty tracking: it is the only mutation
that leaves diff alive. A solve that ends in a non-differentiable termination status also
falls back to discarding diff.

Scope: only differentiation models that natively support MOI.Parameter can be preserved —
with a ParametricOptInterface layer inside diff, a value update would not propagate into the
already-copied problem data. Models opt in through the internal _diff_model_supports_preserve
trait; this PR enables it for NonLinearProgram.Model only (where a parameter update is an
exact value assignment). For every other configuration the attribute has no effect.

Stale sensitivities cannot be served: the re-solve path re-copies x/y/s, clears the input
caches, and the output caches (forw_grad_cache/back_grad_cache) are overwritten by every
differentiation call. A parameter-only re-solve with diff preserved yields results identical
to the full rebuild — the tests assert ==, not .

Tests

  • test_preserve_diff_model_parameter_resolves: 4 rounds of set_parameter_value
    optimize! → reverse on two identical models (one preserved, one stock): gradients == at
    every round, and the preserved model provably reuses the same diff object; forward mode
    checked too.
  • test_preserve_diff_model_parameter_change_without_resolve: changing parameters and
    differentiating without re-solving matches stock (which rebuilds with the new parameter
    values at the old solution).
  • test_preserve_diff_model_structural_change_falls_back: a structural edit discards the
    preserved diff; the rebuilt model matches a fresh one constructed in the final state.
  • Full test suite green locally (Julia 1.12.6), including the 9124-assert parameters.jl set.

Benchmark

Median of 7 (after warmup) of the first reverse_differentiate! after a Parameter-only
re-solve (the solve itself untimed) — the diff-model re-instantiate + copy_to is exactly the
residual this PR removes. Parameterized NLP with P parameters and P variables (script below),
Ipopt, M4 Pro, 1 thread, Julia 1.12.6, both columns on this branch:

P default PreserveDiffModel
10 1.08 ms / 0.5 MiB 0.30 ms / 0.3 MiB
100 2.83 ms / 5.5 MiB 1.29 ms / 4.1 MiB
1000 60.3 ms / 209.0 MiB 49.8 ms / 196.1 MiB

With the attribute enabled, the first backward after a re-solve lands on the level of a
repeated backward on the same solve (0.31 / 1.30 / 45.5 ms in the same run) — i.e. the copy
residual is fully removed; what remains at large P is the dense sensitivity contraction (#369's
target) and the O(P) parameter refresh.

bench_persistent_evaluator.jl
# Reproducer for the two per-backward residuals the DiffOpt follow-on PRs remove
# (upstreaming of task 4.12, companion to jump-dev/DiffOpt.jl#369):
#
#   PR A ("evaluator rebuild"): `_cache_evaluator!` rebuilds the MOI.Nonlinear evaluator
#       (AD tapes, Jacobian/Hessian sparsity, coloring) on EVERY differentiation call,
#       even when the model structure is unchanged. Measured here as `rev_repeat`: a
#       second reverse_differentiate! on the SAME solve — under stock DiffOpt it pays
#       the full rebuild again; with PR A it reuses `model.cache`.
#
#   PR B ("diff-model rebuild"): `MOI.optimize!` unconditionally nulls `Optimizer.diff`,
#       so the first backward after every re-solve re-instantiates the diff model and
#       re-runs a full `MOI.copy_to`. Measured here as `rev_after_resolve`: a
#       Parameter-only re-solve followed by one reverse — with PR B's opt-in
#       `DiffOpt.PreserveDiffModel()` the diff model survives and only the
#       point-dependent state (parameter values + x/y/s) is refreshed.
#
# Run against a DiffOpt checkout (master, or either PR branch):
#
#   julia --project=<env-with-DiffOpt-devved> bench_persistent_evaluator.jl [P ...]
#
# where P is the parameter/variable count (default: 10 100 1000). The script
# auto-detects `DiffOpt.PreserveDiffModel` and adds the preserve column when present.
# Compare the tables across checkouts: the PR-A delta is `rev_repeat` (stock vs PR A
# branch); the PR-B delta is `rev_after_resolve` stock-vs-preserve on the PR B branch.

using JuMP
import DiffOpt
import Ipopt
import MathOptInterface as MOI
using Statistics

const REPS = 7

function build(P; preserve::Bool)
    model = DiffOpt.nonlinear_diff_model(Ipopt.Optimizer)
    set_silent(model)
    @variable(model, p[i = 1:P] in MOI.Parameter(1.0 + 0.2 * i / P))
    @variable(model, x[1:P] >= 0)
    @constraint(model, [i = 1:P], x[i] * (1 + 0.1 * sin(p[i])) >= p[i])
    @constraint(model, sum(x) >= 0.4 * P)
    @objective(
        model,
        Min,
        sum((x[i] - p[i])^2 for i in 1:P) + 0.01 * sum(x[i]^4 for i in 1:P)
    )
    if preserve
        MOI.set(model, DiffOpt.PreserveDiffModel(), true)
    end
    return model, p, x
end

function seed!(model, x, k)
    DiffOpt.empty_input_sensitivities!(model)
    for i in eachindex(x)
        DiffOpt.set_reverse_variable(model, x[i], sin(0.7 * i + k))
    end
    return
end

"""
Median seconds + allocated MiB of `f()` over REPS runs (after one warmup pair).
`setup()` runs untimed before every measurement — the solve lives there, so only the
differentiation call is measured.
"""
function timed(setup, f)
    setup()
    f()
    ts = Float64[]
    bs = Float64[]
    for _ in 1:REPS
        setup()
        GC.gc()
        stats = @timed f()
        push!(ts, stats.time)
        push!(bs, stats.bytes / 2^20)
    end
    return median(ts), median(bs)
end

function bench(P; preserve::Bool)
    model, p, x = build(P; preserve)
    optimize!(model)
    k = Ref(0)
    # first backward after a Parameter-only re-solve (solve itself untimed): pays the
    # diff-model instantiate + copy_to (+ evaluator build) under stock; only the
    # point refresh under PreserveDiffModel
    rev_after_resolve, alloc_resolve = timed(
        () -> begin
            k[] += 1
            for i in 1:P
                set_parameter_value(p[i], 1.0 + 0.2 * i / P + 0.1 * sin(k[]))
            end
            optimize!(model)
            seed!(model, x, k[])
        end,
        () -> DiffOpt.reverse_differentiate!(model),
    )
    # repeated backward on the SAME solve: under stock this still rebuilds the
    # evaluator every time (PR A's target); the diff model itself is reused
    rev_repeat, alloc_repeat = timed(
        () -> begin
            k[] += 1
            seed!(model, x, k[])
        end,
        () -> DiffOpt.reverse_differentiate!(model),
    )
    return rev_after_resolve, alloc_resolve, rev_repeat, alloc_repeat
end

function main()
    Ps = isempty(ARGS) ? [10, 100, 1000] : parse.(Int, ARGS)
    has_preserve = isdefined(DiffOpt, :PreserveDiffModel)
    println("DiffOpt at: ", pkgdir(DiffOpt))
    println("PreserveDiffModel available: ", has_preserve)
    println()
    header = "P      | rev_after_resolve      | rev_repeat"
    has_preserve && (header *= "             | rev_after_resolve (preserve)")
    println(header)
    println("-"^length(header))
    for P in Ps
        r1, a1, r2, a2 = bench(P; preserve = false)
        line = rpad(P, 6) * " | " *
               rpad(string(round(r1 * 1000; digits = 2), " ms  ",
                   round(a1; digits = 1), " MiB"), 22) * " | " *
               rpad(string(round(r2 * 1000; digits = 2), " ms  ",
                   round(a2; digits = 1), " MiB"), 22)
        if has_preserve
            r1p, a1p, _, _ = bench(P; preserve = true)
            line *= " | " * string(round(r1p * 1000; digits = 2), " ms  ",
                round(a1p; digits = 1), " MiB")
        end
        println(line)
    end
    return
end

main()

Context

Same residual-removal series as #369 (single adjoint solve for the nonlinear reverse) and the
evaluator-cache-reuse PR (the two compose: preserving diff skips the copy, reusing the cache
skips the tape rebuild). In a downstream differentiable-ACOPF training loop the combination was
validated bitwise against the rebuild path across a case14→case2000 ladder and yields ×12.6
cumulative per-backward together with #369 at 2000-bus scale.

🤖 Generated with Claude Code

…ly re-solves

MOI.optimize! unconditionally discards Optimizer.diff, so the first
differentiation after every re-solve re-instantiates the diff model and
re-copies the whole problem, even when the only change since the last
instantiation was MOI.Parameter values. Add an opt-in
PreserveDiffModel optimizer attribute (default false, stock behavior
unchanged): when enabled and only parameter values changed - every
structural edit already resets diff to nothing, so a surviving diff is
structurally in sync by construction - optimize! keeps the diff model
and refreshes its point-dependent state instead: input sensitivities
are cleared and the new solution is re-copied after each solve, and
parameter values (their ConstraintSet and their slot in the primal
point) are re-written at the next differentiation call. Results are
identical to the full-rebuild path; a solve that ends in a
non-differentiable status falls back to discarding diff.

Only differentiation models that natively support MOI.Parameter can be
preserved (a ParametricOptInterface layer inside diff would not
propagate value updates); models opt in through
_diff_model_supports_preserve, currently NonLinearProgram.Model.
@codecov

codecov Bot commented Jul 5, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 95.34884% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.92%. Comparing base (f3fdbd8) to head (58a707f).

Files with missing lines Patch % Lines
src/moi_wrapper.jl 94.28% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##           master     #371      +/-   ##
==========================================
+ Coverage   91.88%   91.92%   +0.04%     
==========================================
  Files          17       17              
  Lines        2722     2763      +41     
==========================================
+ Hits         2501     2540      +39     
- Misses        221      223       +2     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@yeabbratz

Copy link
Copy Markdown
Author

Updated for review: applied JuliaFormatter v2 at margin = 80 to green format-check, and added two tests — the non-differentiable-status fallback in MOI.optimize! and the parameter ConstraintSet getter — closing the 6 lines codecov flagged.

Note: the reformatted files are all under docs/src/examples/ and are pre-existing formatting drift on master under JuliaFormatter 2.10 (master's own format-check currently fails for the same reason); this PR's own code needed no reformatting.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Development

Successfully merging this pull request may close these issues.

1 participant