Issue 1:
Description

When using sharding to distribute a large array over several devices, each shard is unable to utilize the entire device's RAM/HBM. In fact, it seems to be limited to about half that size.

I've recreated this both on Cloud TPU (v2-8, having 8 devices of 8GB each) and on the host of the TPU with 8 CPU devices.

I'm using the nifty jax_smi to monitor live memory utilization

Single device allocation

import jax
import jax.numpy as jnp
from jax.experimental import mesh_utils
from jax.sharding import PositionalSharding

from jax_smi import initialise_tracking
initialise_tracking(interval=0.1)

a = jnp.zeros([2_000_000_000])
jax.debug.visualize_array_sharding(a, use_color=False)

Output:

┌───────┐
│ TPU 0 │
└───────┘

Memory usage:

device: Total 7.5GB
7.5GB ( 100%): TPU_0(process=0,(0,0,0,0))

Here I've allocated almost the entire HBM for a single TPU v2 device.

Sharding over 8 devices - success

When sharding, I'm able to allocate half of the above size (combined over all shards).

The following will succeed, with the array size being 1_000_000_000, exactly half of before:

del a # clear previous allocation
devices = mesh_utils.create_device_mesh((jax.device_count()))
a = jax.device_put(
jnp.zeros([1_000_000_000]),
PositionalSharding(devices)
)
jax.debug.visualize_array_sharding(a, use_color=False)

Output:

┌───────┬───────┬───────┬───────┬───────┬───────┬───────┬───────┐
│ TPU 0 │ TPU 1 │ TPU 2 │ TPU 3 │ TPU 6 │ TPU 7 │ TPU 4 │ TPU 5 │
└───────┴───────┴───────┴───────┴───────┴───────┴───────┴───────┘

Memory usage:

device: Total 3.7GB
476.8MB (12.50%): TPU_0(process=0,(0,0,0,0))
476.8MB (12.50%): TPU_1(process=0,(0,0,0,1))
476.8MB (12.50%): TPU_2(process=0,(1,0,0,0))
476.8MB (12.50%): TPU_3(process=0,(1,0,0,1))
476.8MB (12.50%): TPU_4(process=0,(0,1,0,0))
476.8MB (12.50%): TPU_5(process=0,(0,1,0,1))
476.8MB (12.50%): TPU_6(process=0,(1,1,0,0))
476.8MB (12.50%): TPU_7(process=0,(1,1,0,1))

Sharding over 8 devices - fail

When slightly increasing the array size (to 1_100_000_000), we fail with RESOURCE_EXHAUSTED:

del a # clear previous allocation
a = jax.device_put(
jnp.zeros([1_100_000_000]),
PositionalSharding(devices)
)
jax.debug.visualize_array_sharding(a, use_color=False)

Output:

2023-01-19 13:17:02.322393: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2163] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Failed to allocate request for 524.52MiB (550002688B) on device ordinal 0
Output exceeds the [size limit](command:workbench.action.openSettings?[). Open the full output data [in a text editor](command:workbench.action.openLargeOutput?ebf90566-7cb0-46be-a7e1-235621eaddec)
---------------------------------------------------------------------------
XlaRuntimeError Traceback (most recent call last)
/tmp/ipykernel_2457251/453561208.py in <module>
----> 1 a = jax.device_put(
2 jnp.zeros([1_100_000_000]),
3 PositionalSharding(devices)
4 )
5 jax.debug.visualize_array_sharding(a, use_color=False)

...
-> 2136 out_bufs = self.xla_executable.execute_sharded_on_local_devices(
2137 input_bufs)
2138 if dispatch.needs_check_special():

XlaRuntimeError: RESOURCE_EXHAUSTED: Failed to allocate request for 524.52MiB (550002688B) on device ordinal 0

(I've truncated the entire error message, having a long stack)

CPU devices

To verify this is not specific to TPU, I've recreated this on several CPUs on the TPU host, using

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8' # Use 8 CPU devices
import jax
jax.config.update('jax_platform_name', 'cpu')

Having what I believe is 300GB of RAM on the host, I'm able to allocate a single [80_000_000_000] array:

device: Total 298.0GB
298.0GB ( 100%): TFRT_CPU_0

When allocating a sharded [40_000_000_000] array, I can see an intermedia allocation of the entire array on a single device:

device: Total 149.0GB
149.0GB ( 100%): TFRT_CPU_0

and when finalized, it is distributed across all devices:

device: Total 149.0GB
18.6GB (12.50%): TFRT_CPU_0
18.6GB (12.50%): TFRT_CPU_1
18.6GB (12.50%): TFRT_CPU_2
18.6GB (12.50%): TFRT_CPU_3
18.6GB (12.50%): TFRT_CPU_4
18.6GB (12.50%): TFRT_CPU_5
18.6GB (12.50%): TFRT_CPU_6
18.6GB (12.50%): TFRT_CPU_7

However, trying to allocate [45_000_000_000] will fail; this time, with a kernel crash:

Canceled future for execute_request message before replies were done
The Kernel crashed while executing code in the the current cell or a previous cell. Please review the code in the cell(s) to identify a possible cause of the failure. Click [here](https://aka.ms/vscodeJupyterKernelCrash) for more info. View Jupyter [log](command:jupyter.viewOutput) for further details.

The intermediate allocation is once again on a single device; it seems that there is enough memory for both the entire array allocation plus a single shard:

device: Total 167.6GB
167.6GB ( 100%): TFRT_CPU_0

Summary

Is that the intended behavior?

If so - what is the source of this allocation limit?

P.S. - This is part of a larger effort I'm working on.

What jax/jaxlib version are you using?

jax 0.4.1, jaxlib 0.4.1

Which accelerator(s) are you using?

CPU & TPU (v2-8)

Additional system info

Cloud TPU, Ubuntu 20.04.4 LTS

NVIDIA GPU info

No response

Issue 2:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

Code:

import jax
import jaxlib
import jaxlib.xla_extension as xe

jax.random.PRNGKey(0)
print(jax.devices())
print(xe.CpuDevice)

If applicable, include full error messages/tracebacks.

Output:

[GpuDevice(id=0, process_index=0)]
Traceback (most recent call last):
File "~/jax/1_installation_test.py", line 8, in <module>
print(xe.CpuDevice)
AttributeError: module 'jaxlib.xla_extension' has no attribute 'CpuDevice'

Environment details:

Fedora 35 with cuda 11.6, cudnn 8.4

jax 0.3.7

jaxlib 0.3.7+cuda11.cudnn82

python 3.9