Asynchronous dispatch to specific devices / CPU cores

[markus7800]

Hello,

I know that the primary target of JAX is single-program multiple data applications.

However, I would like to kindly ask for advice on whether it is possible to use JAX in multiple-program multiple data applications.

Let's consider following benchmark:

import jax
import time
import jax.numpy as jnp
from jax.random import PRNGKey

devices = jax.devices()

def timer(name, f, x, verbose=False):
   start_wall = time.perf_counter()
   y = f(x)
   s = f"{y}" # this blocks until ready
   if verbose: print(s)
   end_wall = time.perf_counter()
   wall_time = end_wall - start_wall
   print(f"{name}: wall_time:{wall_time:.1f}s")
   
@jax.jit
def step_f(carry, any):
  val, key = carry
  key1, key2 = jax.random.split(key)
  val = val + jax.random.normal(key1)
  return (val,key2), None

@jax.jit
def step_g(carry, any):
  val, key = carry
  key1, key2 = jax.random.split(key)
  val = val + jax.random.uniform(key1, (2,)) - 0.5
  return (val,key2), None
  
@jax.jit
def f(seed):
  return jax.lax.scan(step_f, (jnp.array(0.,float), seed), length=10**6)[0][0]

@jax.jit
def g(seed):
  return jax.lax.scan(step_g, (jnp.array([0.,0.],float), seed), length=10**6)[0][0]

timer("base f", f, jax.device_put(PRNGKey(0), devices[0]))
timer("base g", g, jax.device_put(PRNGKey(1), devices[0]))

def fg_async(x):
    x1, x2 = x
    r1 = f(x1)
    r2 = g(x2)
    return (r1, r2)

timer("async 1 device", fg_async, (jax.device_put(PRNGKey(0), devices[0]),jax.device_put(PRNGKey(1), devices[0])))
timer("async 2 devices", fg_async, (jax.device_put(PRNGKey(0), devices[0]),jax.device_put(PRNGKey(1), devices[1])))

If we run this, for instance, on a Google Colab v2-8 TPU instance, we get following measurements:

base f: wall_time:8.1s
base g: wall_time:6.1s
async 1 device: wall_time:13.7s
async 2 devices: wall_time:7.8s

So far, so expected.

First Question: If we run this on a multi-core CPU (with --xla_force_host_platform_device_count set appropriately), async 2 devices will have the same runtime as async 1 device with only one CPU core being utilised. Is there a way to dispatch mutliple computations asynchronously on specific CPU cores to get back the speed-up?

Note that we cannot use pmap at all because the output shape of f and g are different.

A similar issue arises if we want to split up the computation to specific TPUs. Let's say compute f with the first 4 TPUs and g with the remaining TPUs:

def fg_pmap_async(x):
    x1, x2 = x
    r1 = jax.pmap(f)(x1)
    r2 = jax.pmap(g)(x2)
    return (r1, r2)

timer("pmap async", fg_pmap_async, (jax.random.split(PRNGKey(0),4), jax.random.split(PRNGKey(1),4)))

def fg_pmap_async_devices(x):
    x1, x2 = x
    r1 = jax.pmap(f, devices=devices[:4])(x1)
    r2 = jax.pmap(g, devices=devices[4:])(x2)
    return (r1, r2)

timer("pmap async devices", fg_pmap_async_devices, (jax.random.split(PRNGKey(0),4), jax.random.split(PRNGKey(1),4)))

pmap async: wall_time:14.0s
pmap async devices: wall_time:13.9s

Second Question: Is there a way to dispatch the computations in fg_pmap_async_devices asynchronously and do them in parallel?

Again, we cannot use pmap because the output shape of f and g are different.

General advice on the use case I have in mind is also very much appreciated.

[markus7800]

For CPU, a workaround that I consider would be some sort of multiprocessing setup, where

• we spawn N worker processes
• have one main process (which can compile the functions, serialise them via https://docs.jax.dev/en/latest/export/export.html, and send them to the worker processes)

However, this does not feel like an intended way of using JAX.
In particular, I think this would not work for the second question, as each worker process would try to reserve memory on the TPU.
Maybe there is some way around this with https://docs.jax.dev/en/latest/multi_process.html, but again does not feel like the intended use.