<p dir="auto">Hello Jax team,</p>
<p dir="auto">I have noticed that <code class="notranslate">jax.distributed</code> currently only supports a subset of GPU cores, and I was wondering if it would be possible to extend this support to a subset of TPU cores as well.</p>
<p dir="auto">I have tried running a code snippet with jax.pmap on TPUs using the suggested snippet from <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/skye/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/skye">@skye</a> in <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="1071238862" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/8809" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/8809/hovercard?comment_id=988495501&amp;comment_type=issue_comment" href="https://github.com/google/jax/issues/8809#issuecomment-988495501">#8809 (comment)</a>, but it only recognizes the local TPU device and not other devices, resulting in a failed psum operation.</p>
<p dir="auto">For comparison, I have also included the output of the same code snippet when run on GPUs for reference.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jax.numpy as jnp

jax.distributed.initialize()

print(&quot;process_count&quot;, jax.process_count())
print(&quot;process_index&quot;, jax.process_index())
print(&quot;local_devices&quot;, jax.local_devices())
print(&quot;local_device_count&quot;, jax.local_device_count())
print(&quot;devices&quot;, jax.devices())

def my_function(x):
    return jax.lax.psum(x, 'i')

x = jnp.ones((1, 2))
print(x)
parallel_fn = jax.pmap(my_function, devices=jax.devices(), axis_name='i')
result = parallel_fn(x)
print(result)"><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">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>

<span class="pl-s1">jax</span>.<span class="pl-s1">distributed</span>.<span class="pl-en">initialize</span>()

<span class="pl-en">print</span>(<span class="pl-s">"process_count"</span>, <span class="pl-s1">jax</span>.<span class="pl-en">process_count</span>())
<span class="pl-en">print</span>(<span class="pl-s">"process_index"</span>, <span class="pl-s1">jax</span>.<span class="pl-en">process_index</span>())
<span class="pl-en">print</span>(<span class="pl-s">"local_devices"</span>, <span class="pl-s1">jax</span>.<span class="pl-en">local_devices</span>())
<span class="pl-en">print</span>(<span class="pl-s">"local_device_count"</span>, <span class="pl-s1">jax</span>.<span class="pl-en">local_device_count</span>())
<span class="pl-en">print</span>(<span class="pl-s">"devices"</span>, <span class="pl-s1">jax</span>.<span class="pl-en">devices</span>())

<span class="pl-k">def</span> <span class="pl-en">my_function</span>(<span class="pl-s1">x</span>):
    <span class="pl-k">return</span> <span class="pl-s1">jax</span>.<span class="pl-s1">lax</span>.<span class="pl-en">psum</span>(<span class="pl-s1">x</span>, <span class="pl-s">'i'</span>)

