Saxml을 사용해서 GKE에서 멀티 호스트 TPU를 사용하는 LLM 서빙

표준

이 튜토리얼에서는 효율적이며 확장 가능한 아키텍처를 위해 Saxml을 사용해서 Google Kubernetes Engine(GKE)에서 멀티 호스트 TPU 슬라이스 노드 풀을 사용하는 대규모 언어 모델(LLM)을 배포 및 서빙하는 방법을 보여줍니다.

배경

Saxml은 Paxml, JAX, PyTorch 프레임워크를 지원하는 실험적인 시스템입니다. TPU를 사용해서 이러한 프레임워크에서 데이터 처리를 가속화할 수 있습니다. GKE에서 TPU 배포 방법을 보여주기 위해 이 튜토리얼에서는 175B LmCloudSpmd175B32Test 테스트 모델을 사용합니다. GKE는 각각 4x8 토폴로지를 사용해서 이 테스트 모델을 2개의 v5e TPU 슬라이스 노드 풀에 배포합니다.

테스트 모델을 올바르게 배포하기 위해 TPU 토폴로지는 모델 크기에 따라 정의되었습니다. N 크기(10억 단위)의 16비트 모델에 GB 단위로 약 2배의 메모리(2xN)가 필요하다는 점을 고려할 때 175B LmCloudSpmd175B32Test 모델에는 약 350GB 메모리가 필요합니다. TPU v5e 단일 TPU 칩에는 16GB가 포함됩니다. 350GB를 지원하기 위해서는 GKE에 v5e TPU 칩 21개(350/16= 21)가 필요합니다. TPU 구성 매핑에 따라 이 튜토리얼에 적절한 TPU 구성은 다음과 같습니다.

머신 유형: ct5lp-hightpu-4t
토폴로지: 4x8(TPU 칩 32개)

GKE에서 TPU를 배포할 때는 모델 서빙을 위해 올바른 TPU 토폴로지를 선택하는 것이 중요합니다. 자세한 내용은 TPU 구성 계획을 참조하세요.

목표

이 튜토리얼은 데이터 모델 서빙을 위해 GKE 조정 기능을 사용하려는 MLOps 또는 DevOps 엔지니어 또는 플랫폼 관리자를 대상으로 합니다.

이 튜토리얼은 다음 과정을 다룹니다.

GKE Standard 클러스터로 환경을 준비합니다. 클러스터에는 4x8 토폴로지를 사용하는 2개의 v5e TPU 슬라이스 노드 풀이 포함됩니다.
Saxml을 배포합니다. Saxml에는 관리자 서버, 모델 서버로 작동하는 포드 그룹, 사전 빌드된 HTTP 서버, 부하 분산기가 필요합니다.
Saxml을 사용하여 LLM을 제공합니다.

다음 다이어그램은 다음 튜토리얼이 구현하는 아키텍처를 보여줍니다.

시작하기 전에

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the required API.

Enable the API

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the required API.

Enable the API

Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin
Check for the roles
1. In the Google Cloud console, go to the IAM page.
  Go to IAM
2. Select the project.
3. In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.
4. For all rows that specify or include you, check the Role colunn to see whether the list of roles includes the required roles.
Grant the roles
1. In the Google Cloud console, go to the IAM page.
  IAM으로 이동
2. 프로젝트를 선택합니다.
3. 액세스 권한 부여를 클릭합니다.
4. 새 주 구성원 필드에 사용자 식별자를 입력합니다. 일반적으로 Google 계정의 이메일 주소입니다.
5. 역할 선택 목록에서 역할을 선택합니다.
6. 역할을 추가로 부여하려면 다른 역할 추가를 클릭하고 각 역할을 추가합니다.
7. 저장을 클릭합니다.

GKE에서 Cloud TPU에 대해 프로젝트 할당량이 충분한지 확인합니다.

환경 준비

