Parallelisation with Numba LocalOperators without MPI?

[waleed-sh]

Hi,

I have a question about using NetKet with parallelisation enabled.

So far, I have used Numba based LocalOperators and have used MPI for parallelisation, purely CPU based. This is because I had to, for some reason, modify the get_conn_flattened_kernel function and that was easier to do in the Numba version... I saw now in the latest release that MPI is deprecated, and parallelisation is carried out in Jax.

My questions are, can I still use my Numba operators and make use of parallelisation, or do I have to rewrite my Numba operator in Jax? If I can still use the Numba based operators, are there any performance drawbacks? Also, does the Jax parallelisation run on CPU and GPU, or only GPU?

Thanks in advance!

[PhilipVinc]

Hi,

1a) It is possible to write a JaxOperator and write the implementation in normal python/numba. You 'just' need to use a jax.pure_callback within jax. (If you do try this, It would be nice if we could add a tutorial/example showcasing how to do this). This would not have considerable slowdowns.
1b) Right now NumbaOperators do not work when sharding because of... some issues. If all goes well I should be able to lift the limitation so you would not even have to do the thing above.

Also, does the Jax parallelisation run on CPU and GPU, or only GPU?

2. Easily, only on GPU. However it's possible to instruct jax to use MPI for CPU-to-CPU parallelisation (this is different from the MPI support that we removed). In that case you'd be using jax parallelism, as gpu codes always did, but it would be running on MPI.
   This is very briefly documented at https://netket.readthedocs.io/en/latest/distributed-computing.html#advanced-mpitrampoline-for-multi-cpu-distributed-computing.

You can try to use this package that I experimented with https://github.com/mpi4jax/mpibackend4jax and let me know how it goes. It should allow you to still use MPI.

[waleed-sh]

1a) It is possible to write a JaxOperator and write the implementation in normal python/numba. You 'just' need to use a jax.pure_callback within jax. (If you do try this, It would be nice if we could add a tutorial/example showcasing how to do this). This would not have considerable slowdowns.

Ah, okay. I will give that a shot and come back here for an update in the next days as this seems to be the "easiest" way to go about it for now... Once I get the approach "green lighted" here, I can gladly write a tutorial

1b) Right now NumbaOperators do not work when sharding because of... some issues. If all goes well I should be able to lift the limitation so you would not even have to do the thing above.

Yeah I can only imagine, I saw sharing was removed some versions earlier anyway. That would be super cool if the limitation is lifted but until then, pure callback it is!

2. Easily, only on GPU. However it's possible to instruct jax to use MPI for CPU-to-CPU parallelisation (this is different from the MPI support that we removed). In that case you'd be using jax parallelism, as gpu codes always did, but it would be running on MPI.
   This is very briefly documented at https://netket.readthedocs.io/en/latest/distributed-computing.html#advanced-mpitrampoline-for-multi-cpu-distributed-computing.

You can try to use this package that I experimented with https://github.com/mpi4jax/mpibackend4jax and let me know how it goes. It should allow you to still use MPI.

I will give this a try. In a sense I want to eventually move fully to GPU but the bottleneck is this one specific numba based LocalOperator... I was just curious if to know if I use Jax parallelisation will that break my code or somehow all works out.

In general I don't want to break things for people. As I assume you are doing quite crazy things with your gauge-invariant things, it could be a good playground to check that the callback works fine.
You can get in contact with me over slack, and I can help you sort it out if needed.

Of course, this is understandable!! Yes, indeed I am doing crazy things. My main issue was actually constructing functions of operators, but still maintain the really convenient high level "operator arithmetics" offered by NetKet. This was quite a ride... and it required me to modify the gccf kernels, which I did in numba and currently trying to get working in Jax. I will be working on the pure callback, and if things really get ugly with parallelisation I will get in touch via slack! Thanks for all the help!! Always super helpful!!

[PhilipVinc]

Gccf?

[waleed-sh]

Gccf?

"Get connected components function", typically what I call the get_conn_flattened_kernel and similar functions

[PhilipVinc]

ah yeah, it was the stuff necessary for numba4jax.
Unfortunately I removed because I had no more time to keep maintaining it...
but I think you can still use it within your operators...

[waleed-sh]

Currently I am using NetKet 3.17.1 and it works fine, but I want to keep the codebase always up to date with the latest NetKet release so I am right now trying to get up to speed

[inailuig]

