v6e TPU VM에서 JetStream MaxText 추론

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

시작하기 전에

칩 4개가 있는 TPU v6e를 프로비저닝할 준비를 합니다.

  1. Cloud TPU 환경 설정 가이드를 따라 Cloud TPU를 사용할 수 있는 적절한 액세스 권한이 있는지 확인합니다.

  2. TPU VM의 서비스 ID를 만듭니다.

    gcloud alpha compute tpus tpu-vm service-identity create --zone=zone
  3. TPU 서비스 계정을 만들고 Google Cloud 서비스에 대한 액세스 권한을 부여합니다.

    서비스 계정은 Google Cloud TPU 서비스가 다른 Google Cloud서비스에 액세스하도록 허용합니다. 사용자 관리형 서비스 계정이 권장됩니다. Google Cloud 콘솔 또는 gcloud 명령어를 통해 서비스 계정을 만들 수 있습니다.

    gcloud 명령줄 도구를 사용하여 서비스 계정을 만듭니다.

    gcloud iam service-accounts create your-service-account-name \
    --description="your-sa-description" \
    --display-name="your-sa-display-name"
    export SERVICE_ACCOUNT_NAME=your-service-account-name

    Google Cloud 콘솔에서 서비스 계정을 만듭니다.

    1. Google Cloud Console의 서비스 계정 페이지로 이동합니다.
    2. 서비스 계정 만들기를 클릭합니다.
    3. 서비스 계정 이름을 입력합니다.
    4. (선택사항) 서비스 계정에 대한 설명을 입력합니다.
    5. 만들기를 클릭하고 계속 진행합니다.
    6. 서비스 계정에 부여할 역할을 선택합니다.
    7. 계속을 클릭합니다.
    8. (선택사항) 서비스 계정을 관리할 수 있는 사용자 또는 그룹을 지정합니다.
    9. 완료를 클릭하여 서비스 계정 만들기를 마칩니다.

    서비스 계정을 만든 후 다음 단계에 따라 서비스 계정 역할을 부여합니다.

    다음 역할이 필요합니다.

    • TPU 관리자: TPU를 만드는 데 필요합니다.
    • 스토리지 관리자: Cloud Storage에 액세스하는 데 필요합니다.
    • 로그 작성자
    • 모니터링 측정항목 작성자: Cloud Monitoring에 측정항목을 쓰는 데 필요합니다.

    사용자에게 IAM 역할을 할당하려면 관리자가 나에게 roles/resourcemanager.projectIamAdmin 역할을 부여해야 합니다. 프로젝트 IAM 관리자 roles/resourcemanager.projectIamAdmin 역할이 있는 사용자도 이 역할을 부여할 수 있습니다.

    다음 gcloud 명령어를 사용하여 서비스 계정 역할을 추가합니다.

    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
       --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
       --role roles/tpu.admin
    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
       --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
       --role roles/storage.admin
    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
       --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
       --role roles/logging.logWriter
    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
       --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
       --role roles/monitoring.metricWriter

    Google Cloud 콘솔을 사용하여 역할을 할당할 수도 있습니다.

    Google Cloud 콘솔에서 다음 역할을 선택합니다.

    1. 서비스 계정을 선택하고 주 구성원 추가를 클릭합니다.
    2. 새 주 구성원 필드에 서비스 계정의 이메일 주소를 입력합니다.
    3. 역할 선택 드롭다운에서 역할을 검색하여 (예: 스토리지 관리자) 선택합니다.
    4. 저장을 클릭합니다.
  4. Google Cloud 로 인증하고 Google Cloud CLI의 기본 프로젝트 및 ZONE을 구성합니다.

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

용량 확보

TPU 용량을 확보할 준비가 되면 할당량 페이지에서 Cloud Quotas 시스템에 관해 알아보세요. 용량 확보에 관해 추가 질문이 있는 경우 Cloud TPU 영업팀 또는 계정팀에 문의하세요.

Cloud TPU 환경 프로비저닝

TPU VM은 GKE, GKE 및 XPK 또는 대기열에 추가된 리소스로 프로비저닝할 수 있습니다.

기본 요건

  • 이 튜토리얼은 Python 3.10 이상으로 테스트되었습니다.
  • 프로젝트에Google Cloud 프로젝트 내에서 액세스할 수 있는 최대 칩 수를 지정하는 TPUS_PER_TPU_FAMILY 할당량이 충분한지 확인합니다.
  • 프로젝트에 다음에 충분한 TPU 할당량이 있는지 확인합니다.
    • TPU VM 할당량
    • IP 주소 할당량
    • 하이퍼디스크 균형 할당량
  • 사용자 프로젝트 권한

환경 변수 만들기

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

export NODE_ID=TPU_NODE_ID # TPU name
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=v6e-4
export ZONE=us-central2-b
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=YOUR_SERVICE_ACCOUNT
export QUEUED_RESOURCE_ID=QUEUED_RESOURCE_ID
export VALID_DURATION=VALID_DURATION

# Additional environment variable needed for Multislice:
export NUM_SLICES=NUM_SLICES

# Use a custom network for better performance as well as to avoid having the
# default network becoming overloaded.
export NETWORK_NAME=${PROJECT_ID}-mtu9k
export NETWORK_FW_NAME=${NETWORK_NAME}-fw

명령어 플래그 설명

변수 설명
NODE_ID 큐에 추가된 리소스 요청이 할당될 때 생성되는 TPU의 사용자 할당 ID입니다.
PROJECT_ID Google Cloud 프로젝트 이름 기존 프로젝트를 사용하거나 새 프로젝트를 만듭니다.
ZONE 지원되는 영역에 대해서는 TPU 리전 및 영역 문서를 참조하세요.
ACCELERATOR_TYPE 지원되는 모든 가속기 유형은 가속기 유형 문서를 참고하세요.
RUNTIME_VERSION v2-alpha-tpuv6e
SERVICE_ACCOUNT Google Cloud 콘솔 -> IAM -> 서비스 계정에서 찾을 수 있는 서비스 계정의 이메일 주소입니다.

예: tpu-service-account@<your_project_ID>.iam.gserviceaccount.com.com

NUM_SLICES 생성할 슬라이스 수입니다 (멀티슬라이스에만 필요).
QUEUED_RESOURCE_ID 큐에 추가된 리소스 요청의 사용자 할당 텍스트 ID입니다.
VALID_DURATION 대기열에 추가된 리소스 요청이 유효한 기간입니다.
NETWORK_NAME 사용할 보조 네트워크의 이름입니다.
NETWORK_FW_NAME 사용할 보조 네트워크 방화벽의 이름입니다.

TPU v6e 프로비저닝

gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
    --node-id TPU_NAME \
    --project PROJECT_ID \
    --zone ZONE \
    --accelerator-type v6e-4 \
    --runtime-version v2-alpha-tpuv6e \
    --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

TPU에 연결되면 추론 벤치마크를 실행할 수 있습니다.

Llama2-7B 추론 벤치마크 실행

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

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

삭제

TPU를 삭제합니다.

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