Menyajikan model open source menggunakan TPU di GKE dengan Optimum TPU

Tutorial ini menunjukkan cara menayangkan model open source model bahasa besar (LLM), menggunakan Unit Pemrosesan Tensor (TPU) di Google Kubernetes Engine (GKE) dengan framework penayangan Optimum TPU dari Hugging Face. Dalam tutorial ini, Anda akan mendownload model open source dari Hugging Face dan men-deploy model di cluster Standard GKE menggunakan container yang menjalankan Optimum TPU.

Panduan ini memberikan titik awal jika Anda memerlukan kontrol terperinci, skalabilitas, ketahanan, portabilitas, dan efektivitas biaya Kubernetes terkelola saat men-deploy dan menayangkan beban kerja AI/ML Anda.

Tutorial ini ditujukan untuk pelanggan AI Generatif di ekosistem Hugging Face, pengguna baru atau lama GKE, Engineer ML, engineer MLOps (DevOps), atau administrator platform yang tertarik untuk menggunakan kemampuan orkestrasi penampung Kubernetes untuk menayangkan LLM.

Sebagai pengingat, Anda memiliki beberapa opsi untuk inferensi LLM di Google Cloud—yang mencakup penawaran seperti Vertex AI, GKE, dan Google Compute Engine—tempat Anda dapat menggabungkan library penayangan seperti JetStream, vLLM, dan penawaran partner lainnya. Misalnya, Anda dapat menggunakan JetStream untuk mendapatkan pengoptimalan terbaru dari project. Jika lebih memilih opsi Hugging Face, Anda dapat menggunakan Optimum TPU.

TPU Optimal mendukung fitur berikut:

  • Batch berkelanjutan
  • Streaming token
  • Greedy search dan multinomial sampling menggunakan transformer.

Mendapatkan akses ke model

Anda dapat menggunakan model Gemma 2B atau Llama3 8B. Tutorial ini berfokus pada dua model ini, tetapi Optimum TPU mendukung lebih banyak model.

Gemma 2B

Untuk mendapatkan akses ke model Gemma untuk deployment ke GKE, Anda harus menandatangani perjanjian izin lisensi terlebih dahulu, lalu membuat token akses Hugging Face.

Anda harus menandatangani perjanjian izin untuk menggunakan Gemma. Ikuti petunjuk berikut:

  1. Akses halaman izin model.
  2. Verifikasi izin menggunakan akun Hugging Face Anda.
  3. Setujui persyaratan model.

Membuat token akses

Buat token Hugging Face baru jika Anda belum memilikinya:

  1. Klik Profil Anda > Setelan > Token Akses.
  2. Klik New Token.
  3. Tentukan Nama pilihan Anda dan Peran minimal Read.
  4. Klik Generate a token.
  5. Salin token yang dihasilkan ke papan klip Anda.

Llama3 8B

Anda harus menandatangani perjanjian izin untuk menggunakan Llama3 8b di Hugging Face Repo

Membuat token akses

Buat token Hugging Face baru jika Anda belum memilikinya:

  1. Klik Profil Anda > Setelan > Token Akses.
  2. Pilih New Token.
  3. Tentukan Nama pilihan Anda dan Peran minimal Read.
  4. Pilih Generate a token.
  5. Salin token yang dihasilkan ke papan klip Anda.

Membuat cluster GKE

Buat cluster GKE Standard dengan 1 node CPU:

gcloud container clusters create CLUSTER_NAME \
    --project=PROJECT_ID \
    --num-nodes=1 \
    --location=ZONE

Buat TPU node pool

Buat node pool TPU v5e dengan 1 node dan 8 chip:

gcloud container node-pools create tpunodepool \
    --location=ZONE \
    --num-nodes=1 \
    --machine-type=ct5lp-hightpu-8t \
    --cluster=CLUSTER_NAME

Jika resource TPU tersedia, GKE akan menyediakan node pool. Jika resource TPU tidak tersedia untuk sementara, output akan menampilkan pesan error GCE_STOCKOUT. Untuk memecahkan masalah error kehabisan stok TPU, lihat Sumber daya TPU tidak cukup untuk memenuhi permintaan TPU.

Konfigurasi kubectl untuk berkomunikasi dengan cluster Anda:

gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${ZONE}

Buat container

Jalankan perintah make untuk membangun image

cd optimum-tpu && make tpu-tgi

Kirim image ke Artifact Registry

gcloud artifacts repositories create optimum-tpu --repository-format=docker --location=REGION_NAME && \
gcloud auth configure-docker REGION_NAME-docker.pkg.dev && \
docker image tag huggingface/optimum-tpu REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest && \
docker push REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest

Buat Secret Kubernetes untuk kredensial Hugging Face

