JetStream 및 Pathways를 사용하여 GKE에서 멀티 호스트 TPU를 사용하여 LLM 서빙

이 가이드에서는 여러 노드에서 텐서 프로세싱 단위(TPU)를 사용하여 Google Kubernetes Engine(GKE)에서 Llama 3.1 405B와 같은 최신 대규모 언어 모델(LLM)을 서빙하는 방법을 보여줍니다.

이 가이드에서는 이동식 오픈소스 기술(Kubernetes, JetStream, Pathways on Cloud, LeaderWorkerSet(LWS) API)을 사용하여 GKE의 세밀한 제어, 확장성, 복원력, 이동성, 비용 효율성을 활용하여 GKE에서 AI/ML 워크로드를 배포하고 서빙하는 방법을 설명합니다.

배경

대규모 언어 모델의 크기가 커져 더 이상 단일 호스트 TPU 슬라이스에 수용할 수 없게 되었습니다. ML 추론의 경우 Cloud의 Pathways를 사용하여 여러 상호 연결된 TPU 노드에서 GKE에 대규모 멀티 호스트 추론을 실행할 수 있습니다. 이 가이드에서는 멀티 호스트 TPU 슬라이스로 GKE 클러스터를 프로비저닝하고 Cloud 바이너리의 Pathways를 사용하고, MaxText 프레임워크로 JetStream 서버를 실행하고, 멀티 호스트 추론 요청을 실행하는 방법을 설명합니다.

JetStream, MaxText, Pathways와 함께 GKE에서 TPU를 사용하여 LLM을 서빙하면 비용 효율성, 확장성 및 더 높은 가용성을 비롯하여 관리형 Kubernetes의 모든 이점을 갖춘 강력하고 프로덕션에 즉시 사용 가능한 서빙 솔루션을 빌드할 수 있습니다. 이 섹션에서는 이 튜토리얼에서 사용되는 주요 기술을 설명합니다.

TPU 정보

TPU는 Google이 머신러닝 및 AI 모델을 가속화하는 데 사용하기 위해 TensorFlow, PyTorch 및 JAX와 같은 프레임워크를 활용하여 커스텀 개발한 ASIC(application-specific integrated circuits)입니다.

GKE에서 TPU를 사용하기 전에 다음 학습 과정을 완료하는 것이 좋습니다.

  1. Cloud TPU 시스템 아키텍처를 사용하는 현재 TPU 버전 가용성 알아보기
  2. GKE의 TPU 알아보기

이 튜토리얼에서는 Llama 3.1-405B 모델 서빙을 설명합니다. GKE는 지연 시간이 짧은 프롬프트를 서빙하기 위한 모델 요구사항에 따라 TPU 토폴로지가 구성된 멀티 호스트 TPU v6e 노드에 모델을 배포합니다.

Cloud의 개발자 과정

Pathways는 가속기의 대규모 조정 레이어입니다. Pathways는 현재 모델의 최신 성능을 유지하면서 새로운 시스템과 ML 연구 아이디어를 탐색할 수 있도록 명시적으로 설계되었습니다. Pathways는 단일 JAX 클라이언트 프로세스가 하나 이상의 대규모 TPU 슬라이스에서 계산을 조정할 수 있도록 지원하여 수백 또는 수천 개의 TPU 칩에 걸쳐 있는 ML 계산을 간소화합니다.

JetStream

JetStream은 Google이 개발한 오픈소스 추론 서빙 프레임워크입니다. JetStream은 TPU 및 GPU에서 고성능, 높은 처리량, 메모리 최적화된 추론을 가능하게 합니다. JetStream은 연속 일괄 처리, KV 캐시 최적화, 양자화 기법을 비롯한 고급 성능 최적화를 제공하여 LLM 배포를 용이하게 합니다. JetStream은 성능을 최적화하기 위해 PyTorch/XLA 및 JAX TPU 서빙을 지원합니다.

MaxText

MaxText는 Flax, Orbax 및 Optax와 같은 오픈소스 JAX 라이브러리를 기반으로 구축된 성능이 뛰어나고 확장 가능하며 적응력이 뛰어난 JAX LLM 구현입니다. MaxText의 디코더 전용 LLM 구현은 Python으로 작성됩니다. XLA 컴파일러를 많이 활용하여 커스텀 커널을 빌드할 필요 없이 고성능을 실현합니다.

