{"id":13563803,"url":"https://github.com/ayaka14732/tpu-starter","last_synced_at":"2025-10-25T08:41:01.969Z","repository":{"id":41360323,"uuid":"464360810","full_name":"ayaka14732/tpu-starter","owner":"ayaka14732","description":"Everything you want to know about Google Cloud TPU","archived":false,"fork":false,"pushed_at":"2024-07-16T13:06:58.000Z","size":2129,"stargazers_count":521,"open_issues_count":2,"forks_count":30,"subscribers_count":6,"default_branch":"main","last_synced_at":"2025-03-28T20:06:51.411Z","etag":null,"topics":["cloud-tpu","deep-learning","gcp","google-cloud-platform","jax","machine-learning","tpu"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"cc-by-4.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/ayaka14732.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2022-02-28T06:13:21.000Z","updated_at":"2025-03-25T17:38:03.000Z","dependencies_parsed_at":"2024-01-14T03:50:16.532Z","dependency_job_id":"87640014-bee3-468c-b076-5d51e0d1bf61","html_url":"https://github.com/ayaka14732/tpu-starter","commit_stats":null,"previous_names":[],"tags_count":2,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ayaka14732%2Ftpu-starter","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ayaka14732%2Ftpu-starter/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ayaka14732%2Ftpu-starter/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ayaka14732%2Ftpu-starter/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/ayaka14732","download_url":"https://codeload.github.com/ayaka14732/tpu-starter/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247249525,"owners_count":20908212,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["cloud-tpu","deep-learning","gcp","google-cloud-platform","jax","machine-learning","tpu"],"created_at":"2024-08-01T13:01:23.439Z","updated_at":"2025-10-07T17:21:29.706Z","avatar_url":"https://github.com/ayaka14732.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"# TPU Starter\n\n\u003ch4 align=\"center\"\u003e\n    \u003cp\u003e\n        \u003cb\u003eEnglish\u003c/b\u003e |\n        \u003ca href=\"https://github.com/ayaka14732/tpu-starter/blob/main/README_ko.md\"\u003e한국어\u003c/a\u003e |\n        \u003ca href=\"https://github.com/ayaka14732/tpu-starter/blob/main/README_zh.md\"\u003e中文\u003c/a\u003e\n    \u003cp\u003e\n\u003c/h4\u003e\n\nEverything you want to know about Google Cloud TPU\n\n* [1. Community](#1-community)\n* [2. Introduction to TPU](#2-introduction-to-tpu)\n    * [2.1. Why TPU?](#21-why-tpu)\n    * [2.2. How can I get free access to TPU?](#22-how-can-i-get-free-access-to-tpu)\n    * [2.3. If TPU is so good, why do I rarely see others using it?](#23-if-tpu-is-so-good-why-do-i-rarely-see-others-using-it)\n    * [2.4. I know TPU is great now. Can I touch a TPU?](#24-i-know-tpu-is-great-now-can-i-touch-a-tpu)\n    * [2.5. What does it mean to create a TPU instance? What do I actually get?](#25-what-does-it-mean-to-create-a-tpu-instance-what-do-i-actually-get)\n* [3. Introduction to the TRC Program](#3-introduction-to-the-trc-program)\n    * [3.1. How do I apply for the TRC program?](#31-how-do-i-apply-for-the-trc-program)\n    * [3.2. Is it really free?](#32-is-it-really-free)\n* [4. Using TPU VM](#4-using-tpu-vm)\n    * [4.1. Create a TPU VM](#41-create-a-tpu-vm)\n    * [4.2. Add an SSH public key to Google Cloud](#42-add-an-ssh-public-key-to-google-cloud)\n    * [4.3. SSH into TPU VM](#43-ssh-into-tpu-vm)\n    * [4.4. Verify that TPU VM has TPU](#44-verify-that-tpu-vm-has-tpu)\n    * [4.5. Setting up the development environment in TPU VM](#45-setting-up-the-development-environment-in-tpu-vm)\n    * [4.6. Verify JAX is working properly](#46-verify-jax-is-working-properly)\n    * [4.7. Using Byobu to ensure continuous program execution](#47-using-byobu-to-ensure-continuous-program-execution)\n    * [4.8. Configure VSCode Remote-SSH](#48-configure-vscode-remote-ssh)\n    * [4.9. Using Jupyter Notebook on TPU VM](#49-using-jupyter-notebook-on-tpu-vm)\n* [5. Using TPU Pod](#5-using-tpu-pod)\n    * [5.1. Create a subnet](#51-create-a-subnet)\n    * [5.2. Disable Cloud Logging](#52-disable-cloud-logging)\n    * [5.3. Create TPU Pod](#53-create-tpu-pod)\n    * [5.4. SSH into TPU Pod](#54-ssh-into-tpu-pod)\n    * [5.5. Modify the SSH configuration file on Host 0](#55-modify-the-ssh-configuration-file-on-host-0)\n    * [5.6. Add the SSH public key of Host 0 to all hosts](#56-add-the-ssh-public-key-of-host-0-to-all-hosts)\n    * [5.7. Configure the podrun command](#57-configure-the-podrun-command)\n    * [5.8. Configure NFS](#58-configure-nfs)\n    * [5.9. Setting up the development environment in TPU Pod](#59-setting-up-the-development-environment-in-tpu-pod)\n    * [5.10. Verify JAX is working properly](#510-verify-jax-is-working-properly)\n* [6. TPU Best Practices](#6-tpu-best-practices)\n    * [6.1. Prefer Google Cloud Platform to Google Colab](#61-prefer-google-cloud-platform-to-google-colab)\n    * [6.2. Prefer TPU VM to TPU node](#62-prefer-tpu-vm-to-tpu-node)\n* [7. JAX Best Practices](#7-jax-best-practices)\n    * [7.1. Import convention](#71-import-convention)\n    * [7.2. Manage random keys in JAX](#72-manage-random-keys-in-jax)\n    * [7.3. Conversion between NumPy arrays and JAX arrays](#73-conversion-between-numpy-arrays-and-jax-arrays)\n    * [7.4. Conversion between PyTorch tensors and JAX arrays](#74-conversion-between-pytorch-tensors-and-jax-arrays)\n    * [7.5. Get the shapes of all parameters in a nested dictionary](#75-get-the-shapes-of-all-parameters-in-a-nested-dictionary)\n    * [7.6. The correct way to generate random numbers on CPU](#76-the-correct-way-to-generate-random-numbers-on-cpu)\n    * [7.7. Use optimizers from Optax](#77-use-optimizers-from-optax)\n    * [7.8. Use the cross-entropy loss implementation from Optax](#78-use-the-cross-entropy-loss-implementation-from-optax)\n* [8. How Can I...](#8-how-can-i)\n    * [8.1. Share files across multiple TPU VM instances](#81-share-files-across-multiple-tpu-vm-instances)\n    * [8.2. Monitor TPU usage](#82-monitor-tpu-usage)\n    * [8.3. Start a server on TPU VM](#83-start-a-server-on-tpu-vm)\n    * [8.4. Run separate processes on different TPU cores](#84-run-separate-processes-on-different-tpu-cores)\n* [9. Common Gotchas](#9-common-gotchas)\n    * [9.1. TPU VMs will be rebooted occasionally](#91-tpu-vms-will-be-rebooted-occasionally)\n    * [9.2. One TPU core can only be used by one process at a time](#92-one-tpu-core-can-only-be-used-by-one-process-at-a-time)\n    * [9.3. TCMalloc breaks several programs](#93-tcmalloc-breaks-several-programs)\n    * [9.4. libtpu.so already in used by another process](#94-libtpuso-already-in-used-by-another-process)\n    * [9.5. JAX does not support the multiprocessing fork strategy](#95-jax-does-not-support-the-multiprocessing-fork-strategy)\n\n\u003c!-- Created by https://github.com/ekalinin/github-markdown-toc --\u003e\n\nThis project was inspired by [Cloud Run FAQ](https://github.com/ahmetb/cloud-run-faq), a community-maintained knowledge base about another Google Cloud product.\n\n## 1. Community\n\nGoogle's [official Discord server](https://discord.com/invite/google-dev-community) has established the `#tpu-research-cloud` channel.\n\n## 2. Introduction to TPU\n\n### 2.1. Why TPU?\n\n**TL;DR**: TPU is to GPU as GPU is to CPU.\n\nTPU is hardware specifically designed for machine learning. For performance comparisons, see [Performance Comparison](https://github.com/huggingface/transformers/blob/main/examples/flax/language-modeling/README.md#runtime-evaluation) in Hugging Face Transformers:\n\n![](assets/5.png)\n\nMoreover, Google's [TRC program](https://sites.research.google/trc/about/) offers free TPU resources to researchers. If you've ever wondered what computing resources to use to train a model, you should try the TRC program, as it's the best option I know of. More information about the TRC program is provided below.\n\n### 2.2. How can I get free access to TPU?\n\nResearchers can apply to the [TRC program](https://sites.research.google/trc/about/) to obtain free TPU resources.\n\n### 2.3. If TPU is so good, why do I rarely see others using it?\n\nIf you want to use PyTorch, TPU may not be suitable for you. TPU is poorly supported by PyTorch. In one of my past experiments using PyTorch, a batch took 14 seconds on a CPU but required 4 hours on a TPU. Twitter user @mauricetpunkt also thinks that [PyTorch's performance on TPUs is bad](https://twitter.com/mauricetpunkt/status/1506944350281945090).\n\nIn conclusion, if you want to do deep learning with TPU, you should use JAX as your deep learning framework. In fact, many popular deep learning libraries support JAX. For instance:\n\n- [Many models in Hugging Face Transformers support JAX](https://huggingface.co/docs/transformers/index#supported-frameworks)\n- [Keras supports using JAX as a backend](https://keras.io/keras_core/announcement/)\n- SkyPilot has [examples using Flax](https://github.com/skypilot-org/skypilot/blob/master/examples/tpu/tpuvm_mnist.yaml)\n\nFurthermore, JAX's design is very clean and has been widely appreciated. For instance, JAX is my favorite open-source project. I've tweeted about [how JAX is better than PyTorch](https://twitter.com/ayaka14732/status/1688194164033462272).\n\n### 2.4. I know TPU is great now. Can I touch a TPU?\n\nUnfortunately, we generally can't physically touch a real TPU. TPUs are meant to be accessed via Google Cloud services.\n\nIn some exhibitions, TPUs are [displayed for viewing](https://twitter.com/walkforhours/status/1696654844134822130), which might be the closest you can get to physically touching one.\n\nPerhaps only by becoming a Google Cloud Infrastructure Engineer can one truly feel the touch of a TPU.\n\n### 2.5. What does it mean to create a TPU instance? What do I actually get?\n\nAfter creating a TPU v3-8 instance on [Google Cloud Platform](https://cloud.google.com/tpu), you'll get a cloud server running the Ubuntu system with sudo privileges, 96 CPU cores, 335 GiB memory, and a TPU device with 8 cores (totalling 128 GiB TPU memory).\n\n![](assets/0.png)\n\nIn fact, this is similar to how we use GPUs. Typically, when we use a GPU, we are using a Linux server connected to the GPU. Similarly, when we use a TPU, we're using a server connected to the TPU.\n\n## 3. Introduction to the TRC Program\n\n### 3.1. How do I apply for the TRC program?\n\nApart from the TRC program's [homepage](https://sites.research.google/trc/about/), Shawn wrote a wonderful article about the TRC program on [google/jax#2108](https://github.com/google/jax/issues/2108#issuecomment-866238579). Anyone who is interested in TPU should read it immediately.\n\n### 3.2. Is it really free?\n\nFor the first three months, the TRC program is completely free due to the free trial credit given when registering for Google Cloud. After three months, I spend roughly HK$13.95 (about US$1.78) per month. This expense is for the network traffic of the TPU server, while the TPU device itself is provided for free by the TRC program.\n\n## 4. Using TPU VM\n\n### 4.1. Create a TPU VM\n\nOpen [Google Cloud Platform](https://cloud.google.com/tpu) and navigate to the [TPU Management Page](https://console.cloud.google.com/compute/tpus).\n\n![](assets/1.png)\n\nClick the console button on the top-right corner to activate Cloud Shell.\n\nIn Cloud Shell, type the following command to create a Cloud TPU v3-8 VM:\n\n```sh\nuntil gcloud alpha compute tpus tpu-vm create node-1 --project tpu-develop --zone europe-west4-a --accelerator-type v3-8 --version tpu-vm-base ; do : ; done\n```\n\nHere, `node-1` is the name of the TPU VM you want to create, and `--project` is the name of your Google Cloud project.\n\nThe above command will repeatedly attempt to create the TPU VM until it succeeds.\n\n### 4.2. Add an SSH public key to Google Cloud\n\nFor Google Cloud's servers, if you want to SSH into them, using `ssh-copy-id` is the wrong approach. The correct method is:\n\nFirst, type “SSH keys” into the Google Cloud webpage search box, go to the relevant page, then click edit, and add your computer's SSH public key.\n\nTo view your computer's SSH public key:\n\n```sh\ncat ~/.ssh/id_rsa.pub\n```\n\nIf you haven't created an SSH key pair yet, use the following command to create one, then execute the above command to view:\n\n```sh\nssh-keygen -t rsa -f ~/.ssh/id_rsa -N \"\"\n```\n\nWhen adding an SSH public key to Google Cloud, it's crucial to pay special attention to the value of the username. In the SSH public key string, the part preceding the `@` symbol at the end is the username. When added to Google Cloud, it will create a user with that name on all servers for the current project. For instance, with the string `ayaka@instance-1`, Google Cloud will create a user named `ayaka` on the server. If you wish for Google Cloud to create a different username, you can manually modify this string. Changing the mentioned string to `nixie@instance-1` would lead Google Cloud to create a user named `nixie`. Moreover, making such changes won't affect the functionality of the SSH key.\n\n### 4.3. SSH into TPU VM\n\nCreate or edit your computer's `~/.ssh/config`:\n\n```sh\nnano ~/.ssh/config\n```\n\nAdd the following content:\n\n```\nHost tpuv3-8-1\n    User nixie\n    Hostname 34.141.220.156\n```\n\nHere, `tpuv3-8-1` is an arbitrary name, `User` is the username created in Google Cloud from the previous step, and `Hostname` is the IP address of the TPU VM.\n\nThen, on your own computer, use the following command to SSH into the TPU VM:\n\n```sh\nssh tpuv3-8-1\n```\n\nWhere `tpuv3-8-1` is the name set in `~/.ssh/config`.\n\n### 4.4. Verify that TPU VM has TPU\n\n```sh\nls /dev/accel*\n```\n\nIf the following output appears:\n\n```\n/dev/accel0  /dev/accel1  /dev/accel2  /dev/accel3\n```\n\nThis indicates that the TPU VM indeed has a TPU.\n\n### 4.5. Setting up the development environment in TPU VM\n\nUpdate software packages:\n\n```sh\nsudo apt-get update -y -qq\nsudo apt-get upgrade -y -qq\nsudo apt-get install -y -qq golang neofetch zsh byobu\n```\n\nInstall the latest Python 3.12:\n\n```sh\nsudo apt-get install -y -qq software-properties-common\nsudo add-apt-repository -y ppa:deadsnakes/ppa\nsudo apt-get install -y -qq python3.12-full python3.12-dev\n```\n\nInstall Oh My Zsh:\n\n```sh\nsh -c \"$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)\" \"\" --unattended\nsudo chsh $USER -s /usr/bin/zsh\n```\n\nCreate a virtual environment (venv):\n\n```sh\npython3.12 -m venv ~/venv\n```\n\nActivate the venv:\n\n```sh\n. ~/venv/bin/activate\n```\n\nInstall JAX in the venv:\n\n```sh\npip install -U pip\npip install -U wheel\npip install -U \"jax[tpu]\" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n```\n\n### 4.6. Verify JAX is working properly\n\nAfter activating the venv, use the following command to verify JAX is working:\n\n```sh\npython -c 'import jax; print(jax.devices())'\n```\n\nIf the output contains `TpuDevice`, this means JAX is working as expected.\n\n### 4.7. Using Byobu to ensure continuous program execution\n\nMany tutorials use the method of appending `\u0026` to commands to run them in the background, so they continue executing even after exiting SSH. However, this is a basic method. The correct approach is to use a window manager like Byobu.\n\nTo run Byobu, simply use the `byobu` command. Then, execute commands within the opened window. To close the window, you can forcefully close the current window on your computer. Byobu will continue running on the server. The next time you connect to the server, you can retrieve the previous window using the `byobu` command.\n\nByobu has many advanced features. You can learn them by watching the official video [Learn Byobu while listening to Mozart](https://youtu.be/NawuGmcvKus).\n\n### 4.8. Configure VSCode Remote-SSH\n\nOpen VSCode, access the Extensions panel on the left, search and install Remote - SSH.\n\nPress \u003ckbd\u003eF1\u003c/kbd\u003e to open the command palette. Type ssh, click \"Remote-SSH: Connect to Host...\", then click on the server name set in `~/.ssh/config` (e.g., `tpuv3-8-1`). Once VSCode completes the setup on the server, you can develop directly on the server with VSCode.\n\n![](assets/3.png)\n\nOn your computer, you can use the following command to quickly open a directory on the server:\n\n```sh\ncode --remote ssh-remote+tpuv3-8-1 /home/ayaka/tpu-starter\n```\n\nThis command will open the directory `/home/ayaka/tpu-starter` on `tpuv3-8-1` using VSCode.\n\n### 4.9. Using Jupyter Notebook on TPU VM\n\nAfter configuring VSCode with Remote-SSH, you can use Jupyter Notebook within VSCode. The result is as follows:\n\n![](assets/6.png)\n\nThere are two things to note here: First, in the top-right corner of the Jupyter Notebook interface, you should select the Kernel from `venv`, which refers to the `~/venv/bin/python` we created in the previous steps. Second, the first time you run it, you'll be prompted to install the Jupyter extension for VSCode and to install `ipykernel` within `venv`. You'll need to confirm these operations.\n\n## 5. Using TPU Pod\n\n### 5.1. Create a subnet\n\nTo create a TPU Pod, you first need to create a new VPC network and then create a subnet in the corresponding area of that network (e.g., `europe-west4-a`).\n\nTODO: Purpose?\n\n### 5.2. Disable Cloud Logging\n\nTODO: Reason? Steps?\n\n### 5.3. Create TPU Pod\n\nOpen Cloud Shell using the method described earlier for creating the TPU VM and use the following command to create a TPU v3-32 Pod:\n\n```sh\nuntil gcloud alpha compute tpus tpu-vm create node-1 --project tpu-advanced-research --zone europe-west4-a --accelerator-type v3-32 --version v2-alpha-pod --network advanced --subnetwork advanced-subnet-for-europe-west4 ; do : ; done\n```\n\nWhere `node-1` is the name you want for the TPU VM, `--project` is the name of your Google Cloud project, and `--network` and `--subnetwork` are the names of the network and subnet created in the previous step.\n\n### 5.4. SSH into TPU Pod\n\nSince the TPU Pod consists of multiple hosts, we need to choose one host, designate it as Host 0, and then SSH into Host 0 to execute commands. Given that the SSH public key added on the Google Cloud web page will be propagated to all hosts, every host can be directly connected through the SSH key, allowing us to designate any host as Host 0. The method to SSH into Host 0 is the same as for the aforementioned TPU VM.\n\n### 5.5. Modify the SSH configuration file on Host 0\n\nAfter SSH-ing into Host 0, the following configurations need to be made:\n\n```sh\nnano ~/.ssh/config\n```\n\nAdd the following content:\n\n```\nHost 172.21.12.* 127.0.0.1\n    StrictHostKeyChecking no\n    UserKnownHostsFile /dev/null\n    LogLevel ERROR\n```\n\nHere, `172.21.12.*` is determined by the IP address range of the subnet created in the previous steps. We use `172.21.12.*` because the IP address range specified when creating the subnet was 172.21.12.0/24.\n\nWe need to do so because the `known_hosts` in ssh is created for preventing man-in-the-middle attacks. Since we are using an internal network environment here, we don't need to prevent such attacks or require this file, so we direct it to `/dev/null`. Additionally, having `known_hosts` requires manually confirming the server's fingerprint during the first connection, which is unnecessary in an internal network environment and is not conducive to automation.\n\nThen, run the following command to modify the permissions of this configuration file. If the permissions are not modified, the configuration file will not take effect:\n\n```sh\nchmod 600 ~/.ssh/config\n```\n\n### 5.6. Add the SSH public key of Host 0 to all hosts\n\nGenerate a key pair on Host 0:\n\n```sh\nssh-keygen -t rsa -f ~/.ssh/id_rsa -N \"\"\n```\n\nView the generated SSH public key:\n\n```sh\ncat ~/.ssh/id_rsa.pub\n```\n\nAdd this public key to the SSH keys in Google Cloud. This key will be automatically propagated to all hosts.\n\n### 5.7. Configure the `podrun` command\n\nThe `podrun` command is a tool under development. When executed on Host 0, it can run commands on all hosts via SSH.\n\nDownload `podrun`:\n\n```sh\nwget https://raw.githubusercontent.com/ayaka14732/llama-2-jax/18e9625f7316271e4c0ad9dea233cfe23c400c9b/podrun\nchmod +x podrun\n```\n\nEdit `~/podips.txt` using:\n\n```sh\nnano ~/podips.txt\n```\n\nSave the internal IP addresses of the other hosts in `~/podips.txt`, one per line. For example:\n\n```sh\n172.21.12.86\n172.21.12.87\n172.21.12.83\n```\n\nA TPU v3-32 includes 4 hosts. Excluding Host 0, there are 3 more hosts. Hence, the `~/podips.txt` for TPU v3-32 should contain 3 IP addresses.\n\nInstall Fabric using the system pip3:\n\n```sh\npip3 install fabric\n```\n\nUse `podrun` to make all hosts purr like a kitty:\n\n```sh\n./podrun -iw -- echo meow\n```\n\n### 5.8. Configure NFS\n\nInstall the NFS server and client:\n\n```sh\n./podrun -i -- sudo apt-get update -y -qq\n./podrun -i -- sudo apt-get upgrade -y -qq\n./podrun -- sudo apt-get install -y -qq nfs-common\nsudo apt-get install -y -qq nfs-kernel-server\nsudo mkdir -p /nfs_share\nsudo chown -R nobody:nogroup /nfs_share\nsudo chmod 777 /nfs_share\n```\n\nModify `/etc/exports`:\n\n```sh\nsudo nano /etc/exports\n```\n\nAdd:\n\n```\n/nfs_share  172.21.12.0/24(rw,sync,no_subtree_check)\n```\n\nExecute:\n\n```sh\nsudo exportfs -a\nsudo systemctl restart nfs-kernel-server\n\n./podrun -- sudo mkdir -p /nfs_share\n./podrun -- sudo mount 172.21.12.2:/nfs_share /nfs_share\n./podrun -i -- ln -sf /nfs_share ~/nfs_share\n\ntouch ~/nfs_share/meow\n./podrun -i -- ls -la ~/nfs_share/meow\n```\n\nReplace `172.21.12.2` with the actual internal IP address of Host 0.\n\n### 5.9. Setting up the development environment in TPU Pod\n\nSave to `~/nfs_share/setup.sh`:\n\n```sh\n#!/bin/bash\n\nexport DEBIAN_FRONTEND=noninteractive\n\nsudo apt-get update -y -qq\nsudo apt-get upgrade -y -qq\nsudo apt-get install -y -qq golang neofetch zsh byobu\n\nsudo apt-get install -y -qq software-properties-common\nsudo add-apt-repository -y ppa:deadsnakes/ppa\nsudo apt-get install -y -qq python3.12-full python3.12-dev\n\nsh -c \"$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)\" \"\" --unattended\nsudo chsh $USER -s /usr/bin/zsh\n\npython3.12 -m venv ~/venv\n\n. ~/venv/bin/activate\n\npip install -U pip\npip install -U wheel\npip install -U \"jax[tpu]\" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n```\n\nThen execute:\n\n```sh\nchmod +x ~/nfs_share/setup.sh\n./podrun -i ~/nfs_share/setup.sh\n```\n\n### 5.10. Verify JAX is working properly\n\n```sh\n./podrun -ic -- ~/venv/bin/python -c 'import jax; jax.distributed.initialize(); jax.process_index() == 0 and print(jax.devices())'\n```\n\nIf the output contains `TpuDevice`, this means JAX is working as expected.\n\n## 6. TPU Best Practices\n\n### 6.1. Prefer Google Cloud Platform to Google Colab\n\n[Google Colab](https://colab.research.google.com/) only provides TPU v2-8 devices, while on [Google Cloud Platform](https://cloud.google.com/tpu) you can select TPU v2-8 and TPU v3-8.\n\nBesides, on Google Colab you can only use TPU through the Jupyter Notebook interface. Even if you [log in into the Colab server via SSH](https://ayaka.shn.hk/colab/), it is a docker image and you don't have root access. On Google Cloud Platform, however, you have full access to the TPU VM.\n\nIf you really want to use TPU on Google Colab, you need to run [the following script](01-basics/setup_colab_tpu.py) to set up TPU:\n\n```python\nimport jax\nfrom jax.tools.colab_tpu import setup_tpu\n\nsetup_tpu()\n\ndevices = jax.devices()\nprint(devices)  # should print TpuDevice\n```\n\n### 6.2. Prefer TPU VM to TPU node\n\nWhen you are creating a TPU instance, you need to choose between TPU VM and TPU node. Always prefer TPU VM because it is the new architecture in which TPU devices are connected to the host VM directly. This will make it easier to set up the TPU device.\n\n## 7. JAX Best Practices\n\n### 7.1. Import convention\n\nYou may see two different kind of import conventions. One is to import `jax.numpy` as `np` and import the original numpy as `onp`. Another one is to import `jax.numpy` as `jnp` and leave original numpy as `np`.\n\nOn 16 Jan 2019, Colin Raffel wrote in [a blog article](https://colinraffel.com/blog/you-don-t-know-jax.html) that the convention at that time was to import original numpy as `onp`.\n\nOn 5 Nov 2020, Niru Maheswaranathan said in [a tweet](https://twitter.com/niru_m/status/1324078070546882560) that he thinks the convention at that time was to import `jax.numpy` as `jnp` and to leave original numpy as `np`.\n\nWe can conclude that the new convention is to import `jax.numpy` as `jnp`.\n\n### 7.2. Manage random keys in JAX\n\nThe regular way is this:\n\n```python\nkey, *subkey = rand.split(key, num=4)\nprint(subkey[0])\nprint(subkey[1])\nprint(subkey[2])\n```\n\n### 7.3. Conversion between NumPy arrays and JAX arrays\n\nUse [`np.asarray`](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.asarray.html) and [`onp.asarray`](https://numpy.org/doc/stable/reference/generated/numpy.asarray.html).\n\n```python\nimport jax.numpy as np\nimport numpy as onp\n\na = np.array([1, 2, 3])  # JAX array\nb = onp.asarray(a)  # converted to NumPy array\n\nc = onp.array([1, 2, 3])  # NumPy array\nd = np.asarray(c)  # converted to JAX array\n```\n\n### 7.4. Conversion between PyTorch tensors and JAX arrays\n\nConvert a PyTorch tensor to a JAX array:\n\n```python\nimport jax.numpy as np\nimport torch\n\na = torch.rand(2, 2)  # PyTorch tensor\nb = np.asarray(a.numpy())  # JAX array\n```\n\nConvert a JAX array to a PyTorch tensor:\n\n```python\nimport jax.numpy as np\nimport numpy as onp\nimport torch\n\na = np.zeros((2, 2))  # JAX array\nb = torch.from_numpy(onp.asarray(a))  # PyTorch tensor\n```\n\nThis will result in a warning:\n\n```\nUserWarning: The given NumPy array is not writable, and PyTorch does not support non-writable tensors. This means writing to this tensor will result in undefined behavior. You may want to copy the array to protect its data or make it writable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  ../torch/csrc/utils/tensor_numpy.cpp:178.)\n```\n\nIf you need writable tensors, you can use `onp.array` instead of `onp.asarray` to make a copy of the original array.\n\n### 7.5. Get the shapes of all parameters in a nested dictionary\n\n```python\njax.tree_map(lambda x: x.shape, params)\n```\n\n### 7.6. The correct way to generate random numbers on CPU\n\nUse the [jax.default_device()](https://jax.readthedocs.io/en/latest/_autosummary/jax.default_device.html) context manager:\n\n```python\nimport jax\nimport jax.random as rand\n\ndevice_cpu = jax.devices('cpu')[0]\nwith jax.default_device(device_cpu):\n    key = rand.PRNGKey(42)\n    a = rand.poisson(key, 3, shape=(1000,))\n    print(a.device())  # TFRT_CPU_0\n```\n\nSee \u003chttps://github.com/google/jax/discussions/9691#discussioncomment-3650311\u003e.\n\n### 7.7. Use optimizers from Optax\n\n### 7.8. Use the cross-entropy loss implementation from Optax\n\n`optax.softmax_cross_entropy_with_integer_labels`\n\n## 8. How Can I...\n\n### 8.1. Share files across multiple TPU VM instances\n\nTPU VM instances in the same zone are connected with internal IPs, so you can [create a shared file system using NFS](https://tecadmin.net/how-to-install-and-configure-an-nfs-server-on-ubuntu-20-04/).\n\n### 8.2. Monitor TPU usage\n\n[jax-smi](https://github.com/ayaka14732/jax-smi)\n\n### 8.3. Start a server on TPU VM\n\nExample: Tensorboard\n\nAlthough every TPU VM is allocated with a public IP, in most cases you should expose a server to the Internet because it is insecure.\n\nPort forwarding via SSH\n\n```\nssh -C -N -L 127.0.0.1:6006:127.0.0.1:6006 tpu1\n```\n\n### 8.4. Run separate processes on different TPU cores\n\nhttps://gist.github.com/skye/f82ba45d2445bb19d53545538754f9a3\n\n## 9. Common Gotchas\n\n### 9.1. TPU VMs will be rebooted occasionally\n\nAs of 24 Oct 2022, the TPU VMs will be rebooted occasionally if there is a maintenance event.\n\nThe following things will happen:\n\n1. All the running processes will be terminated\n2. The external IP address will be changed\n\nWe can save the model parameters, optimiser states and other useful data occasionally, so that the model training can be easily resumed after termination.\n\nWe should use `gcloud` command instead of connect directly to it with SSH. If we have to use SSH (e.g. if we want to use VSCode, SSH is the only choice), we need to manually change the target IP address.\n\n### 9.2. One TPU core can only be used by one process at a time\n\nSee also: §10.5.\n\nUnlike GPU, you will get an error if you run two processes on TPU at a time:\n\n```\nI0000 00:00:1648534265.148743  625905 tpu_initializer_helper.cc:94] libtpu.so already in use by another process. Run \"$ sudo lsof -w /dev/accel0\" to figure out which process is using the TPU. Not attempting to load libtpu.so in this process.\n```\n\n### 9.3. TCMalloc breaks several programs\n\n[TCMalloc](https://github.com/google/tcmalloc) is Google's customized memory allocation library. On TPU VM, `LD_PRELOAD` is set to use TCMalloc by default:\n\n```sh\n$ echo LD_PRELOAD\n/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4\n```\n\nHowever, using TCMalloc in this manner may break several programs like gsutil:\n\n```sh\n$ gsutil --help\n/snap/google-cloud-sdk/232/platform/bundledpythonunix/bin/python3: /snap/google-cloud-sdk/232/platform/bundledpythonunix/bin/../../../lib/x86_64-linux-gnu/libm.so.6: version `GLIBC_2.29' not found (required by /usr/lib/x86_64-linux-gnu/libtcmalloc.so.4)\n```\n\nThe [homepage of TCMalloc](http://goog-perftools.sourceforge.net/doc/tcmalloc.html) also indicates that `LD_PRELOAD` is tricky and this mode of usage is not recommended.\n\nIf you encounter problems related to TCMalloc, you can disable it in the current shell using the command:\n\n```sh\nunset LD_PRELOAD\n```\n\n### 9.4. `libtpu.so` already in used by another process\n\n```sh\nif ! pgrep -a -u $USER python ; then\n    killall -q -w -s SIGKILL ~/.venv311/bin/python\nfi\nrm -rf /tmp/libtpu_lockfile /tmp/tpu_logs\n```\n\nSee also \u003chttps://github.com/google/jax/issues/9220#issuecomment-1015940320\u003e.\n\n### 9.5. JAX does not support the multiprocessing `fork` strategy\n\nUse the `spawn` or `forkserver` strategies.\n\nSee \u003chttps://github.com/google/jax/issues/1805#issuecomment-561244991\u003e.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fayaka14732%2Ftpu-starter","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fayaka14732%2Ftpu-starter","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fayaka14732%2Ftpu-starter/lists"}