Google Cloud 콘솔에서 Cloud Shell 인스턴스를 시작합니다.
Cloud Shell 열기
기본 환경 변수를 설정합니다.
```
  gcloud config set project PROJECT_ID
  export PROJECT_ID=$(gcloud config get project)
  export REGION=COMPUTE_REGION
  export ZONE=COMPUTE_ZONE
  export GSBUCKET=PROJECT_ID-gke-bucket
```
다음 값을 바꿉니다.
- PROJECT_ID: Google Cloud 프로젝트 ID
- COMPUTE_REGION: Compute Engine 리전
- COMPUTE_ZONE: ct5lp-hightpu-4t를 사용할 수 있는 영역

GKE Standard 클러스터 만들기

Cloud Shell을 사용하여 다음을 수행합니다.

GKE용 워크로드 아이덴티티 제휴를 사용하는 표준 클러스터를 만듭니다.
```
gcloud container clusters create saxml \
    --zone=${ZONE} \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --cluster-version=VERSION \
    --num-nodes=4
```
VERSION을 GKE 버전 번호로 바꿉니다. GKE는 버전 1.27.2-gke.2100 이상의 TPU v5e를 지원합니다. 자세한 내용은 GKE의 TPU 가용성을 참조하세요.

클러스터 만들기는 몇 분 정도 걸릴 수 있습니다.

tpu1이라는 첫 번째 노드 풀을 만듭니다.

gcloud container node-pools create tpu1 \
    --zone=${ZONE} \
    --num-nodes=8 \
    --machine-type=ct5lp-hightpu-4t \
    --tpu-topology=4x8 \
    --cluster=saxml

tpu2라는 두 번째 노드 풀을 만듭니다.

gcloud container node-pools create tpu2 \
    --zone=${ZONE} \
    --num-nodes=8 \
    --machine-type=ct5lp-hightpu-4t \
    --tpu-topology=4x8 \
    --cluster=saxml

다음 리소스를 만들었습니다.

CPU 노드가 4개 있는 표준 클러스터
4x8 토폴로지를 사용하는 2개의 v5e TPU 슬라이스 노드 풀. 각 노드 풀에 TPU 슬라이스 노드 8개가 포함되고 이러한 노드에는 개당 TPU 칩 4개가 있습니다.

175B 모델은 최소한 4x8 토폴로지 슬라이스(v5e TPU 칩 32개)를 사용하는 멀티 호스트 v5e TPU 슬라이스로 제공되어야 합니다.

Cloud Storage 버킷 만들기

Saxml 관리자 서버 구성을 저장할 Cloud Storage 버킷을 만듭니다. 실행 중인 관리자 서버는 서버 상태 및 게시된 모델의 세부정보를 주기적으로 저장합니다.

Cloud Shell에서 다음을 실행합니다.

gcloud storage buckets create gs://${GSBUCKET}

GKE용 워크로드 아이덴티티 제휴를 사용하여 워크로드 액세스 구성

Kubernetes 서비스 계정을 애플리케이션에 할당하고 IAM 서비스 계정으로 작동하도록 Kubernetes 서비스 계정을 구성합니다.

클러스터와 통신하도록 kubectl을 구성합니다.

gcloud container clusters get-credentials saxml --zone=${ZONE}

애플리케이션이 사용할 Kubernetes 서비스 계정을 만듭니다.
```
kubectl create serviceaccount sax-sa --namespace default
```
애플리케이션의 IAM 서비스 계정을 만듭니다.
```
gcloud iam service-accounts create sax-iam-sa
```

IAM 서비스 계정이 Cloud Storage에서 읽고 쓸 수 있도록 IAM 정책 바인딩을 추가합니다.

gcloud projects add-iam-policy-binding ${PROJECT_ID} \
  --member "serviceAccount:sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com" \
  --role roles/storage.admin

