Repository navigation
determinism api - #1687
determinism api#1687
Conversation
f63a701 to
a8469a9
Compare
a8469a9 to
b451d45
Compare
eric-heiden
left a comment
eric-heiden
left a comment
There was a problem hiding this comment.
Thanks for adding the configurable determinism path. I reviewed b451d45 and found two correctness issues when ATOMICS is enabled, plus a smaller regression in analyzer suppression; the reproductions and suggested fixes are inline.
Validation on an NVIDIA L4 with Warp 1.15.0: the 50 API/analyzer tests and seven existing factor/solve tests passed, while focused reproductions exposed the issues below. The GPU failures also reproduce through the public solve_m and fwd_actuation APIs. I treated the contact, constraint, and island sorting hooks as the explicitly deferred work described in the PR.
| """Fused sparse backsubstitution: UP + diag + DOWN in one kernel.""" | ||
| module_options = {"enable_backward": False} | ||
| if deterministic: | ||
| module_options["deterministic"] = wp.DeterministicMode.RUN_TO_RUN |
There was a problem hiding this comment.
[P1] Preserve the substitution dependencies in the deterministic path
Could we give this fused solver a separate deterministic implementation? Warp's RUN_TO_RUN mode records these atomic contributions and applies the reductions after the kernel finishes. This kernel needs the updated x_out values between substitution levels, before diagonal scaling, and during backward substitution, so _syncthreads() cannot preserve those dependencies under this mode.
On the L4 with Warp 1.15.0, a minimal two-DOF kernel reproduction with y=[2, 3], D=[1, 1], and an off-diagonal L coefficient of 0.5 returns [0.5, 2.0] instead of [0.5, 2.75]. I also reproduced incorrect results through solve_m on a 66-DOF model, which takes the sparse path.
Splitting the dependent stages/levels into separate launches, or using a fixed-order gather reduction that completes before the next stage, would preserve the solve. A numerical parity test with ATOMICS enabled would help cover this.
| def _qfrc_actuator(deterministic: bool = False): | ||
| module_options = {"enable_backward": False} | ||
| if deterministic: | ||
| module_options["deterministic"] = wp.DeterministicMode.RUN_TO_RUN |
There was a problem hiding this comment.
[P1] Size deterministic buffers for repeated atomics in dynamic loops
Could we pass a sufficient deterministic_max_records bound here and audit the other factories with dynamic loops? This kernel emits one contribution per moment nonzero, but the default code-generated record bound does not account for every iteration of range(rownnz).
I reproduced this through fwd_actuation with 1,024 worlds and one ball-joint motor (gear="1 1 1 0 0 0", ctrl=1). The normal path returns [1, 1, 1] in every world; enabling ATOMICS raises RuntimeError: Deterministic scatter buffer overflow. Capturing and replaying the same actuator kernel is more concerning: it returns without that exception, and only 1,024 of the 3,072 expected unit-valued entries are correct.
A model-derived upper bound on records per thread, or restructuring the launch to emit a bounded number of contributions per thread, should address this. A batched test that also exercises graph replay would catch the silent loss.
| start_line = iss.node.lineno | ||
| end_line = iss.node.end_lineno or start_line | ||
| matched_rules: Set[str] = set() | ||
| for l in range(start_line, end_line + 1): | ||
| if l in ignore_lines: | ||
| matched_rules.update(ignore_lines[l]) |
There was a problem hiding this comment.
[P2] Keep inline ignores local for function-level diagnostics
Could we limit the expanded ignore range to the statement or diagnostic it belongs to? MissingModuleUnique is attached to a FunctionDef, whose end_lineno includes the entire kernel body. Collecting every ignore within that span means an unrelated inline ignore now suppresses the missing-module warning for the enclosing kernel.
For example, this reports MissingModuleUnique until the assignment's comment is added:
def factory():
@wp.kernel
def kernel(arr_out: wp.array[float]):
arr_out[0] = 1.0 # kernel_analyzer: ignoreKeeping function-level suppression anchored to the function declaration, while allowing multiline atomic-call suppression across that call's own span, would retain the new functionality without weakening the existing check.
b451d45 to
94ad52b
Compare
|
@eric-heiden thank you for the review! 1.1
|
94ad52b to
796113d
Compare
796113d to
bd12638
Compare
There was a problem hiding this comment.
Thanks for the updates! I re-reviewed bd12638 and confirmed that all three findings from the previous review are addressed. The separate sparse-solve stages preserve the substitution dependencies, the dynamic atomic loops now have record bounds, and inline analyzer ignores no longer suppress unrelated function-level diagnostics. The earlier public-API reproductions also pass now.
There was a problem hiding this comment.
Approving this as the API and deterministic-atomics foundation, with contact, constraint-row, and island sorting remaining the explicitly described follow-up work.
I also exercised this PR through Newton's SolverMuJoCo at 18abe5585, using an NVIDIA L4 and MuJoCo 3.16.1.dev999309006:
- On Warp 1.17.0, a 67-DOF branching articulation with contacts disabled was bitwise repeatable with this PR's flag enabled: 16 worlds, 64 steps, three rollouts from identical initial state, comparing body/joint poses, velocities, and accelerations.
- With both Newton's
deterministic=wp.DeterministicMode.RUN_TO_RUNandsolver.mjw_model.opt.deterministic = mjw.DeterminismType.ALL, a six-box stack using deterministic Newton contacts and a driven quadruped also produced repeatable results. The stack/Newton-contact and quadruped/MuJoCo-contact CUDA graph trajectories matched ordinary stepping bitwise, with finite states and no overflow flags. These are checks of the combined configuration, rather than a general contact-determinism guarantee from this PR alone.
Non-blocking FYI for Warp 1.19 compatibility: Newton's lockfile at that commit uses 1.19.0.dev20261020. On that build, enabling this PR's flag causes both the articulation and quadruped to fail on the first step with:
WarpDeterminismError: Deterministic mode cannot run kernel '_subtree_com_acc__locals__kernel'.
Cause: argument 'subtree_com_out' receives deterministic atomic updates, but it overlaps
argument 'subtree_com_in', which is read directly by the kernel.
The launch in smooth.py passes d.subtree_com as both arguments, which Warp 1.19's runtime alias validation rejects. The normal-mode articulation control runs successfully on the same build. Flagging this for a compatibility follow-up; it does not block my approval of the scoped API work here.
bd12638 to
c5bd793
Compare
…erministic in create_render_context, and remove DeterminismType.ISLANDS
c5bd793 to
bae3042
Compare
|
thanks @eric-heiden! made a few additional changes and updated the pr description. 1. Remove
|
90bb78e to
9c9ec5c
Compare
|
Thanks for catching this on Warp 1.19! Changes in
|
eric-heiden
left a comment
eric-heiden
left a comment
There was a problem hiding this comment.
Thanks for the fix, Taylor! Retested 9c9ec5c with Newton on Warp 1.17 and 1.19 dev: the reported _subtree_com_acc aliasing error is resolved, and the articulation and Newton-contact quadruped are bitwise repeatable with both eager stepping and CUDA graphs (16 worlds, 64 steps, 3 repeats). Looks good to me!
Resolve the conflicts with main's determinism API (google-deepmind#1687), fast-math line search (google-deepmind#1803) and the main/deepmind sync (google-deepmind#1834). main turned the Hessian metric kernels (_add_tendon_metric_dense, _add_actuator_metric_dense, _add_efmK_metric_dense, _add_flexcon_metric_dense) and _update_gradient_JTCJ_dense into @cache_kernel factories with deterministic_max_records sizing. The float64 scalar type now goes into those factories as an htype argument (diagonal too, for the tendon and actuator metrics), placed before the trailing deterministic=False that determinism_test requires; the PR's *_typed variants of these kernels are gone. _JTDACJ_sparse and _update_gradient_h_incremental_sparse take htype between max_records and deterministic, and their call sites (including main's deterministic re-creation) pass it positionally, since cache_kernel takes no keywords. JTCJ keeps main's wp.atomic_add, in htype. Kernels main did not touch keep the PR design (factories over htype, module-level names are the float32 instance).
Fill the CONTACTS and CONSTRAINT hooks introduced by google-deepmind#1687 without changing atomic reductions or the NONE path. Canonicalize contact features and geometry, permute constraint rows with their sparse Jacobian metadata, and repair contact and JTDAJ references independently of incoming contact order. Add permutation, CPU parity, interleaved-world, mixed-dimension, overflow, flex, SDF, and CUDA capture regressions. Credit Mark Yang (@mar-yan24) for the original joint determinism work that motivated this follow-up.
* Implement canonical contact and constraint sorting Fill the CONTACTS and CONSTRAINT hooks introduced by #1687 without changing atomic reductions or the NONE path. Canonicalize contact features and geometry, permute constraint rows with their sparse Jacobian metadata, and repair contact and JTDAJ references independently of incoming contact order. Add permutation, CPU parity, interleaved-world, mixed-dimension, overflow, flex, SDF, and CUDA capture regressions. Credit Mark Yang (@mar-yan24) for the original joint determinism work that motivated this follow-up. * Preserve equality force semantics during constraint sorting Order equality rows by subtype and ID while preserving each connect/weld group's semantic component offsets. Keep physical Jacobian tie-breaks for repeated flex IDs and leave force reconstruction unchanged. Add CPU-parity regressions for solved-state sorting and full forward with asymmetric loads, interleaved equality types, dense/sparse Jacobians, heterogeneous worlds, reversed groups, and JTDAJ block support. * Make sorting regression snapshots portable to CPU Copy retained NumPy snapshots before mutating Warp storage so CPU assertions compare independent data rather than aliased views. Use moderate binary-exact equality loads and velocities without relaxing the standard parity tolerance or changing production sorting. * Refine sorting kernels and regression coverage after review * Initialize contact addresses and imported collision IDs * Address the second review round of the sorting kernels and tests Compare contacts lazily through a `_ContactKeys` struct instead of materializing a 43-column key buffer, evaluate the comparison once per merge step, test raw bits before computing float keys, gather rows as `uint32` words, and name kernel parameters after their Model and Data fields. `_remap_blocks` reads the head row of each block directly and `constraints` always ranks contacts. Call `sort.constraints` from `make_constraint` instead of going through `_sort_constraints`. Clear only the contact address components that assembly does not write, so the addresses of active contacts are no longer written twice. Move the contact-address and imported-contact regressions into `sort_test.py`, replace attribute reflection and random permutations with explicit fields and fixed permutations, and replay captured contact and constraint sorts across input orders and active counts. JTDAJ blocks are asserted to tile every row except the DOF friction rows, which have no block on current main. `test_contact_permutation` still fails at this commit because main now stores `contact.surfacevel`; the next commit adapts sorting to the merged main. * Follow upstream contact and constraint references through sorting Merging main brought three things that sorting has to carry. `contact.surfacevel` is computed before the contact sort, so it is permuted with the other contact fields and used as a tie breaker. `dof_frictionadr` holds the row of each DOF friction constraint for the sparse Newton Hessian diagonal, so it is remapped through the row permutation. Flex strain rows are now allocated as one JTDAJ block per cell, so they keep their mode order within the cell like connect and weld components, which keeps the head row of each block first. Add regressions for the friction addresses and the flex strain blocks, and a surface velocity tie-break case. * Sort contacts and constraint rows with fewer launches The deterministic sort launched one kernel per merge pass and one clone and one gather per contact or constraint field: 80 sorting launches per step on the humanoid benchmark with 8192 worlds. Contacts are now counted per world, grouped and placed by rank, moving every field in the same launch. The second contact sort of the constraint rows is gone: one launch checks whether the contact buffer is already in canonical order, and only ranks the contacts if it is not. The row sort starts from a counting launch over blocks of rows, and the row fields and the Jacobian move through one packed copy that skips the rows that are already in place. Buffers larger than the static limits fall back to the merge passes. Sorting launches per step go from 80 to 15 (humanoid) and from 99 to 23 (three humanoids). With 8192 worlds, CONTACTS|CONSTRAINT costs 18.4% of the steps per second of the humanoid benchmark instead of 33.6%, and 14.9% of three humanoids instead of 29.1%. The sorted output is bitwise identical to the previous implementation. * Address review: sort tests, typed sort structs, stale Jqvel sort_test.py: test_contact_rank_paths and test_sparse_offsets_paths build the unsorted state once, snapshot Data and restore it before each path, so both paths sort the same input instead of the output of two forward calls, whose efc.Ma differs by one ULP. sort.py: the fields of _ContactRows and _ConstraintRows are assigned directly, the contact copy is allocated as a _ContactRows instance by _empty_contact_rows, and _rank and _merge take the type of their keys from compare.input_types["keys"]. This removes getattr, setattr, SimpleNamespace, Any, _CONTACT_FIELDS and _CONSTRAINT_FIELDS. test_contact_bucket_ties clones the contact into a types.Contact. sort_test.py: test_permutation_blocks, test_contact_paths and test_sparse_move use modular index arithmetic instead of np.random.default_rng. collision_driver.py: collision calls sort.contacts(d) and _sort_contacts is removed. sort_test.py and the spies of determinism_test.py follow. collision_flex.py: the duplicate reset of contact_efc_address_out is removed, so the file matches main. sort.py: _constraint_primary zeroes efc.Jqvel of the rows that are not contact rows. Assembly never writes it for them, so sorted rows carried the values of earlier slots. test_noncontact_jqvel covers it.
This PR introduces the API and foundation for deterministic simulation in MuJoCo Warp, along with deterministic kernel factories across all simulation subsystems and kernel analyzer enforcement for atomic operations. Non-deterministic execution in GPU physics arises from non-deterministic floating-point atomic reduction order, arbitrary contact generation order, and uncanonical constraint row ordering (island discovery and DOF/constraint island mappings are already deterministic following #1541). This change establishes the configuration bitmask and pipeline integration hooks to control these behaviors selectively while keeping the default non-deterministic path zero-overhead.
Determinism is configured via a new
DeterminismTypeIntFlagenum (NONE=0,CONTACTS=1<<0,CONSTRAINT=1<<1,ATOMICS=1<<2,ALL=7) exposed onm.opt.deterministicand top-levelmujoco_warp.DeterminismType. The option mirrors thewarn_overflowproperty pattern, accepting boolean inputs (Truemaps toALL,FalsetoNONE), bitmask combinations, or string expressions throughoverride_model(e.g.["opt.deterministic=CONSTRAINT|ATOMICS"]).Key changes included in this PR:
DeterminismTypeenum (NONE,CONTACTS,CONSTRAINT,ATOMICS,ALL) andOption.deterministicproperty/setter intypes.py, withput_modelinitialization and string/bitmask parsing registered inio.py._sort_contactsincollision_driver.pyand_sort_constraintsinconstraint.py, placed immediately above their callers and gated byDeterminismType.CONTACTSandDeterminismType.CONSTRAINT.@cache_kernelfactories that conditionally compile withwp.DeterministicMode.RUN_TO_RUNwhenDeterminismType.ATOMICSis active, maintainingenable_backward=Falseby default and passing0fordeterministic_max_recordsdimension arguments on the default path to avoid redundant cache compilations. Subsystems covered include smooth dynamics, tree accumulations, momentum, tendon dynamics, transmissions, and sparse Cholesky backsubstitution (smooth.py), passive forces and flex body dynamics/bending (passive.py), sensors and contact matching (sensor.py), analytical derivatives (derivative.py), actuator forces (forward.py), solver linesearches, Hessians, and gradient reductions (solver.py), dense contact Jacobian velocity tile reductions and flex equality constraints (constraint.py), awake body/DOF compaction (sleep.py), flex BVH bounds and vertex normal accumulations (bvh.py,render_util.py), and model subtree mass accumulations (set_const.py).@cache_kernelby(func.__module__, func.__name__)inwarp_util.pyso identically named kernel factories across modules do not collide.contrib/kernel_analyzerto version 0.9.0 with AST checks (MissingDeterministicFactory,MissingDeterministicMaxRecords,AtomicTargetConflict) enforcing deterministic factories, dynamic-loopdeterministic_max_recordsbounds, and atomic target read/write separation on kernels and device functions containing atomic operations, enabled via--check_atomic=Truein the CLI,.pre-commit-config.yaml, and VS Code extension (kernel-analyzer-0.9.0.vsix).determinism_test.pyandast_analyzer_test.pycovering allDeterminismTypebitmask operations, hook preconditions, CUDA graph capture, static AST verification of every deterministic kernel factory, and bitwise reproducibility across all subsystems.Future work:
_sort_contactsforDeterminismType.CONTACTS._sort_constraintsforDeterminismType.CONSTRAINT.