Repository: jax-ml/jax
Pull request #28529 (2025-05-05 22:04:23)

PR #28157: Warn on excessive captured constants
PR #28157: Warn on excessive captured constants

Imported from GitHub PR https://github.com/jax-ml/jax/pull/28157

One of the most common modes by which users may have extended lowering times is by unintentionally capturing, rather than tracing, a large number of constants, e.g. capturing the model weights as part of the computation. 

To address this we introduce two new flags

- `jax_captured_constants_warn_bytes` defaults to `2 * 10 ** 9` (2GB).  The number of total bytes of captured constants before warning is issued. (Note that the maximum size binary XLA can serialize for the compilation cache is 2GB.)
- `jax_captured_constants_report_frames` defaults to `0`. If a warning is issued, how many stack frames to report for each constant. Defaults to 0, which means we don't generate the report by default. Reports all frames if set to `-1`.

Both message are returned by using `warnings.warn`. The envisioned workflow for debugging captured constants is as follows:

1. The user is alerted to the problem by the initial (by default). This looks something like:
```
UserWarning: A large amount of constants were captured during lowering (125.00GB total). 
If this is intentional, disable this warning by setting JAX_CAPTURED_CONSTANTS_WARN_BYTES=-1. 
To obtain a report of where these constants were encountered, set JAX_CAPTURED_CONSTANTS_REPORT_FRAMES=-1.
```
2. The user sets the `JAX_CAPTURED_CONSTANTS_REPORT_FRAMES=-1` to obtain debugging information to locate the source of the constants. Setting `JAX_CAPTURED_CONSTANTS_REPORT_FRAMES` to a small positive integer will return a suffix of the number of captured frames. Upon rerunning the code, the report generated will be something like this:

```
UserWarning: A large amount of constants were captured during lowering (125.00GB total). 
If this is intentional, disable this warning by setting JAX_CAPTURED_CONSTANTS_WARN_BYTES=-1. 
The subsequent report may be disabled by setting JAX_CAPTURED_CONSTANTS_REPORT_FRAMES=0.

Largest 5 allocation(s):
Constant <class 'numpy.ndarray'>, float32[1439,721,720], 5.00GB captured at:
  /home/user/project/main.py:193 (<module>)
  /home/user/project/main.py:156 (run_export)
  /home/user/project/main.py:147 (run_forward_exported)
  /home/user/project/.venv/lib/python3.10/site-packages/dinosaur/pytree_utils.py:98 (tree_map_over_nonscalars)
  /home/user/project/.venv/lib/python3.10/site-packages/dinosaur/pytree_utils.py:97 (g)
  /home/user/project/.venv/lib/python3.10/site-packages/dinosaur/spherical_harmonic.py:661 (g)
  /home/user/project/.venv/lib/python3.10/site-packages/dinosaur/spherical_harmonic.py:266 (inverse_transform)
  
Constant <class 'numpy.ndarray'>, float32[1439,721,720], 5.00GB captured at:
  /home/user/project/main.py:193 (<module>)
  /home/user/project/main.py:156 (run_export)
  /home/user/project/main.py:147 (run_forward_exported)
  /home/user/project/.venv/lib/python3.10/site-packages/dinosaur/pytree_utils.py:98 (tree_map_over_nonscalars)
  /home/user/project/.venv/lib/python3.10/site-packages/dinosaur/pytree_utils.py:97 (g)
  /home/user/project/.venv/lib/python3.10/site-packages/dinosaur/spherical_harmonic.py:661 (g)
  /home/user/project/.venv/lib/python3.10/site-packages/dinosaur/spherical_harmonic.py:266 (inverse_transform)
```
and so fourth. The report is hard coded to only report the the top 5 (ordered by `nbytes`) largest constants. 

3. The user can then set a breakpoint at last line listed, and inspect the operands to find the problematic constant by the shape and type we report. 

