Optimum TPU is the interface between the Hugging Face Transformers library and Google Cloud TPU devices.
Project description
Tensor Processing Units (TPU) are AI accelerator made by Google to optimize performance and cost from AI training to inference.
This repository exposes an interface similar to what Hugging Face transformers library provides to interact with a magnitude of models developed by research labs, institutions and the community.
We aim at providing our user the best possible performances targeting Google Cloud TPUs for both training and inference working closely with Google and Google Cloud to make this a reality.
Supported Model and Tasks
We currently support a few LLM models targeting text generation scenarios:
- 💎 Gemma (2b, 7b)
- 🦙 Llama2 (7b) and Llama3 (8b)
- 💨 Mistral (7b)
Installation
optimum-tpu
comes with an handy PyPi released package compatible with your classical python dependency management tool.
pip install optimum-tpu -f https://storage.googleapis.com/libtpu-releases/index.html
Inference
optimum-tpu
provides a set of dedicated tools and integrations in order to leverage Cloud TPUs for inference, especially
on the latest TPU version v5e
.
Other TPU versions will be supported along the way.
Text-Generation-Inference
As part of the integration, we do support a text-generation-inference (TGI) backend allowing to deploy and serve incoming HTTP requests and execute them on Cloud TPUs.
Please see the TGI specific documentation on how to get started
Training
Fine-tuning is supported and tested on the TPU v5e
. We have tested so far:
- 🦙 Llama-2 7B and Llama-3 8B
- 💎 Gemma 2B and 7B
You can check the examples:
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for optimum_tpu-0.1.5-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | e12e154a974ba41ca90df63daa963632b2dfa6642139f281a9a1e8a1ca721b3f |
|
MD5 | 2833072aa7a53e82dc0df5ac81755dc2 |
|
BLAKE2b-256 | 4ce60960fe4c878f7ae4a50ec880b33b415e246493e6063cc04133dfd19c5eb7 |