Activity

1. Jammy2211 commented on Jul 9, 2026

   

   Jammy2211

   on Jul 9, 2026

   CollaboratorAuthor

   More actions

   Phase A checkpoint — measured shares, go/no-go, and the Delaunay verdict

   All numbers are post-#368 main (89a5fbd), local CPU fp64, HST production configs (MGE-60 lens light; rect 39×39 RectangularAdaptImage+Constant, n=1581; Delaunay 1500-vertex Hilbert+ConstantSplit, n=1560). Probes ran on the real extracted inversion matrices, not synthetic systems. Working dir: autolens_profiling/scratch/nnls_speedup/ (standalone, untracked).

   Per-eval breakdown (likelihood_breakdown, CPU)

   Step	Rect	Delaunay
   Curvature matrix (F)	3.01s (50.6%)	2.99s (51.2%)
   Inversion setup (relocate→map→PSF-convolve)	1.58s (26.5%)	1.76s (30.2%)
   NNLS reconstruction	1.21s (20.3%)	0.97s (16.6%)
   Everything else	~2.5%	~2%

   F's dominance is CPU-only: it is one 15361×n GEMM (~4ms on A100 fp64). The NNLS solve is sequential — it cannot be bought back with GPU parallelism — so its A100 share of the rect eval will be much larger than 20% (plausibly the dominant block of the 52.7ms/32.9ms rect rows; RAL re-profile will confirm).

   PDIP probe on the real matrices (probe_pdip.py)

   cond(Q_raw) ≈ 4.1e10 both meshes (Jacobi rescale does not reduce it: 1.7e11). Per iteration = one fresh n² Cholesky (factorize_kkt) + 2 triangular-solve pairs; iteration cost fully explains solve time (rect: 21 × 64.5ms; Delaunay: 19 × 59.8ms).

   Variant	Rect iters	Rect speedup	Delaunay iters	Delaunay speedup	rel Δobjective
   baseline (tol=n·eps·5e3 ≈ 1.7e-9)	21	1.00×	19	1.00×	—
   tol 1e-8	20	0.95×*	18	1.04×	≤ 6e-15
   tol 1e-6	18	1.06×	15	1.18×	≤ 3.8e-13
   tol 1e-4	15	1.28×	13	1.32×	≤ 2.1e-10

   * timing noise; iteration counts are the signal.

   Because the log-det terms of the evidence are reconstruction-independent, Δlog_evidence from a looser solve equals the Δ of the NNLS quadratic objective — i.e. ≤ ~1e-8 at tol 1e-6 and ~1e-6 at tol 1e-4 for these systems. Both are far below sampler sensitivity; the Phase C parity gate will verify end-to-end.

   vmap (8 lanes, real noise-covariance perturbations of q): confirmed max-lane effect — rect lanes spread 18–22 iterations and the whole batch pays 22. Tolerance loosening shifts the whole spread down (~uniform −3), so the saving survives batching. Per-lane early exit is not structurally available under vmap'd while_loop.

   Note: the rect fit's saved NumPy-path reconstruction (fnnls_cholesky) differs from PDIP's by up to 0.27 (rel max) with matching objectives — active-set degeneracy of the near-empty rect outskirts, not a solver bug; Delaunay agrees to 1.7e-5.

   Verdict

   GO — scoped. Phase B = expose nnls_solver_tol and nnls_max_iter as config keys (defaults exactly reproduce current behavior: tol=n·eps·5e3, max_iter=50), with the measured recommendation that searches run at 1e-6 (rect −14%, Delaunay −21% of solve; ~5–15% of a whole A100 rect eval if NNLS's GPU share is as expected). jaxnnls does not expose these parameters, so the loop is vendored into autoarray/util/ (reusing jaxnnls.pdip.pdip_pc_step/initialize as imports, only the while_loop driver re-implemented — smallest possible fork surface) unless upstream takes a params PR quickly.

   Parked (documented, pending A100 re-profile): cheaper PDIP iterations (PCG-on-fixed-factorization; the varying z/s diagonal degrades any fixed preconditioner — research-grade risk), alternative first-order solvers (cond 1e10 makes them iteration-hungry), fp32/mixed-precision PDIP (cond 1e10 vs fp32 eps).

   The Delaunay question — is the mesh/scipy machinery still the target?

   No. Post-#368 the qhull callback is ~9ms host-side (was ~53ms) and the Delaunay per-eval budget is now dominated by the same blocks as rect: F on CPU, NNLS (+ inversion setup) on GPU. Two Delaunay-specific residuals remain:

   1. the callback still runs vmap_method="sequential", so under batched sampling it serializes per lane — this, not single-eval cost, is where Delaunay still pays vs rect (expected A100 vmap lane ~40–50ms vs rect 32.9ms; RAL re-profile will quantify). The fix (batch-parallel/threaded callback, plan item 3 of the 2026-07-09 review) stays parked until those numbers exist;
   2. nothing else — mapping/interpolation moved to JAX in refactor: qhull-only Delaunay callback, exact JAX visibility-walk point location #368 and regularization H is ~2%.

   So the right next Delaunay target = the same NNLS levers as rect (this task), then the vmap-callback question after the A100 re-profile.
2. Jammy2211 commented on Jul 9, 2026

   

   Jammy2211

   on Jul 9, 2026

   CollaboratorAuthor

   More actions

   Interim update — A100 run queued, Phase B implemented + locally validated

   A100 (RAL restored): all five HPC checkouts pulled to current main (PyAutoArray at the #368 merge), probe package + real extracted systems rsynced over, and SLURM job 330046 (submit_probe_nnls_a100) submitted: PDIP probe on both real systems + A100 re-runs of the two likelihood_breakdown scripts (also clears the outstanding post-#368 A100 re-profile from the Delaunay review). Currently pending — euclid-ral-gpu-[1-2] are still down/drained; it will start when the GPU nodes return.

   Phase B (in the feature/nnls-solver-optimization worktree, unshipped):

   • autoarray/util/jax_nnls.py — vendored while_loop driver with solver_tol / max_iter knobs; reuses all jaxnnls building blocks + the relaxed-KKT custom-vjp backward; knobs are static closure params behind an lru_cache so JIT tracing caches hit.
   • inversion_util.reconstruction_positive_only_from reads general.yaml -> inversion -> nnls_solver_tol / nnls_max_iter (try/except fallbacks, same pattern as the existing nnls keys); packaged defaults reproduce upstream exactly.
   • Validated on the real rect+Delaunay systems: defaults bit-identical to upstream jaxnnls (solutions and gradients), tol 1e-6 saves 3-4 iterations with rel drift ~1e-6, knob-cache stable. New NumPy-only unit tests pass (module-level JAX-import guard, packaged config defaults, numpy-path positive-only solve); full suite run in progress.

   Ship gate: A100 numbers land in the PR benchmark table first.
3. Jammy2211 commented on Jul 9, 2026

   

   Jammy2211

   on Jul 9, 2026

   CollaboratorAuthor

   More actions

   Laptop-GPU probe (RTX 2060 Max-Q 6GB, WSL2) — consumer feasibility datapoint

   Full rect pixelization+MGE HST likelihood, real production configs, post-#368:

   	Laptop CPU fp64	RTX 2060 fp64	RTX 2060 mp	GPU/CPU
   Curvature matrix (F)	3.01s	0.397s	0.395s	7.6x
   Inversion setup	1.58s	0.125s	0.077s	12.6x
   NNLS reconstruction	1.21s	0.298s	0.294s	4.1x
   whole eval	5.96s	0.881s	0.820s	6.8x

   NNLS solver probe on the real systems (identical iteration counts to CPU, as expected): rect solve 336ms (3.4x CPU), Delaunay 329ms (2.7x); vmap-8 amortizes rect to 267ms/lane. Tolerance knob worth more on GPU for Delaunay (1e-6 -> 1.43x, 1e-4 -> 1.60x solve speedup).

   Takeaways

   • A consumer laptop GPU runs the full production likelihood ~7x faster than the same laptop's CPU — sub-second/eval — so laptop-GPU modeling is genuinely feasible for patient users.
   • NNLS share grows 20% (CPU) -> 34% (2060) — confirms the solver as the right accelerator-side target; on A100 it should dominate outright (F's GEMM gets ~100x cheaper there, the sequential PDIP chain doesn't).
   • Mixed precision is NOT the consumer rescue: F and the NNLS solve stay fp64 (cond(Q)~1e10), together 84% of the GPU eval. Log-evidence parity passed in both modes.
   • UX trap for laptop users: jax silently picks CPU unless BOTH JAX_PLATFORM_NAME and JAX_PLATFORMS are unset/overridden — worth a docs note.
   • mp run OOM'd at XLA_PYTHON_CLIENT_MEM_FRACTION=0.5 on 6GB (eager FFT mapping-matrix convolution, 398MB alloc); fine at 0.85.

   A100 job 330046 still queued (euclid-ral-gpu-[1-2] down, "Not responding").
4. Jammy2211 commented on Jul 9, 2026

   

   Jammy2211

   on Jul 9, 2026

   CollaboratorAuthor

   More actions

   Deeper-lever experiment: Gondzio multiple centrality corrections — tried, near-wash, parked

   Anatomy first (RTX 2060, real rect system): the fresh Cholesky is the PDIP iteration (~14ms step; one extra KKT solve ~2ms) — so reusing the factorization for extra centrality-corrector solves (Gondzio MCC) had a real ceiling. Implemented in scratch/nnls_speedup/probe_gondzio.py (composite directions via KKT linearity; branch-free acceptance; correct: solutions match baseline to ~1e-8, all variants converge).

   Result: the iteration reduction is real, the wall-time win is not.

   (RTX 2060)	rect iters	rect single	rect vmap-8/lane	delaunay iters	delaunay single	delaunay vmap-8/lane
   baseline	21	298ms	284ms	19	258ms	211ms
   Gondzio K=2	18	1.09x	288ms (wash)	15	0.90x	222ms (worse)
   tol 1e-6 alone	18	1.17x	247ms	15	1.29x	202ms
   K=2 + tol 1e-6	16	1.23x	275ms	13	1.34x	196ms

   The extra solves + line searches eat the saved factorizations, and under vmap (the production mode) every lane pays the corrector cost every iteration while acceptance diverges across lanes — net zero to negative. Stacking on the tolerance knob adds only ~3-6% over the knob alone.

   Verdict: ship the tolerance/max-iter knobs (Phase B, done); park Gondzio with these numbers. Solver dead-end ledger for this problem class (cond(Q)1e10, n1600): warm starts (hurt), fp32/iterative refinement (cond x eps32 ~ 600, diverges), Gondzio MCC (iterations -15% but wall-time wash). One-command re-bench on A100 if its factorization:solve ratio turns out very different: python3 scratch/nnls_speedup/probe_gondzio.py --mesh rectangular.
5. Jammy2211 commented on Jul 9, 2026

   

   Jammy2211

   on Jul 9, 2026

   CollaboratorAuthor

   More actions

   Outcome: findings shipped, implementation parked (user decision 2026-07-09)

   The measured optimization (solver tolerance/iteration knobs, ~15-20% of the NNLS solve) is not being merged — the risk of a new solver code path was judged not worth the size of the win. The validated implementation is preserved on branch feature/nnls-solver-optimization (d8a1c84: vendored driver + config keys + tests, defaults bit-identical to upstream jaxnnls in solutions and gradients on the real production systems, full suite 876-passed) should the trade-off change — e.g. if the A100 numbers show a larger production share than expected, or the follow-up experiment (below) promotes a new solver anyway.

   What this task established (the durable knowledge)

   Per-eval breakdowns, production HST pixelization+MGE configs, post-#368:

   step	laptop CPU fp64	RTX 2060 fp64	RTX 2060 mp
   Curvature matrix (F)	3.01s (51%)	0.397s (45%)	0.395s (48%)
   Inversion setup	1.58s (27%)	0.125s (14%)	0.077s (9%)
   NNLS reconstruction	1.21s (20%)	0.298s (34%)	0.294s (36%)
   whole eval	5.96s	0.881s	0.820s

   • Consumer-laptop modeling is feasible: 6.8x over the same laptop's CPU, sub-second/eval, evidence parity passing. Mixed precision is not the consumer rescue (F + NNLS stay fp64 at cond(Q)~1e10 = 84% of the eval).
   • NNLS share grows on accelerators (20% CPU -> 34% GPU) and the solve is latency-bound, not flops-bound (2060 does an iteration in ~14ms despite 1/100th the A100's fp64 flops).
   • Delaunay verdict: post-refactor: qhull-only Delaunay callback, exact JAX visibility-walk point location #368 the mesh/scipy machinery is no longer the target — remaining Delaunay-specific cost is the callback's vmap serialization only (quantified by the queued A100 job).

   The interior-point ledger is complete — measured on the real systems:

   lever	result
   tolerance 1e-6 / 1e-4	-3 to -6 of ~20 iters, Δlog_ev ~1e-8/1e-6 (parked branch exposes it)
   warm start from unconstrained solve	HURTS (un-centers; cond ~1e10-1e16)
   fp32 / iterative refinement	diverges (cond x eps32 ~ 600; KKT worsens near convergence)
   Gondzio multiple centrality corrections	iters -15% but wall-time wash; worse under vmap
   per-iteration anatomy	the fresh Cholesky IS the iteration (~14ms vs ~2ms/extra solve, RTX 2060)

   Follow-up

   Different-algorithm candidates (block principal pivoting; ADMM with fixed factorization) are filed as PyAutoMind/feature/autolens_profiling/nnls_bpp_admm_solver_experiment.md with hard promotion gates (>=2x at evidence parity, vmap-clean, PDIP fallback). Scratch harness + real extracted systems in autolens_profiling/scratch/nnls_speedup/ are the shared testbed.

   Pending addendum

   A100 job 330046 remains queued (euclid-ral-gpu-[1-2] down). When it runs it delivers: A100 per-step shares (the true production NNLS share), the post-#368 Delaunay re-profile, and on-A100 PDIP probe numbers. Results will be posted here; suggest closing this issue after that addendum.
6. Jammy2211 commented on Jul 9, 2026

   

   Jammy2211

   on Jul 9, 2026

   CollaboratorAuthor

   More actions

   Ledger update from #370: BPP and ADMM both measured and negative (BPP machine-exact but 36-53 factorizations vs PDIP's 19-21; ADMM plateaus at rel dObj ~3e-6, four orders short of the evidence envelope). The dead-end ledger is now: warm starts, fp32/IR, Gondzio, BPP, ADMM. PDIP ~20 iters x Cholesky is near-optimal for this problem class — the remaining routes to a faster positive-only solve are problem-level (fewer source pixels, conditioning via regularization design, or relaxing positivity), not algorithmic.
7. Jammy2211

   added a commit that references this issue on Jul 9, 2026

   feat: NNLS solver knobs via Settings (nnls_solver_tol / nnls_max_iter…

   ...

   8771f29
8. Jammy2211 commented on Jul 9, 2026

   

   Jammy2211

   on Jul 9, 2026

   CollaboratorAuthor

   More actions

   User decision update: the knobs are now shipping after all — reworked as per-fit Settings options (default off) per request, PR #371. The A100 addendum (job 330046) will be posted here when the GPU nodes return. Closing; the durable record is autolens_profiling/results/notes/nnls_solver_ledger.md.
9. Jammy2211

   closed this as completedon Jul 9, 2026