SAX를 사용하여 TPU v5e에서 Llama 2 모델 배포 및 서빙

이 튜토리얼은 다음을 안내합니다.

  • 다양한 크기(7B, 13B 또는 70B)로 사용 가능한 대규모 언어 모델(LLM)의 Llama 2 계열을 배포하는 Cloud TPU VM 생성
  • 모델의 체크포인트 준비 및 SAX에 배포
  • HTTP 엔드포인트를 통해 모델과 상호작용

AGI 실험(SAX) 서빙은 추론을 위해 Paxml, JAX, PyTorch 모델을 제공하는 실험용 시스템입니다. SAX 코드 및 문서는 Saxml Git 저장소에 있습니다. TPU v5e가 지원되는 현재 안정화 버전은 v1.1.0입니다.

SAX 셀 정보

SAX 셀(또는 클러스터)은 모델 서빙을 위한 핵심 단위입니다. 다음 두 가지 주요 부분으로 구성됩니다.

  • 관리 서버: 이 서버는 모델 서버를 추적하고, 모델 서버에 모델을 할당하고, 클라이언트가 상호작용할 올바른 모델 서버를 찾을 수 있도록 지원합니다.
  • 모델 서버: 이 서버는 모델을 실행합니다. 들어오는 요청을 처리하고 응답을 생성하는 역할을 합니다.

다음 다이어그램은 SAX 셀의 다이어그램을 보여줍니다.

관리 서버 및 모델 서버가 있는 SAX 셀

그림 1. 관리 서버 및 모델 서버가 있는 SAX 셀.

Python, C++ 또는 Go로 작성된 클라이언트를 사용하거나 HTTP 서버를 통해 직접 SAX 셀과 상호작용할 수 있습니다. 다음 다이어그램은 외부 클라이언트가 SAX 셀과 상호작용하는 방법을 보여줍니다.

SAX 셀과 상호작용하는 외부 클라이언트 및 HTTP 엔드포인트

그림 2. SAX 셀과 상호작용하는 외부 클라이언트의 런타임 아키텍처

목표

  • 서빙을 위한 TPU 리소스 설정
  • SAX 클러스터 만들기
  • Llama 2 모델 게시
  • 모델과 상호작용

비용

이 문서에서는 비용이 청구될 수 있는 다음과 같은 Google Cloud 구성요소를 사용합니다.

  • Cloud TPU
  • Compute Engine
  • Cloud Storage

프로젝트 사용량을 기준으로 예상 비용을 산출하려면 가격 계산기를 사용하세요. Google Cloud를 처음 사용하는 사용자는 무료 체험판을 사용할 수 있습니다.

시작하기 전에

Cloud TPU 환경 설정의 안내에 따라 Google Cloud 프로젝트를 설정하고 Cloud TPU API를 활성화하고 서비스 계정을 만듭니다.

TPU 만들기

다음 단계에서는 모델을 서빙할 TPU VM을 만드는 방법을 보여줍니다.

  1. 환경 변수를 만듭니다.

    export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export ZONE=ZONE
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=SERVICE_ACCOUNT
export TPU_NAME=TPU_NAME
export QUEUED_RESOURCE_ID=QUEUED_RESOURCE_ID

    환경 변수 설명

    PROJECT_ID
    Google Cloud 프로젝트의 ID입니다.
    ACCELERATOR_TYPE
    가속기 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. Llama 2 모델 크기에 따라 TPU 크기 요구사항이 다릅니다.
    • 7B: v5litepod-4 이상
    • 13B: v5litepod-8 이상
    • 70B: v5litepod-16 이상
    ZONE
    Cloud TPU를 만들려는 영역입니다.
    SERVICE_ACCOUNT
    Cloud TPU에 연결할 서비스 계정입니다.
    TPU_NAME
    Cloud TPU의 이름입니다.
    QUEUED_RESOURCE_ID
    큐에 추가된 리소스 요청의 식별자입니다.

  2. 활성 Google Cloud CLI 구성에서 프로젝트 ID 및 영역을 설정합니다.

    gcloud config set project $PROJECT_ID  && gcloud config set compute/zone $ZONE

  3. TPU VM 만들기

    gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
  --node-id ${TPU_NAME} \
  --project ${PROJECT_ID} \
  --zone ${ZONE} \
  --accelerator-type ${ACCELERATOR_TYPE} \
  --runtime-version ${RUNTIME_VERSION} \
  --service-account ${SERVICE_ACCOUNT}

  4. TPU가 활성 상태인지 확인합니다.

    gcloud compute tpus queued-resources list --project $PROJECT_ID --zone $ZONE

