Google TPUs documentation

πŸ€— Optimum TPU

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

πŸ€— Optimum TPU

Optimum TPU provides all the necessary machinery to leverage and optimize AI workloads running on Google Cloud TPU devices.

The API provides the overall same user-experience as Hugging Face transformers with the minimum amount of changes required to target performance for inference.

Training support is underway, stay tuned! πŸš€

Installation

Optimum TPU is meant to reduce as much as possible the friction in order to leverage Google Cloud TPU accelerators. As such, we provide a pip installable package to make sure everyone can get easily started.

Run Cloud TPU with pip

pip install optimum-tpu

Run Cloud TPU within Docker container

PyTorch

export TPUVM_IMAGE_URL=us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla
export TPUVM_IMAGE_VERSION=8f1dcd5b03f993e4da5c20d17c77aff6a5f22d5455f8eb042d2e4b16ac460526
docker pull
docker run -ti --rm --privileged --network=host ${TPUVM_IMAGE_URL}@sha256:${TPUVM_IMAGE_VERSION} bash

From there you can install optimum-tpu through the pip instructions above.