Bins-sharded multi-device fits (--nDevices) - #154
Open
davidwalter2 wants to merge 12 commits into
Open
Conversation
…vice New --minimizerMethod tf-trust-exact: a TensorFlow port of scipy's trust-exact (More-Sorensen) where the Hessian, the H + lambda*I builds, the Cholesky factorizations and the triangular solves stay on the TF device, so the n x n matrix never round-trips to LAPACK per lambda trial. The python outer loop mirrors scipy's _minimize_trust_region and plugs into the existing callback / early-stopping / restart / result plumbing unchanged, including the preconditioner (applied at the same internal-coordinate boundary as for the scipy methods). Differences from a line-for-line port, all deliberate: - No custom op for the potrf failure index: tf.linalg.cholesky signals a non-PD input by filling the factor with NaNs on every backend, and a failed factorization already proves lambda_current is a valid lower bound, so the safeguarded bracket update converges without the index-accelerated bound (a few extra cheap factorizations at worst). - The Hessian closure runs only on accepted steps; scipy's subproblem constructor consumes Hessian norms and therefore pays a Hessian per proposal. - lambda_new is clamped to >= 0 (the interior-case Newton correction can otherwise push the bracket negative) and the matrix norms are computed as matrix norms -- tf.norm's default axis=None flattens the tensor and its max|H_ij| can sit below |lambda_min|, silently invalidating the lambda_ub bracket. - The rare hard-case refinement downloads the factor once and reuses scipy's host-side smallest-singular-value estimate. Validated against scipy on random PD/indefinite subproblems (worst-case model reduction 96.7% of the exact optimum vs scipy's 96.0%, ~1.6 factorizations per solve), on Rosenbrock and an ill-conditioned quadratic (identical minima to ~1e-8), and through a full Fitter fit on the test tensor against the scipy path. Without a visible GPU a warning points users back to scipy trust-exact: TF's CPU Cholesky kernel is single-threaded Eigen and measured ~7x slower at n=2000. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY
The matrix-free counterpart of tf-trust-exact: the same truncated-CG trust-region subproblem as scipy's trust-ncg (the practical stand-in for trust-krylov, which solves the identical subproblem via GLTR), but the entire CG iteration runs inside a single tf.function while_loop. One solve costs one graph dispatch instead of one python round trip -- x assignment, numpy conversion both ways, forced device sync -- per Hessian-vector product, which is what the scipy callback path pays and what dominates when individual HVPs are fast. Exact improvements over scipy's loop, with identical iterates: - Hz is carried through the CG recurrence, so the boundary and negative-curvature exits price their steps from dot products where scipy spends two extra HVPs; - the model value of the returned step falls out of the same bookkeeping and is handed to the outer loop, which otherwise needs one more HVP per proposal. The preconditioner gains tf_transforms(): graph versions of the block T/T^T applications (gather -> dense matvec -> scatter), because for an HVP subproblem the reparameterisation runs inside the compiled loop, once per CG iteration, where the numpy path cannot. The subproblem also re-pins the fitter's parameter state before each solve, since the outer loop's objective evaluations at proposed points move it in between. Validated: steps agree with scipy's CGSteihaugSubproblem to float precision across PD/indefinite models and radii (same algorithm, so bitwise-deterministic, unlike the nearly-exact solver); Rosenbrock matches scipy trust-ncg; full Fitter fits match the trust-krylov reference with and without preconditioning. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY
The actual trust-krylov algorithm (Gould-Lucidi-Roma-Toint, the method behind trlib/scipy trust-krylov), completing the native set. Where Steihaug-CG stops at the first boundary crossing, GLTR keeps expanding the Krylov subspace and returns the step that is optimal within it: the Lanczos basis tridiagonalizes H, the projected subproblem min gamma0*e1.h + h.T_k h/2, ||h|| <= Delta is solved per iteration by eigendecomposition plus a safeguarded secular solve (hard case included), and the full-space residual comes free as gamma_k |h_k|. Division of labour: HVPs, the Lanczos recurrence and full CGS2 reorthogonalization run on the TF device through one compiled step function against a fixed-shape [kmax, n] basis buffer (zero rows make masking unnecessary, and fixed shapes mean no retracing); the k x k tridiagonal solves run in numpy where they cost microseconds against HVPs costing milliseconds. Boundary-phase residuals use a 10x tighter tolerance than the interior forcing sequence, as in trlib. The Krylov data is radius-independent, so the outer loop's rejected-step path -- re-solving at a shrunken radius -- reuses it and usually costs zero new HVPs, verified by a test counting them. Validated: subproblem reaches >= 99.9% of the exact-optimum model reduction on problems the subspace exhausts (PD and indefinite, all radii) where truncated CG has no such guarantee; explicit hard-case KKT check; Rosenbrock matches scipy; full Fitter fits match the trust-krylov reference with and without preconditioning (32 tests). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY
Fixes the two issues the first GPU run surfaced: - The interior branch estimated the smallest singular value with the host-side Cline et al. recurrence, downloading the full factor per call -- ~1 s per outer iteration at n=4000 over PCIe, erasing the cheap on-device factorizations. Replaced by inverse iteration on L L^T: a handful of O(n^2) triangular solves on device, with only two scalars and one n-vector crossing to the host. Inverse iteration converges fastest exactly in the near-singular regime the hard case lives in; far from singularity it only gives an upper bound on sigma_min, which every use is safe against (lambda_lb just gets looser, and acceptance is guarded, below). - The hard-case acceptance trusted the s_min estimate outright, so an inaccurate estimate could accept a corrected step that raises the model; the outer loop then rejects it, which showed up as a doubled outer iteration count under GPU rounding at n=2000. Both hard-case exits now verify the candidate against the model value (one matvec) and fall through to the normal lambda updates when it fails. Subproblem quality scan is unchanged (worst case 96.6% of the exact optimum over 60 PD/indefinite cases vs scipy's 96.0%, ~1.5 factorizations per solve); new unit test pins the device estimator to the true sigma_min in the near-singular regime. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY
A proposal whose objective overflows to NaN produced a NaN rho, and IEEE comparisons on NaN are all False: the radius was neither shrunk nor the step accepted, so the outer loop spun at a fixed radius until early stopping gave up far from the minimum. First observed on the first preconditioned run at scale, where an internal-coordinate step of norm 1 is an enormous physical step whose exponentials overflow. A non-finite rho now counts as a hard rejection (radius shrinks), with a regression test whose objective is NaN outside a box around the minimum. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY
The NLL is a sum over bins plus parameter-level terms, and every large tensor in the fit (logk, norm, nobs, sumw/sumw2, beta) is bins-proportional, so sharding the bins axis over N devices splits both memory and compute while only the parameter vector out and one [nparams] partial per shard back ever cross devices. The primary motivation is memory -- models whose logk exceeds a single GPU now load host-side (FitInputData(host_memory=True)) and only per-device slices are materialized on GPUs -- with compute scaling as a bonus in exactly that regime (each shard stays large enough to saturate its device). Two structural rules, established on a 4-GPU proof of concept (worst-case numbers in rabbit/sharding.py): one GradientTape per shard differentiating w.r.t. a device-local copy of x, because a single tape over the combined loss misplaces the backward ops and drags every shard's logk across the bus per call (135 ms vs 4.5 ms at 4 shards); and per-shard jit compilation with a plain-graph combiner, since XLA clusters are single-device. No duplication of the likelihood mathematics: each shard is a duck-typed evaluator (rabbit/sharding.py) onto which the Fitter binds its own unbound methods, so both interpolation forms, both systematic types and the BinByBinStat machinery (instantiated per shard from a sliced indata view; kstat/beta are per-bin) are shared with the single-device path by construction. The minimizers, the preconditioner and the outer fit loop are untouched -- they consume the same loss_val/loss_val_grad/loss_val_grad_hessp/loss_val_grad_hess interface, now assembled from per-shard partials. The dense Hessian shards too, which also lets the preconditioner's reference matrix build fit in memory where a single device could not. Not supported (rejected at construction or call time): sparse mode, --covarianceFit, regularizers' full=True path, fwdrev HVP (falls back to revrev), profile=False Hessians. Tests run the sharded machinery on CPU replicas (select_devices falls back when no GPU is visible) and hold it to near-bitwise agreement with the single-device path: loss/grad/HVP/Hessian at a randomized point, profiled beta, full fits via scipy and native minimizers, and deepcopy (toys) rebuilding the shard machinery. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY (cherry picked from commit 295b8c1)
By default TF creates a CUDA context on -- and reserves the memory of -- every visible GPU, so a single-device fit on a shared 4-GPU machine blocked all four while computing on one, and always landed on GPU:0 even when another process was already using it. rabbit_fit now restricts the visible devices to --nDevices before the TF runtime initializes, choosing the least-occupied GPUs by current memory use (nvidia-smi query, first-N fallback); --devices gives explicit indices for full control. Verified on a 4-GPU node: a single-device fit reserves exactly one GPU leaving the others untouched (4 MiB, no CUDA context), and with a process occupying GPU:0 the fit lands on GPU:1. Under slurm with --gres this is all transparent, since CUDA_VISIBLE_DEVICES already hides other jobs' GPUs. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY (cherry picked from commit aa85a4b)
The device layout is fixed at initialization and never changes during a fit, so the single-device / multi-device split is a static choice best expressed as a class: rabbit.sharding.MultiDeviceFitter subclasses Fitter and owns all sharding -- shard construction, the per-device tf.function assembly, the set_nobs propagation -- while the base Fitter carries none of it. make_fitter() picks the class from --nDevices, so call sites stay generic. The base-class diff reduces to making the deepcopy strip-list an extensible class attribute: since __deepcopy__ already re-instantiates via type(self) and ends in _make_tf_functions(), and init_fit_parms already ends in _make_tf_functions(), the subclass rebuilds its shards by overriding that one method -- no hooks, no if-shards branches, and toys preserve the subclass automatically. The likelihood mathematics remains shared by construction: shard evaluators borrow the base class's own unbound methods, so this class is orchestration only. Behavior is unchanged: the full sharded and single-device test suites pass identically (101 tests). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY (cherry picked from commit 30bbfa3)
The factory introduced by the MultiDeviceFitter refactor accepted only do_blinding and dropped globalImpactsFromJVP, which rabbit_fit passes -- a TypeError on every rabbit_fit invocation. make_fitter and MultiDeviceFitter.__init__ now forward **kwargs verbatim so the wrapper cannot drift from Fitter.__init__'s signature, the tests construct through the factory with rabbit_fit's exact kwargs, and rabbit_fit runs end-to-end on both fitter classes. Reported from a parallel session reviewing the branch. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY (cherry picked from commit a0076df)
A penalty on the parameters alone -- a bound on polynomial coefficients, a smoothness prior -- does not need the predicted yields; one that shapes the prediction does. Nothing expressed that, so the fitter had to assume the worst for every regularizer. The distinction is not cosmetic. It decides whether the likelihood has to build full=True yields on every call, and it is what makes a penalty usable at all in a bins-sharded multi-device fit, where no single device holds the full yield vector. The default is True because it is the safe answer: a subclass that wrongly inherited False would be handed None and fail loudly rather than silently penalise the wrong thing.
tape.jacobian vectorises over every parameter at once, holding one [nbins, 9] intermediate each: [nparams, nbins, 9] doubles, 43 GB for a 6538-parameter 92144-bin model and ~70 GB per device once sharded. No GPU has that, so the reference Hessian could not be formed and the fit fell back to running unpreconditioned -- easy to miss, because the fit then proceeds normally and only the wall clock gives it away. Assemble it a batch of columns at a time from HVPs instead. tf.vectorized_map is a pfor of width k rather than the while_loop tape.jacobian's non-pfor path builds, so XLA compiles it without the unrolling that made chunking unusable, and memory scales with k instead of nparams: k = 256 costs a few GB where vectorising over all parameters costs 43. The batch halves itself on ResourceExhaustedError, so --hvpBatch is an upper bound rather than a number that has to be right. These are the same HVPs the minimiser already uses, so on a sharded fit they stay distributed. Also let a regularizer declare that its penalty does not read the predicted yields. _compute_nll_components forced full=True yields on every likelihood call whenever any regularizer was attached; now that cost is only paid for penalties that actually use them, and each regularizer is handed None rather than a truncated vector when it declared it does not. Drop three experiments that did not work out: - preconditioner refresh on curvature drift (--preconditionRefresh and friends). 1.5-2.5x faster but landed 3-48 above the minimum: a rebuild rescales the gradient scipy's gtol test is applied to, so an unconverged point can look converged. It needs a convergence test in a fixed metric before it is worth having, and it was off by default meanwhile. - --hessianParallelIterations. Fixed the memory but replaced the vectorised pfor with a while_loop XLA unrolls, so compilation never finished. - --preconditionOnHost. Kept the fast pfor path but reached 161 GB of RSS on a 187 GB node before producing anything. (cherry picked from commit ea172ba)
Regularizers, reduced_nll, global impacts, toys and the full-NLL Hessian were all inherited from the base class and computed over all bins on one device, which is what the sharding exists to avoid. A penalty on the parameters alone has no reason to be unsupported here -- every device holds the whole parameter vector -- so it is evaluated in gnll_local, the only global term, which gvg and gvgp differentiate, giving the gradient and HVP for free. gnll_local previously summed the constraint and external-likelihood pieces only, so a penalty was accepted and then contributed nothing: the fit minimised the unregularized likelihood while the command line said otherwise, and its loss compared favourably with a single-device fit's for lack of a positive term. Penalties that read the yields are now refused, in arm_regularizers, which runs when regularizers exist and before any loss is evaluated. The graphs are rebuilt there, since __init__ traces them before --regularizer has been parsed. reduced_nll is the same quantity the sharded loss_val computes, so it uses that. The remaining four all run after the minimiser, where an allocation failure would take the completed result down with it; they now refuse at the point of use, before the fit starts, and --fullNll says why it cannot work rather than running out of memory discovering it. Drop the sharded jacobian Hessian and the shard/global graph entries that fed it. It needed ~70 GB per device, so it could not run at all, and it is what the postfit covariance called -- the failure would have landed at the end of the fit. The batched-HVP Hessian replaces it unconditionally. (cherry picked from commit 76b5f3d)
davidwalter2
force-pushed
the
260828_multiDeviceFitter
branch
from
August 30, 2026 20:23
a04bf4e to
aeaa83d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Depends on #153 (native TensorFlow trust-region minimizers). This branch is
based on
260826_tfMinimizer, so until #153 merges the diff below also showsits commits; it reduces to the seven listed at the end once #153 lands. Review
from "Add multi-device fits" onward.
What this adds
--nDevices Nsplits the likelihood's bins across N devices. Each shard is aduck-typed evaluator exposing the attributes the existing
Fittermethodsread, and the Fitter's own unbound methods are bound onto it, so both
interpolation forms, both systematic types and the bin-by-bin-stat machinery
are shared with the single-device path by construction rather than duplicated.
Shard graphs are jit-compiled individually (XLA clusters are single-device) and
the thin combiner stays a plain graph.
--devicespicks the GPUs explicitly;otherwise the least-occupied ones are chosen, so a fit started on a shared node
lands next to nobody.
On a 92144-bin, 6538-parameter W+Z fit across 4 V100s this preconditions all 71
parameter blocks in under 7 minutes and then iterates in seconds.
The dense Hessian
Both the preconditioner's reference matrix and the postfit covariance need a
dense
[npar, npar]Hessian. It is assembled a batch of columns at a time fromHessian-vector products, which are the same HVPs the minimiser already
evaluates each iteration, so the work stays distributed across the shards.
tf.vectorized_mapis a pfor of width k, so XLA compiles it without theunrolling that a
while_loopof the same shape produces, and peak memory isone
[k, npar]batch rather than the[npar, nbins, 9]that vectorisingtape.jacobianover every parameter at once would need (43 GB for this model,~70 GB per device once sharded).
--hvpBatch(default 256) sets k, and thebatch halves itself on
ResourceExhaustedError, so it is an upper bound ratherthan a number that has to be tuned.
Regularizers
Regularizer.needs_observablesdeclares whether a penalty reads the predictedyields. Penalties that depend only on the parameters are supported here: every
device holds the whole parameter vector, so the penalty is evaluated once in
the global term
gnll_local, whichgvgandgvgpdifferentiate — gradientand HVP come for free. Penalties that read the yields are rejected in
arm_regularizers, which runs when regularizers are present and before anyloss is evaluated; the graphs are rebuilt there, since
__init__traces thembefore
--regularizerhas been parsed.The flag also pays off single-device:
_compute_nll_componentsonly buildsfull=Trueyields when some regularizer actually needs them, and eachregularizer is handed
Nonerather than a truncated vector when it hasdeclared it does not.
Not supported in sharded mode
Checked at construction: sparse tensors,
--covarianceFit.Rejected at the point of use, before the fit starts: regularizers with
needs_observables=True,--fullNll, global impacts (--doImpacts), and toys(
-t > 0). These are global over all bins on one device, which is what thesharding exists to avoid; each raises
NotImplementedErrornaming the flag andthe single-device remedy. The fwdrev HVP falls back to revrev.
Testing
tests/test_sharded_fit.pychecks a 2-device fit against the single-device oneon parameters, uncertainties, NLL and EDM; that a parameter-only penalty shifts
the sharded loss by exactly the amount it shifts the single-device loss, and
reaches the gradient; and that each unsupported postfit step raises rather than
running. 114 tests pass on this branch.
Commits
--nDevicesMultiDeviceFittersubclassmake_fitter🤖 Generated with Claude Code
https://claude.ai/code/session_01HrsrmERvJF2Vafn7wE1B3v