MaxText에서 지원하는 최신 모델 및 매개변수 크기에 관한 자세한 내용은 MaxText 프로젝트 저장소를 참고하세요.

Llama 3.1 405B

Llama 3.1 405B는 텍스트 생성, 번역, 질의 응답을 포함한 다양한 자연어 처리 태스크를 위해 설계된 Meta의 대규모 언어 모델입니다. GKE는 이러한 규모의 모델의 분산 학습 및 서빙 니즈를 지원하는 데 필요한 인프라를 제공합니다.

자세한 내용은 Llama 문서를 참고하세요.

아키텍처

이 섹션에서는 이 튜토리얼에서 사용되는 GKE 아키텍처를 설명합니다. 이 아키텍처는 TPU를 프로비저닝하고 모델을 배포하고 서빙하기 위해 JetStream 및 Pathways 구성요소를 호스팅하는 GKE Standard 클러스터를 포함합니다.

다음은 이 아키텍처의 구성요소를 보여주는 다이어그램입니다.

JetStream 및 Pathways 구성요소가 포함된 멀티 호스트 TPU 노드 풀이 있는 GKE 클러스터의 아키텍처

이 아키텍처에는 다음 구성요소가 포함됩니다.

  • GKE Standard 리전 클러스터
  • JetStream 배포 및 개발자 과정 구성요소를 호스팅하는 멀티 호스트 TPU 슬라이스 노드 풀입니다.
  • Pathways resource manager는 가속기 리소스를 관리하고 사용자 작업에 대한 가속기 할당을 조정합니다.
  • Pathways client는 Pathways resource manager와 협력하여 컴파일된 프로그램이 실행될 위치를 결정합니다.
  • Pathways worker는 가속기 머신에서 계산을 실행하고 IFRT 프록시 서버를 통해 데이터를 워크로드로 다시 전송합니다.
  • IFRT proxy client는 OSS 임시 프레임워크 런타임 (IFRT) API를 구현하고 워크로드와 개발자 과정 구성요소 간의 통신 브리지 역할을 합니다.
  • IFRT proxy server는 IFRT proxy client에서 요청을 수신하고 Pathways client에 전달하여 작업을 배포합니다.
  • JetStream-Pathways 컨테이너는 추론 요청을 수신하고 실행 프로세스를 Pathways workers에 위임하는 JAX 기반 추론 서버를 제공합니다.
  • 서비스 구성요소는 수신 트래픽을 모든 JetStream HTTP 복제본으로 전파합니다.
  • JetStream HTTP는 요청을 JetStream의 필수 형식에 관한 래퍼로 수락하고 JetStream의 GRPC 클라이언트로 전송하는 HTTP 서버입니다.

시작하기 전에

  • 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, roles/resourcemanager.projectIamAdmin

    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 column 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. 저장을 클릭합니다.
  • TPU v6e PodSlice Lite 칩 16개에 대해 충분한 할당량이 있는지 확인합니다. 이 튜토리얼에서는 온디맨드 인스턴스를 사용합니다.
  • Google Cloud 프로젝트가 개발자 과정의 허용 목록에 포함되어 있는지 확인합니다.

모델 액세스 권한 얻기

GKE에 배포하기 위한 Meta Llama 3.1-405B 체크포인트에 액세스하려면 다음 단계를 따르세요.

  1. 라이선스 동의 계약에 서명합니다.
  2. Meta Llama 다운로드 페이지에 액세스합니다.
  3. 모델 이용약관을 검토하고 동의한 후 모델을 다운로드하는 데 필요한 URL을 가져옵니다.
  4. 모델 체크포인트를 다운로드하려면 적절한 모델의 모델 ID를 찾습니다. 지원되는 모델 및 ID 목록은 llama CLI 문서를 참고하세요. 예를 들어 Llama 3.1-405B 모델의 경우 Llama 3.1-405B-Instruct:bf16-mp16을 사용합니다.

환경 준비

이 튜토리얼에서는 Cloud Shell을 사용하여Google Cloud에서 호스팅되는 리소스를 관리합니다. Cloud Shell에는 kubectl 및 gcloud CLI 등 이 튜토리얼에 필요한 소프트웨어가 사전 설치되어 있습니다.

