One Paragraph of project description goes here
These instructions will get you a copy of the project up and running on your local machine for development and testing purposes. See deployment for notes on how to deploy the project on a live system.
Boris Dayma borisdayma
skye
/ tpu_topology_env_vars.py
Last active
April 13, 2025 18:33
You can use these environment variables to run a Python process on a subset of the TPU cores on a Cloud TPU VM. This allows running multiple TPU processes at the same time, since only one process can access a given TPU chip at a time. Note that on TPU v2 and v3, 1 TPU chip = 2 TpuDevice as reported by `jax.devices()` (8 devices total). On v4, 1 …
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # ==== Non-communicating processes | |
| # 4x 1 chip per process: | |
| os.environ["TPU_CHIPS_PER_PROCESS_BOUNDS"] = "1,1,1" | |
| os.environ["TPU_PROCESS_BOUNDS"] = "1,1,1" | |
| # Different per process: | |
| os.environ["TPU_VISIBLE_DEVICES"] = "0" # "1", "2", "3" | |
| # 1-liner for bash: TPU_CHIPS_PER_PROCESS_BOUNDS=1,1,1 TPU_PROCESS_BOUNDS=1,1,1 TPU_VISIBLE_DEVICES=0 TPU_MESH_CONTROLLER_ADDRESS=localhost:8476 TPU_MESH_CONTROLLER_PORT=8476 |