The reason for the two stage process is it is very easy and cheap to check the number of bytes each time we lower a function. However generating the report involves a traversal of the Jaxpr, and so is not. 

For Googlers [b/403532544](http://b/403532544#comment26)

Merging this change closes #28157



Patch:
diff --git a/jax/BUILD b/jax/BUILD
index 6358cce78..99cf508fa 100644
--- a/jax/BUILD
+++ b/jax/BUILD
@@ -650,6 +650,7 @@ pytype_strict_library(
         ":core",
         ":dtypes",
         ":effects",
+        ":jaxpr_util",
         ":layout",
         ":mesh",
         ":op_shardings",
diff --git a/jax/_src/config.py b/jax/_src/config.py
index 2abb457ef..e79993958 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -970,6 +970,28 @@ check_tracer_leaks = bool_state(
           'to disable any debuggers while leak checking is enabled.'))
 checking_leaks = functools.partial(check_tracer_leaks, True)
 
+
+captured_constants_warn_bytes = int_state(
+    name='jax_captured_constants_warn_bytes',
+    default=2 * 10 ** 9,
+    help=('The number of bytes of parameters that may be captured as constants '
+          'before a warning is issued. Defaults to approximately 2GB. '
+          'Set to -1 to disable issuing a warning.'
+    )
+)
+
+captured_constants_report_frames = int_state(
+    name='jax_captured_constants_report_frames',
+    default=0,
+    help=('The number of stack frames reported for each captured constant '
+          'indicating the file and operation where the constant was captured. '
+          'Set to -1 to print the complete set of frames, or 0 to disable. '
+          'N.b. the report is only generated if the total amount of captured '
+          'constants exceeds `jax_captured_constants_warn_bytes`, as it is expensive'
+          'to generate the report.'
+    )
+)
+
 debug_nans = bool_state(
     name='jax_debug_nans',
     default=False,
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index ccbe84bcf..e9deb8d3f 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -16,11 +16,12 @@
 from __future__ import annotations
 
 import collections
-import contextlib
 from collections.abc import Callable, Iterable, Iterator, Sequence
+import contextlib
 import dataclasses
 import functools
 from functools import partial
+import heapq
 import io
 import itertools
 import operator
@@ -31,14 +32,13 @@ import typing
 from typing import Any, NamedTuple, Protocol, Union, cast as type_cast
 import warnings
 
-import numpy as np
-
 from jax._src import ad_util
 from jax._src import api_util
 from jax._src import config
 from jax._src import core
 from jax._src import dtypes
 from jax._src import effects as effects_lib
+from jax._src import jaxpr_util
 from jax._src import linear_util as lu
 from jax._src import path
 from jax._src import sharding_impls
@@ -48,19 +48,20 @@ from jax._src import xla_bridge as xb
 from jax._src.interpreters import partial_eval as pe
 from jax._src.interpreters import xla
 from jax._src.layout import AutoLayout, DeviceLocalLayout
-from jax._src.partition_spec import PartitionSpec
-from jax._src.mesh import AxisType
-from jax._src.sharding import Sharding as JSharding
-from jax._src.sharding_impls import (AUTO, NamedSharding,
-                                     modify_sdy_sharding_wrt_axis_types,
-                                     SdyArraySharding, SdyArrayShardingList)
-from jax._src.util import foreach
-from jax._src.lib import xla_client as xc
 from jax._src.lib import _jax
+from jax._src.lib import xla_client as xc
 from jax._src.lib.mlir import dialects, ir, passmanager
-from jax._src.lib.mlir.dialects import func as func_dialect, hlo
 from jax._src.lib.mlir import register_jax_dialects
+from jax._src.lib.mlir.dialects import func as func_dialect, hlo
+from jax._src.mesh import AxisType
+from jax._src.partition_spec import PartitionSpec
+from jax._src.sharding import Sharding as JSharding
+from jax._src.sharding_impls import ( AUTO, NamedSharding,
+                                     SdyArraySharding, SdyArrayShardingList,
+                                     modify_sdy_sharding_wrt_axis_types)
 from jax._src.state.types import AbstractRef
+from jax._src.util import foreach
+import numpy as np
 
 # mypy: ignore-errors
 
@@ -1104,6 +1105,51 @@ def _get_unconstrained_variants(s, aval) -> UnconstrainedVariants:
       unconstrained_dims=unconstrained_dims)
 
 
+def check_jaxpr_constants(closed_jaxpr: core.ClosedJaxpr):
+  """Check if a JAXPR contains an excessive amount of constants, if so, report where they were captured"""
+  if (threshold := config.captured_constants_warn_bytes.value) == -1:
+    return
+
+  # need the unaesthetic getter here as some of the consts in the test suite are arbitrary objects
+  total_iter, nbytes_iter = itertools.tee(
+      map(lambda c: getattr(c, "nbytes", 0), closed_jaxpr.consts)
+  )
+
+  if (total_bytes := sum(total_iter)) < threshold:
+    return
+
+  message = (
+      "A large amount of constants were captured during lowering"
+      f" ({util.pprint_bytes(total_bytes)} total). If this is intentional,"
+      " disable this warning by setting JAX_CAPTURED_CONSTANTS_WARN_BYTES=-1. "
+  )
+
+  if not (num_frames := config.captured_constants_report_frames.value):
+    message += (
+        "To obtain a report of where these constants were encountered, "
+        "set JAX_CAPTURED_CONSTANTS_REPORT_FRAMES=-1."
+    )
+    warnings.warn(message)
+    return
+
+  message += (
+      "The subsequent report may be disabled by setting JAX_CAPTURED_CONSTANTS_REPORT_FRAMES=0.\n\n"
+      f"Largest {min(num_frames, len(closed_jaxpr.consts))} allocation(s):\n"
+  )
+  try:
+    nbytes_var_const = zip(nbytes_iter, closed_jaxpr.jaxpr.constvars, closed_jaxpr.consts)
+    for nbytes, var, const in heapq.nlargest(5, nbytes_var_const, key=operator.itemgetter(0)):
+      message += f"  Constant {type(const)}, {var.aval.str_short()}, {util.pprint_bytes(nbytes)} captured at:\n"
+
+      for eqn in jaxpr_util.eqns_using_var(closed_jaxpr.jaxpr, var):
+        call_frame_source_info = source_info_util.summarize(eqn.source_info, num_frames)
+        message += "  " * 2 + call_frame_source_info.replace("\n", "\n" + "  " * 2) + "\n\n"
+
+    warnings.warn(message)
+  except Exception as exc:
+    warnings.warn(message + f" Exception raised while generating report: {exc}")
+
+
 def lower_jaxpr_to_module(
     module_name: str,
     jaxpr: core.ClosedJaxpr,
@@ -1441,6 +1487,8 @@ def lower_jaxpr_to_fun(
     MLIR func op
   """
   util.test_event("lower_jaxpr_to_fun", name)
+  check_jaxpr_constants(jaxpr)
+
   # The first dimension variable may be the platform index
   num_dim_vars = len(ctx.shape_poly_state.dim_vars)
   dim_var_avals = [core.ShapedArray((), dtypes.canonicalize_dtype(np.int64))] * num_dim_vars
diff --git a/jax/_src/jaxpr_util.py b/jax/_src/jaxpr_util.py
index ab72634d3..a6c93c8c1 100644
--- a/jax/_src/jaxpr_util.py
+++ b/jax/_src/jaxpr_util.py
@@ -209,3 +209,35 @@ def pprof_equation_profile(jaxpr: core.Jaxpr) -> bytes:
       for _, eqn in all_eqns(jaxpr)
   )
   return _pprof_profile(d)
+
+def eqns_using_var_with_invar_index(jaxpr: core.Jaxpr, invar: core.Var) -> Iterator[tuple[core.JaxprEqn, int]]:
+  """Find all the equations which use invar and the positional index of its binder"""
+  for eqn in jaxpr.eqns:
+    for invar_index, eqn_var in enumerate(eqn.invars):
+      if eqn_var == invar:
+        yield eqn, invar_index
+        break # we found the var, no need to keep looking in this eqn
+
+def jaxpr_and_binder_in_params(params, index: int) -> Iterator[tuple[core.Jaxpr, core.Var]]:
+  for val in params.values():
+    vals = val if isinstance(val, tuple) else (val,)
+    for v in vals:
+      if isinstance(v, core.Jaxpr):
+        if index >= len(v.invars):
+          raise RuntimeError(f"Failed to find index {index} in jaxpr.invars while building report")
+        yield v, v.invars[index]
+      elif isinstance(v, core.ClosedJaxpr):
+        if index >= len(v.jaxpr.invars):
+          raise RuntimeError(f"Failed to find index {index} in jaxpr.invars while building report")
+        yield v.jaxpr, v.jaxpr.invars[index]
+
+def eqns_using_var(jaxpr: core.Jaxpr, invar: core.Var) -> Iterator[core.JaxprEqn]:
+  """Find the leaf equations using a variable"""
+  # The complexity of this call is becauase the invar might originate from a nested jaxpr
+  for eqn, invar_index in eqns_using_var_with_invar_index(jaxpr, invar):
+    if (child_jaxprs_and_vars := tuple(jaxpr_and_binder_in_params(eqn.params, invar_index))):
+      for (jaxpr, invar) in child_jaxprs_and_vars:
+        yield from eqns_using_var(jaxpr, invar)
+    else:
+      # if the previous condition fails, there is no deeper jaxpr to explore =(
+      yield eqn
diff --git a/jax/_src/util.py b/jax/_src/util.py
index 94b6380b1..34f748544 100644
--- a/jax/_src/util.py
+++ b/jax/_src/util.py
@@ -20,6 +20,7 @@ import functools
 from functools import partial
 import itertools as it
 import logging
+import math
 import operator
 from typing import (Any, Generic, SupportsIndex, TypeVar, overload, TYPE_CHECKING, cast)
 import weakref
@@ -675,3 +676,12 @@ def test_event(name: str, *args) -> None:
 
 if hasattr(jaxlib_utils, "Mutex"):
   Mutex = jaxlib_utils.Mutex
+
+
+def pprint_bytes(num_bytes: int | float) -> str:
+  prefixes = ("", "K", "M", "G", "T")
+  if num_bytes <= 0:
+    return "0.00B"
+  exponent = min(math.floor(math.log(num_bytes, 1000)), len(prefixes) - 1)
+  scaled_value = num_bytes / (1000**exponent)
+  return f"{scaled_value:.2f}{prefixes[exponent]}B"


Test patch:
diff --git a/tests/jax_jit_test.py b/tests/jax_jit_test.py
index 5946d557d..cbf7c710f 100644
--- a/tests/jax_jit_test.py
+++ b/tests/jax_jit_test.py
@@ -227,6 +227,25 @@ class JaxJitTest(jtu.JaxTestCase):
     self.assertArraysEqual(v1, v1_expected)
     self.assertArraysEqual(v2, v2_expected)
 
+  def test_check_for_large_number_of_constants(self):
+    y = jnp.ones((128, 128))
+    x = jnp.zeros((128,))
+
+    def jit_maker(): # need to ensure we lower at each test
+      def func(x):
+        return x @ y
+      return jax.jit(func)
+
+    with self.assertWarnsRegex(UserWarning, "A large amount of constants were captured during lowering"):
+      with config.captured_constants_warn_bytes(y.nbytes):
+        jit_maker()(x)
+
+    with self.assertNoWarnings():
+      with config.captured_constants_warn_bytes(y.nbytes + 1):
+        jit_maker()(x)
+
+      with config.captured_constants_warn_bytes(-1):
+        jit_maker()(x)
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())