Add Google TPU to the mix (#1797)

Co-authored-by: regisss <15324346+regisss@users.noreply.github.com>
huggingface · Apr 8, 2024 · 5d194dd · 5d194dd
1 parent 2c06aa2
commit 5d194dd
Show file tree

Hide file tree

Showing 5 changed files with 68 additions and 3 deletions.
diff --git a/.github/workflows/build_main_documentation.yml b/.github/workflows/build_main_documentation.yml
@@ -49,6 +49,11 @@ jobs:
           repository: 'huggingface/optimum-amd'
           path: optimum-amd
 
+      - uses: actions/checkout@v2
+        with:
+          repository: 'huggingface/optimum-tpu'
+          path: optimum-tpu
+
       - name: Free disk space
         run: |
           df -h
@@ -150,6 +155,16 @@ jobs:
           mv furiosa-doc-build ../optimum
           cd ..
 
+      - name: Make TPU documentation
+        run: |
+          sudo docker system prune -a -f
+          cd optimum-tpu
+          pip install -U pip
+          pip install .
+          doc-builder build optimum.tpu docs/source/ --build_dir tpu-doc-build --version pr_$PR_NUMBER --version_tag_suffix "" --html --clean
+          mv tpu-doc-build ../optimum
+          cd ..
+
       - name: Make AMD documentation
         run: |
           sudo docker system prune -a -f
@@ -171,7 +186,7 @@ jobs:
       - name: Combine subpackage documentation
         run: |
           cd optimum
-          sudo python docs/combine_docs.py --subpackages nvidia amd intel neuron habana furiosa --version ${{ env.VERSION }}
+          sudo python docs/combine_docs.py --subpackages nvidia amd intel neuron tpu habana furiosa --version ${{ env.VERSION }}
           cd ..
 
       - name: Push to repositories

diff --git a/.github/workflows/build_pr_documentation.yml b/.github/workflows/build_pr_documentation.yml
@@ -53,6 +53,11 @@ jobs:
           repository: 'huggingface/optimum-amd'
           path: optimum-amd
 
+      - uses: actions/checkout@v2
+        with:
+          repository: 'huggingface/optimum-tpu'
+          path: optimum-tpu
+
       - name: Setup environment
         run: |
           pip uninstall -y doc-builder
@@ -91,6 +96,16 @@ jobs:
           sudo mv amd-doc-build ../optimum
           cd ..
 
+      - name: Make TPU documentation
+        run: |
+          sudo docker system prune -a -f
+          cd optimum-tpu
+          pip install -U pip
+          pip install .
+          doc-builder build optimum.tpu docs/source/ --build_dir tpu-doc-build --version pr_$PR_NUMBER --version_tag_suffix "" --html --clean
+          mv tpu-doc-build ../optimum
+          cd ..
+
       - name: Make Optimum documentation
         run: |
           sudo docker system prune -a -f
@@ -101,7 +116,7 @@ jobs:
       - name: Combine subpackage documentation
         run: |
           cd optimum
-          sudo python docs/combine_docs.py --subpackages nvidia amd intel neuron habana furiosa --version pr_$PR_NUMBER
+          sudo python docs/combine_docs.py --subpackages nvidia amd intel neuron tpu habana furiosa --version pr_$PR_NUMBER
           sudo mv optimum-doc-build ../
           cd ..
 

diff --git a/docs/combine_docs.py b/docs/combine_docs.py
@@ -108,6 +108,31 @@ def add_neuron_doc(base_toc: List):
     )
 
 
+def add_tpu_doc(base_toc: List):
+    """
+    Extends the table of content with a section about Optimum TPU.
+
+    Args:
+        base_toc (List): table of content for the doc of Optimum.
+    """
+    # Update optimum table of contents
+    base_toc.insert(
+        SUBPACKAGE_TOC_INSERT_INDEX,
+        {
+            "sections": [
+                {
+                    # Ideally this should directly point at https://huggingface.co/docs/optimum-tpu/index
+                    # Current hacky solution is to have a redirection in _redirects.yml
+                    "local": "docs/optimum-tpu/index",
+                    "title": "🤗 Optimum-TPU",
+                }
+            ],
+            "title": "Google TPUs",
+            "isExpanded": False,
+        },
+    )
+
+
 def main():
     args = parser.parse_args()
     optimum_path = Path("optimum-doc-build")
