Best way to debug "Fetching value for jax.Array that spans non-addressable devices is not possible" (ideally on single machine)

[dlwh]

Hi,

I'm running into a non-addressable-data error when running distributed (on TPU) and trying to debug where I went wrong.

I've seen this before, and I've usually been able to figure it out by looking at the stack trace, but today it's not terribly helpful (it's erroring out inside JAX itself). I was wondering if there's an easy way to reliably emulate the non-addressable data condition on a laptop, since it's a lot easier to debug there than on TPUs...

FWIW, this is the stack trace

   File "/home/dlwh/levanter/src/levanter/main/train_lm.py", line 347, in <module>
    main()
  File "/home/dlwh/levanter/src/levanter/config.py", line 233, in wrapper_inner
    response = fn(cfg, *args, **kwargs)
  File "/home/dlwh/levanter/src/levanter/main/train_lm.py", line 327, in main
    jax_step_loss, model, opt_state = train_step(model, opt_state, example, example_keys)
  File "/home/dlwh/venv310/lib/python3.10/site-packages/haliax/partitioning.py", line 296, in f
    return cached_pjitted_fun(dynamic_donated, dynamic_reserved, static)
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/pjit.py", line 208, in cache_miss
    outs, out_flat, out_tree, args_flat = _python_pjit_helper(
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/pjit.py", line 155, in _python_pjit_helper
    out_flat = pjit_p.bind(*args_flat, **params)
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/core.py", line 2633, in bind
    return self.bind_with_trace(top_trace, args, params)
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/core.py", line 383, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/core.py", line 790, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/pjit.py", line 1085, in _pjit_call_impl
    compiled = _pjit_lower(
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/pjit.py", line 1177, in _pjit_lower
    return _pjit_lower_cached(jaxpr, in_shardings, out_shardings, *args, **kwargs)
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/pjit.py", line 1237, in _pjit_lower_cached
    return pxla.lower_sharding_computation(
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 2072, in lower_sharding_computation
    nreps, tuple_args) = _cached_lowering_to_hlo(
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 1919, in _cached_lowering_to_hlo
    lowering_result = mlir.lower_jaxpr_to_module(
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py", line 617, in lower_jaxpr_to_module
    lower_jaxpr_to_fun(
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py", line 933, in lower_jaxpr_to_fun
    jaxpr.jaxpr, tokens_in, map(ir_constants, jaxpr.consts),
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py", line 217, in ir_constants
    out = handler(val, canonicalize_types)
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/array.py", line 679, in _array_mlir_constant_handler
    return mlir.ir_constants(val._value,
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/home/dlwh/venv310/lib/python3.10/site-packages/jax/_src/array.py", line 523, in _value
    raise RuntimeError("Fetching value for `jax.Array` that spans "

[yashk2810]

This means that the constant you are closing over is a jax.Array which is not fully addressable.

I would suggest passing it as an argument to the jitted function instead and then it should work.

[dlwh]

hrm, i'm not doing that on purpose. is there a good way to see what it is?

[dlwh]

found it. thank you for the hint!