v6e TPU VM에서 JetStream을 사용한 PyTorch 추론

이 튜토리얼에서는 JetStream을 사용하여 TPU v6e에서 PyTorch 모델을 서빙하는 방법을 설명합니다. JetStream은 XLA 기기(TPU)에서 대규모 언어 모델(LLM) 추론을 위한 처리량 및 메모리 최적화 엔진입니다. 이 튜토리얼에서는 Llama2-7B 모델에 대해 추론 벤치마크를 실행합니다.

시작하기 전에

TPU v6e(4칩 구성) 프로비저닝을 준비하세요.

  1. Cloud TPU 환경 설정 가이드를 따라 Google Cloud 프로젝트 설정, Google Cloud CLI 구성, Cloud TPU API 사용 설정, Cloud TPU 사용 권한 확보 작작업을 수행합니다.

  2. Google Cloud 에 인증을 수행하고 Google Cloud CLI의 기본 프로젝트 및 영역을 구성합니다.

    gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE

용량 확보

TPU 용량을 확보할 준비가 되면, Cloud TPU 할당량에 대한 자세한 내용을 확인하세요. 추가로 용량 확보에 대해 궁금한 점이 있으면 Cloud TPU 영업팀 또는 계정팀에 문의하세요.

Cloud TPU 환경 프로비저닝

TPU VM은 GKE, GKE와 XPK, 큐에 추가된 리소스를 사용하여 프로비저닝할 수 있습니다.

기본 요건

  • Google Cloud 프로젝트 내에서 액세스할 수 있는 최대 칩 수를 지정하는 TPUS_PER_TPU_FAMILY 할당량이 충분한지 확인합니다.
  • 프로젝트의 TPU 할당량이 충분한지 확인합니다.
    • TPU VM 할당량
    • IP 주소 할당량
    • Hyperdisk Balanced 할당량
  • 사용자 프로젝트 권한

환경 변수 만들기

Cloud Shell에서 다음 환경 변수를 만듭니다.

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-central2-b
export ACCELERATOR_TYPE=v6e-4
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id

환경 변수 설명

변수 설명
PROJECT_ID Google Cloud 프로젝트 ID입니다. 기존 프로젝트를 사용하거나 새 프로젝트를 만듭니다.
TPU_NAME TPU의 이름입니다.
ZONE TPU VM을 만들 영역입니다. 지원되는 영역에 대한 자세한 내용은 TPU 리전 및 영역을 참조하세요.
ACCELERATOR_TYPE 가속기 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. 각 TPU 버전에서 지원되는 가속기 유형에 대한 자세한 내용은 TPU 버전을 참조하세요.
RUNTIME_VERSION Cloud TPU 소프트웨어 버전입니다.
SERVICE_ACCOUNT 서비스 계정의 이메일 주소입니다. Google Cloud 콘솔의 서비스 계정 페이지로 이동하여 찾을 수 있습니다.

예를 들면 tpu-service-account@PROJECT_ID.iam.gserviceaccount.com입니다.
QUEUED_RESOURCE_ID 큐에 추가된 리소스 요청의 사용자 할당 텍스트 ID입니다.

TPU v6e 프로비저닝

    gcloud alpha 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}

list 또는 describe 명령어를 사용하여 큐에 추가된 리소스의 상태를 쿼리합니다.

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project ${PROJECT_ID} --zone ${ZONE}

큐에 추가된 리소스 요청 상태의 전체 목록을 보려면 큐에 추가된 리소스 문서를 참조하세요.

SSH를 사용하여 TPU에 연결

  gcloud compute tpus tpu-vm ssh ${TPU_NAME}

JetStream PyTorch Llama2-7B 벤치마크 실행

JetStream-PyTorch를 설정하고, 모델 체크포인트를 변환하고, 추론 벤치마크를 실행하려면 GitHub 저장소의 안내를 따르세요.

추론 벤치마크가 완료되면 TPU 리소스를 삭제해야 합니다.

삭제

TPU를 삭제합니다.

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
      --project ${PROJECT_ID} \
      --zone ${ZONE} \
      --force \
      --async