@@ -121,6 +146,9 @@ def main():
         if subpackage == "neuron":
             # Neuron has its own doc so it is managed differently
             add_neuron_doc(base_toc)
+        elif subpackage == "tpu":
+            # Optimum TPU has its own doc so it is managed differently
+            add_tpu_doc(base_toc)
         elif subpackage == "nvidia":
             # At the moment, Optimum Nvidia's doc is the README of the GitHub repo
             # It is linked to in optimum/docs/source/nvidia_overview.mdx

diff --git a/docs/source/_redirects.yml b/docs/source/_redirects.yml
@@ -28,3 +28,6 @@ intel_trainer: intel/reference_inc
 
 # Optimum Neuron
 docs/optimum-neuron/index: /docs/optimum-neuron/index
+
+# Optimum TPU
+docs/optimum-tpu/index: /docs/optimum-tpu/index
diff --git a/docs/source/index.mdx b/docs/source/index.mdx
@@ -25,7 +25,7 @@ As such, Optimum enables developers to efficiently use any of these platforms wi
 The packages below enable you to get the best of the 🤗 Hugging Face ecosystem on various types of devices.
 
 <div class="mt-10">
-  <div class="w-full flex flex-col space-y-4 md:space-y-0 md:grid md:grid-cols-3 md:gap-y-4 md:gap-x-5">
+  <div class="w-full flex flex-col space-y-4 md:space-y-0 md:grid md:grid-cols-4 md:gap-y-4 md:gap-x-5">
     <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://github.com/huggingface/optimum-nvidia"
       ><div class="w-full text-center bg-gradient-to-br from-green-600 to-green-600 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">NVIDIA</div>
       <p class="text-gray-700">Accelerate inference with NVIDIA TensorRT-LLM on the <span class="underline" onclick="event.preventDefault(); window.open('https://developer.nvidia.com/blog/nvidia-tensorrt-llm-supercharges-large-language-model-inference-on-nvidia-h100-gpus/', '_blank');">NVIDIA platform</span></p>
@@ -42,6 +42,10 @@ The packages below enable you to get the best of the 🤗 Hugging Face ecosystem
       ><div class="w-full text-center bg-gradient-to-br from-orange-400 to-orange-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">AWS Trainium/Inferentia</div>
       <p class="text-gray-700">Accelerate your training and inference workflows with <span class="underline" onclick="event.preventDefault(); window.open('https://aws.amazon.com/machine-learning/trainium/', '_blank');">AWS Trainium</span> and <span class="underline" onclick="event.preventDefault(); window.open('https://aws.amazon.com/machine-learning/inferentia/', '_blank');">AWS Inferentia</span></p>
     </a>
+    <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="https://huggingface.co/docs/optimum-tpu/index"
+      ><div class="w-full text-center bg-gradient-to-tr from-blue-200 to-blue-600 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Google TPUs</div>
+      <p class="text-gray-700">Accelerate your training and inference workflows with <span class="underline" onclick="event.preventDefault(); window.open('https://cloud.google.com/tpu', '_blank');">Google TPUs</span></p>
+    </a>
     <a class="!no-underline border dark:border-gray-700 p-5 rounded-lg shadow hover:shadow-lg" href="./habana/index"
       ><div class="w-full text-center bg-gradient-to-br from-indigo-400 to-indigo-500 rounded-lg py-1.5 font-semibold mb-5 text-white text-lg leading-relaxed">Habana</div>
       <p class="text-gray-700">Maximize training throughput and efficiency with <span class="underline" onclick="event.preventDefault(); window.open('https://docs.habana.ai/en/latest/Gaudi_Overview/Gaudi_Architecture.html', '_blank');">Habana's Gaudi processor</span></p>