Cloud Shell로 환경을 설정하려면 다음 단계를 따르세요.

  1. Google Cloud 콘솔의 경우 Google Cloud 콘솔에서 Cloud Shell 활성화 아이콘 Cloud Shell 활성화를 클릭하여 Cloud Shell 세션을 시작합니다. 그러면 Google Cloud 콘솔 하단 창에서 세션이 실행됩니다.

  2. 기본 환경 변수를 설정합니다.

    gcloud config set project PROJECT_ID
export PROJECT_ID=$(gcloud config get project)
export CLUSTER_NAME=CLUSTER_NAME
export BUCKET_NAME=BUCKET_NAME
export REGION=REGION
export LOCATION=LOCATION
export CLUSTER_VERSION=CLUSTER_VERSION
export MACHINE_TYPE=ct6e-standard-4t
export TPU_TYPE=v6e
export TOPOLOGY=4x4
export WORKERS_PER_SLICE=4

    다음 값을 바꿉니다.

    • PROJECT_ID: Google Cloud 프로젝트 ID
    • CLUSTER_NAME: GKE 클러스터의 이름
    • BUCKET_NAME: Cloud Storage 버킷 이름. gs:// 프리픽스를 지정할 필요가 없습니다.
    • REGION: GKE 클러스터, Cloud Storage 버킷, TPU 노드가 위치한 리전. 리전에는 TPU v6e 머신 유형을 사용할 수 있는 영역이 포함됩니다(예: us-east1, us-east5, europe-west4, asia-northeast1, 또는 us-south1).
    • LOCATION: TPU 리소스를 사용할 수 있는 영역(예: us-east1-d)
    • CLUSTER_VERSION: 사용하려는 머신 유형을 지원해야 하는 GKE 버전. 기본 GKE 버전은 대상 TPU에 대한 가용성이 없을 수 있습니다. TPU 머신 유형별로 사용 가능한 최소 GKE 버전 목록은 GKE의 TPU 가용성을 참조하세요.
    • MACHINE_TYPE: v6e 머신 유형
    • TPU_TYPE: 노드 풀 이름 지정 시 사용되는 접두사(v6e)
    • TOPOLOGY: TPU v6e 토폴로지
    • WORKERS_PER_SLICE: 노드 풀 또는 TPU 슬라이스당 노드 수

Google Cloud 리소스 만들기 및 구성

필요한 리소스를 만들려면 다음 안내를 따르세요.

GKE 클러스터 만들기

  1. 지역 GKE Standard 클러스터를 만듭니다.

    gcloud container clusters create CLUSTER_NAME \
    --project=PROJECT_ID \
    --cluster-version=CLUSTER_VERSION \
    --region=REGION \
    --scopes=cloud-platform \
    --machine-type=n2-standard-32

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

    CLUSTER_VERSION을 적절한 클러스터 버전으로 바꿉니다.

  2. 4x4 토폴로지와 노드 4개가 있는 TPU v6e 노드 풀을 하나 만듭니다.

    gcloud container node-pools create multihost-np \
--project=PROJECT_ID \
--region=REGION \
--node-locations=LOCATION \
--cluster=CLUSTER_NAME \
--machine-type=MACHINE_TYPE \
--num-nodes=WORKERS_PER_SLICE \
--tpu-topology=TOPOLOGY \
--scopes cloud-platform \
--placement-type=COMPACT \
--workload-metadata=GCE_METADATA

스토리지 객체 액세스를 위한 서비스 계정 구성

Kubernetes 서비스 계정이 IAM 서비스 계정으로 작동하도록 구성합니다.

  1. 애플리케이션의 IAM 서비스 계정을 만듭니다.

    gcloud iam service-accounts create jetstream-pathways

  2. IAM 서비스 계정이 Cloud Storage를 관리할 수 있도록 IAM 정책 바인딩을 추가합니다. 이렇게 하면 IAM 서비스 계정이 체크포인트가 저장될 스토리지 버킷에 액세스할 수 있습니다.

    gcloud projects add-iam-policy-binding ${PROJECT} \
  --member "serviceAccount:jetstream-pathways@${PROJECT}.iam.gserviceaccount.com" \
  --role roles/storage.objectUser