<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>((<span class="pl-c1">1</span>, <span class="pl-c1">2</span>))
<span class="pl-en">print</span>(<span class="pl-s1">x</span>)
<span class="pl-s1">parallel_fn</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">pmap</span>(<span class="pl-s1">my_function</span>, <span class="pl-s1">devices</span><span class="pl-c1">=</span><span class="pl-s1">jax</span>.<span class="pl-en">devices</span>(), <span class="pl-s1">axis_name</span><span class="pl-c1">=</span><span class="pl-s">'i'</span>)
<span class="pl-s1">result</span> <span class="pl-c1">=</span> <span class="pl-en">parallel_fn</span>(<span class="pl-s1">x</span>)
<span class="pl-en">print</span>(<span class="pl-s1">result</span>)</pre></div>
<p dir="auto">On GPUs, when running with</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="CUDA_VISIBLE_DEVICES=&quot;0&quot; SLURM_JOB_ID=26017 SLURM_STEP_NODELIST=localhost SLURM_NTASKS=2 SLURM_PROCID=0 SLURM_LOCALID=0 SLURM_STEP_NUM_NODES=2  python t.py
CUDA_VISIBLE_DEVICES=&quot;1&quot; SLURM_JOB_ID=26017 SLURM_STEP_NODELIST=localhost SLURM_NTASKS=2 SLURM_PROCID=1 SLURM_LOCALID=0 SLURM_STEP_NUM_NODES=2  python t.py"><pre class="notranslate"><code class="notranslate">CUDA_VISIBLE_DEVICES="0" SLURM_JOB_ID=26017 SLURM_STEP_NODELIST=localhost SLURM_NTASKS=2 SLURM_PROCID=0 SLURM_LOCALID=0 SLURM_STEP_NUM_NODES=2  python t.py
CUDA_VISIBLE_DEVICES="1" SLURM_JOB_ID=26017 SLURM_STEP_NODELIST=localhost SLURM_NTASKS=2 SLURM_PROCID=1 SLURM_LOCALID=0 SLURM_STEP_NUM_NODES=2  python t.py
</code></pre></div>
<p dir="auto">produces</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="process_count 2
process_index 1
local_devices [StreamExecutorGpuDevice(id=1, process_index=1, slice_index=0)]
local_device_count 1
devices [StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=1, process_index=1, slice_index=0)]
[[1. 1.]]
[[2. 2.]]"><pre class="notranslate"><code class="notranslate">process_count 2
process_index 1
local_devices [StreamExecutorGpuDevice(id=1, process_index=1, slice_index=0)]
local_device_count 1
devices [StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=1, process_index=1, slice_index=0)]
[[1. 1.]]
[[2. 2.]]
</code></pre></div>
<p dir="auto">on TPUs, when running with the snippet suggested by <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/skye/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/skye">@skye</a> at <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="1071238862" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/8809" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/8809/hovercard?comment_id=988495501&amp;comment_type=issue_comment" href="https://github.com/google/jax/issues/8809#issuecomment-988495501">#8809 (comment)</a></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="TPU_CHIPS_PER_HOST_BOUNDS=1,2,1 TPU_HOST_BOUNDS=1,1,1 TPU_VISIBLE_DEVICES=0 TPU_MESH_CONTROLLER_ADDRESS=localhost:8476 TPU_MESH_CONTROLLER_PORT=8476 SLURM_JOB_ID=26017 SLURM_STEP_NODELIST=localhost SLURM_NTASKS=2 SLURM_PROCID=0 SLURM_LOCALID=0 SLURM_STEP_NUM_NODES=2 python t.py
TPU_CHIPS_PER_HOST_BOUNDS=1,2,1 TPU_HOST_BOUNDS=1,1,1 TPU_VISIBLE_DEVICES=1 TPU_MESH_CONTROLLER_ADDRESS=localhost:8476 TPU_MESH_CONTROLLER_PORT=8476 SLURM_JOB_ID=26017 SLURM_STEP_NODELIST=localhost SLURM_NTASKS=2 SLURM_PROCID=1 SLURM_LOCALID=0 SLURM_STEP_NUM_NODES=2 python t.py"><pre class="notranslate"><code class="notranslate">TPU_CHIPS_PER_HOST_BOUNDS=1,2,1 TPU_HOST_BOUNDS=1,1,1 TPU_VISIBLE_DEVICES=0 TPU_MESH_CONTROLLER_ADDRESS=localhost:8476 TPU_MESH_CONTROLLER_PORT=8476 SLURM_JOB_ID=26017 SLURM_STEP_NODELIST=localhost SLURM_NTASKS=2 SLURM_PROCID=0 SLURM_LOCALID=0 SLURM_STEP_NUM_NODES=2 python t.py
TPU_CHIPS_PER_HOST_BOUNDS=1,2,1 TPU_HOST_BOUNDS=1,1,1 TPU_VISIBLE_DEVICES=1 TPU_MESH_CONTROLLER_ADDRESS=localhost:8476 TPU_MESH_CONTROLLER_PORT=8476 SLURM_JOB_ID=26017 SLURM_STEP_NODELIST=localhost SLURM_NTASKS=2 SLURM_PROCID=1 SLURM_LOCALID=0 SLURM_STEP_NUM_NODES=2 python t.py
</code></pre></div>
<p dir="auto">we get the following.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="process_count 1
process_index 0
local_devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)]
local_device_count 1
devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)]
[[1. 1.]]
[[1. 1.]]"><pre class="notranslate"><code class="notranslate">process_count 1
process_index 0
local_devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)]
local_device_count 1
devices [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)]
[[1. 1.]]
[[1. 1.]]
</code></pre></div>
<p dir="auto">Note that the TPU example does not recognize other devices, thus making the <code class="notranslate">psum</code> fails.</p>
<h2 dir="auto">Use cases</h2>
<p dir="auto">We are using <code class="notranslate">jax.distributed</code> this way in our new distributed RL platform <a href="https://github.com/vwxyzjn/cleanba">Cleanba</a>, which supports different ways to leverage accelerators. Specifically, the users can only use 1 accelerator core for the actor, multiple accelerator cores for the learner, and replicate the training process through <code class="notranslate">jax.distributed</code>. Below is a figure showing our results</p>
<p dir="auto"><a target="_blank" rel="noopener noreferrer nofollow" href="https://user-images.githubusercontent.com/5555347/225036444-522fcb39-f70b-41b0-bd97-297a109fae39.png"><img width="753" alt="image" src="https://user-images.githubusercontent.com/5555347/225036444-522fcb39-f70b-41b0-bd97-297a109fae39.png" style="max-width: 100%;"></a></p>
<p dir="auto">The format <code class="notranslate">a0_l1 2 3 4_d1</code> means the actor runs on GPU 0 (<code class="notranslate">a0</code>), the learner runs on GPU 1 2 3 4 (<code class="notranslate">l1 2 3 4</code>), and the computation is replicated/distributed one time (<code class="notranslate">d1</code>). The blue baseline curve corresponds to <code class="notranslate">a0_l1 2 3_d2</code>.</p>
<p dir="auto">TPUv4 is fast, and the results on <code class="notranslate">a0_l1 2_d1</code> is pretty great. However, because of the problem listed in this issue, I couldn't test out the <code class="notranslate">a0_l1_d2</code> to fully utilize 4 chips on TPUv4. If <code class="notranslate">a0_l1_d2</code> were to work, I can expect it to be 10-20% faster than our current fastest 8 A100 setting, which would be very exciting 😬</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 requests.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Describe your goal, and if possible provide a code snippet with a motivating example.</li>
</ul>
<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"> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp
print(jnp.arange(10))
# [0 1 2 3 4 5 6 7 8 9]"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>
<span class="pl-en">print</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">arange</span>(<span class="pl-c1">10</span>))
<span class="pl-c"># [0 1 2 3 4 5 6 7 8 9]</span></pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> If applicable, include full error messages/tracebacks.</li>
</ul>