Does pjit have partition rules?

[coderlemon17]

Hi, I'm trying to understand what pjit is actually doing following the official docs, and I choose the simple identical mapping lambda x:x for the experiment as follows.

import jax
from jax.config import config
from jax.experimental import maps
from jax.experimental import PartitionSpec
from jax.experimental.pjit import pjit

# [GpuDevice(id=0, process_index=0), GpuDevice(id=1, process_index=0)]
print(jax.devices())

devices = np.asarray(jax.devices()).reshape(2, 1)
mesh = maps.Mesh(devices, ('x', 'y'))

# array([[0, 1, 2, 3],
#       [4, 5, 6, 7]])
input_data = np.arange(8).reshape(2, 4)

def dummy_f(x):
  return x

f = pjit(
    dummy_f,
    in_axis_resources=PartitionSpec('x', None),
    out_axis_resources=None
  )

with maps.Mesh(mesh.devices, ('x', 'y')):
    data = f(input_data)

print(data,  data.device_buffers)

What I expect is:

array([[0, 1, 2, 3],        =(in_axis)=>    device0: [[0, 1, 2, 3]]    
       [4, 5, 6, 7]])         =(in_axis)=>   device1: [[4, 5, 6, 7]] 

=(out_axis)=> device0: [[0, 1, 2, 3], [4, 5, 6, 7]]
=(out_axis)=> device1: [[0, 1, 2, 3], [4, 5, 6, 7]]

However, I receive an NCCL error:

2022-07-31 19:06:04.092649: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2130] Execution of replica 0 failed: INTERNAL: external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nccl_utils.cc:245: NCCL operation ncclCommInitRank(comm.get(), nranks, id, rank) failed: unhandled cuda error
2022-07-31 19:06:04.092780: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2130] Execution of replica 0 failed: INTERNAL: external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nccl_utils.cc:245: NCCL operation ncclCommInitRank(comm.get(), nranks, id, rank) failed: unhandled cuda error

Note: If I change out_axis_resources=None to out_axis_resources=PartitionSpec('x', None), the code can run normally, so I think the hardware should be fine.

I have a few concerns:

• In official docs, it says One way to do data redistribution cross-host is to use an identity pjit with in_axis_resources different from out_axis_resources. XLA will do the global data reordering for you via pjit..

  • So I assume with different in_axis_resources and out_axis_resources, I can move tensors across devices, is that right?

• I understand that in_axis_resources and out_axis_resources define how tensor is sharded. But does out_axis_resources shard each tensor device-wise or global-wise? I.e. does out_axis_resources act on each tensor per device? or on the concatenation of tensors from all devices?

[sudhakarsingh27]

Which jax version are you using?

I can run this code successfully on my end with following stdout (I'm using jax 0.3.15):

[GpuDevice(id=0, process_index=0), GpuDevice(id=1, process_index=0)]
[[0 1 2 3]
 [4 5 6 7]] [DeviceArray([[0, 1, 2, 3],
             [4, 5, 6, 7]], dtype=int32), DeviceArray([[0, 1, 2, 3],
             [4, 5, 6, 7]], dtype=int32)]
# pip list | grep jax
jax                           0.3.15
jaxlib                        0.3.15

[coderlemon17]

Thanks for the reply, here is my jax version

~/Workspace/Jax » pip list | grep jax 
jax                  0.3.15
jaxlib               0.3.15+cuda11.cudnn82

And I'm using NVIDIA Driver 510.47.03, here is my cuda version:

~/Workspace/Jax » nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2020 NVIDIA Corporation
Built on Tue_Sep_15_19:10:02_PDT_2020
Cuda compilation tools, release 11.1, V11.1.74
Build cuda_11.1.TC455_06.29069683_0

[coderlemon17]

Thanks, it's due to Jax and CUDA version mismatch, I upgrade CUDA from 11.1 to 11.4 and it works.
Another solution is to install "jax[cuda11_cudnn805]".

Also, I was thinking maybe we can specify CUDA requirement to make the installation instruction clearer by changing:

pip install --upgrade pip
# Installs the wheel compatible with CUDA 11 and cuDNN 8.2 or newer.
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

to

pip install --upgrade pip
# Installs the wheel compatible with CUDA 11.4 and cuDNN 8.2 or newer.
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html