gcloud projects add-iam-policy-binding ${PROJECT} \
  --member "serviceAccount:jetstream-pathways@${PROJECT}.iam.gserviceaccount.com" \
  --role roles/storage.insightsCollectorService

  3. Kubernetes 서비스 계정에 IAM 서비스 계정의 이메일 주소를 주석으로 추가합니다. 

    kubectl annotate serviceaccount default \
iam.gke.io/gcp-service-account=jetstream-pathways@${PROJECT}.iam.gserviceaccount.com

Artifact Registry에 인증하도록 Docker 구성

허용 목록에 추가된 개발자 과정 이미지를 가져올 수 있도록 Artifact Registry에 인증하도록 Docker를 구성합니다.

gcloud auth login
gcloud auth configure-docker

체크포인트 변환

Meta Llama 3.1-405B 체크포인트를 MaxText 호환 int8 추론 체크포인트로 변환하려면 Llama3.1-405B를 사용한 체크포인트 변환의 단계를 완료합니다. 배포에서 load_parameters_path 플래그가 있는 체크포인트를 사용합니다.

Pathways 임시 파일을 저장할 Cloud Storage 버킷 만들기

컴파일 캐시와 같은 Pathways 임시 파일을 저장할 Cloud Storage 버킷을 만듭니다.

export PATHWAYS_BUCKET=PATHWAYS_BUCKET
gcloud storage buckets create gs://$PATHWAYS_BUCKET

JetStream-MaxText 및 Pathways 배포

JetStream-MaxText 및 Pathways 모델 서버를 배포합니다.

GKE 클러스터에 연결

gcloud container clusters get-credentials "${CLUSTER}" --project "${PROJECT}" --location "${ZONE}"

LeaderWorkerSet(LWS) API 배포

LWS는 특히 리더-워커 아키텍처가 있는 스테이트풀 분산 애플리케이션을 배포하고 관리하도록 설계된 맞춤 리소스입니다. 대규모 모델이 여러 노드의 여러 기기에서 샤딩되고 서빙되는 AI/ML 워크로드에 특히 적합합니다.

VERSION=v0.6.1
kubectl apply --server-side -f https://github.com/kubernetes-sigs/lws/releases/download/$VERSION/manifests.yaml

LeaderWorkerSet 컨트롤러를 완전히 사용할 수 있을 때까지 기다립니다.

kubectl wait deploy/lws-controller-manager -n lws-system --for=condition=available --timeout=5m

출력은 다음과 비슷하게 표시됩니다.

deployment.apps/lws-controller-manager condition met

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

kubectl get pod -n lws-system

출력은 다음과 비슷하게 표시됩니다.

NAME                          READY   STATUS    RESTARTS    AGE
lws-controller-manager-abcd   1/1     Running   0           40s
lws-controller-manager-efgh   1/1     Running   0           40s

워크로드 매니페스트 배포

  1. 다음 매니페스트를 jetstream-pathways-llama-3-1-405b-4x4.yaml로 저장합니다.

    apiVersion: leaderworkerset.x-k8s.io/v1
