colab_gcp_tpu_tutorial.md

Tips and tricks on using Colab, Google Cloud Platform and TPUs

Authors: murphyk, mjsML, gerdm, summer 2021.

Links

• Colab tutorial

• Using Google Cloud Storage from Colab

• Using TPU VMs from Colab

• Accessing colab machine via ssh

• Hugging face tutorial page on TPU VMs

• Skye's tutorial video on TPU VMs

• Using GCP/ TPU VMs via SSH + VSCode

• Using tensorboard

Random stuff

• Screenshot of how to split your v3-8 TPU VM into two 4-core machines. This lets you run two separate JAX processes at once, in different jupyter notebooks (eg if working on two projects simultaneously). The magic incantation is from this gist written by Skye Wanderman-Milne and is reproduced below.

# 2x 2 chips (4 cores) per process:
os.environ["TPU_CHIPS_PER_HOST_BOUNDS"] = "1,2,1"
os.environ["TPU_HOST_BOUNDS"] = "1,1,1"
# Different per process:
os.environ["TPU_VISIBLE_DEVICES"] = "0,1" # Change to "2,3" for the second machine

• Python multiprocessing library

• Jax on CPUs

• pmap on CPUs

• pip install git+https://github.com/blackjax-devs/blackjax