Parallel signaling

[imoneoi]

Is there any parallel signaling method in JAX? I want to run parallel optimization on multiple nodes, and stop when all nodes have reached the stopping criterion.

One approach is use psum() every step to check the stopping criterion like below.
However, this will block nodes every step until all other nodes finished stepping.

@partial(jax.pmap, axis_name="num_devices")
def _opt_runner(state):
    # Evaluate
    def _parallel_stop_cond(state):
        is_finish = state.may_stop
        # Stop when all nodes have finished
        # However, psum will block until all nodes finished a step
        is_finish = jax.lax.psum(is_finish, "num_devices") == device_count
        return is_finish

    state = jax.lax.while_loop(_parallel_stop_cond, _optimize_a_step, state)

    # some other logic ...

Is there any non-blocking signaling approach, that every node notifies others when it finished, and stops when received notifications from all other nodes? (like code below)

@partial(jax.pmap, axis_name="num_devices")
def _opt_runner(state):
    # Evaluate
    def _parallel_stop_cond(state):
        # Stop when all nodes have finished
        is_finish = state.may_stop
        # Notify all other devices that this node have finished
        pnotify(is_finish)
        # Stop if it get enough notifications
        return pnotify_count() == device_count

    state = jax.lax.while_loop(_parallel_stop_cond, _optimize_a_step, state)