Some options for writing a wrapper for the numba oparator mentioned in point 1 above:

I) use jax.pure_callback inside of shard_map

• works under jit, you can use the same code as for the jax operators
• you need unflatten and pad the gcf to a known, common to all ranks max_conn_size inside of the callback
• callback overhead; if you only call it once on all samples together it's probably negligible
• See https://github.com/inailuig/netket/blob/1fe9d5ea534ec653ec597b351db4e71f7e5a3775/netket/jax/distributed.py#L11 (assumes shape is the same on all ranks)

II) just run it outside of jit

• convert to numpy, run numba, do an allgather with jax to find the common maxconnsize, pad in numpy, convert to a global device array with the correct sharding
• some overhead, recompilation if max_conn_size changes
• this is essentially what we were doing previously with MPI, modulo some extra paddding to make it the same on all ranks
• see

  netket/netket/jax/sharding.py

  Line 60 in c4f1748

  def replicate_sharding_decorator_for_get_conn_padded(f):

  (needs to be updated for NamedSharding)

III) run numba inside of jax directly
- e.g. using something like numba4jax (can probably be updated/revived using ffi_call for newer versions of jax)
- works under jit, should have little to no overhead, can be called for small batches at a time
- you would need to know the max_conn_size a priori, and pad inside of the function

[PhilipVinc]

See https://github.com/inailuig/netket/blob/1fe9d5ea534ec653ec597b351db4e71f7e5a3775/netket/jax/distributed.py#L11 (assumes shape is the same on all ranks)

we should merge this in netket. It's super useful. We could do it in jax 0.7 which I'll merge later today, for numba operators.

[PhilipVinc]

2. will still work. It would be easy to ensure the sharding is correct.

[waleed-sh]

This was very helpful! I am getting a first version ready which follows jax.pure_callback inside of a shard_map. Still some way to go, but this was very helpful. I will look into (II) afterwards and once I have everything running I will see any performance drawbacks and see if there is any need to do (III). Hopefully some day I can write the get_conn_padded in purely Jax but until then, thanks for all the feedback!!

See https://github.com/inailuig/netket/blob/1fe9d5ea534ec653ec597b351db4e71f7e5a3775/netket/jax/distributed.py#L11 (assumes shape is the same on all ranks)

we should merge this in netket. It's super useful. We could do it in jax 0.7 which I'll merge later today, for numba operators.

I think this would be helpful for perhaps the very small number of people like me who are still using Numba operators for some reason

[waleed-sh]

Here is my attempt, can someone please give me some feedback? I would really appreciate it! If it all checks out, I would be happy to write a tutorial for it or perhaps a "Jax local operator wrapper" class if there is actually a use case for it other than myself...

P.S. please dont yell at me if it is really horrible... also sorry for the long upcoming discussion

What I did: I went with suggestion (1a) by @PhilipVinc (jax.pure_callback) and using the decorator given by @inailuig

1 - The sharding decorator

I used the code provided by @inailuig, modified a tiny bit (made the out_specs=(P("d"), P("d")) instead of out_specs=(P("d"),) since I thought the Numba get_conn_padded() returns two Jax arrays), which is here below

from functools import partial

import jax
from jax import ShapeDtypeStruct, pure_callback
from jax.experimental.shard_map import shard_map
from jax.sharding import Mesh, PartitionSpec as P


def _replicate_shmap_callback(f_py, x):
    """
    This is a decorator which is used in Jax-compatible custom LocalOperator types. It is designed
    to call the .get_conn_padded() of the Numba version of the operator and return its result.

    So it executes a python callback `f_py` (returns xp, mels) on every host with its local shard of
    `x`, then stitch results together.
    """

    mesh = Mesh(jax.devices(), axis_names=("d",))

    @partial(shard_map, mesh=mesh, in_specs=(P("d"),), out_specs=(P("d"), P("d")))
    def _per_shard(x_shard):

        local_batch, n_sites = x_shard.shape

        xp_spec = ShapeDtypeStruct((local_batch, f_py.max_conn, n_sites), x_shard.dtype)

        mels_spec = ShapeDtypeStruct((local_batch, f_py.max_conn), f_py.mel_dtype)

        xp_loc, mels_loc = pure_callback(
            f_py,
            (xp_spec, mels_spec),
            x_shard,
            vmap_method="legacy_vectorized",
        )

        return xp_loc, mels_loc

    return _per_shard(x)