kind: LeaderWorkerSet
metadata:
  name: jetstream-pathways
  annotations:
    leaderworkerset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  replicas: 1
  leaderWorkerTemplate:
    leaderTemplate:
      metadata:
        labels:
          app: jetstream-pathways
      spec:
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 4x4
        tolerations:
        - key: "google.com/tpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: pathways-proxy
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
          args:
          imagePullPolicy: Always
          ports:
          - containerPort: 38681
        - name: pathways-rm
          env:
          - name: HOST_ADDRESS
            value: "$(LWS_LEADER_ADDRESS)"
          - name: TPU_SKIP_MDS_QUERY
            value: "true"
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
          args:
          - --server_port=38677
          - --gcs_scratch_location=PATHWAYS_BUCKET
          - --node_type=resource_manager
          - --instance_count=1
          - --instance_type=tpuv6e:4x4
          imagePullPolicy: Always
          ports:
          - containerPort: 38677
        - name: jax-tpu
          image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
          env:
          - name: LOG_LEVEL
            value: "INFO"
          args:
          - MaxText/configs/v5e/inference/llama3_405b_v5e-64.yml
          - model_name=llama3.1-405b
          - load_parameters_path=CHECKPOINT_PATH
          - max_prefill_predict_length=1024
          - max_target_length=2048
          - async_checkpointing=false
          - steps=1
          - ici_fsdp_parallelism=1
          - ici_autoregressive_parallelism=2
          - ici_tensor_parallelism=8
          - scan_layers=false
          - weight_dtype=bfloat16
          - per_device_batch_size=6
          - enable_single_controller=true
          - quantization=int8
          - quantize_kvcache=true
          - checkpoint_is_quantized=true
          - enable_model_warmup=true
          imagePullPolicy: Always
          ports:
          - containerPort: 9000
          startupProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 1
            initialDelaySeconds: 600
            failureThreshold: 10000
          livenessProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 60
            failureThreshold: 10
          readinessProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 60
            failureThreshold: 10
        - name: jetstream-http
          image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
          imagePullPolicy: Always
          ports:
          - containerPort: 8000
    size: 5
    workerTemplate:
      spec:
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 4x4
        tolerations:
        - key: "google.com/tpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: worker
          args:
          - --server_port=38679
          - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
          - --gcs_scratch_location=PATHWAYS_BUCKET
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
          imagePullPolicy: Always
          ports:
          - containerPort: 38679
          resources:
            limits:
              google.com/tpu: "4"
--- 
apiVersion: v1
kind: Service
metadata:
  name: jetstream-svc
spec:
  selector:
    app: jetstream-pathways
  ports:
  - protocol: TCP
    name: jetstream-http
    port: 8000
    targetPort: 8000

  2. load_parameters_path 필드의 값을 체크포인트 변환 프로세스에서 생성된 체크포인트 경로로 설정합니다.

    • bf16 체크포인트의 경우 경로는 gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items와 유사해야 합니다.
    • int8 체크포인트의 경우 gs://OUTPUT_BUCKET_DIRECTORY/int8와 유사해야 합니다.

    gcs_scratch_location 필드의 값을 이전에 만든 Pathways 버킷으로 설정합니다.

    perl -pi -e 's|CHECKPOINT_PATH|gs://OUTPUT_BUCKET_DIRECTORY/int8|g' jetstream-pathways-llama-3-1-405b-4x4.yaml
perl -pi -e 's|PATHWAYS_BUCKET|gs://PATHWAYS_BUCKET|g' jetstream-pathways-llama-3-1-405b-4x4.yaml

배포 매니페스트 적용

매니페스트를 적용하여 서버를 배포합니다.

kubectl apply -f jetstream-pathways-llama-3-1-405b-4x4.yaml

모델 서버가 시작됩니다.

모델 서버 시작 확인

405B 모델의 경우 체크포인트를 복원하는 데 10~20분 정도 걸릴 수 있습니다. enable_model_warmup 플래그를 사용 설정한 경우 모델 워밍업 중에 추가로 기다릴 수도 있습니다.

kubectl logs -f jetstream-pathways-0 -c jax-tpu

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

2025-03-02 02:15:07,682 - JetstreamLogger - INFO - Initializing the driver with 1 prefill engines and 1 generate engines in interleaved mode
2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up prefill thread 0.
2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up transfer thread 0.
2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up generate thread 0.
2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up detokenize thread 0.
2025-03-02 02:15:07,685 - JetstreamLogger - INFO - Driver initialized.
...
...
...
INFO:     Started server process [7]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:9999 (Press CTRL+C to quit)

Llama 3.1-405b 서빙

Llama 3.1-405b 모델을 서빙하려면 포트 전달을 설정합니다.

kubectl port-forward svc/jetstream-svc 8000:8000

포트 전달을 사용하면 클러스터 외부에서 서비스에 액세스할 수 있습니다. GKE의 ClusterIP 서비스를 통해 JetStream-Pathways 배포에 액세스할 수 있습니다. ClusterIP 서비스는 클러스터 내에서만 연결할 수 있습니다.

모델과 상호작용

새 터미널에서 다음 명령어를 실행합니다.

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8000/generate \
--data \
'{
    "prompt": "What are the top 5 programming languages",
    "max_tokens": 200
}'

모델 준비로 인해 초기 요청이 완료되는 데 몇 초 정도 걸릴 수 있습니다. 출력은 다음과 비슷하게 표시됩니다.

{
    "response": " for web development?\nThe top 5 programming languages for web development are:\n1. **JavaScript**: JavaScript is the most popular language for web development, used by over 90% of websites for client-side scripting. It's also popular for server-side programming with technologies like Node.js.\n2. **HTML/CSS**: HTML (Hypertext Markup Language) and CSS (Cascading Style Sheets) are not programming languages, but are essential for building websites. HTML is used for structuring content, while CSS is used for styling and layout.\n3. **Python**: Python is a popular language for web development, especially with frameworks like Django and Flask. It's known for its simplicity, flexibility, and large community of developers.\n4. **Java**: Java is a popular language for building enterprise-level web applications, especially with frameworks like Spring and Hibernate. It's known for its platform independence, strong security features, and large community of developers.\n5. **PHP**: PHP is a mature language for web"
}

지금까지 다음 작업을 완료했습니다.

  1. TPU를 사용하여 GKE에 MaxText 및 Pathways가 포함된 JetStream 모델 서버를 배포했습니다.
  2. gs://BUCKET_NAME에 Llama 3.1-405B int8 체크포인트를 만들었습니다.
  3. 모델을 서빙하고 상호작용했습니다.

분할 서빙

분할 서빙은 미리 채우기 및 디코딩 단계를 여러 호스트로 분할하면서 LLM을 서빙하는 기법입니다. 이 접근 방식은 리소스 사용률을 최적화하여 처리량과 지연 시간을 개선할 수 있습니다.

  • 미리 채우기: 키-값 캐시를 초기화하기 위해 입력 프롬프트의 전방 패스입니다.

  • Decode: 출력 토큰을 점진적으로 생성하는 절차로, 단계당 토큰 1개, 반복당 KV 캐시 값 1개를 생성합니다.

  1. 기본 환경 변수를 설정합니다.

    export NODE_POOL_NAME=dis-v6e-8
export NODE_POOL_SIZE=2
export MACHINE_TYPE=ct6e-standard-4t
export TOPOLOGY=2x4
export WORKERS_PER_SLICE=2

  2. v6e-8 노드를 사용하는 노드 풀 2개를 만듭니다.

    for i in $(seq 1 NODE_POOL_SIZE); do
  gcloud container node-pools create NODE_POOL_NAME-${i}-np \
  --project=PROJECT \
  --zone=ZONE \
  --cluster=CLUSTER_NAME \
  --machine-type=MACHINE_TYPE \
  --num-nodes=WORKERS_PER_SLICE \
  --tpu-topology=TOPOLOGY \
  --scopes=cloud-platform \
  --workload-metadata=GCE_METADATA
done

체크포인트 변환

Meta Llama 2-70B 체크포인트를 MaxText 호환 int8 추론 체크포인트로 변환하려면 Llama2-70B를 사용한 체크포인트 변환의 단계를 완료합니다. Meta 이용약관에 동의할 때 Llama2-70B를 모델로 선택합니다. 배포에서 load_parameters_path 플래그가 있는 체크포인트를 사용합니다.

checkpoint-job.yaml 파일에서 다음 매개변수를 바꿉니다.

- --meta_url=META_URL
- --model_name=llama-2
- --model_path=Llama-2-70b-chat
- --output_directory=gs://BUCKET_NAME/maxtext/llama-2-70b

체크포인트는 load_parameters_path 플래그와 함께 배포에 사용됩니다.

분할 서빙을 사용하여 JetStream 개발자 과정 배포

  1. 다음 매니페스트를 jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml로 저장합니다.

    apiVersion: leaderworkerset.x-k8s.io/v1
kind: LeaderWorkerSet
metadata:
  name: jetstream-pathways
  annotations:
    leaderworkerset.sigs.k8s.io/subgroup-exclusive-topology: cloud.google.com/gke-nodepool
spec:
  replicas: 1
  leaderWorkerTemplate:
    subGroupPolicy:
      subGroupSize: 2
    leaderTemplate:
      metadata:
        labels:
          app: jetstream-pathways
      spec:
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 2x4
        tolerations:
        - key: "google.com/tpu"
          operator: "Exists"
          effect: "NoSchedule"
        containers:
        - name: pathways-proxy
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
          args:
          - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
          - --server_port=38681
          - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
          - --xla_jf_auto_cross_replica_sharding=false
          - --xla_tpu_enable_windowed_einsum_for_reduce_scatter=false
          - --xla_tpu_enable_windowed_einsum_for_all_gather=false
          - --xla_tpu_prefer_latch_optimized_rhs_layouts=true
          - --xla_tpu_enable_experimental_fusion_cost_model=false
          - --xla_tpu_dot_dot_fusion_duplicated=false
          - --xla_tpu_dot_dot_fusion=true
          - --xla_jf_conv_input_fusion=true
          - --xla_jf_conv_output_fusion=true
          - --xla_tpu_rwb_fusion=false
          - --xla_tpu_copy_fusion_pad_unpad_ratio=0
          - --xla_tpu_licm_size_inflation_ratio=1
          - --xla_tpu_copy_elision_analysis_allowance=150000
          - --xla_tpu_copy_insertion_use_region_analysis_limit=10000
          - --xla_tpu_order_dot_after_layout=true
          - --xla_jf_rematerialization_percent_shared_memory_limit=100
          - --xla_tpu_use_repeated_instance_for_preferred_prefetch_time=true
          - --xla_tpu_enforce_prefetch_fifo_order=false
          - --xla_tpu_prefetch_interval_picker_size_override=6000000
          - --xla_tpu_async_copy_bandwidth_scaling_factor=1
          - --xla_tpu_nd_short_transfer_max_chunks=-1
          - --xla_tpu_enable_aggressive_broadcast_priority_update=true
          - --xla_tpu_alternate_memory_benefit_scaling_factor_for_large_buffers=SQRT
          - --xla_tpu_memory_bound_loop_optimizer_options=enabled:true
          - --xla_tpu_enable_copy_fusion=true
          - --xla_tpu_enable_cross_program_prefetch_freeing=false
          - --xla_tpu_enable_dot_strength_reduction=true
          - --xla_tpu_layout_use_dot_grouping=false
          - --xla_tpu_msa_inefficient_use_to_copy_ratio=0.5
          - --xla_tpu_reduce_loop_fusion_dup_with_unfusable_user=false
          - --xla_tpu_vector_load_fusion_window=1024
          - --xla_tpu_vector_store_fusion_window=256
          - --xla_jf_conv_reshape_fusion=false
          - --xla_tpu_input_conv_multi_users=false
          - --xla_tpu_enable_multi_level_input_dot_dot_fusion=false
          - --xla_tpu_enable_multi_level_output_dot_dot_fusion=false
          - --xla_tpu_dot_dot_fusion_separable_convs_only=false
          - --xla_tpu_enable_multi_level_nested_loop_fusion=true
          - --xla_tpu_nested_dot_fusion=true
          - --xla_tpu_enable_multi_level_nested_dot_fusion=false
          - --xla_jf_enable_multi_output_fusion=true
          - --xla_tpu_use_lp_llo_scheduler_for_dot_dot_fusions=false
          - --xla_tpu_enable_flash_attention=true
          imagePullPolicy: Always
          ports:
          - containerPort: 38681
        - name: pathways-rm
          env:       
          - name: HOST_ADDRESS
            value: "$(LWS_LEADER_ADDRESS)"
          - name: TPU_SKIP_MDS_QUERY
            value: "true"
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
          args:
          - --server_port=38677
          - --gcs_scratch_location=PATHWAYS_BUCKET
          - --node_type=resource_manager
          - --instance_count=2
          - --instance_type=tpuv6e:2x4
          imagePullPolicy: Always
          ports:
          - containerPort: 38677
        - name: jax-tpu
          image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
          args:
          - MaxText/configs/base.yml
          - tokenizer_path=assets/tokenizer.llama2
          - load_parameters_path=CHECKPOINT_PATH
          - max_prefill_predict_length=1024
          - max_target_length=2048
          - model_name=llama2-70b
          - ici_fsdp_parallelism=1
          - ici_autoregressive_parallelism=1
          - ici_tensor_parallelism=-1
          - scan_layers=false
          - weight_dtype=bfloat16
          - per_device_batch_size=27
          - checkpoint_is_quantized=true 
          - quantization=int8
          - quantize_kvcache=true
          - compute_axis_order=0,2,1,3
          - ar_cache_axis_order=0,2,1,3
          - stack_prefill_result_cache=True
          - inference_server=ExperimentalMaxtextDisaggregatedServer_8
          - inference_benchmark_test=True
          - enable_model_warmup=True
          env:
          - name: LOG_LEVEL
            value: "INFO"
          imagePullPolicy: Always
          securityContext:
            capabilities:
              add: ["SYS_PTRACE", "NET_ADMIN", "SYS_TIME"]
          ports: 
          - containerPort: 9000
          startupProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 1
            initialDelaySeconds: 240
            failureThreshold: 10000
          livenessProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 60
            failureThreshold: 100
          readinessProbe:
            httpGet:
              path: /healthcheck
              port: 8000
              scheme: HTTP
            periodSeconds: 60
            failureThreshold: 100
        - name: jetstream-http
          image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
          imagePullPolicy: Always
          ports:
          - containerPort: 8000
    size: 5
    workerTemplate:
      spec:
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 2x4
        containers:
        - name: worker
          args:
          - --server_port=38679
          - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
          - --gcs_scratch_location=PATHWAYS_BUCKET
          image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
          imagePullPolicy: Always
          ports:
          - containerPort: 38679
          resources:
            limits:
              google.com/tpu: "4"
--- 
apiVersion: v1
kind: Service
metadata:
  name: jetstream-svc
spec:
  selector:
    app: jetstream-pathways
  ports:
  - protocol: TCP
    name: jetstream-http
    port: 8000
    targetPort: 8000

  2. load_parameters_path 필드의 값을 체크포인트 변환 프로세스에서 생성된 체크포인트 경로로 설정합니다.

    • bf16 체크포인트의 경우 경로는 gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items와 유사해야 합니다.
    • int8 체크포인트의 경우 gs://OUTPUT_BUCKET_DIRECTORY/int8와 유사해야 합니다.

    gcs_scratch_location 필드의 값을 이전에 만든 Pathways 버킷으로 설정합니다.

    perl -pi -e 's|CHECKPOINT_PATH|BUCKET_NAME/maxtext/llama-2-70b/int8|g' jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
perl -pi -e 's|PATHWAYS_BUCKET|gs://PATHWAYS_BUCKET|g' jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml

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

    kubectl apply -f jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml

    체크포인트 크기에 따라 모델 서버에서 체크포인트를 복원하는 데 다소 시간이 걸릴 수 있습니다. 70B 모델의 경우 모델 준비를 포함하여 체크포인트를 복원하는 데 약 8분이 걸릴 수 있습니다. 로그를 계속 관찰하여 모델 서버 시작을 확인하여 준비 상태를 확인하고 모델과 상호작용할 수 있도록 포트 전달을 설정하여 모델을 서빙할 수 있습니다.

지금까지 다음 작업을 완료했습니다.

  1. TPU 및 분할 서빙을 사용하여 GKE에 MaxText 및 Pathways가 포함된 JetStream 모델 서버를 배포했습니다.
  2. gs://BUCKET_NAME에 Llama 2-70B int8 체크포인트를 만들었습니다.
  3. 모델을 서빙하고 상호작용했습니다.

문제해결

  • Empty reply from server 메시지가 표시되면 컨테이너에서 모델 데이터 다운로드가 완료되지 않았을 수 있습니다. 모델을 서빙할 준비가 되었음을 나타내는 Connected 메시지가 있는지 포드의 로그를 다시 확인합니다.
  • Connection refused 메시지가 표시되면 포트 전달이 활성 상태인지 확인합니다.

삭제

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

배포된 리소스 삭제

이 가이드에서 만든 리소스의 비용이 Google Cloud 계정에 청구되지 않도록 하려면 다음 명령어를 실행하고 프롬프트를 따르세요.

gcloud container clusters delete CLUSTER_NAME --region=REGION

gcloud iam service-accounts delete jetstream-pathways@PROJECT_ID.iam.gserviceaccount.com

gcloud storage rm --recursive gs://BUCKET_NAME

다음 단계