체크포인트 변환 노드 설정

SAX 클러스터에서 LLama 모델을 실행하려면 원래 Llama 체크포인트를 SAX 호환 형식으로 변환해야 합니다.

변환에는 모델 크기에 따라 상당한 메모리 리소스가 필요합니다.

모델 머신 유형
7B 50~60GB 메모리
13B 120GB 메모리
70B 500~600GB 메모리(N2 또는 M1 머신 유형)

7B 및 13B 모델의 경우 TPU VM에서 변환을 실행할 수 있습니다. 70B 모델의 경우 약 1TB의 디스크 공간을 갖는 Compute Engine 인스턴스를 만들어야 합니다.

gcloud compute instances create INSTANCE_NAME --project=$PROJECT_ID --zone=$ZONE \
  --machine-type=n2-highmem-128 \
  --network-interface=network-tier=PREMIUM,stack-type=IPV4_ONLY,subnet=default \
  --maintenance-policy=MIGRATE --provisioning-model=STANDARD \
  --service-account=$SERVICE_ACCOUNT \
  --scopes=https://www.googleapis.com/auth/cloud-platform \
  --tags=http-server,https-server \
  --create-disk=auto-delete=yes,boot=yes,device-name=bk-workday-dlvm,image=projects/ml-images/global/images/c0-deeplearning-common-cpu-v20230925-debian-10,mode=rw,size=500,type=projects/$PROJECT_ID/zones/$ZONE/diskTypes/pd-balanced \
  --no-shielded-secure-boot \
  --shielded-vtpm \
  --shielded-integrity-monitoring \
  --labels=goog-ec-src=vm_add-gcloud \
  --reservation-affinity=any

TPU 또는 Compute Engine 인스턴스를 변환 서버로 사용하는 경우 Llama 2 체크포인트를 변환하도록 서버를 설정합니다.

  1. 7B 및 13B 모델의 경우 서버 이름 환경 변수를 TPU 이름으로 설정합니다.

    export CONV_SERVER_NAME=$TPU_NAME

    70B 모델의 경우 서버 이름 환경 변수를 Compute Engine 인스턴스 이름으로 설정합니다.

    export CONV_SERVER_NAME=INSTANCE_NAME

  2. SSH를 사용해서 변환 노드에 연결합니다.

    변환 노드가 TPU이면 TPU에 연결합니다.

    gcloud compute tpus tpu-vm ssh $CONV_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE

    변환 노드가 Compute Engine 인스턴스이면 Compute Engine VM에 연결합니다.

    gcloud compute ssh $CONV_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE

  3. 변환 노드에 필수 패키지를 설치합니다.

    sudo apt update
sudo apt-get install python3-pip
sudo apt-get install git-all

pip3 install paxml==1.1.0
pip3 install torch
pip3 install jaxlib==0.4.14

  4. Llama 체크포인트 변환 스크립트를 다운로드합니다.

    gcloud storage cp gs://cloud-tpu-inference-public/sax-tokenizers/llama/convert_llama_ckpt.py .

Llama 2 가중치 다운로드

모델을 변환하기 전에 Llama 2 가중치를 다운로드해야 합니다. 이 튜토리얼에서는 변환된 가중치가 아닌 원래 Llama 2 가중치 (예: meta-llama/Llama-2-7b)를 사용해야 합니다. Hugging Face Transformer 형식(예: meta-llama/Llama-2-7b-hf)을 사용합니다.

이미 Llama 2 가중치가 있으면 가중치 변환으로 건너뜁니다.

Hugging Face 허브에서 가중치를 다운로드하려면 사용자 액세스 토큰을 설정하고 Llama 2 모델에 대한 액세스를 요청해야 합니다. 액세스를 요청하려면 사용하려는 모델의 Hugug Face 페이지에 있는 안내를 따르세요(예: meta-llama/Llama-2-7b).

  1. 가중치 디렉터리를 만듭니다.

    sudo mkdir WEIGHTS_DIRECTORY

  2. Hugging Face 허브에서 Llama2 가중치를 가져옵니다.

    1. Hugging Face 허브 CLI를 설치합니다.

      pip install -U "huggingface_hub[cli]"

    2. 가중치 디렉터리로 변경합니다.

      cd WEIGHTS_DIRECTORY

    3. Llama 2 파일을 다운로드합니다.

      python3
from huggingface_hub import login
login()
from huggingface_hub import hf_hub_download, snapshot_download
import os
PATH=os.getcwd()
snapshot_download(repo_id="meta-llama/LLAMA2_REPO", local_dir_use_symlinks=False, local_dir=PATH)

      LLAMA2_REPO를 다운로드하려는 Hugging Face 저장소 이름(Llama-2-7b, Llama-2-13b 또는 Llama-2-70b)으로 바꿉니다.

가중치 변환

변환 스크립트를 수정한 후 변환 스크립트를 실행하여 모델 가중치를 변환합니다.

  1. 변환된 가중치를 저장할 디렉터리를 만듭니다.

    sudo mkdir CONVERTED_WEIGHTS

  2. 읽기, 쓰기, 실행 권한이 있는 디렉터리에 Saxml GitHub 저장소를 클론합니다.

      git clone https://github.com/google/saxml.git -b r1.1.0

  3. saxml 디렉터리로 변경합니다.

    cd saxml

  4. saxml/tools/convert_llama_ckpt.py 파일을 엽니다.

  5. saxml/tools/convert_llama_ckpt.py 파일에서 169행을 다음과 같이 변경합니다.

    'scale': pytorch_vars[0]['layers.%d.attention_norm.weight' % (layer_idx)].type(torch.float16).numpy()

    변경 후:

    'scale': pytorch_vars[0]['norm.weight'].type(torch.float16).numpy()

  6. saxml/tools/init_cloud_vm.sh 스크립트를 실행합니다.

    saxml/tools/init_cloud_vm.sh

  7. 70B만 해당: 테스트 모드를 사용 중지합니다.

    1. saxml/server/pax/lm/params/lm_cloud.py 파일을 엽니다.

    2. saxml/server/pax/lm/params/lm_cloud.py 파일에서 344행을 다음과 같이 변경합니다.

      return True

      변경 후:

      return False

  8. 가중치를 변환합니다.

    python3 saxml/tools/convert_llama_ckpt.py --base-model-path WEIGHTS_DIRECTORY \
  --pax-model-path CONVERTED_WEIGHTS \
  --model-size MODEL_SIZE

    다음을 바꿉니다.

    • WEIGHTS_DIRECTORY: 원래 가중치의 디렉터리입니다.
    • CONVERTED_WEIGHTS: 변환된 가중치의 대상 경로입니다.
    • MODEL_SIZE: 7b, 13b 또는 70b

체크포인트 디렉터리 준비

체크포인트를 변환한 후에는 체크포인트 디렉터리 구조가 다음과 같습니다.

checkpoint_00000000
  metadata/
      metadata
    state/
        mdl_vars.params.lm*/
        ...
        ...
        step/

이름이 commit_success.txt인 빈 파일을 만들고 복사본을 checkpoint_00000000, metadata, state 디렉터리에 저장합니다. 이를 통해 SAX는 이 체크포인트가 완전히 변환되어 로드할 준비가 되었는지 알 수 있습니다.

  1. 체크포인트 디렉터리로 변경합니다.

    cd CONVERTED_WEIGHTS/checkpoint_00000000

  2. 이름이 commit_success.txt인 빈 파일을 만듭니다.

    touch commit_success.txt

  3. 메타데이터 디렉터리로 변경하고 commit_success.txt라는 빈 파일을 만듭니다.

    cd metadata && touch commit_success.txt

  4. 상태 디렉터리로 변경하고 commit_success.txt라는 빈 파일을 만듭니다.

    cd .. && cd state && touch commit_success.txt

이제 체크포인트 디렉터리 구조는 다음과 같습니다.

checkpoint_00000000
    commit_success.txt
metadata/
    commit_success.txt
    metadata
    state/
        commit_success.txt
        mdl_vars.params.lm*/
        ...
        ...
        step/

Cloud Storage 버킷 만들기

모델을 게시할 때 사용할 수 있도록 변환된 체크포인트를 Cloud Storage 버킷에 저장해야 합니다.

  1. Cloud Storage 버킷 이름의 환경 변수를 설정합니다.

    export GSBUCKET=BUCKET_NAME

  2. 버킷 만들기

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

  3. 변환된 체크포인트 파일을 버킷에 복사합니다.

    gcloud storage cp -r CONVERTED_WEIGHTS/checkpoint_00000000  gs://$GSBUCKET/sax_models/llama2/SAX_LLAMA2_DIR/

    SAX_LLAMA2_DIR을 적절한 값으로 바꿉니다.

    • 7B: saxml_llama27b
    • 13B: saxml_llama213b
    • 70B: saxml_llama270b

SAX 클러스터 만들기

SAX 클러스터를 만들려면 다음을 수행해야 합니다.

일반적인 배포에서는 Compute Engine 인스턴스에서 관리 서버를, TPU 또는 GPU에서 모델 서버를 실행합니다. 이 튜토리얼에서는 관리 서버와 모델 서버를 동일한 TPU v5e 인스턴스에 배포합니다.

관리 서버 만들기

관리 서버 Docker 컨테이너를 만듭니다.

  1. 변환 서버에 Docker를 설치합니다.

    sudo apt-get update
sudo apt-get install docker.io

  2. 관리 서버 Docker 컨테이너를 시작합니다.

    sudo docker run --name sax-admin-server \
  -it \
  -d \
  --rm \
  --network host \
  --env GSBUCKET=${GSBUCKET} us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.1.0

-d 옵션 없이 docker run 명령어를 실행하여 로그를 보고 관리 서버가 올바르게 시작되는지 확인할 수 있습니다.

모델 서버 만들기

다음 섹션에서는 모델 서버를 만드는 방법을 보여줍니다.

7b 모델

모델 서버 Docker 컨테이너를 시작합니다.

sudo docker run --privileged \
  -it \
  -d \
  --rm \
  --network host \
  --name "sax-model-server" \
  --env SAX_ROOT=gs://${GSBUCKET}/sax-root us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \
  --sax_cell="/sax/test" \
  --port=10001 \
  --platform_chip=tpuv5e \
  --platform_topology='4'

13b 모델

LLaMA13BFP16TPUv5e의 구성이 lm_cloud.py에 누락되었습니다. 다음 단계에서는 lm_cloud.py를 업데이트하고 새 Docker 이미지를 커밋하는 방법을 보여줍니다.

  1. 모델 서버를 시작합니다.

    sudo docker run --privileged \
  -it \
  -d \
  --rm \
  --network host \
  --name "sax-model-server" \
  --env SAX_ROOT=gs://${GSBUCKET}/sax-root \
  us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \
  --sax_cell="/sax/test" \
  --port=10001 \
  --platform_chip=tpuv5e \
  --platform_topology='8'

  2. SSH를 사용하여 Docker 컨테이너에 연결합니다.

    sudo docker exec -it sax-model-server bash

  3. Docker 이미지에 Vim을 설치합니다.

    $ apt update
$ apt install vim

  4. saxml/server/pax/lm/params/lm_cloud.py 파일을 엽니다. LLaMA13B를 찾습니다. 다음 코드가 표시됩니다.

    @servable_model_registry.register
@quantization.for_transformer(quantize_on_the_fly=False)
class LLaMA13B(BaseLLaMA):
  """13B model on a A100-40GB.

  April 12, 2023
  Latency = 5.06s with 128 decoded tokens. 38ms per output token.
  """

  NUM_LAYERS = 40
  VOCAB_SIZE = 32000
  DIMS_PER_HEAD = 128
  NUM_HEADS = 40
  MODEL_DIMS = 5120
  HIDDEN_DIMS = 13824
  ICI_MESH_SHAPE = [1, 1, 1]

  @property
  def test_mode(self) -> bool:
    return True

  5. @quantization으로 시작하는 줄을 주석 처리하거나 삭제합니다. 변경 후에는 파일이 다음과 같이 표시됩니다.

    @servable_model_registry.register
class LLaMA13B(BaseLLaMA):
  """13B model on a A100-40GB.

  April 12, 2023
  Latency = 5.06s with 128 decoded tokens. 38ms per output token.
  """

  NUM_LAYERS = 40
  VOCAB_SIZE = 32000
  DIMS_PER_HEAD = 128
  NUM_HEADS = 40
  MODEL_DIMS = 5120
  HIDDEN_DIMS = 13824
  ICI_MESH_SHAPE = [1, 1, 1]

  @property
  def test_mode(self) -> bool:
    return True

  6. TPU 구성을 지원하려면 다음 코드를 추가합니다.

    @servable_model_registry.register
class LLaMA13BFP16TPUv5e(LLaMA13B):
"""13B model on TPU v5e-8.

"""

BATCH_SIZE = [1]
BUCKET_KEYS = [128]
MAX_DECODE_STEPS = [32]
ENABLE_GENERATE_STREAM = False

ICI_MESH_SHAPE = [1, 1, 8]

@property
def test_mode(self) -> bool:
  return False

  7. Docker 컨테이너 SSH 세션을 종료합니다.

    exit

  8. 새 Docker 이미지에 대해 변경사항을 커밋합니다.

    sudo docker commit sax-model-server sax-model-server:v1.1.0-mod

  9. 새 Docker 이미지가 생성되었는지 확인합니다.

    sudo docker images

    프로젝트의 Artifact Registry에 Docker 이미지를 게시할 수 있지만 이 튜토리얼은 로컬 이미지를 사용합니다.

  10. 모델 서버를 중지합니다. 이 튜토리얼의 나머지 부분에서는 업데이트된 모델 서버가 사용됩니다.

    sudo docker stop sax-model-server

  11. 업데이트된 Docker 이미지를 사용하여 모델 서버를 시작합니다. 업데이트된 이미지 이름 sax-model-server:v1.1.0-mod를 지정해야 합니다.

    sudo docker run --privileged \
  -it \
  -d \
  --rm \
  --network host \
  --name "sax-model-server" \
  --env SAX_ROOT=gs://${GSBUCKET}/sax-root \
  sax-model-server:v1.1.0-mod \
  --sax_cell="/sax/test" \
  --port=10001 \
  --platform_chip=tpuv5e \
  --platform_topology='8'

70B 모델

SSH를 사용하여 TPU에 연결하고 모델 서버를 시작합니다.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
 --project ${PROJECT_ID} \
 --zone ${ZONE} \
 --worker=all \
 --command="
   gcloud auth configure-docker \
     us-docker.pkg.dev
   # Pull SAX model server image
   sudo docker pull us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0
   # Run model server
   sudo docker run \
     --privileged  \
     -it \
     -d \
     --rm \
     --network host \
     --name "sax-model-server"  \
     --env SAX_ROOT=gs://${GSBUCKET}/sax-root \
     us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \
     --sax_cell="/sax/test" \
     --port=10001 \
     --platform_chip=tpuv5e \
     --platform_topology='16'
"

로그 확인

모델 서버 로그를 확인하여 모델 서버가 올바르게 시작되었는지 확인합니다.

docker logs -f sax-model-server

모델 서버가 시작되지 않았다면 문제 해결 섹션에서 자세한 내용을 참조하세요.

70B 모델의 경우 각 TPU VM에 대해 다음 단계를 반복합니다.

  1. SSH를 사용하여 TPU에 연결합니다.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
  --project ${PROJECT_ID} \
  --zone ${ZONE} \
  --worker=WORKER_NUMBER

    WORKER_NUMBER는 0부터 시작하는 색인이며, 연결할 TPU VM을 나타냅니다.

  2. 로그 확인하기:

    sudo docker logs -f sax-model-server

    3개의 TPU VM이 다른 인스턴스에 연결된 것으로 표시됩니다.

    I1117 00:16:07.196594 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.3:10001
I1117 00:16:07.197484 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.87:10001
I1117 00:16:07.199437 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.13:10001

    TPU VM 중 하나에는 모델 서버가 시작됨을 보여주는 로그가 있어야 합니다.

    I1115 04:01:29.479170 139974275995200 model_service_base.py:867] Started joining SAX cell /sax/test
ERROR: logging before flag.Parse: I1115 04:01:31.479794       1 location.go:141] Calling Join due to address update
ERROR: logging before flag.Parse: I1115 04:01:31.814721       1 location.go:155] Joined 10.182.0.44:10000

모델 게시

SAX에는 SAX 모델 서버와의 상호작용을 쉽게 수행할 수 있게 해주는 saxutil이라는 명령줄 도구가 제공됩니다. 이 튜토리얼에서는 saxutil을 사용하여 모델을 게시합니다. saxutil 명령어의 전체 목록은 Saxml README 파일을 참조하세요.

  1. Saxml GitHub 저장소를 복제한 디렉터리로 변경합니다.

    cd  saxml

  2. 70B 모델의 경우 변환 서버에 연결합니다.

    gcloud compute ssh ${CONV_SERVER_NAME} \
  --project ${PROJECT_ID} \
  --zone ${ZONE}

  3. Bazel 설치

    sudo apt-get install bazel

  4. Cloud Storage 버킷으로 saxutil을 실행하기 위한 별칭을 설정합니다.

    alias saxutil='bazel run saxml/bin:saxutil -- --sax_root=gs://${GSBUCKET}/sax-root'

  5. saxutil을 사용하여 모델을 게시합니다. TPU v5litepod-8에서 약 10분 정도 걸립니다.

    saxutil --sax_root=gs://${GSBUCKET}/sax-root publish '/sax/test/MODEL' \
    saxml.server.pax.lm.params.lm_cloud.PARAMETERS \
    gs://${GSBUCKET}/sax_models/llama2/SAX_LLAMA2_DIR/checkpoint_00000000/ \
    1

    다음 변수를 바꿉니다.

    모델 크기 값
    7B MODEL: llama27b
    PARAMETERS: saxml.server.pax.lm.params.lm_cloud.LLaMA7BFP16TPUv5e
    SAX_LLAMA2_DIR: saxml_llama27b
    13B MODEL: llama213b
    PARAMETERS: saxml.server.pax.lm.params.lm_cloud.LLaMA13BFP16TPUv5e
    SAX_LLAMA2_DIR: saxml_llama213b
    70B MODEL: llama270b
    PARAMETERS: saxml.server.pax.lm.params.lm_cloud.LLaMA70BFP16TPUv5e
    SAX_LLAMA2_DIR: saxml_llama270b

테스트 배포

배포가 성공했는지 확인하려면 saxutil ls 명령어를 사용합니다.

saxutil ls /sax/test/MODEL

배포에 성공하면 0보다 큰 복제본 수를 포함해야 하며 다음과 비슷하게 표시됩니다.

INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root is /sax/test/1lama27b

+----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+
| MODEL    | MODEL PATH                                            | CHECKPOINT PATH                                                       | # OF REPLICAS | (SELECTED) REPLICAADDRESS |
+----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+
| llama27b | saxml.server.pax.lm.params.lm_cloud.LLaMA7BFP16TPUv5e | gs://${MODEL_BUCKET}/sax_models/llama2/7b/pax_7B/checkpoint_00000000/ | 1             | 10.182.0.28:10001         |
+----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+

모델 서버의 Docker 로그는 다음과 유사합니다.

I1114 17:31:03.586631 140003787142720 model_service_base.py:532] Successfully loaded model for key: /sax/test/llama27b

INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root is /sax/test/1lama27b

문제 해결

배포에 실패하면 모델 서버 로그를 확인합니다.

sudo docker logs -f sax-model-server

배포에 성공하면 다음 출력이 표시됩니다.

Successfully loaded model for key: /sax/test/llama27b

로그에 모델이 배포되었다고 표시되지 않으면 모델 구성과 모델 체크포인트의 경로를 확인합니다.

응답 생성

saxutil 도구를 사용하여 프롬프트에 대한 응답을 생성할 수 있습니다.

질문에 대한 응답을 생성합니다.

saxutil lm.generate -extra="temperature:0.2"  /sax/test/MODEL "Q: Who is Harry Potter's mother? A:"

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

INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root' lm.generate /sax/test/llama27b 'Q: Who is Harry Potter's mother? A: `
+-------------------------------+------------+
| GENERATE                      | SCORE      |
+-------------------------------+------------+
| 1. Harry Potter's mother is   | -20.214787 |
| Lily Evans. 2. Harry Potter's |            |
| mother is Petunia Evans       |            |
| (Dursley).                    |            |
+-------------------------------+------------+

클라이언트에서 모델과 상호작용

SAX 저장소에는 SAX 셀과 상호작용하는 데 사용할 수 있는 클라이언트가 포함됩니다. 클라이언트는 C++, Python, Go로 제공됩니다. 다음 예시에서는 Python 클라이언트를 빌드하는 방법을 보여줍니다.

  1. Python 클라이언트를 빌드합니다.

    bazel build saxml/client/python:sax.cc --compile_one_dependency

  2. 클라이언트를 PYTHONPATH에 추가합니다. 이 예시에서는 홈 디렉터리에 saxml이 있다고 가정합니다.

    export PYTHONPATH=${PYTHONPATH}:$HOME/saxml/bazel-bin/saxml/client/python/

  3. Python 셸에서 SAX와 상호작용합니다.

    $ python3
Python 3.10.12 (main, Jun 11 2023, 05:26:28) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import sax
>>>

HTTP 엔드포인트에서 모델과 상호작용

HTTP 엔드포인트에서 모델과 상호작용하려면 HTTP 클라이언트를 만듭니다.

  1. Compute Engine VM 만들기

    export PROJECT_ID=PROJECT_ID
export ZONE=ZONE
export HTTP_SERVER_NAME=HTTP_SERVER_NAME
export SERVICE_ACCOUNT=SERVICE_ACCOUNT
export MACHINE_TYPE=e2-standard-8
gcloud compute instances create $HTTP_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE \
  --machine-type=$MACHINE_TYPE \
  --network-interface=network-tier=PREMIUM,stack-type=IPV4_ONLY,subnet=default \
  --maintenance-policy=MIGRATE --provisioning-model=STANDARD \
  --service-account=$SERVICE_ACCOUNT \
  --scopes=https://www.googleapis.com/auth/cloud-platform \
  --tags=http-server,https-server \
  --create-disk=auto-delete=yes,boot=yes,device-name=$HTTP_SERVER_NAME,image=projects/ml-images/global/images/c0-deeplearning-common-cpu-v20230925-debian-10,mode=rw,size=500,type=projects/$PROJECT_ID/zones/$ZONE/diskTypes/pd-balanced \
  --no-shielded-secure-boot \
  --shielded-vtpm \
  --shielded-integrity-monitoring \
  --labels=goog-ec-src=vm_add-gcloud \
  --reservation-affinity=any

  2. SSH를 사용하여 Compute Engine VM에 연결합니다.

    gcloud compute ssh $HTTP_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE

  3. GKE 기반 AI GitHub 저장소를 클론합니다.

    git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git

  4. HTTP 서버 디렉터리로 변경합니다.

    cd ai-on-gke/tools/saxml-on-gke/httpserver

  5. Docker 파일을 빌드합니다.

    docker build -f Dockerfile -t sax-http .

  6. HTTP 서버 실행

    docker run -e SAX_ROOT=gs://${GSBUCKET}/sax-root -p 8888:8888 -it sax-http

다음 명령어를 사용하여 로컬 머신 또는 포트 8888에 액세스할 수 있는 다른 서버에서 엔드포인트를 테스트합니다.

  1. 서버 IP 주소 및 포트에 대한 환경 변수를 내보냅니다.

    export LB_IP=HTTP_SERVER_EXTERNAL_IP
export PORT=8888

  2. 모델 및 쿼리가 포함된 JSON 페이로드를 설정합니다.

    json_payload=$(cat  << EOF
{
  "model": "/sax/test/MODEL",
  "query": "Example query"
}
EOF
)

  3. 요청을 전송합니다.

    curl --request POST --header "Content-type: application/json" -s $LB_IP:$PORT/generate --data "$json_payload"

삭제

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

이 튜토리얼을 마쳤으면 다음 단계에 따라 리소스를 삭제하세요.

  1. Cloud TPU를 삭제합니다.

    $ gcloud compute tpus tpu-vm delete $TPU_NAME  --zone $ZONE

  2. Compute Engine 인스턴스를 만든 경우 삭제합니다.

    gcloud compute instances delete INSTANCE_NAME

  3. Cloud Storage 버킷과 콘텐츠를 삭제합니다.

    gcloud storage rm --recursive gs://BUCKET_NAME

다음 단계