Pytorch/xla SPMD on Cloud TPU pod

abdelkrim_zitouni: <p>I’m currently trying to train a Vision Transformer on a Cloud TPU Pod, using Pytorch/xla SPMD, I followed the later’s <a href="https://github.com/pytorch/xla/blob/master/docs/spmd.md" rel="noopener nofollow ugc">documentation</a>, and my code is mostly inspired by their <a href="https://github.com/pytorch/xla/blob/master/test/spmd/test_train_spmd_imagenet.py" rel="noopener nofollow ugc">example</a> of running Resnet with SPMD on a single TPU device.</p>
<p>However, as quoted from their docs below:</p>
<blockquote>
<p>There is no code change required to go from single TPU host to TPU Pod if you construct your mesh and partition spec based on the number of devices instead of some hardcode constant. To run the PyTorch/XLA workload on TPU Pod.</p>
</blockquote>
<p>I truly couldn’t understand how it is possible to run the exact same code and distribute it over all TPU pod VMs only by invoking the same gcloud command with the <code>--worker=all</code> argument, does it work out of the box? how can I check if the data and weights are sharded across all hosts in the pod?</p>
<p>I’m finding it difficult to use Pytorch/xla with a lack of documentation and examples on TPU Pod.</p>

Yeounoh_Chung: <p>Hi Abdelkrim, the best way to monitor your workload (and see if it’s utilizing all the devices) would be to profile, <a href="https://cloud.google.com/tpu/docs/pytorch-xla-performance-profiling-tpu-vm" class="inline-onebox" rel="noopener nofollow ugc">PyTorch XLA performance profiling &nbsp;|&nbsp; Cloud TPU &nbsp;|&nbsp; Google Cloud</a> – and check the profiled traces.</p>
<p>You can also enable some loggings, to see if your model is being executed across multiple devices, e.g., <code>TF_CPP_MIN_LOG_LEVEL=0 TF_CPP_VMODULE=xla_graph_executor=5 python test/spmd/test_train_spmd_imagenet.py --fake_data --num_epoch=1 --sharding batch</code></p>
<p>This should print out useful debugging info, like</p>
<pre><code class="lang-auto">2023-11-21 01:49:24.349482: I torch_xla/csrc/xla_graph_executor.cpp:994] Executing IR graph hash 4b2ecd7c298db70445084a8ad90e9fd9 on devices: TPU:0,TPU:1,TPU:2,TPU:3 done!
2023-11-21 01:49:24.352578: I torch_xla/csrc/xla_graph_executor.cpp:603] waiting barrier for device SPMD:0 done
</code></pre>
<p>on each host. Here, <code>Executing IR graph hash 4b2ecd7c298db70445084a8ad90e9fd9 on devices: TPU:0,TPU:1,TPU:2,TPU:3 done!</code> means that the execution on this host was done with all 4 local devices.</p>
<p>When you run your training on a pod slice, with <code>--worker=all --command={your command}</code>, could you try adding <code>{your command} 2&gt;&amp;1 | tee output.txt}</code> and check the logs afterwards on each host?</p>

abdelkrim_zitouni: <p>Hi <a href="/u/Yeounoh_Chung">Yeounoh</a>, Thank you very much, and I apologize for the late response as I paused the project to study for exams.</p>
<p>I just tried setting the environment variables as per your instructions and I got a huge file containing debugging info,  the screenshot below contains some information similar to your example,<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/0/c0fdb895c051f9cba8c74d90e0573df625979508.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/c0fdb895c051f9cba8c74d90e0573df625979508" title="Screenshot 2023-11-28 at 3.39.50 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/0/c0fdb895c051f9cba8c74d90e0573df625979508_2_690x333.jpeg" alt="Screenshot 2023-11-28 at 3.39.50 PM" data-base62-sha1="rxhgGyODlqCjOfbatgJ5d2ZQxo4" width="690" height="333" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/0/c0fdb895c051f9cba8c74d90e0573df625979508_2_690x333.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/c/0/c0fdb895c051f9cba8c74d90e0573df625979508_2_1035x499.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/c/0/c0fdb895c051f9cba8c74d90e0573df625979508_2_1380x666.jpeg 2x" data-dominant-color="404040"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2023-11-28 at 3.39.50 PM</span><span class="informations">1920×927 382 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>