Warp Debug Gradients

Debugging Gradients in Warp

Gradient bugs in Warp are almost never math bugs. The forward simulation looks perfectly healthy while the backward pass silently reads clobbered values, skips arrays, or double-counts adjoints. Users routinely burn days tuning physics knobs, loss functions, and assets when the real cause is a two-line taping-pattern fix. Your job is to find that fix with evidence, not intuition.

The single most important discipline: measure before hypothesizing. It is cheap for you to run a shrunk reproduction and compare autodiff against finite differences. The way the gradient is wrong (its signature) prunes the hypothesis space far faster than reading code ever will. Do not start proposing fixes from code reading alone — plausible-looking diagnoses of differentiability bugs are very often wrong, and an unverified "fix" that happens to perturb the numbers wastes everyone's time.

When to Use This Skill

Anything trained, optimized, calibrated, or fit through Warp kernels flows through wp.Tape gradients — so when such a workflow misbehaves, gradients are the prime suspect even if the user never says the word. Activate on the symptoms users actually report: training that diverges, NaNs, or does nothing; loss that stalls or plateaus above where it should; fits that converge to a wrong or biased answer or are worse than a reference implementation; pipelines that work at small scale but fail at production scale or fail a QA recheck. Also activate on explicit gradient symptoms — exploding, NaN/inf, zero, or subtly wrong gradients, wp.autograd.gradcheck failures, suspected wp.Tape/backward issues — and when the user asks whether their gradients can be trusted.

Do not activate for forward-only Warp work (kernel authoring, rendering, performance tuning), Warp build or installation problems, autograd questions in other frameworks with no Warp involvement, or pure performance work on a backward pass whose gradients the user has already validated.

The canonical background is Warp's own documentation — consult the relevant section before diagnosing in its territory (online at https://nvidia.github.io/warp/stable/; in a Warp source checkout the same content is under docs/user_guide/; pip installs do not include it):

• The "Differentiability" guide — especially "Array Overwrites", "Debugging Gradients", "Array Overwrite Tracking", and "Limitations and Workarounds" (in-place math, component assignment, dynamic loops).
• The FAQ, section "Differentiation and Interoperability" — what state a tape does and does not preserve, and checkpointing.

Prerequisites

Executing this skill assumes all of the following; if one is missing, surface that to the user instead of improvising around it:

• The user's script (or a faithful reproduction) is available in the workspace, runnable, and modifiable — diagnosis executes it repeatedly and edits it to apply fixes.
• This skill's references/ files (quick-checks.md, verification.md, custom-gradients.md, case-studies.md) accompany it and are consulted at the steps that cite them.

Instructions

1. Note the user's Warp version first (wp.__version__ or the banner Warp prints at init). Several verification behaviors changed in Warp 1.17 — copy-adjoint accumulation, overwrite-warning call sites, read-flag lifetime, gradcheck's restore_inputs — and the references mark each with a version caveat. On Warp < 1.17, a whole bug class exists that later versions fixed (quick-checks §1's version caveat), and some tools need workarounds.

2. Reproduce and shrink. Get the user's script running, then cut it down: fewer particles/elements, fewer time steps, fewer optimizer iterations, CPU device if the sim allows. You need a repro that runs in seconds, because you will run it many times. Keep the structure (number of kernels, the taping pattern, buffer reuse) intact — that is where the bug lives. Shrinking the physics is fine; restructuring the dataflow is not.

   If the script cannot be made to run (missing dependencies, broken code), report the blocking issue as the deliverable and stop — do not proceed to verify a program that never ran.

3. Instrument and establish ground truth (details and templates in references/verification.md):

   • Set wp.config.verify_autograd_array_access = True before module load and rerun under an active tape. Capture every warning. This catches the single most common bug class (write-after-read overwrites) nearly for free. Know its blind spots: it needs a tape, it cannot see arrays stored inside Warp structs, and it disables kernel caching (expect a kernel rebuild — JIT module recompilation only, not a rebuild of the native library). If the tracker runs clean but gradients are still wrong, specifically check for in-place mutations of arrays held inside Warp structs (quick-checks §1 and Limitations) before trusting the clean result.
   • Run one end-to-end finite-difference check: wrap the full forward pass (sim steps + loss) in a Python callable and hand it to wp.autograd.gradcheck with the true optimization inputs — it compares the autodiff gradient against central differences, restoring array inputs between evaluations (Warp 1.17+) so in-place-mutating forwards are checked from pristine state; on older Warp use the manual harness in references/verification.md. The reference is the user's actual objective over the full horizon, compared against the gradient the optimizer actually consumes — never a narrower window (see references/verification.md). This confirms gradients are actually wrong (users are sometimes wrong about this — report "gradients are correct" findings honestly) and yields the error signature. The template in references/verification.md fixes the eps/tolerance choices and the seed-pinning a stochastic forward needs — do not eyeball pass/fail against floating-point or sampling noise. If the backward pass runs out of memory while establishing ground truth, apply the checkpointing pattern from "Edge case: out of memory" below before proceeding.

4. Match the signature against the table below to rank hypotheses.

5. Scan the code against the known-pattern checklist (references/quick-checks.md). This is fast for you — do it in the same pass, but let the signature decide which findings are plausible causes versus incidental smells.

6. Localize if still ambiguous. Binary-search the pipeline: truncate to K steps and find where FD and autodiff first diverge; run wp.autograd.gradcheck_tape to test each recorded launch in isolation. Remember gradcheck_tape validates kernels individually — it is structurally blind to inter-kernel overwrites, so a clean per-kernel pass plus a wrong end-to-end gradient points at the taping pattern, not the kernels. It also silently skips kernels compiled with enable_backward=False (see Limitations) — if any kernel in the pipeline sets that, a clean pass says nothing about it; verify it separately.

7. Fix minimally, then re-verify with the exact same FD harness that established the failure. A gradient fix without a before/after FD comparison is not a fix. Verify the exact program you are shipping — the fixed file as it stands, every line included — never a re-implementation of it in a diagnostic script: a rebuilt pipeline silently drops whatever you believed was irrelevant, and if that belief is wrong the verification passes while the shipped code stays broken. Mechanically: the harness must import the fixed module (or execute the fixed file) and call into it — the only code that may live outside the shipped program is the FD driver itself. Also rerun the overwrite tracker to confirm the warnings are gone. "Minimally" applies to the code diff, not the diagnosis: when the root cause is structural (e.g., accidental gradient truncation, quick-checks §8), the minimal correct fix is the restructure — do not substitute a smaller change that only silences the surface symptom.

8. Close the loop on the user's original complaint. Rerun their actual workflow (their script, their printed metrics). The job is done when the symptom they reported is resolved — an optimization that was "exploding" should now demonstrably improve its objective, not merely avoid NaN. If gradients verify correct at the full horizon but training still fails, that is a new signature-table entry, not a victory; keep diagnosing (or report the verified gradients and the remaining non-gradient cause, e.g. learning rate).