Kubernetes 서비스 계정과 IAM 서비스 계정 사이에 IAM 정책 바인딩을 추가하여 Kubernetes 서비스 계정이 IAM 서비스 계정을 가장하도록 허용합니다. 이 바인딩을 통해 Kubernetes 서비스 계정이 IAM 서비스 계정 역할을 하므로 Kubernetes 서비스 계정이 Cloud Storage에서 읽고 쓸 수 있습니다.
```
gcloud iam service-accounts add-iam-policy-binding sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com \
  --role roles/iam.workloadIdentityUser \
  --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/sax-sa]"
```
Kubernetes 서비스 계정에 IAM 서비스 계정의 이메일 주소를 주석으로 추가합니다. 이를 통해 샘플 앱은 Google Cloud 서비스에 액세스하는 데 사용할 서비스 계정을 알 수 있습니다. 따라서 앱이 표준 Google API 클라이언트 라이브러리를 사용하여 Google Cloud 서비스에 액세스하면 해당 IAM 서비스 계정을 사용합니다.
```
kubectl annotate serviceaccount sax-sa \
  iam.gke.io/gcp-service-account=sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com
```

Saxml 배포

이 섹션에서는 Saxml 관리자 서버 및 Saxml 모델 서버를 배포합니다.

Saxml 관리자 서버 배포

다음 sax-admin-server.yaml 매니페스트를 만듭니다.

apiVersion: apps/v1
kind: Deployment
metadata:
  name: sax-admin-server
spec:
  replicas: 1
  selector:
    matchLabels:
      app: sax-admin-server
  template:
    metadata:
      labels:
        app: sax-admin-server
    spec:
      hostNetwork: false
      serviceAccountName: sax-sa
      containers:
      - name: sax-admin-server
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.1.0
        securityContext:
          privileged: true
        ports:
        - containerPort: 10000
        env:
        - name: GSBUCKET
          value: BUCKET_NAME

BUCKET_NAME을 Cloud Storage 버킷 이름으로 바꿉니다.

매니페스트를 적용합니다.
```
kubectl apply -f sax-admin-server.yaml
```

관리자 서버 포드가 작동되어 실행되는지 확인합니다.

kubectl get deployment

출력은 다음과 비슷합니다.

NAME               READY   UP-TO-DATE   AVAILABLE   AGE
sax-admin-server   1/1     1            1           52s

Saxml 모델 서버 배포

멀티 호스트 TPU 슬라이스에서 실행되는 워크로드는 동일한 TPU 슬라이스에 있는 피어를 검색하기 위해 각 포드에 대해 안정적인 네트워크 식별자가 필요합니다. 이러한 식별자를 정의하려면 IndexedJob, 헤드리스 서비스가 포함된 StatefulSet, JobSet에 속하는 모든 작업에 대해 헤드리스 서비스를 자동으로 만드는 JobSet를 사용합니다. 다음 섹션에서는 JobSet로 여러 모델 서버 포드 그룹을 관리하는 방법을 보여줍니다.

v0.2.3 이상의 JobSet을 설치합니다.
```
kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/JOBSET_VERSION/manifests.yaml
```
JOBSET_VERSION을 JobSet 버전으로 바꿉니다. 예를 들면 v0.2.3입니다.

JobSet 컨트롤러가 jobset-system 네임스페이스에서 실행 중인지 확인합니다.

kubectl get pod -n jobset-system

출력은 다음과 비슷합니다.

NAME                                        READY   STATUS    RESTARTS   AGE
jobset-controller-manager-69449d86bc-hp5r6   2/2     Running   0          2m15s

2개의 TPU 슬라이스 노드 풀에 2개의 모델 서버를 배포합니다. 다음 sax-model-server-set 매니페스트를 저장합니다.

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: sax-model-server-set
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: sax-model-server
      replicas: 2
      template:
        spec:
          parallelism: 8
          completions: 8
          backoffLimit: 0
          template:
            spec:
              serviceAccountName: sax-sa
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                cloud.google.com/gke-tpu-topology: 4x8
              containers:
              - name: sax-model-server
                image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0
                args: ["--port=10001","--sax_cell=/sax/test", "--platform_chip=tpuv5e"]
                ports:
                - containerPort: 10001
                - containerPort: 8471
                securityContext:
                  privileged: true
                env:
                - name: SAX_ROOT
                  value: "gs://BUCKET_NAME/sax-root"
                - name: MEGASCALE_NUM_SLICES
                  value: ""
                resources:
                  requests:
                    google.com/tpu: 4
                  limits:
                    google.com/tpu: 4