Buat Secret Kubernetes yang berisi token Hugging Face:

kubectl create secret generic hf-secret \
  --from-literal=hf_api_token=${HF_TOKEN} \
  --dry-run=client -o yaml | kubectl apply -f -

Men-deploy TPU Optimum

Untuk men-deploy Optimum TPU, tutorial ini menggunakan Deployment Kubernetes. Deployment adalah objek Kubernetes API yang memungkinkan Anda menjalankan beberapa replika Pod yang didistribusikan di antara node dalam cluster.

Gemma 2B

  1. Simpan manifes Deployment berikut sebagai optimum-tpu-gemma-2b-2x4.yaml:

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: tgi-tpu
spec:
  replicas: 1
  selector:
    matchLabels:
      app: tgi-tpu
  template:
    metadata:
      labels:
        app: tgi-tpu
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      containers:
      - name: tgi-tpu
        image: REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
        args:
        - --model-id=google/gemma-2b
        - --max-concurrent-requests=4
        - --max-input-length=8191
        - --max-total-tokens=8192
        - --max-batch-prefill-tokens=32768
        - --max-batch-size=16
        securityContext:
            privileged: true
        env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
        ports:
        - containerPort: 80
        resources:
          limits:
            google.com/tpu: 8
        livenessProbe:
          httpGet:
            path: /health
            port: 80
          initialDelaySeconds: 300
          periodSeconds: 120

---
apiVersion: v1
kind: Service
metadata:
  name: service
spec:
  selector:
    app: tgi-tpu
  ports:
    - name: http
      protocol: TCP
      port: 8080
      targetPort: 80

    Manifes ini menjelaskan deployment TPU Optimal dengan load balancer internal pada port TCP 8080.

  2. Terapkan manifes

    kubectl apply -f optimum-tpu-gemma-2b-2x4.yaml

Llama3 8B

  1. Simpan manifes berikut sebagai optimum-tpu-llama3-8b-2x4.yaml:

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: tgi-tpu
spec:
  replicas: 1
  selector:
    matchLabels:
      app: tgi-tpu
  template:
    metadata:
      labels:
        app: tgi-tpu
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      containers:
      - name: tgi-tpu
        image: REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
        args:
        - --model-id=meta-llama/Meta-Llama-3-8B
        - --max-concurrent-requests=4
        - --max-input-length=8191
        - --max-total-tokens=8192
        - --max-batch-prefill-tokens=32768
        - --max-batch-size=16
        env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
        ports:
        - containerPort: 80
        resources:
          limits:
            google.com/tpu: 8
        livenessProbe:
          httpGet:
            path: /health
            port: 80
          initialDelaySeconds: 300
          periodSeconds: 120
---
apiVersion: v1
kind: Service
metadata:
  name: service
spec:
  selector:
    app: tgi-tpu
  ports:
    - name: http
      protocol: TCP
      port: 8080
      targetPort: 80

    Manifes ini menjelaskan deployment TPU Optimal dengan load balancer internal pada port TCP 8080.

  2. Terapkan manifes

    kubectl apply -f optimum-tpu-llama3-8b-2x4.yaml

Lihat log dari Deployment yang sedang berjalan:

kubectl logs -f -l app=tgi-tpu

Outputnya akan mirip dengan berikut ini:

2024-07-09T22:39:34.365472Z  WARN text_generation_router: router/src/main.rs:295: no pipeline tag found for model google/gemma-2b
2024-07-09T22:40:47.851405Z  INFO text_generation_router: router/src/main.rs:314: Warming up model
2024-07-09T22:40:54.559269Z  INFO text_generation_router: router/src/main.rs:351: Setting max batch total tokens to 64
2024-07-09T22:40:54.559291Z  INFO text_generation_router: router/src/main.rs:352: Connected
2024-07-09T22:40:54.559295Z  WARN text_generation_router: router/src/main.rs:366: Invalid hostname, defaulting to 0.0.0.0

Pastikan model didownload sepenuhnya sebelum melanjutkan ke bagian berikutnya.

Menyajikan model

Siapkan penerusan port ke model:

kubectl port-forward svc/service 8080:8080

Berinteraksi dengan server model menggunakan curl

Verifikasi model yang di-deploy:

Dalam sesi terminal baru, gunakan curl untuk memulai percakapan dengan model:

curl 127.0.0.1:8080/generate     -X POST     -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":40}}'     -H 'Content-Type: application/json'

Outputnya akan mirip dengan berikut ini:

{"generated_text":"\n\nDeep learning is a subset of machine learning that uses artificial neural networks to learn from data.\n\nArtificial neural networks are inspired by the way the human brain works. They are made up of multiple layers"}