From f92066c3f2b5bbe41dbfa58f447eb5a128a7bd9b Mon Sep 17 00:00:00 2001 From: Funtowicz Morgan Date: Mon, 8 Apr 2024 11:49:42 +0200 Subject: [PATCH] Add documentation to the repository (#13) * Initial commit for the documentation * Initial doc-build workflow * Initial how to deploy instance * Add the gcloud alpha requirement * Address general comments * Add PR doc workflows * Fix workflow * Add _toctree * Fix docs/source/howto/overview --------- Co-authored-by: regisss <15324346+regisss@users.noreply.github.com> --- .github/workflows/doc-build.yml | 57 +++++++++++++ .github/workflows/doc-pr-build.yml | 52 ++++++++++++ .github/workflows/upload_pr_documentation.yml | 16 ++++ docs/source/_toctree.yml | 17 ++++ docs/source/howto/deploy.mdx | 82 +++++++++++++++++++ docs/source/howto/overview.mdx | 26 ++++++ docs/source/howto/serving.mdx | 0 docs/source/index.mdx | 69 ++++++++++++++++ docs/source/tutorials/overview.mdx | 0 pyproject.toml | 2 +- 10 files changed, 320 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/doc-build.yml create mode 100644 .github/workflows/doc-pr-build.yml create mode 100644 .github/workflows/upload_pr_documentation.yml create mode 100644 docs/source/_toctree.yml create mode 100644 docs/source/howto/deploy.mdx create mode 100644 docs/source/howto/overview.mdx create mode 100644 docs/source/howto/serving.mdx create mode 100644 docs/source/index.mdx create mode 100644 docs/source/tutorials/overview.mdx diff --git a/.github/workflows/doc-build.yml b/.github/workflows/doc-build.yml new file mode 100644 index 00000000..112df125 --- /dev/null +++ b/.github/workflows/doc-build.yml @@ -0,0 +1,57 @@ +name: Build documentation + +on: + push: + branches: + - main + tags: + - 'v[0-9]+.[0-9]+.[0-9]+' + + paths: + - 'docs/source/**' + - 'docs/assets/**' + - 'optimum/**' + - '.github/workflows/doc-build.yml' + workflow_dispatch: + +jobs: + build_documentation: + runs-on: ubuntu-latest + env: + COMMIT_SHA: ${{ github.event.pull_request.head.sha }} + PR_NUMBER: ${{ github.event.number }} + EVENT_CONTEXT: ${{ toJSON(github.event) }} + PR_CLONE_URL: ${{ github.event.pull_request.head.repo.clone_url }} + + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-node@v4 + with: + node-version: '18' + cache-dependency-path: "kit/package-lock.json" + + - name: Set environment variables + run: | + cd optimum + version=`echo "$(grep '^__version__ =' tpu/version.py | cut -d '=' -f 2- | xargs)"` + + if [[ $version == *.dev0 ]] + then + echo "VERSION=main" >> $GITHUB_ENV + else + echo "VERSION=v$version" >> $GITHUB_ENV + fi + + cd .. + + - name: Setup environment + run: | + pip install ".[quality]" + + - name: Make documentation + shell: bash + run: | + doc-builder build optimum.tpu docs/source/ --repo_name optimum-tpu --build_dir tpu-doc-build/ --version ${{ env.VERSION }} --version_tag_suffix "" --html --clean + cd tpu-doc-build/ + mv optimum.tpu optimum-tpu + doc-builder push optimum-tpu --doc_build_repo_id "hf-doc-build/doc-build" --token "${{ secrets.HF_DOC_BUILD_PUSH }}" --commit_msg "Updated with commit $COMMIT_SHA See: https://github.com/huggingface/optimum-tpu/commit/$COMMIT_SHA" --n_retries 5 diff --git a/.github/workflows/doc-pr-build.yml b/.github/workflows/doc-pr-build.yml new file mode 100644 index 00000000..e951b80a --- /dev/null +++ b/.github/workflows/doc-pr-build.yml @@ -0,0 +1,52 @@ +name: Build PR Documentation + +on: + pull_request: + branches: [ main ] + paths: + - 'docs/source/**' + - 'docs/assets/**' + - 'optimum/**' + - '.github/workflows/doc-pr-build.yml' + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} + cancel-in-progress: true + +jobs: + build_documentation: + runs-on: ubuntu-latest + env: + COMMIT_SHA: ${{ github.event.pull_request.head.sha }} + PR_NUMBER: ${{ github.event.number }} + EVENT_CONTEXT: ${{ toJSON(github.event) }} + PR_CLONE_URL: ${{ github.event.pull_request.head.repo.clone_url }} + + steps: + - uses: actions/checkout@v3 + - uses: actions/setup-node@v3 + with: + node-version: '18' + cache-dependency-path: "kit/package-lock.json" + + - name: Setup environment + run: | + pip install -U pip + pip install ".[quality]" + + - name: Make documentation + shell: bash + run: | + doc-builder build optimum.tpu docs/source/ --repo_name optimum-tpu --build_dir tpu-doc-build/ --version pr_${{ env.PR_NUMBER }} --version_tag_suffix "" --html --clean + + - name: Save commit_sha & pr_number + run: | + cd tpu-doc-build/ + mv optimum.tpu optimum-tpu + echo ${{ env.COMMIT_SHA }} > ./commit_sha + echo ${{ env.PR_NUMBER }} > ./pr_number + + - uses: actions/upload-artifact@v3 + with: + name: doc-build-artifact + path: tpu-doc-build/ diff --git a/.github/workflows/upload_pr_documentation.yml b/.github/workflows/upload_pr_documentation.yml new file mode 100644 index 00000000..18d9e3ec --- /dev/null +++ b/.github/workflows/upload_pr_documentation.yml @@ -0,0 +1,16 @@ +name: Upload PR Documentation + +on: + workflow_run: + workflows: ["Build PR Documentation"] + types: + - completed + +jobs: + build: + uses: huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml@main + with: + package_name: optimum-tpu + secrets: + hf_token: ${{ secrets.HF_DOC_BUILD_PUSH }} + comment_bot_token: ${{ secrets.COMMENT_BOT_TOKEN }} diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml new file mode 100644 index 00000000..898ff81e --- /dev/null +++ b/docs/source/_toctree.yml @@ -0,0 +1,17 @@ +- sections: + - local: index + title: 🤗 Optimum-TPU + - sections: + - local: tutorials/overview + title: Overview + title: Tutorials + - sections: + - local: howto/overview + title: Overview + - local: howto/deploy + title: Deploying a Google Cloud TPU instance + - local: howto/serving + title: Deploying a TGI server on a Google Cloud TPU instance + title: How-To Guides + title: Optimum-TPU + isExpanded: true diff --git a/docs/source/howto/deploy.mdx b/docs/source/howto/deploy.mdx new file mode 100644 index 00000000..a44b1d04 --- /dev/null +++ b/docs/source/howto/deploy.mdx @@ -0,0 +1,82 @@ +# Deploying a Google TPU instance on Google Cloud Platform (GCP) + + +## Context + +We assume the reader has already created a Google Cloud Platform (GCP) user or organisation account and an +associated project. + +We also assume the reader to have the Google Cloud CLI installed. If not please follow the links right after to +[install](https://cloud.google.com/sdk/docs/install) and [setup](https://cloud.google.com/sdk/docs/initializing). + +## Creating the initial TPUVM on GCP + +In order to create your initial TPU instance, you will need to provide some information: + +- The zone in GCP you would like to see the instance being deployed (close to the reader for development purpose, close to the end user for production for instance) +- Which kind of TPU you would like to target +- Which version of the TPU runtime you want to leverage on the instance +- Custom instance name to quickly skim and refer back to the instance + +Overall the end command looks like this: + +```bash +gcloud compute tpus tpu-vm create \ +--zone= \ +--accelerator-type= \ +--version= +``` + +### Deploying a TPU v5litepod-8 instance + +In our case we will be deploying a `v5litepod-8` instance name `optimum-tpu-get-started` +in the GCP region `us-west4-a` using the latest `v2-alpha-tpuv5-lite` runtime version. + +Of course, feel free to adjust all these parameters to the one that match with your usage and quotas. + +Before creating the instance, please make sure to install `gcloud alpha component` as it is required to be able to +target TPUv5 VMs: `gcloud components install alpha` + +```bash +gcloud alpha compute tpus tpu-vm create optimum-tpu-get-started \ +--zone=us-west4-a \ +--accelerator-type=v5litepod-8 \ +--version=v2-alpha-tpuv5 +``` + +## Connecting to the instance + +```bash +gcloud compute tpus tpu-vm ssh --zone= +$ > +``` + +In the example above deploying v5litepod-8 it would be something like: + +```bash +gcloud compute tpus tpu-vm ssh optimum-tpu-get-started --zone=us-west4-a +$ > +``` + +## Setting up the instance to run AI workloads on TPUs + +### Optimum-TPU with PyTorch/XLA + +If you want to leverage PyTorch/XLA through Optimum-TPU, it should be as simple as + +```bash +$ python3 -m pip install optimum-tpu +$ export PJRT_DEVICE=TPU +``` + +Now you can validate the installation with the following command which should print `xla:0` as we do have a single +TPU device bound to this instance. + +```bash +$ python -c "import torch_xla.core.xla_model as xm; print(xm.xla_device())" +xla:0 +``` + +### Optimum-TPU with JAX + +JAX is coming very soon - stay tuned! \ No newline at end of file diff --git a/docs/source/howto/overview.mdx b/docs/source/howto/overview.mdx new file mode 100644 index 00000000..226b4ce3 --- /dev/null +++ b/docs/source/howto/overview.mdx @@ -0,0 +1,26 @@ +# Optimum-TPU How To + +This page gives you access to handful of walkthrough scenarios to leverage Google TPUs for your use case. + +## Looking for? + + \ No newline at end of file diff --git a/docs/source/howto/serving.mdx b/docs/source/howto/serving.mdx new file mode 100644 index 00000000..e69de29b diff --git a/docs/source/index.mdx b/docs/source/index.mdx new file mode 100644 index 00000000..602dbcca --- /dev/null +++ b/docs/source/index.mdx @@ -0,0 +1,69 @@ + + +# 🤗 Optimum TPU + +Optimum TPU provides all the necessary machinery to leverage and optimize AI workloads runningo on [Google Cloud TPU devices](https://cloud.google.com/tpu/docs). + +The API provides the overall same user-experience as Hugging Face transformers with the minimum amount of changes required to target performance for both 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 +```bash +pip install optimum-tpu +``` + +### Run Cloud TPU within Docker container + +### PyTorch +```bash +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. + + + \ No newline at end of file diff --git a/docs/source/tutorials/overview.mdx b/docs/source/tutorials/overview.mdx new file mode 100644 index 00000000..e69de29b diff --git a/pyproject.toml b/pyproject.toml index f5957a18..4d5370e1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,7 +50,7 @@ dependencies = [ [project.optional-dependencies] tests = ["pytest", "safetensors"] -quality = ["black", "ruff", "isort",] +quality = ["black", "ruff", "isort", "hf_doc_builder @ git+https://github.com/huggingface/doc-builder.git"] [project.urls] Homepage = "https://hf.co/hardware"