BUCKET_NAME을 Cloud Storage 버킷 이름으로 바꿉니다.

이 매니페스트에서 각 항목은 다음을 수행합니다.

replicas: 2는 작업 복제본의 수입니다. 각 작업은 모델 서버를 나타냅니다. 따라서 8개 포드가 포함된 그룹입니다.
parallelism: 8 및 completions: 8는 각 노드 풀에 있는 노드 수와 같습니다.
포드가 실패할 때 작업을 실패로 표시하려면 backoffLimit: 0이 0이어야 합니다.
ports.containerPort: 8471은 VM 커뮤니케이션을 위한 기본 포트입니다.
name: MEGASCALE_NUM_SLICES는 GKE가 멀티슬라이스 학습을 실행 중이 아니기 때문에 환경 변수를 설정 해제합니다.

매니페스트를 적용합니다.

kubectl apply -f sax-model-server-set.yaml

Saxml 관리 서버 및 모델 서버 포드의 상태를 확인합니다.

kubectl get pods

출력은 다음과 비슷합니다.

NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server-557c85f488-lnd5d                 1/1     Running   0          35h
sax-model-server-set-sax-model-server-0-0-nj4sm   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-1-sl8w4   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-2-hb4rk   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-3-qv67g   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-4-pzqz6   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-5-nm7mz   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-6-7br2x   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-7-4pw6z   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-0-8mlf5   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-1-h6z6w   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-2-jggtv   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-3-9v8kj   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-4-6vlb2   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-5-h689p   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-6-bgv5k   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-7-cd6gv   1/1     Running   0          24m

이 예시에는 16개의 모델 서버 컨테이너가 있습니다. sax-model-server-set-sax-model-server-0-0-nj4sm 및 sax-model-server-set-sax-model-server-1-0-8mlf5는 각 그룹에 있는 2개의 기본 모델 서버입니다.

Saxml 클러스터에는 각각 4x8 토폴로지가 사용되는 2개의 v5e TPU 슬라이스 노드 풀에 배포된 2개의 모델 서버가 있습니다.

Saxml HTTP 서버 및 부하 분산기 배포

다음과 같은 사전 빌드된 HTTP 서버 이미지를 사용합니다. 다음 sax-http.yaml 매니페스트를 저장합니다.

apiVersion: apps/v1
kind: Deployment
metadata:
  name: sax-http
spec:
  replicas: 1
  selector:
    matchLabels:
      app: sax-http
  template:
    metadata:
      labels:
        app: sax-http
    spec:
      hostNetwork: false
      serviceAccountName: sax-sa
      containers:
      - name: sax-http
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-http:v1.0.0
        ports:
        - containerPort: 8888
        env:
        - name: SAX_ROOT
          value: "gs://BUCKET_NAME/sax-root"
---
apiVersion: v1
kind: Service
metadata:
  name: sax-http-lb
spec:
  selector:
    app: sax-http
  ports:
  - protocol: TCP
    port: 8888
    targetPort: 8888
  type: LoadBalancer

BUCKET_NAME을 Cloud Storage 버킷 이름으로 바꿉니다.

sax-http.yaml 매니페스트를 적용합니다.
```
kubectl apply -f sax-http.yaml
```
참고: Saxml에 대해 빌드된 자체 HTTP 서버를 사용할 수 있습니다. 자세한 내용은 Saxml 및 HTTP 서버를 사용한 추론을 참조하세요.

HTTP 서버 컨테이너 생성이 완료될 때까지 기다립니다.

kubectl get pods

출력은 다음과 비슷합니다.

NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server-557c85f488-lnd5d                 1/1     Running   0          35h
sax-http-65d478d987-6q7zd                         1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-0-nj4sm   1/1     Running   0          24m
...

서비스에 외부 IP 주소가 할당될 때까지 기다립니다.