def replicate_sharding(f_py):
    """
    A decorator to replicate the sharding mechanism for jax.pure_callback of get_conn_padded()
    of Numba LocalOperators from their Jax-compatible types
    """
    return partial(_replicate_shmap_callback, f_py)

2 - Jax LocalOperator Wrapper

I then created a wrapper essentially which acts as a LocalOperatorJax. It does the jax.pure_callback (wrapped with the sharding decorator above) to the Numba version of the operator, there is no more any computation done by the Jax native _local_operator_kernel_jax(). However, I had to modify both the tree_flatten() and tree_unflatten(). Maybe this was unnecessary, but I kept getting a JaxOperatorNotConvertibleToNumba causing the callback to fail. Also, I did not test converting it to Pauli strings so I just marked it as not implemented for now.

Here is the implementation

from typing import TYPE_CHECKING

from jax.tree_util import register_pytree_node_class

from netket.errors import JaxOperatorNotConvertibleToNumba
from netket.operator._discrete_operator_jax import DiscreteJaxOperator

from netket.operator._local_operator.base import LocalOperatorBase

if TYPE_CHECKING:
    from netket.operator._local_operator.numba import LocalOperator


@register_pytree_node_class
class LocalOperatorJaxWrapper(LocalOperatorBase, DiscreteJaxOperator):
    """
    A Jax-compatible wrapper LocalOperator which works (?) with Jax sharding. It does not implement
    the .get_conn_padded() function in Jax, but it executes a jax.pure_callback to the Numba version
    and that is sharded across devices (?).
    """

    # can still go back to Numba
    _convertible: bool = True

    def _numba_operator(self) -> "LocalOperator":
        """
        Cached Numba version of the operator
        """

        if not hasattr(self, "_cached_numba_op"):

            # if we dont have a cached version, construct the operator once and cache it
            from netket.operator._local_operator.numba import LocalOperator

            if self._convertible is False:
                # check if convertible or not
                raise JaxOperatorNotConvertibleToNumba(self)

            # construct the concrete LocalOperator and cache
            self._cached_numba_op = self.to_numba_operator()

        # return the cache
        return self._cached_numba_op

    def to_numba_operator(self) -> "LocalOperator":
        """
        Convert the Jax-compatible operator to a Numba version
        """

        if self._convertible is False:
            raise JaxOperatorNotConvertibleToNumba(self)

        from netket.operator._local_operator.numba import LocalOperator

        return LocalOperator(
            self.hilbert,
            self.operators,
            self.acting_on,
            self.constant,
            dtype = self.dtype,
            mel_cutoff = self.mel_cutoff,
        )

    def _setup(self, force: bool = False):
        """
        Setup up the Jax-compatible LocalOperator wrapper. This is done once-per-object and the
        _setup() of the Numba based implementation is called instead.
        """

        if force or not getattr(self, "_initialized", False):

            # get the Numba operator
            op_nb = self._numba_operator()

            # setup the Numba operator
            op_nb._setup()

            # store some things for tracing
            self._max_conn_size = int(op_nb._max_conn_size)
            self._nonzero_diagonal = bool(op_nb._nonzero_diagonal)

            # set as initialised to avoid doing this again
            self._initialized = True

    def _get_conn_padded(self, x):
        """
        Host‑side Numba lookup wrapped in a sharded jax.pure_callback.
        """

        # setup first, will call Numba's setup if not already initialised
        self._setup()

        # get the maximum number of connections and the expected dtype of the matrix elements
        max_conn = self._max_conn_size

        # match the dtype
        mel_dtype = self.dtype

        # python callback executed on host
        def _py_kernel(x_np: np.ndarray):
            """
            The kernel function for the jax.pure_callback
            """

            # call the Numba .get_conn_padded()
            xp_np, mels_np = self._numba_operator().get_conn_padded(x_np)

            # return the values
            return xp_np, mels_np

        # attach static attributes used by wrapper
        _py_kernel.max_conn = max_conn
        _py_kernel.mel_dtype = mel_dtype

        # wrap the python kernel so it is called once per host shard
        xp, mels = replicate_sharding(_py_kernel)(x)

        # on‑device post‑processing
        # count the number of nonzero connections (nonzero means > the mels cutoff)
        n_conn = (jax.numpy.abs(mels) > self._mel_cutoff).sum(axis = -1, dtype = jax.numpy.int32)

        # return the values
        return xp, mels, n_conn

    def get_conn_padded(self, x):
        xp, mels, _ = self._get_conn_padded(x)
        return xp, mels

    def n_conn(self, x):
        _, _, n_conn = self._get_conn_padded(x)
        return n_conn

    # stay object opaque for flatten/unflatten
    @classmethod
    def tree_flatten(self):
        return (), self

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        return aux_data

    def to_pauli_strings(self):
        raise NotImplementedError(
            "This operator cannot be converted into Pauli strings type"
        )

