{"id":22090354,"url":"https://github.com/rajeshthallam/fastertransformer-converter","last_synced_at":"2026-05-19T06:01:31.321Z","repository":{"id":94021948,"uuid":"601306718","full_name":"RajeshThallam/fastertransformer-converter","owner":"RajeshThallam","description":"This repository is a code sample to serve Large Language Models (LLM) on a Google Kubernetes Engine (GKE) cluster with GPUs running NVIDIA Triton Inference Server with FasterTransformer backend.","archived":false,"fork":false,"pushed_at":"2023-04-05T23:52:59.000Z","size":142,"stargazers_count":0,"open_issues_count":0,"forks_count":0,"subscribers_count":3,"default_branch":"main","last_synced_at":"2025-01-29T06:30:42.081Z","etag":null,"topics":["fastertransformer","gke","googlecloudplatform","inference","large-scale-machine-learning","llm","triton-inference-server"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/RajeshThallam.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-02-13T19:46:28.000Z","updated_at":"2023-09-05T22:41:13.000Z","dependencies_parsed_at":"2025-01-29T06:29:08.150Z","dependency_job_id":"f6deaaac-0833-476b-8338-420f5b926ee3","html_url":"https://github.com/RajeshThallam/fastertransformer-converter","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/RajeshThallam%2Ffastertransformer-converter","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/RajeshThallam%2Ffastertransformer-converter/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/RajeshThallam%2Ffastertransformer-converter/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/RajeshThallam%2Ffastertransformer-converter/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/RajeshThallam","download_url":"https://codeload.github.com/RajeshThallam/fastertransformer-converter/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":245182789,"owners_count":20573916,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["fastertransformer","gke","googlecloudplatform","inference","large-scale-machine-learning","llm","triton-inference-server"],"created_at":"2024-12-01T02:16:12.323Z","updated_at":"2026-05-19T06:01:26.258Z","avatar_url":"https://github.com/RajeshThallam.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Serving Large Language Models on Google Kubernetes Engine (GKE) using NVIDIA Triton Inference Server with FasterTransformer\n\nThis repository compiles prescriptive guidance and code sample to serve Large Language Models (LLM) such as [Unified Language Learner (UL2)](https://ai.googleblog.com/2022/10/ul2-20b-open-source-unified-language.html) on a Google Kubernetes Engine (GKE) cluster with GPUs running NVIDIA Triton Inference Server with FasterTransformer backend.\n\n- [NVIDIA Triton Inference Server](https://developer.nvidia.com/nvidia-triton-inference-server) is an open-source inference serving solution from NVIDIA to simplify and standardize the inference serving process supporting multiple frameworks and optimizations for both CPUs and GPUs.\n- [NVIDIA FasterTransformer](https://github.com/NVIDIA/FasterTransformer/) library implements an accelerated engine for the inference of transformer-based models, spanning multiple GPUs and nodes in a distributed manner.\n\nThe solution provides a Terraform standardized template to [deploy Triton inference server on GKE](https://github.com/jarokaz/triton-on-gke-sandbox) and integrate with other Google Cloud Managed Services.\n\n## Table of Contents\n---\n  - [High Level Flow](#high-level-flow)\n  - [Checkpoints](#checkpoints)\n  - [Getting started](#getting-started)\n  - [Repository Structure](#repository-structure)\n  - [Getting help](#getting-help)\n  - [Citations](#citations)\n---\n\n\n## High Level Flow\n\nThe solution demonstrates deploying the UL2 (20B parameter) model on a GKE cluster with GPUs. Assuming, JAX based checkpoints of a [pre-trained](https://github.com/google-research/google-research/tree/master/ul2#checkpoints) or fine-tuned UL2 model are available, the workflow has the following steps:\n\n1. Set up the environment running [Triton server on GKE cluster](https://github.com/jarokaz/triton-on-gke-sandbox).\n2. Convert JAX checkpoint to FasterTransformer checkpoint\n3. Serve the resultant model on GPUs using NVIDIA Triton Inference server with FasterTransformer backend\n4. Run evaluation script with test instances to compute model eval metrics\n\n![flow](/images/flow.png)\n\n## Checkpoints\n\nYou have following ways to access JAX based checkpoints for UL2:\n\n- You can find pre-trained UL2 checkpoints [here](https://github.com/google-research/google-research/tree/master/ul2#checkpoints).\n- We fined-tuned UL2 model with [XSum dataset](https://www.tensorflow.org/datasets/catalog/xsum) and made checkpoints available on Google Cloud Storage bucket at `gs://se-checkpoints/ul2-xsum/`.\n\n---\n\n**NOTE:** You can refer to the following [solution accelerator](https://github.com/GoogleCloudPlatform/t5x-on-vertex-ai) to fine-tune UL2 model with custom datasets using [T5X framework](https://github.com/google-research/t5x) that creates JAX based checkpoints. \n\n---\n\n## Getting Started\n\n### 1. Clone the GitHub repo\n\n``` bash\ncd ~\ngit clone https://github.com/RajeshThallam/fastertransformer-converter\ncd ~/fastertransformer-converter\n```\n\n### 2. Environment setup\n\nFollow the environment setup guide [here](https://github.com/jarokaz/triton-on-gke-sandbox) to create a GKE cluster running NVIDIA Triton on GPU node pool using Terraform standardized template. The setup performs the following steps:\n\n- [ ] Enable APIs\n- [ ] Run Terraform to provision the required resources\n- [ ] Deploy Ingress Gateway\n- [ ] Deploy NVIDIA GPU drivers\n- [ ] Configure and deploy Triton Inference Server\n- [ ] Run health check to validate the Triton deployment\n\nAs part of the environment setup, you configure following environment variables:\n\n``` bash\nexport PROJECT_ID=my-project-id\nexport REGION=us-central1\nexport ZONE=us-central1-a\nexport NETWORK_NAME=my-gke-network\nexport SUBNET_NAME=my-gke-subnet\nexport GCS_BUCKET_NAME=my-triton-repository\nexport GKE_CLUSTER_NAME=my-ft-gke\nexport TRITON_SA_NAME=triton-sa\nexport TRITON_NAMESPACE=triton\n```\n\n\u003cdiv class=\"alert alert-block alert-warning\"\u003e \n    \u003cp\u003e\n        \u003cstrong\u003e⚠️ NOTE: \u003c/strong\u003e\nBefore running through the setup, ensure the GPU node pool in the GKE cluster is configured with minimum 1 NVIDIA A100 GPU by setting the following variables. This is prerequisite for deploying large model, such as UL2, which is deployed with `bfloat16` (BF16) activation in this tutorial.\n    \u003c/p\u003e\n\u003c/div\u003e\n\n``` bash\nexport MACHINE_TYPE=a2-highgpu-1g\nexport ACCELERATOR_TYPE=nvidia-tesla-a100\nexport ACCELERATOR_COUNT=1\n```\n\nAfter provisioning the GKE cluster, configure access to the cluster:\n\n``` bash\ngcloud container clusters get-credentials ${GKE_CLUSTER_NAME} --project ${PROJECT_ID} --zone ${ZONE} \n\nkubectl create clusterrolebinding cluster-admin-binding --clusterrole cluster-admin --user \"$(gcloud config get-value account)\"\n```\n\n### 3. Convert JAX checkpoint to FasterTransformer checkpoint\n\nThe checkpoint format conversion from JAX to NVIDIA FasterTransformer is run on GKE cluster as a kubernetes Job on the GPU node pool. The source code for conversion script is located [here](./converter).\n\n- Create Docker repository in Google Artifact Registry to manage images\n\n``` bash\n# Configure parameters \nexport DOCKER_ARTIFACT_REPO=llm-inference     # \u003c-- Change to your repo name\n\n# Enable API\ngcloud services enable artifactregistry.googleapis.com\n\n# Create repository\ngcloud artifacts repositories create ${DOCKER_ARTIFACT_REPO} \\\n    --repository-format=docker \\\n    --location={REGION} \\\n    --description=\"Triton Docker repository\"\n\n# Authenticate to repository\ngcloud auth configure-docker ${REGION}-docker.pkg.dev --quiet\n```\n\n- Build container image\n\n``` bash\n# Configure container image name\nexport JAX_TO_FT_IMAGE_NAME=\"jax-to-fastertransformer\"\nexport JAX_TO_FT_IMAGE_URI=${REGION}\"-docker.pkg.dev/\"${PROJECT_ID}\"/\"${DOCKER_ARTIFACT_REPO}\"/\"${JAX_TO_FT_IMAGE_NAME}\n\n# Change directory\ncd ~/fastertransformer-converter\n\n# Run Cloud Build job to build the container image\nexport FILE_LOCATION=\"./converter\"\ngcloud builds submit \\\n  --region ${REGION} \\\n  --config converter/cloudbuild.yaml \\\n  --substitutions _IMAGE_URI=${JAX_TO_FT_IMAGE_URI},_FILE_LOCATION=${FILE_LOCATION} \\\n  --timeout \"2h\" \\\n  --machine-type=e2-highcpu-32 \\\n  --quiet\n```\n\n- Copy UL2 checkpoints to your Google Cloud Storage bucket\n\n``` bash\ngcloud storage cp -r gs://se-checkpoints/ul2-xsum/ gs://${GCS_BUCKET_NAME}/models/\n```\n\n- Run conversion job start with configuring job parameters\n\n``` bash\ncd ~/fastertransformer-converter/env-setup/kustomize\n  \ncat \u003c\u003c EOF \u003e ~/fastertransformer-converter/converter/configs.env\nksa=triton-sa\nconverter_image_uri=$JAX_TO_FT_IMAGE_URI\naccelerator_count=1\nmodel_name=ul2\ngcs_jax_ckpt=gs://${GCS_BUCKET_NAME}/models/ul2-xsum/\ngcs_ft_ckpt=gs://${GCS_BUCKET_NAME}/triton_model_repository/ul2-ft/\nEOF\n```\n\n- Deploy the configuration\n\n``` bash\nkubectl kustomize ./\n```\n\n- Run the job\n\n``` bash\nkubectl apply -k ./\n```\n\n- Monitor the job\n\n``` bash\nkubectl logs\n```\n\n### 4. Serve FasterTransformer checkpoint with NVIDIA Triton\n\n- Pull NVIDIA NeMo Inference container with NVIDIA Triton and FasterTransformer\n  - Sign in to [NGC](https://ngc.nvidia.com/signin) and select organization as `ea-participants`\n  - Get [API key from NGC](https://ngc.nvidia.com/setup/api-key ) to authorize docker client to access NGC Private Registry\n  - Authorize docker\n\n``` bash\ndocker login nvcr.io\n\nUsername: $oauthtoken\nPassword: Your key\n```\n-  \n  - Push the container to your Docker repository in Cloud Artifact Registry \n\n``` bash\nexport TRITON_FT_IMAGE_URI=${REGION}\"-docker.pkg.dev/\"${PROJECT_ID}\"/\"${DOCKER_ARTIFACT_REPO}\"/bignlp-inference:22.08-py3\"\n\ndocker pull nvcr.io/ea-bignlp/bignlp-inference:22.08-py3\n\ndocker tag nvcr.io/ea-bignlp/bignlp-inference:22.08-py3 ${TRITON_FT_IMAGE_URI}\n\ndocker push ${TRITON_FT_IMAGE_URI}\n```\n\n- Configure NVIDIA Triton Deployment parameters\n\n``` bash\ncd ~/triton-on-gke-sandbox/env-setup/kustomize\n\ncat \u003c\u003c EOF \u003e ~/triton-on-gke-sandbox/env-setup/kustomize/configs.env\nmodel_repository=gs://${GCS_BUCKET_NAME}/triton_model_repository/ul2-ft/\nksa=${TRITON_SA_NAME}\nEOF\n```\n\n- Update NVIDIA Triton container image\n\n``` bash\nkustomize edit set image \"nvcr.io/nvidia/tritonserver:22.01-py3=\"${TRITON_FT_IMAGE_URI}\n```\n\n- Deploy the configuration\n\n``` bash\nkubectl kustomize ./\n```\n\n- Deploy to the cluster\n\n``` bash\nkubectl apply -k ./\n```\n\n- Run health check\n\nTo validate that NVIDIA Triton Inference Server has been deployed successfully try to access the server's health check API. You will access the server through Istio Ingress Gateway. Start by getting the external IP address of the  `istio-ingressgateway` service.\n\n```bash\nkubectl get services -n $TRITON_NAMESPACE\n\n# Invoke the health check API\nISTIO_GATEWAY_IP_ADDRESS=$(kubectl get services -n $TRITON_NAMESPACE \\\n   -o=jsonpath='{.items[?(@.metadata.name==\"istio-ingressgateway\")].status.loadBalancer.ingress[0].ip}')\n\ncurl -v ${ISTIO_GATEWAY_IP_ADDRESS}/v2/health/ready\n```\n\nIf the returned status is `200OK` the server is up and accessible through the gateway.\n\n### 5. Run evaluation \n\n\n\n## Repository Structure\n\n```\n.\n├── converter\n├── evaluator\n└── README.md\n```\n\n- `/converter`: Source for converting JAX checkpoints to FasterTransformer checkpoints\n- `/evaluator`: Source for running model evaluation on validation dataset with model hosted on NVIDIA Triton\n\n## Getting Help\nIf you have any questions or if you found any problems with this repository, please report through GitHub issues.\n\n## Citations\n- XSum\n    ```\n    @InProceedings{xsum-emnlp,\n    author =      \"Shashi Narayan and Shay B. Cohen and Mirella Lapata\",\n    title =       \"Don't Give Me the Details, Just the Summary! {T}opic-Aware Convolutional Neural Networks for Extreme Summarization\",\n    booktitle =   \"Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing \",\n    year =        \"2018\",\n    address =     \"Brussels, Belgium\",\n    }\n    ```\n- UL2\n    ```\n    @article{tay2022unifying,\n    title={Unifying Language Learning Paradigms},\n    author={Yi Tay*, Mostafa Dehghani*, Vinh Q. Tran, Xavier Garcia, Dara Bahr, Tal Schuster, Huaixiu Steven Zheng, Neil Houlsby, and Donald Metzler},\n    year={2022}\n    }\n    ```","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Frajeshthallam%2Ffastertransformer-converter","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Frajeshthallam%2Ffastertransformer-converter","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Frajeshthallam%2Ffastertransformer-converter/lists"}