Abstract:

This tutorial provides a comprehensive overview of tools and techniques for debugging JAX programs, led by a JAX engineering manager. The session covers fundamental debugging methods, distinguishing between standard Python tools that operate at trace-time and specialized JAX APIs (jax.debug.print, jax.debug.breakpoint) designed for runtime inspection within JIT-compiled functions. The speaker emphasizes the performance implications of these debug tools, which can prevent XLA optimizations by materializing intermediate values. The tutorial details various JAX configurations for identifying issues such as frequent recompilations (JAX_LOG_COMPILES), NaN/Inf generation (jax_debug_nans), unintended data transfers (jax_transfer_guard), and memory leaks (jax.live_arrays). It also provides guidance on interpreting and reporting hard crashes like segmentation faults and check failures, which are typically bugs within the JAX/XLA framework itself. Finally, the session addresses the complexities of debugging in distributed, multi-host environments, offering strategies for identifying out-of-sync processes or host failures.

Navigating JAX Debugging: Tools and Techniques

• 00:01:53 - Standard vs. JAX Debugging: Standard Python print() and breakpoints only execute once during the initial "tracing" of a JIT-compiled function. To inspect values at runtime on every execution, you must use specialized APIs like jax.debug.print() and jax.debug.breakpoint().
• 00:06:41 - Performance Impact of Debugging: Using jax.debug APIs can negatively affect performance. They force XLA to materialize intermediate values, which can prevent compiler optimizations like operator fusion and increase memory pressure. These tools are best for functional debugging, not performance analysis.
• 00:16:16 - Tracking Unwanted Recompilations: To debug slow performance caused by excessive recompilation, you can enable logging. JAX_LOG_COMPILES will log every time a function is recompiled, and JAX_EXPLAIN_CACHE_MISSES provides detailed reasons why a cached function wasn't used.
• 00:25:48 - Automatically Finding NaNs and Infs: If your JIT-compiled function produces NaNs or Infs, setting the JAX_DEBUG_NANS or JAX_DEBUG_INFS flags will cause JAX to run the code in an eager, unoptimized mode to pinpoint the exact operation that created the invalid value.
• 00:27:20 - Globally Disabling JIT: For easier debugging with standard Python tools (like an IDE's debugger), you can set JAX_DISABLE_JIT=True. This effectively makes JAX run in an eager mode, as if the @jit decorators were commented out.
• 00:30:28 - Finding Memory Leaks: The jax.live_arrays() function returns a list of all JAX arrays currently allocated on the accelerator. This is useful for identifying memory leaks, especially in interactive environments like notebooks where references to arrays can be unintentionally kept alive.
• 00:32:30 - Explicit Error Checking with checkify: The jax.experimental.checkify transform allows you to add explicit, runtime error checks inside your JIT-compiled code (e.g., asserting that a value is positive). It modifies the function to return an additional output indicating if an error occurred.
• 00:36:50 - Interpreting Hard Crashes: If your program aborts with a segmentation fault (SIGSEGV) or a "Check failed" error, it is almost always a bug in the JAX or XLA C++ backend, not your code. These should be reported as issues to the JAX team.
• 00:40:00 - Debugging Multi-Host Failures: In a distributed setting, crashes are often caused by processes getting out of sync (e.g., due to uneven data loading) or one host failing. The key to debugging is to collect logs from all hosts and find the one with a unique, root-cause error message, as others will typically just report that a peer has died.
• 00:45:17 - Inspecting Gradients: A useful technique to print or inspect gradients during the backward pass is to use jax.custom_vjp. You can define a custom backward pass for a function that simply prints the incoming gradients before passing them along.