<h3 dir="auto">Description</h3>
<p dir="auto">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.<br>
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.</p>
<p dir="auto"><em>I'm using the nifty <a href="https://github.com/ayaka14732/jax-smi">jax_smi</a> to monitor live memory utilization</em></p>
<h2 dir="auto">Single device allocation</h2>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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)"><pre class="notranslate"><code class="notranslate">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)
</code></pre></div>
<p dir="auto">Output:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="┌───────┐
│ TPU 0 │
└───────┘"><pre class="notranslate"><code class="notranslate">┌───────┐
│ TPU 0 │
└───────┘
</code></pre></div>
<p dir="auto">Memory usage:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content=" device: Total 7.5GB
         7.5GB (  100%): TPU_0(process=0,(0,0,0,0))"><pre class="notranslate"><code class="notranslate"> device: Total 7.5GB
         7.5GB (  100%): TPU_0(process=0,(0,0,0,0))
</code></pre></div>
<p dir="auto">Here I've allocated almost the entire HBM for a single TPU v2 device.</p>
<h2 dir="auto">Sharding over 8 devices - success</h2>
<p dir="auto">When sharding, I'm able to allocate half of the above size (combined over all shards).<br>
The following will succeed, with the array size being 1_000_000_000, exactly half of before:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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)"><pre class="notranslate"><code class="notranslate">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)
</code></pre></div>
<p dir="auto">Output:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="┌───────┬───────┬───────┬───────┬───────┬───────┬───────┬───────┐
│ TPU 0 │ TPU 1 │ TPU 2 │ TPU 3 │ TPU 6 │ TPU 7 │ TPU 4 │ TPU 5 │
└───────┴───────┴───────┴───────┴───────┴───────┴───────┴───────┘"><pre class="notranslate"><code class="notranslate">┌───────┬───────┬───────┬───────┬───────┬───────┬───────┬───────┐
│ TPU 0 │ TPU 1 │ TPU 2 │ TPU 3 │ TPU 6 │ TPU 7 │ TPU 4 │ TPU 5 │
└───────┴───────┴───────┴───────┴───────┴───────┴───────┴───────┘
</code></pre></div>
<p dir="auto">Memory usage:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content=" 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))"><pre class="notranslate"><code class="notranslate"> 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))
</code></pre></div>
<h2 dir="auto">Sharding over 8 devices - fail</h2>
<p dir="auto">When slightly increasing the array size (to 1_100_000_000), we fail with RESOURCE_EXHAUSTED:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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)"><pre class="notranslate"><code class="notranslate">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)
</code></pre></div>
<p dir="auto">Output:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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 &lt;module&gt;
----&gt; 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)

...
-&gt; 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"><pre class="notranslate"><code class="notranslate">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 &lt;module&gt;
----&gt; 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)

...
-&gt; 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
</code></pre></div>
<p dir="auto">(<em>I've truncated the entire error message, having a long stack</em>)</p>
<h2 dir="auto">CPU devices</h2>
<p dir="auto">To verify this is not specific to TPU, I've recreated this on several CPUs on the TPU host, using</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import os
os.environ[&quot;XLA_FLAGS&quot;] = '--xla_force_host_platform_device_count=8' # Use 8 CPU devices
import jax
jax.config.update('jax_platform_name', 'cpu')"><pre class="notranslate"><code class="notranslate">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')
</code></pre></div>
<p dir="auto">Having what I believe is 300GB of RAM on the host, I'm able to allocate a single [80_000_000_000] array:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content=" device: Total 298.0GB
         298.0GB (  100%): TFRT_CPU_0"><pre class="notranslate"><code class="notranslate"> device: Total 298.0GB
         298.0GB (  100%): TFRT_CPU_0
</code></pre></div>
<p dir="auto">When allocating a sharded [40_000_000_000] array, I can see an intermedia allocation of the entire array on a single device:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content=" device: Total 149.0GB
         149.0GB (  100%): TFRT_CPU_0"><pre class="notranslate"><code class="notranslate"> device: Total 149.0GB
         149.0GB (  100%): TFRT_CPU_0
</code></pre></div>
<p dir="auto">and when finalized, it is distributed across all devices:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content=" 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"><pre class="notranslate"><code class="notranslate"> 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
</code></pre></div>
<p dir="auto">However, trying to allocate [45_000_000_000] will fail; this time, with a kernel crash:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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."><pre class="notranslate"><code class="notranslate">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.
</code></pre></div>
<p dir="auto">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:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content=" device: Total 167.6GB
         167.6GB (  100%): TFRT_CPU_0"><pre class="notranslate"><code class="notranslate"> device: Total 167.6GB
         167.6GB (  100%): TFRT_CPU_0
</code></pre></div>
<h2 dir="auto">Summary</h2>
<p dir="auto">Is that the intended behavior?<br>
If so - what is the source of this allocation limit?</p>
<p dir="auto">P.S. - This is part of a <a href="https://github.com/google/jax/discussions/13842" data-hovercard-type="discussion" data-hovercard-url="/google/jax/discussions/13842/hovercard">larger effort </a>I'm working on.</p>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">jax 0.4.1, jaxlib 0.4.1</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">CPU &amp; TPU (v2-8)</p>
<h3 dir="auto">Additional system info</h3>
<p dir="auto">Cloud TPU, Ubuntu 20.04.4 LTS</p>
<h3 dir="auto">NVIDIA GPU info</h3>
<p dir="auto"><em>No response</em></p>
<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<p dir="auto">Code:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jaxlib
import jaxlib.xla_extension as xe

jax.random.PRNGKey(0)
print(jax.devices())
print(xe.CpuDevice)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">jaxlib</span>
<span class="pl-k">import</span> <span class="pl-s1">jaxlib</span>.<span class="pl-s1">xla_extension</span> <span class="pl-k">as</span> <span class="pl-s1">xe</span>

<span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">0</span>)
<span class="pl-en">print</span>(<span class="pl-s1">jax</span>.<span class="pl-en">devices</span>())
<span class="pl-en">print</span>(<span class="pl-s1">xe</span>.<span class="pl-v">CpuDevice</span>)</pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>
<p dir="auto">Output:</p>
<div class="highlight highlight-source-shell notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="[GpuDevice(id=0, process_index=0)]
Traceback (most recent call last):
  File &quot;~/jax/1_installation_test.py&quot;, line 8, in &lt;module&gt;
    print(xe.CpuDevice)
AttributeError: module 'jaxlib.xla_extension' has no attribute 'CpuDevice'"><pre class="notranslate">[GpuDevice(id=0, process_index=0)]
Traceback (most recent call last):
  File <span class="pl-s"><span class="pl-pds">"</span>~/jax/1_installation_test.py<span class="pl-pds">"</span></span>, line 8, <span class="pl-k">in</span> <span class="pl-k">&lt;</span>module<span class="pl-k">&gt;</span>
    print(xe.CpuDevice)
AttributeError: module <span class="pl-s"><span class="pl-pds">'</span>jaxlib.xla_extension<span class="pl-pds">'</span></span> has no attribute <span class="pl-s"><span class="pl-pds">'</span>CpuDevice<span class="pl-pds">'</span></span></pre></div>
<p dir="auto">Environment details:</p>
<ul dir="auto">
<li>Fedora 35 with cuda 11.6, cudnn 8.4</li>
<li>jax 0.3.7</li>
<li>jaxlib 0.3.7+cuda11.cudnn82</li>
<li>python 3.9</li>
</ul>