Guard tile Cholesky against float32-indefinite Hessians with a pivot floor - #1488
Guard tile Cholesky against float32-indefinite Hessians with a pivot floor#1488elliot-at-liminalnook wants to merge 2 commits into
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
d6ee1c2 to
657d17d
Compare
|
A few notes on how this PR relates to the discussion on #1415, since the thread had converged on a tradeoff this design tries to escape: On the cuSolverDx constraint (@adenzler-nvidia): this PR deliberately does not modify On the performance concern with always-on diagonal regularization: because detection is post-hoc, the healthy-path cost is unchanged — the well-conditioned parity test asserts the factor is bit-identical pre/post-patch, and On matching Just pushed a strengthening (d79f0b4): detection now triggers on finite-but-below-floor pivots, not only non-finite values — rank deficiency can factor to finite garbage without a NaN, and this also covers backends whose indefinite-input behavior is finite garbage rather than NaN. New regression test included where the unguarded factor is fully finite but a pivot sits far below the floor. One question we can't answer without CUDA hardware (validation here is CPU/LLVM only, as flagged in the PR body): what does the cuSolverDx-backed |
…floor The Newton solver Hessian H = M + J'DJ is SPD in exact arithmetic, but on stiff-contact states the float32 assembly error (~||H||*eps32) swamps the smallest eigenvalues and the stored matrix goes numerically indefinite; the unguarded tile Cholesky then square-roots a negative pivot and NaNs propagate through qfrc_constraint into qacc and qpos. Reference MuJoCo is protected by the mindiag pivot floor in mju_cholFactor, which the Warp port dropped. The fast path is untouched: after the existing unguarded factorization, scan the factor diagonal for non-finite values, and only then rebuild the block and refactor element-wise with a scale-relative pivot floor eps32*max|diag|. Applied to both call sites, the blocked factorize-solve (nv > 32) and the single-tile _update_gradient_cholesky (nv <= 32). One deliberate deviation from mju_cholFactor: its floor-then-divide treatment of a failing pivot amplifies consecutive floored pivots quadratically and overflows float32, so floored pivots are instead treated as rank deficient - the pivot becomes sqrt(mindiag) and the rest of the row is zeroed, decoupling the direction so it responds like a spring of stiffness mindiag. Tests build the Hessian from issue google-deepmind#1415 (float64-SPD, float32-indefinite), assert the pre-fix path still NaNs on it, and check both guarded paths stay finite and bounded while well-conditioned matrices keep matching a float64 reference to ~1e-7.
…NaNs Rank deficiency can factor to finite garbage pivots without producing a NaN, and backends may return finite garbage rather than NaN on indefinite input (cuSolverDx behavior on the CUDA path). Replace the non-finite diagonal scan with a health check: every factor pivot must be finite and >= sqrt(mindiag), with mindiag hoisted so the fast path cost is unchanged. Adds a regression test where the unguarded factorization is finite but a pivot sits far below the mju_cholFactor-style floor.
d79f0b4 to
792db87
Compare
|
Update: we got A100 time and can answer our own cuSolverDx question. On sm_80 (warp 1.14, CUDA 12.9), the cuSolverDx-backed tile factorization produces NaNs on indefinite input — the in-test unguarded assertion passes on GPU, so the failure mode matches CPU/LLVM. All 4 tests in |
|
Thanks - I should have written in the comments that cuSolverDx uses NaN as a signaling mechanism. Without going too deep into the code yet things look correct on the surface, but I wonder how much this impacts performance. Given that you got your hands on an A100, are you able to run the benchmarks? Otherwise I can also give generate some blackwell numbers next week. |
|
Benchmarks run on the A100 (80GB PCIe, warp 1.14.0, CUDA 12.9, mujoco 3.10.0),
Rep spread within each config is ≤0.2%, so the deltas are real. The cost is the always-on part — the If ~1% is more than this should cost, two easy directions: (a) check the diagonal only after the final solver iteration rather than every iteration (a NaN'd factor still gets caught before its solve is consumed), or (b) gate the check behind an option so it's opt-in for stiff-contact workloads. Happy to implement either — and Blackwell numbers would still be a welcome cross-check! |
|
We ran into #1415 in an RL training workload (large batched manipulation scenes with stiff contact materials, elliptic cone, impratio=10, solver=newton) and tried this fix. First off, thanks — it does exactly what it promises, the NaNs are gone. But a heads-up on the perf assumption: in our setup the float32 indefiniteness isn't an occasional event, it's chronic. With light bodies and stiff contacts (efc_D up in the 1e5–1e6 range, impratio inflating the friction rows on top) the assembled Hessian fails the health check on basically every solver iteration, so the element-wise rebuild+refactor becomes the steady-state path rather than a rare fallback. Microbenchmark of the nv<=32 single-tile path (8192 worlds, nv=32, RTX 5090, warp 1.15 dev): on healthy matrices the guarded kernel is ~1.05x the unguarded one, which matches the ~1% number reported here. On indefinite matrices it's ~11.6x per launch. End to end, our training throughput dropped ~8x on a pyramidal-cone config and ~30x with elliptic+impratio=10 (roughly 1k env-steps/s where we normally see tens of thousands). I suspect part of the end-to-end gap beyond the raw kernel cost is that chronically floored, rank-deficient factors give poor search directions, so the Newton solver stops early-exiting and burns its full iteration budget. One caveat: we measured this on a backport of the patch to the 3.8.1 solver structure (the lower-triangular split factorize/solve), not this branch as-is. The guard structure is the same though — unguarded fast path, then element-wise rebuild and refactor inside the branch — so I'd expect the repair cost to carry over. Might be worth parallelizing the repair across the block's threads instead of the redundant element-wise loop, or folding the pivot floor into the primary factorization so there's no second pass, at least for workloads where indefiniteness is the common case rather than the exception. For now we've switched these scenes to solver=cg, which avoids the factorization entirely. |
3 participants
Fixes #1415.
The Newton solver builds the constraint Hessian H = M + JᵀDJ and factors it with the tile Cholesky in float32, with no pivot floor. On stiff-contact states the float32 assembly error (~‖H‖·eps32) swamps H's smallest eigenvalues, so a Hessian that is SPD in float64 becomes numerically indefinite in float32 (eigmin -4.28 on the test matrix here). The unguarded factorization then square-roots a negative pivot and the NaNs propagate through qfrc_constraint into qacc and qpos. Reference MuJoCo avoids this via the mindiag pivot floor in mju_cholFactor, which the Warp port dropped.
The design keeps the fast path untouched: the existing unguarded tile factorization runs first, and only if its diagonal comes back non-finite do we rebuild the block and refactor element-wise with a scale-relative pivot floor eps32·max|diag|. Both call sites are guarded, the blocked factorize-solve for nv > 32 and the single-tile
_update_gradient_choleskyfor nv <= 32.One numerical finding worth stating explicitly, because it drove a deliberate deviation from mju_cholFactor: implementing C's literal "floor the pivot, then divide the row by it" overflows float32. Consecutive floored pivots amplify the row entries quadratically, and on a badly indefinite float32 matrix that runs away to inf well before the factor completes. So instead of dividing by a floored pivot, this treats a floored pivot as rank deficient: the pivot becomes sqrt(mindiag), the rest of its row is zeroed, and the direction is decoupled from the trailing blocks. Numerically that direction then behaves like a spring of stiffness mindiag rather than a hard constraint, which is a reasonable degradation for a mode that float32 cannot resolve anyway. On well-conditioned matrices the guard never fires and parity with a float64 reference solve is ~1.4e-7; the test also asserts that the pre-fix path still NaNs on the indefinite matrix, so the guard doesn't quietly become dead code.
One caveat I want to be honest about: I've only been able to validate this on the CPU/LLVM Warp backend. The repair branch writes the tile element-wise and is only correct because every lane computes identical values in identical order; I'd appreciate a reviewer with CUDA sanity-checking that the per-lane tile writes in the repair path behave as intended on GPU. The fast path is unchanged, so this only concerns the newly added repair branch.
I prepared this change with AI assistance (Claude); per AGENTS.md I have not added it as a commit co-author (it cannot sign the CLA), and I've reviewed the change and take responsibility for it.