3 - Sanity checks

I solved the Heisenberg model from the NetKet tutorials just for sanity. After optimisation, all of the expectation values of the following matched (to the correct value):

• the NetKet Numba operator (called ha = nk.operator.Heisenberg(hilbert=hi, graph=g))
• the NetKet Jax operator (called ha_jax_nk and used NetKet's native ha.to_jax_operator())
• the modified Jax wrapper (called ha_cp_jax, obtained by replacing the NetKet .to_jax_operator() with one which calls the LocalOperatorJaxWrapper instead)

Also a simple %timeit on all three gave the following

%timeit vs.expect(ha)
>>> 7.2 ms ± 1.2 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit vs.expect(ha_jax_nk)
>>> 4.88 ms ± 735 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit vs.expect(ha_cp_jax)
>>> 7.21 ms ± 111 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

So it does not seem that there is a (noticeable) drawback in compute time compared to the Numba version. I have not done any GPU testing yet! I wanted some feedback first if possible. Which brings me to the point below...

4 - Concerns/questions

Now my questions/concerns are:

1. Is the implementation correct or did I do something really stupid somewhere? Locally, everything seems to work. But I dont know if I have sharded everything correctly or not (I am not the most experienced with Jax) or if any issues arise anywhere because of this.
2. How bad is it that I modified the .tree_flatten() (original here) and .tree_unflatten() (original here). At least as of now, I did not find a workaround (but I just wanted something that works right now so I can still think of something). But I dont know if this will cause any issues when running on GPUs or performance drawbacks.

Thanks a lot for any feedback, this was of tremendous help.

[waleed-sh]

To address it I think you can try to call òp_numba.get_conn_padded(x, pad=op.max_conn_size)?

I tried doing xp_np, mels_np = self._numba_operator().get_conn_padded(x_np, pad = max_conn) but that gave me a TypeError: DiscreteOperator.get_conn_padded() got an unexpected keyword argument 'pad'. However, using this py_kernel instead

def _py_kernel(x_np: np.ndarray):
    """
    The kernel function for the jax.pure_callback
    """

    # call the Numba .get_conn_padded()
    xp_np, mels_np = self._numba_operator().get_conn_padded(x_np)

    # pad to common max_conn, in case the Numba version did not do a good job
    if xp_np.shape[1] < max_conn:
        pad = max_conn - xp_np.shape[1]
        xp_pad = np.repeat(x_np[:, None, :], pad, 1)
        mels_pad = np.zeros((x_np.shape[0], pad), dtype = mels_np.dtype)
        xp_np = np.concatenate([xp_np,   xp_pad], axis = 1)
        mels_np = np.concatenate([mels_np, mels_pad], axis = 1)

    # return the values
    return xp_np, mels_np

should do the job. I also kept vmap_method="legacy_vectorized", and optimising with respect to ha_cp_jax should then run normally and produce the correct results.

[waleed-sh]

it's a bug in Shardy. It would be good to reduce a MWE and share it with jax devs...
If I take this script https://gist.github.com/PhilipVinc/bd3e08b5b500070995761e001ff73207
and try to run it with JAX_USE_SHARDY_PARTITIONER=0 ipython -I example2.py it works fine. So it's the new shady partitioner that is broken when you use a callback...

(I'm working with netket master and jax 0.7)

I will try later tonight to install from the master branch and rerun everything. For some reason I got some issues with the dependencies, using Python 3.12.1 and 3.11.7 but I am probably messing up something somewhere. I also might have an idea for resolving this issue but I need to do this on the new code.

Either way if it is indeed a bug in Jax that is being hit then it's worth sharing it with them.

[waleed-sh]

it's a bug in Shardy. It would be good to reduce a MWE and share it with jax devs...
If I take this script https://gist.github.com/PhilipVinc/bd3e08b5b500070995761e001ff73207
and try to run it with JAX_USE_SHARDY_PARTITIONER=0 ipython -I example2.py it works fine. So it's the new shady partitioner that is broken when you use a callback...

Can confirm this too, after finally getting the master branch running locally. I will try to do a MWE which is just pure Jax that can be easily shared with the Jax devs if that would be helpful.

- A better implementation would be

   # stay object opaque for flatten/unflatten
   @classmethod
   def tree_flatten(self):
       self._setup()
       data, static = DiscreteJaxOperator.tree_flatten(self)
       static['_cached_numba_op'] = self. _cached_numba_op 
       return data, static
   

   @classmethod
   def tree_unflatten(cls, static, data):
        # copy paste from DiscreteJaxOperator
        op. _cached_numba_op = static['_cached_numba_op']
        return op

This is somehow still giving me some errors/headache. I think if the issue is with in-place modification of operators, then in my use case, that will never happen, and perhaps it is safe to use as is with the not-so-kosher version of the flatten/unflatten until figured out...

I really appreciate all of the feedback given here! This was extremely helpful!!

[waleed-sh]

It would be good to reduce a MWE and share it with jax devs...

Maybe not the most sensible code, but here is a MWE with just Jax that reproduces the bug. Didn't test with djaxrun, just python locally, but the issue should still be the same.

[PhilipVinc]

I tried doing xp_np, mels_np = self._numba_operator().get_conn_padded(x_np, pad = max_conn) but that gave me a TypeError: DiscreteOperator.get_conn_padded()

sorry, I thought it would work but it's not implemented anywhere. You are right,
Your solution is the good one.

[PhilipVinc]

This is somehow still giving me some errors/headache

what's the error?

[waleed-sh]

I basically copy/pasted the code for the flatten/unflatten from the original implementation and modified them to add the _cached_numba_op as suggested. Now they look like this

    def tree_flatten(self):
        self._setup()

        data = (
            self._acting_on,
            self._n_conns,
            self._diag_mels,
            self._x_prime,
            self._mels,
            self._basis,
            self._constant,
        )

        metadata = {
            "hilbert": self.hilbert,
            "dtype": self.dtype,
            "nonzero_diagonal": self._nonzero_diagonal,
            "max_conn_size": self._max_conn_size,
            "_cached_numba_op": self. _cached_numba_op
        }

        return data, metadata

    @classmethod
    def tree_unflatten(cls, metadata, data):

        hi = metadata["hilbert"]
        dtype = metadata["dtype"]

        op = cls(hi, dtype=dtype)

        op._nonzero_diagonal = metadata["nonzero_diagonal"]
        op._max_conn_size = metadata["max_conn_size"]
        (
            op._acting_on,
            op._n_conns,
            op._diag_mels,
            op._x_prime,
            op._mels,
            op._basis,
            op._constant,
        ) = data

        op._cached_numba_op = metadata["_cached_numba_op"]

        op._initialized = True
        op._convertible = False

        return op

I have also modified the _setup() to do this

    def _setup(self, force: bool = False):
        """
        Setup up the Jax-compatible MarkedLocalOperator. This is done once-per-object and the
        _setup() of the Numba based implementation is called instead.
        """

        # get the Numba operator
        op_nb = self._numba_operator()

        # setup the Numba operator
        op_nb._setup()

        self._acting_on = op_nb._acting_on
        self._n_conns = op_nb._n_conns
        self._diag_mels = op_nb._diag_mels
        self._x_prime = op_nb._x_prime
        self._mels = op_nb._mels
        self._basis = op_nb._basis

        # store some things for tracing
        self._max_conn_size = int(op_nb._max_conn_size)
        self._nonzero_diagonal = bool(op_nb._nonzero_diagonal)

        # set as initialised to avoid doing this again
        self._initialized = True

Now,

1. If I keep op._convertible = False as per the original implementation, I get a JaxOperatorNotConvertibleToNumba error arising when taking expectation values because the _setup() will try to do this op_nb = self._numba_operator()
2. If I set it to True, I get a TracerArrayConversionError because again, the Numba operator will be constructed and in its __init__ this line self._constant = np.array(constant, dtype=dtype) triggers the issue
3. And if I (just for fun) remove the constant from the flatten/unflatten, the expectation value is just incorrect (just 0.000e+00 ± nan [σ²=0.000e+00] so clearly something went wrong somewhere).

It's probably something I am overlooking/doing incorrectly, but it's late so I will try again in the morning...