kubectl get svc

출력은 다음과 비슷합니다.

NAME           TYPE           CLUSTER-IP    EXTERNAL-IP   PORT(S)          AGE
sax-http-lb    LoadBalancer   10.48.11.80   10.182.0.87   8888:32674/TCP   7m36s

Saxml 사용

v5e TPU 멀티호스트 슬라이스에서 Saxml로 모델을 로드, 배포, 제공합니다.

모델 로드

Saxml에 대한 부하 분산기 IP 주소를 검색합니다.

LB_IP=$(kubectl get svc sax-http-lb -o jsonpath='{.status.loadBalancer.ingress[*].ip}')
PORT="8888"

v5e TPU 슬라이스 노드 풀 두 개에 LmCloudSpmd175B 테스트 모델을 로드합니다.

curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/publish --data \
'{
    "model": "/sax/test/spmd",
    "model_path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test",
    "checkpoint": "None",
    "replicas": 2
}'

테스트 모델에는 미세 조정된 체크포인트가 포함되지 않으며, 가중치가 무작위로 생성됩니다. 모델을 로드하는 데 최대 10분이 걸릴 수 있습니다.

출력은 다음과 비슷합니다.

{
    "model": "/sax/test/spmd",
    "path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test",
    "checkpoint": "None",
    "replicas": 2
}

모델 준비 상태를 확인합니다.

kubectl logs sax-model-server-set-sax-model-server-0-0-nj4sm

출력은 다음과 비슷합니다.

...
loading completed.
Successfully loaded model for key: /sax/test/spmd

모델이 완전히 로드되었습니다.

모델 정보를 가져옵니다.

curl --request GET \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/listcell --data \
'{
    "model": "/sax/test/spmd"
}'

출력은 다음과 비슷합니다.

{
"model": "/sax/test/spmd",
"model_path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test",
"checkpoint": "None",
"max_replicas": 2,
"active_replicas": 2
}

모델 제공

프롬프트 요청을 처리합니다.

curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/generate --data \
'{
  "model": "/sax/test/spmd",
  "query": "How many days are in a week?"
}'

출력에는 모델 응답 예시가 표시됩니다. 테스트 모델에 무작위 가중치가 포함되기 때문에 이 응답은 의미가 없을 수 있습니다.

모델 게시 취소

다음 명령어를 사용하여 모델을 게시 취소합니다.

curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/unpublish --data \
'{
    "model": "/sax/test/spmd"
}'

출력은 다음과 비슷합니다.

{
  "model": "/sax/test/spmd"
}

삭제

이 튜토리얼에서 사용된 리소스 비용이 Google Cloud 계정에 청구되지 않도록 하려면 리소스가 포함된 프로젝트를 삭제하거나 프로젝트를 유지하고 개별 리소스를 삭제하세요.

배포된 리소스 삭제

이 튜토리얼에서 만든 클러스터를 삭제합니다.
```
gcloud container clusters delete saxml --zone ${ZONE}
```

서비스 계정을 삭제합니다.

gcloud iam service-accounts delete sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com

Cloud Storage 버킷을 삭제합니다.
```
gcloud storage rm -r gs://${GSBUCKET}
```

다음 단계

Cloud TPU 시스템 아키텍처를 사용하는 현재 TPU 버전 알아보기
GKE의 TPU 자세히 알아보기

Saxml을 사용해서 GKE에서 멀티 호스트 TPU를 사용하는 LLM 서빙

배경

목표

시작하기 전에

Check for the roles

Grant the roles

환경 준비

GKE Standard 클러스터 만들기

Cloud Storage 버킷 만들기

GKE용 워크로드 아이덴티티 제휴를 사용하여 워크로드 액세스 구성

Saxml 배포

Saxml 관리자 서버 배포

Saxml 모델 서버 배포

Saxml HTTP 서버 및 부하 분산기 배포

Saxml 사용

모델 로드

모델 제공

모델 게시 취소

삭제

배포된 리소스 삭제

다음 단계