[pastewka]

Why

In the 512³ MI300A scaling runs, the two solves of each Hessian-vector product are about 80% of the run at 16 ranks. They take about 98 CG iterations each, against 47 for the forward solves, because they run at a fixed hessian_tolerance of 1e-3 while the forward solves start at 1e-2. Ten times tighter costs about twice the iterations. Early on that accuracy is wasted: the Steihaug CG that asks for the product solves its subproblem only to min(0.5, sqrt(|g|/|g0|)).

Changes

• Adaptive Hessian tolerance (default). With NuMPI's new hessp_tol (BUG: Scale-invariant Steihaug forcing in tr_newton_bounded; pass its tolerance to hessp; changelog for v0.18.0 NuMPI#105), the trust region passes the Steihaug tolerance to hessp, and the solves use a tenth of it, clipped to [1e-4, 1e-2] (hessian_rtol). So the first subproblems get 1e-2, and the tolerance tightens only as the outer iteration converges. optimizer.trust_region.hessian_tolerance now defaults to null for this; a number still fixes the tolerance. The log shows the tolerance of the last product in each outer iteration (hv-cg=… (n Hv, rtol=…)).
• No warm-start check in cold Hessian solves. They passed warm_start=True even without hessian_warm_start, after zeroing their start vector, so solve_rhs applied the operator once per solve only to reject a zero warm start.
• Thinner CG fallback snapshots. The callback copied the full iterate on every improving iteration (about 2% of a solve at 2 ranks). It now refreshes the snapshot when the residual has halved, so a solve that falls back returns an iterate within a factor two of the best residual.

Depends on

muSpectre/NuMPI#105 and a NuMPI 0.18.0 release, which pyproject.toml now requires. Until that release, CI can't install the dependencies.

Tests

• New: test_hessian_rtol_follows_the_steihaug_tolerance, and test_tr_adapts_the_hessian_tolerance, which checks that the first products are requested at 1e-2, that all stay within the bounds, that the objective goes down, and that a fixed hv_rtol is used as given.
• With NuMPI from #105, the suite passes locally serially (249) and on 2 ranks (243 per rank).

A comparison of the optimization itself (3D 32³, 4 ranks, up to 60 outer iterations) is running: today's defaults, fixed 1e-2, the NuMPI fix alone, and NuMPI fix plus adaptive tolerance. I'll post the results here.

🤖 Generated with Claude Code