Question about pjit

[zhao-han]

I wrote a simple demo to test pjit, the demo code tries to split the graph across 8 devices.

import jax
from jax import numpy as jnp
from jax.experimental import maps
from jax.experimental import PartitionSpec as P
from jax.experimental.pjit import pjit
import numpy as np

print("Backend: {}, device_count: {}".format(jax.default_backend(), jax.device_count()))

x = np.random.normal(size=(4, 8))
w0 = np.random.normal(size=(8, 8))
b0 = np.random.normal(size=(4, 8))
w1 = np.random.normal(size=(8, 8))
b1 = np.random.normal(size=(4, 8))

def simple_graph(x, w0, b0, w1, b1):
  output0 = jnp.matmul(x, w0) + b0
  act0 = jnp.tanh(output0)

  output1 = jnp.matmul(act0, w1) + b1
  act1 = jnp.tanh(output1)
  return act1

mesh_shape = (1, 8)
devices = np.asarray(jax.devices()).reshape(*mesh_shape)
mesh = maps.Mesh(devices, ('x', 'y'))
print(mesh)

f = pjit(
  simple_graph,
  in_axis_resources=(P('x', 'y'), P('y', 'x'), P('x', 'y'), P('y', 'x'), P('x', 'y')),
  out_axis_resources=P('x', 'y'))

with maps.mesh(mesh.devices, mesh.axis_names):
 result = f(x, w0, b0, w1, b1)

print(result)

It works fine.
But after I modified input variables to below code:

x = np.random.normal(size=(4, 8))
w0 = np.random.normal(size=(8, 5))
b0 = np.random.normal(size=(4, 5))
w1 = np.random.normal(size=(5, 7))
b1 = np.random.normal(size=(4, 7))

I got error:

ValueError: One of pjit arguments was given the resource assignment of PartitionSpec('x', 'y'), which implies that the size of its dimension 1 should be divisible by 8, but it is equal to 5

Seems like pjit can't handle the uneven partition. But from the paper, unevenly partition is supported by GSPMD.
My question is: how to modify my code to split that uneven graph on 8 devices? Look forward to your feedback!

[apaszke]

Hi! At the moment pjit only supports even partitioning on the function boundary. with_sharding_constraint doesn't have the same limitation, but it can only be applied in the body of a pjitted function.

I don't think this is a fundamental restriction, but it would certainly require some care in handling the boundary conditions that we haven't gotten to just yet!

[zhao-han]

Thanks for your reply, really helpful!