NCCL4JAX - Using JIT on the allreduce

[sophia-estrela]

Hi! I am currently building an NCCL library for JAX, more specifically I am finishing writing the all reduce operation right now.

I already finished the XLA part and I am now trying to make the function work with JIT.

def capsule_allReduce(sendbuffer, count, dataType, op):
    res = allreduce(sendbuffer, count, dataType, op, comm=comm)
    return res
jit_all_reduce = jit(capsule_allReduce, device=jax.devices()[device_id])

I have something like above where count, dataType, and op, are integers, sendbuffer is an array (that will be put on device), and comm is the communicator (object of the nccl communicator class).

But when I try to run I get the following error:

Traceback (most recent call last):
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/interpreters/mlir.py", line 1156, in cached_lowering
func = ctx.module_context.cached_primitive_lowerings.get(key)
TypeError: unhashable type: 'DynamicJaxprTracer'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/scratch/app/python/3.9.12/lib/python3.9/runpy.py", line 188, in _run_module_as_main
mod_name, mod_spec, code = _get_module_details(mod_name, _Error)
File "/scratch/app/python/3.9.12/lib/python3.9/runpy.py", line 111, in _get_module_details
import(pkg_name)
File "", line 1007, in _find_and_load
File "", line 986, in _find_and_load_unlocked
File "", line 680, in _load_unlocked
File "", line 850, in exec_module
File "", line 228, in _call_with_frames_removed
File "/petrobr/parceirosbr/brcluster/sophia.estrela/nccl4jax/nccl4jax_allreduce.py", line 49, in
recv_buffer, token = jit_all_reduce(send_buffer, count, dataType, op)
File "/petrobr/parceirosbr/brcluster/sophia.estrela/nccl4jax/nccl4jax_allreduce.py", line 24, in capsule_allReduce
res = allreduce(send_buffer_device, count, dataType, op, comm=comm)
File "/petrobr/parceirosbr/brcluster/sophia.estrela/nccl4jax/all_reduce_nccl.py", line 28, in allreduce
return tuple(nccl_allreduce_primitive.bind(sendBuf, token, count = count, dataType = dataType, op = op, comm=comm))
jax._src.source_info_util.JaxStackTraceBeforeTransformation: jax._src.errors.TracerArrayConversionError: The numpy.ndarray conversion method array() was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
This Tracer was created on line /petrobr/parceirosbr/brcluster/sophia.estrela/nccl4jax/nccl4jax_allreduce.py:49 ()
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

---

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/petrobr/parceirosbr/brcluster/sophia.estrela/nccl4jax/nccl4jax_allreduce.py", line 49, in
recv_buffer, token = jit_all_reduce(send_buffer, count, dataType, op)
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/src/api.py", line 473, in cache_miss
out_flat = xla.xla_call(
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/core.py", line 1765, in bind
return call_bind(self, fun, *args, **params)
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/core.py", line 1781, in call_bind
outs = top_trace.process_call(primitive, fun, tracers, params)
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/core.py", line 678, in process_call
return primitive.impl(f, *tracers, **params)
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/_src/dispatch.py", line 182, in _xla_call_impl
compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/linear_util.py", line 285, in memoized_fun
ans = call(fun, *args)
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/_src/dispatch.py", line 230, in _xla_callable_uncached
return lower_xla_callable(fun, device, backend, name, donated_invars, False,
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/_src/profiler.py", line 206, in wrapper
return func(*args, **kwargs)
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/_src/dispatch.py", line 340, in lower_xla_callable
module, keepalive = mlir.lower_jaxpr_to_module(
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/interpreters/mlir.py", line 556, in lower_jaxpr_to_module
lower_jaxpr_to_fun(
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/interpreters/mlir.py", line 810, in lower_jaxpr_to_fun
out_vals, tokens_out = jaxpr_subcomp(ctx.replace(name_stack=callee_name_stack),
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/interpreters/mlir.py", line 929, in jaxpr_subcomp
ans = rule(rule_ctx, *map(_unwrap_singleton_ir_values, in_nodes),
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/interpreters/mlir.py", line 1161, in cached_lowering
return f(ctx, *args, **params)
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/interpreters/mlir.py", line 1219, in fallback
xla_computation = xla.primitive_subcomputation(
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/interpreters/xla.py", line 303, in primitive_subcomputation
ans = rule(ctx, avals_in, avals_out, *xla_args, **params)
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/interpreters/xla.py", line 531, in wrapped
ans = f(ctx.builder, *args, **kw)
File "/petrobr/parceirosbr/brcluster/sophia.estrela/nccl4jax/all_reduce_nccl.py", line 47, in nccl_allreduce_xla_encode_gpu
_np.intc(count),
File "/petrobr/parceirosbr/brcluster/sophia.estrela/cupy/lib/python3.9/site-packages/jax/core.py", line 534, in array
raise TracerArrayConversionError(self)
jax._src.traceback_util.UnfilteredStackTrace: jax._src.errors.TracerArrayConversionError: The numpy.ndarray conversion method array() was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
This Tracer was created on line /petrobr/parceirosbr/brcluster/sophia.estrela/nccl4jax/nccl4jax_allreduce.py:49 ()
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

---

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/scratch/app/python/3.9.12/lib/python3.9/runpy.py", line 188, in _run_module_as_main
mod_name, mod_spec, code = _get_module_details(mod_name, _Error)
File "/scratch/app/python/3.9.12/lib/python3.9/runpy.py", line 111, in _get_module_details
import(pkg_name)
File "/petrobr/parceirosbr/brcluster/sophia.estrela/nccl4jax/nccl4jax_allreduce.py", line 49, in
recv_buffer, token = jit_all_reduce(send_buffer, count, dataType, op)
File "/petrobr/parceirosbr/brcluster/sophia.estrela/nccl4jax/all_reduce_nccl.py", line 47, in nccl_allreduce_xla_encode_gpu
_np.intc(count),
jax._src.errors.TracerArrayConversionError: The numpy.ndarray conversion method array() was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
This Tracer was created on line /petrobr/parceirosbr/brcluster/sophia.estrela/nccl4jax/nccl4jax_allreduce.py:49 ()
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

When I remove the variables count, dataType, and op it works. I do not understand why this does not work, since they are simple integers at start, and the conversion to DynamicJaxprTracer is done by the jit itself.
Can someone explain to me what is wrong or why this is returning an error?

Github with the code

Thank you for the help!

[YouJiacheng]

JAX JIT will replace all non-static arguments with Tracers, and Tracers cannot be converted into numpy.ndarray.
I suggest you use sendbuffer.size as count and sendbuffer.dtype as dataType, just like mpi4jax.

[sophia-estrela]

Hi! I changed what you said, it was really helpful, but I still can not understand why the op cannot be passed, I do not use the op as numpy.ndarray at any point of my code ?
Thank you for your help!

[YouJiacheng]

JAX might recognize your op as int...

[hawkinsp]

This isn't an answer to your question, but are you aware that JAX already supports using NCCL on GPU, including across hosts? See jax.pmap and operations like psum.