From 0ec38ca2fac54c0002dbd0e76409b93d9391ecc9 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Thu, 29 Aug 2024 15:03:12 +0000
Subject: [PATCH] doc(README): add mention on how to install and enable
 Pytorch/Jetstream

---
 README.md | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/README.md b/README.md
index 9dc0ee0f..c998b955 100644
--- a/README.md
+++ b/README.md
@@ -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: