skye
/ tpu_topology_env_vars.py
Last active
September 4, 2025 15:32
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 |
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
| #!/bin/bash | |
| # | |
| # script to extract ImageNet dataset | |
| # ILSVRC2012_img_train.tar (about 138 GB) | |
| # ILSVRC2012_img_val.tar (about 6.3 GB) | |
| # make sure ILSVRC2012_img_train.tar & ILSVRC2012_img_val.tar in your current directory | |
| # | |
| # https://github.com/facebook/fb.resnet.torch/blob/master/INSTALL.md | |
| # | |
| # train/ |