Enable multiprocessing on pytorch XLA for TPU vm

username: adhamali450
cooked: <p>I’m fairly new to this and have little to no experience.<br>
I had a notebook running PyTorch that I wanted to run a Google Cloud TPU VM.<br>
Machine specs:</p>
<pre><code class="lang-auto">* Ubuntu
* TPU v2-8
* pt-2.0
</code></pre>
<p>I should have 8 cores. Correct me if I’m wrong.</p>
<p>So, I followed the guidelines for making the notebook TPU-compatible via XLA. I did the following:</p>
<pre><code class="lang-auto">os.environ['PJRT_DEVICE'] = 'TPU'
import torch_xla.core.xla_model as xm
import torch_xla.distributed.parallel_loader as pl
import torch_xla.distributed.xla_multiprocessing as xmp
device = xm.xla_device()

print(device)
</code></pre>
<blockquote>
<p>It prints <code>xla:0</code>.</p>
</blockquote>
<p>*Models was sent to the device via the <code>model.to(device)</code> function.<br>
*Dataloaders was wrapped in a <code>pl.MpDeviceLoader(loader, device)</code><br>
*Optimizer is stepped via the <code>xm.optimizer_step(optimizer)</code> function</p>
<ul>
<li>As far as I know, this is how to enable multiprocessing:</li>
</ul>
<pre><code class="lang-auto">def _mp_fn(index):
     # models creation
     # data preparation
     # training loop 
if __name__ == '__main__':
  xmp.spawn(_mp_fn, args=())
</code></pre>
<p>I get the error:</p>
<blockquote>
<p>BrokenProcessPool: A process in the process pool was terminated abruptly while the future was running or pending.</p>
</blockquote>
<p>I could be all totally wrong about this. So, I’m sorry for that.<br>
If you need a further look at the code, I can share the notebook if you want.<br>
When I follow the guidelines for single-core processing, and I don’t use <code>xmp.spawn</code> , I get <code>1.2 iterations/sec</code> which can be significantly increased if used all cores.</p>

username: alanwaketan
cooked: <p>To do multi-process in PyTorch/XLA, all xla related code such as <code>device = xm.xla_device()</code> need to be done in the _mp_fn function. You can’t do it in global scope.</p>

username: adhamali450
cooked: <p>I followed an article which led me to write another code but It still doesn’t work.<br>
I state what I did on details here: <a href="https://stackoverflow.com/questions/76039948/use-8-cores-on-pytorch-xla-for-tpu-vm" rel="noopener nofollow ugc">use-8-cores-on-pytorch-xla-for-tpu-vm</a>.<br>
Can you please check it out?<br>
Thank you so much</p>

username: alanwaketan
cooked: <p>Hi Adham,</p>
<p>First, thanks for using PyTorch/XLA. I have looked into your code. One thing you can try is to replace your ParallelLoader with MpDeviceLoader. You can find the usage examples here: <a href="https://github.com/pytorch/xla/blob/dc97329679fb1baa2d4ea5e9f112e8f58dc9b1b4/test/test_train_mp_mnist.py#L188" class="inline-onebox" rel="noopener nofollow ugc">xla/test_train_mp_mnist.py at dc97329679fb1baa2d4ea5e9f112e8f58dc9b1b4 · pytorch/xla · GitHub</a></p>
<p>Will love to diagnose further if the above suggestion doesn’t work. BTW, can you also provide a code that I can easily run locally? Your sample code doesn’t run.</p>

username: adhamali450
cooked: <p><a class="mention" href="/u/alanwaketan">@alanwaketan</a> Thank you for responding. I did that and now here’s what’s happening:<br>
When using <code>nprocs=8</code> I get an exception:</p>
<blockquote>
<p>A process in the process pool was terminated abruptly while the future was running or pending</p>
</blockquote>
<p>I looked it up and found this <a href="https://github.com/facebookresearch/demucs/issues/344" rel="noopener nofollow ugc">GitHub issue</a> which says make it 1 for debugging.  It WORKED but it’s <code>1.44it/s</code>.<br>
It would be great if you helped be to use all the 8 workers.</p>
<p>Here’s my notebook as per your request <a href="https://colab.research.google.com/drive/1RdhTl6ZlT-BZ-_DDYQtEzKqnVdOp9cRH?usp=sharing" rel="noopener nofollow ugc">Colab notebook</a><br>
The code is the cell below <code>Main</code>. You have an edit access</p>

username: alanwaketan
cooked: <p>Just want to confirm that you are using V2-8 not V3-8?</p>
<p>P.S. PJRT doesn’t work on colab. We recommend to use kaggle instead. Here are examples how to use kaggle with pytorch/xla: <a href="https://github.com/pytorch/xla/tree/master/contrib/kaggle" class="inline-onebox" rel="noopener nofollow ugc">xla/contrib/kaggle at master · pytorch/xla · GitHub</a>.</p>

username: hrbigelow
cooked: <p>Hi Jiewen,</p>
<p>I just saw this - actually, I have been having other issues (not using PJRT) on Colab TPU.  I opened an <a href="https://github.com/pytorch/xla/issues/5096" rel="noopener nofollow ugc">issue</a>, would you be able to take a look?</p>
<p>It would be nice to be able to use Colab since saving checkpoints to my google drive is easier and more reliable.  (Kaggle has sometimes evaporated my checkpoints even after selecting ‘Persistence → Variables and Files’.)</p>
<p>Best,</p>
<p>Henry</p>