Skip to content

Commit

Permalink
doc(README): add mention on how to install and enable Pytorch/Jetstream
Browse files Browse the repository at this point in the history
  • Loading branch information
tengomucho committed Aug 30, 2024
1 parent 4cd1d9a commit 0ec38ca
Showing 1 changed file with 11 additions and 0 deletions.
11 changes: 11 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,17 @@ incoming HTTP requests and execute them on Cloud TPUs.

Please see the [TGI specific documentation](text-generation-inference) on how to get started.

### JetStream Pytorch Engine

`optimum-tpu` provides an optional support of JetStream Pytorch engine inside of TGI. This support can be installed using the dedicated command:

```shell
pip install "optimum-tpu[jetstream-pt]" -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
-f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
```

To enable the support, export the environment variable `JETSTREAM_PT=1`.

## Training

Fine-tuning is supported and tested on the TPU `v5e`. We have tested so far:
Expand Down

0 comments on commit 0ec38ca

Please sign in to comment.