BLOOM 🌸 Inference in JAX

Structure

CPU Host: as defined in TPU manager

TPU Host: as defined in Host worker

ray: distributes load from CPU host -> TPU hosts

Example usage: run.py

Setting Up a TPU-Manager

The TPU hosts are managed by a single TPU manager. This TPU manager takes the form of a single CPU device.

First, create a CPU VM in the same region as that of the TPU pod. This is important to enable the TPU manager to communicate with the TPU hosts. A suitable device config is as follows:

1. Region & Zone: TO MATCH TPU ZONE
2. Machine type: c2-standard-8
3. CPU platform: Intel Cascade Lake
4. Boot disk: 256GB balanced persistent disk

SSH into the CPU and set-up a Python environment with the same Python version as that of the TPUs. The default TPU Python version is 3.8.10. You should ensure the Python version of th CPU matches this.

python3.8 -m venv /path/to/venv

If the above does not work, run the following and then repeat:

sudo apt-get update
sudo apt-get install python3-venv

Activate Python env:

source /path/to/venv/bin/activate

Check Python version is 3.8.10:

python --version

Clone the repository and install requirements:

git clone https://github.com/huggingface/bloom-jax-inference.git
cd bloom-jax-inference
pip install -r requirements.txt

Authenticate gcloud, which will require copy-and-pasting a command into a terminal window on a machine with a browser installed:

gcloud auth login

Now SSH into one of the workers. This will generate an SSH key:

gcloud alpha compute tpus tpu-vm ssh patrick-tpu-v3-32 --zone europe-west4-a --worker 0

Logout of the TPU worker:

logout

You should now be back in the CPU host.