Saxml을 통해 Vertex AI Prediction에서 TPU를 사용하여 Gemma 개방형 모델 제공

이 가이드에서는 Saxml을 사용해서 Vertex AI Prediction에서 Tensor Processing Unit(TPU)을 사용하여 Gemma 개방형 모델 대규모 언어 모델(LLM)을 제공하는 방법을 보여줍니다. 이 가이드에서는 2B 및 7B 매개변수 안내를 통해 조정된 Gemma 모델을 Cloud Storage에 다운로드하고 TPU에서 Saxml을 실행하는 Vertex AI Prediction에 배포합니다.

배경

Saxml을 통해 Vertex AI Prediction에서 TPU를 사용하여 Gemma를 제공합니다. 하위 수준의 인프라를 관리하고 LLM을 제공하는 비용 효율적인 방법을 제공하는 관리형 AI 솔루션을 활용할 수 있습니다. 이 섹션에서는 이 튜토리얼에서 사용되는 주요 기술을 설명합니다.

Gemma

Gemma는 오픈 라이선스로 출시된 공개적으로 사용 가능한 가벼운 생성형 인공지능(AI) 모델 집합입니다. 이러한 AI 모델은 애플리케이션, 하드웨어, 휴대기기 또는 호스팅된 서비스에서 실행할 수 있습니다. 텍스트 생성에 Gemma 모델을 사용할 수 있지만 특수한 태스크를 위해 이러한 모델을 조정할 수도 있습니다.

자세한 내용은 Gemma 문서를 참조하세요.

Saxml

Saxml은 추론을 위해 Paxml, JAX, PyTorch 모델을 제공하는 실험용 시스템입니다. 이 튜토리얼에서는 Saxml에 더 비용 효율적인 TPU에 Gemma를 제공하는 방법을 설명합니다. GPU 설정도 비슷합니다. Saxml은 이 튜토리얼에서 사용할 Vertex AI Prediction용 컨테이너를 빌드하는 스크립트를 제공합니다.

TPU

TPU는 Google에서 커스텀 개발한 ASIC(Application-Specific Integrated Circuits)로서 TensorFlow, PyTorch, JAX와 같은 데이터 처리 프레임워크를 가속화하는 데 사용됩니다.

이 튜토리얼은 Gemma 2B 및 Gemma 7B 모델을 사용합니다. Vertex AI Prediction은 다음 단일 호스트 TPU v5e 노드 풀에서 이러한 모델을 호스팅합니다.

Gemma 2B: 하나의 TPU 칩을 나타내는 1x1 토폴로지가 있는 TPU v5e 노드 풀에서 호스팅됩니다. 노드의 머신 유형은 ct5lp-hightpu-1t입니다.
Gemma 7B: 4개의 TPU 칩을 나타내는 2x2 토폴로지가 있는 TPU v5e 노드 풀에서 호스팅됩니다. 노드의 머신 유형은 ct5lp-hightpu-4t입니다.

시작하기 전에

Google Cloud 계정에 로그인합니다. Google Cloud를 처음 사용하는 경우 계정을 만들고 Google 제품의 실제 성능을 평가해 보세요. 신규 고객에게는 워크로드를 실행, 테스트, 배포하는 데 사용할 수 있는 $300의 무료 크레딧이 제공됩니다.

Google Cloud Console의 프로젝트 선택기 페이지에서 Google Cloud 프로젝트를 선택하거나 만듭니다.

프로젝트 선택기로 이동

Google Cloud 프로젝트에 결제가 사용 설정되어 있는지 확인합니다.

API Vertex AI API and Artifact Registry API 사용 설정

API 사용 설정

Google Cloud Console의 프로젝트 선택기 페이지에서 Google Cloud 프로젝트를 선택하거나 만듭니다.

프로젝트 선택기로 이동

Google Cloud 프로젝트에 결제가 사용 설정되어 있는지 확인합니다.

API Vertex AI API and Artifact Registry API 사용 설정

API 사용 설정

Google Cloud 콘솔에서 Cloud Shell을 활성화합니다.

Cloud Shell 활성화

Google Cloud 콘솔 하단에서 Cloud Shell 세션이 시작되고 명령줄 프롬프트가 표시됩니다. Cloud Shell은 Google Cloud CLI가 사전 설치된 셸 환경으로, 현재 프로젝트의 값이 이미 설정되어 있습니다. 세션이 초기화되는 데 몇 초 정도 걸릴 수 있습니다.

이 튜토리얼에서는 Google Cloud와 상호작용하기 위해 Cloud Shell을 사용한다고 가정합니다. Cloud Shell 대신 다른 셸을 사용하려면 다음 추가 구성을 수행하세요.

Google Cloud CLI를 설치합니다.
gcloud CLI를 초기화하려면 다음 명령어를 실행합니다.
```
gcloud init
```

Artifact Registry 문서에 따라 Docker를 설치합니다.
Vertex AI Prediction용 TPU v5e 칩 5개에 충분한 할당량이 있는지 확인합니다.
Kaggle 계정이 아직 없다면 만듭니다.

모델 액세스 권한 얻기

Vertex AI Prediction에 배포하기 위해 Gemma 모델에 액세스하려면 Kaggle 플랫폼에 로그인하고 라이선스 동의 계약에 서명하고 Kaggle API 토큰을 가져와야 합니다. 이 튜토리얼에서는 Kaggle 사용자 인증 정보에 Kubernetes 보안 비밀을 사용합니다.

Gemma를 사용하려면 동의 계약에 서명해야 합니다. 다음 안내를 따르세요.

Kaggle.com의 모델 동의 페이지에 액세스합니다.
아직 로그인하지 않았다면 Kaggle에 로그인합니다.
액세스 요청을 클릭합니다.
동의를 위한 계정 선택 섹션에서 Kaggle 계정을 통해 인증을 선택하여 동의를 위해 Kaggle 계정을 사용합니다.
모델 이용약관에 동의합니다.

액세스 토큰 생성

Kaggle을 통해 모델에 액세스하려면 Kaggle API 토큰이 필요합니다.

아직 토큰이 없으면 다음 단계에 따라 새 토큰을 생성합니다.

브라우저에서 Kaggle 설정으로 이동합니다.
API 섹션에서 새 토큰 만들기를 클릭합니다.

kaggle.json이라는 파일이 다운로드됩니다.

Cloud Shell에 액세스 토큰 업로드

Cloud Shell에서 Kaggle API 토큰을 Google Cloud 프로젝트에 업로드할 수 있습니다.

Cloud Shell에서 더보기 > 업로드를 클릭합니다.
파일을 선택하고 파일 선택을 클릭합니다.
kaggle.json 파일을 엽니다.
업로드를 클릭합니다.

Cloud Storage 버킷 만들기

모델 체크포인트를 저장할 Cloud Storage 버킷을 만듭니다.

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

gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME

CHECKPOINTS_BUCKET_NAME을 모델 체크포인트를 저장하는 Cloud Storage 버킷의 이름으로 바꿉니다.

Cloud Storage 버킷에 모델 복사

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

pip install kaggle --break-system-packages

# For Gemma 2B
mkdir -p /data/gemma_2b-it
kaggle models instances versions download google/gemma/pax/2b-it/1 --untar -p /data/gemma_2b-it
gsutil -m cp -R /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/

# For Gemma 7B
mkdir -p /data/gemma_7b-it
kaggle models instances versions download google/gemma/pax/7b-it/1 --untar -p /data/gemma_7b-it
gsutil -m cp -R /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/

Artifact Registry 저장소 만들기

다음 섹션에서 만들 컨테이너 이미지를 저장할 Artifact Registry 저장소를 만듭니다.

프로젝트에 Artifact Registry API 서비스를 사용 설정하세요.

gcloud services enable artifactregistry.googleapis.com

셸에서 다음 명령어를 실행하여 Artifact Registry 저장소를 만듭니다.

gcloud artifacts repositories create saxml \
 --repository-format=docker \
 --location=LOCATION \
 --description="Saxml Docker repository"

LOCATION을 Artifact Registry의 컨테이너 이미지 저장 리전으로 바꿉니다. 나중에 이 리전과 일치하는 리전 엔드포인트에서 Vertex AI 모델 리소스를 만들어야 하므로, TPU의 경우 us-west1과 같이 Vertex AI에 리전 엔드포인트가 있는 리전을 선택합니다.

컨테이너 이미지를 Artifact Registry에 푸시

사전 빌드된 Saxml 컨테이너는 us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest에서 사용할 수 있습니다. 이를 Artifact Registry에 복사합니다. Artifact Registry에 액세스하도록 Docker를 구성합니다. 그런 후 컨테이너 이미지를 Artifact Registry 저장소에 푸시합니다.

선택한 리전의 Artifact Registry로 푸시할 수 있도록 로컬 Docker에 설치 권한을 제공하려면 셸에서 다음 명령어를 실행합니다.
```
gcloud auth configure-docker LOCATION-docker.pkg.dev
```
- LOCATION을 저장소를 만든 리전으로 바꿉니다.
Artifact Registry에 바로 전에 빌드한 컨테이너 이미지를 복사하려면 셸에서 다음 명령어를 실행합니다.
```
docker tag us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest
```
Artifact Registry에 바로 전에 빌드한 컨테이너 이미지를 푸시하려면 셸에서 다음 명령어를 실행합니다.
```
docker push LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest
```
이전 섹션에서 한 것처럼 다음을 바꿉니다.
- LOCATION: Artifact Registry 저장소의 리전입니다.
- PROJECT_ID: Google Cloud 프로젝트의 ID

모델 배포

모델 업로드

Saxml 컨테이너를 사용하는 Model 리소스를 업로드하려면 다음 gcloud ai models upload 명령어를 실행합니다.

Gemma 2B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma2BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

Gemma 7B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma7BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

다음을 바꿉니다.

PROJECT_ID: Google Cloud 프로젝트의 ID
LOCATION_ID: Vertex AI를 사용하는 리전 TPU는 us-west1에서만 사용 가능합니다.
DEPLOYED_MODEL_NAME: DeployedModel의 이름. DeployedModel의 Model 표시 이름도 사용할 수 있습니다.

엔드포인트 만들기

온라인 예측 제공을 위해 모델을 사용하려면 먼저 모델을 엔드포인트에 배포해야 합니다. 기존 엔드포인트에 모델을 배포하는 경우 이 단계를 건너뛸 수 있습니다. 다음 예시에서는 gcloud ai endpoints create 명령어를 사용합니다.

gcloud ai endpoints create \
  --region=LOCATION \
  --display-name=ENDPOINT_NAME

다음을 바꿉니다.

LOCATION_ID: Vertex AI를 사용하는 리전
ENDPOINT_NAME: 엔드포인트의 표시 이름

Google Cloud CLI 도구가 엔드포인트를 만드는 데 몇 초 정도 걸릴 수 있습니다.

엔드포인트에 모델 배포

엔드포인트가 준비되면 엔드포인트에 모델을 배포합니다.

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

MODEL_ID=$(gcloud ai models list \
   --region=LOCATION \
   --filter=display_name=DEPLOYED_MODEL_NAME \
   --format="value(name)")

gcloud ai endpoints deploy-model $ENDPOINT_ID \
  --region=LOCATION \
  --model=$MODEL_ID \
  --display-name=DEPLOYED_MODEL_NAME \
  --machine-type=ct5lp-hightpu-4t \
  --traffic-split=0=100

다음을 바꿉니다.

LOCATION_ID: Vertex AI를 사용하는 리전
ENDPOINT_NAME: 엔드포인트의 표시 이름
DEPLOYED_MODEL_NAME: DeployedModel의 이름. DeployedModel의 Model 표시 이름도 사용할 수 있습니다.

Gemma 2B는 더 작은 ct5lp-hightpu-1t 머신에 배포할 수 있습니다. 이 경우 모델을 업로드할 때 --platform_topology=1x1을 지정해야 합니다.

Google Cloud CLI 도구로 모델을 엔드포인트에 배포하려면 몇 분 정도 걸릴 수 있습니다. 모델이 성공적으로 배포되면 이 명령어가 다음 출력을 표시합니다.

  Deployed a model to the endpoint xxxxx. Id of the deployed model: xxxxx.

배포된 모델에서 온라인 예측 가져오기

Vertex AI Prediction 엔드포인트를 통해 모델을 호출하려면 표준 추론 요청 JSON 객체 를 사용하여 예측 요청 형식을 지정합니다.

다음 예시에서는 gcloud ai endpoints predict 명령어를 사용합니다.

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

gcloud ai endpoints predict $ENDPOINT_ID \
  --region=LOCATION \
  --http-headers=Content-Type=application/json \
  --json-request instances.json

다음을 바꿉니다.

LOCATION_ID: Vertex AI를 사용하는 리전
ENDPOINT_NAME: 엔드포인트의 표시 이름
instances.json의 형식은 {"instances": [{"text_batch": "<your prompt>"},{...}]}입니다.

삭제

Vertex AI 요금 및 Artifact Registry 비용이 계속 부과되지 않도록 하려면 이 튜토리얼에서 만든 Google Cloud 리소스를 삭제하세요.

엔드포인트에서 모델을 배포 해제하고 엔드포인트를 삭제하려면 셸에서 다음 명령어를 실행합니다.

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

DEPLOYED_MODEL_ID=$(gcloud ai endpoints describe $ENDPOINT_ID \
   --region=LOCATION \
   --format="value(deployedModels.id)")

gcloud ai endpoints undeploy-model $ENDPOINT_ID \
  --region=LOCATION \
  --deployed-model-id=$DEPLOYED_MODEL_ID

gcloud ai endpoints delete $ENDPOINT_ID \
   --region=LOCATION \
   --quiet

LOCATION을 이전 섹션에서 모델을 만든 리전으로 바꿉니다.

모델을 삭제하려면 셸에서 다음 명령어를 실행합니다.

MODEL_ID=$(gcloud ai models list \
   --region=LOCATION \
   --filter=display_name=DEPLOYED_MODEL_NAME \
   --format="value(name)")

gcloud ai models delete $MODEL_ID \
   --region=LOCATION \
   --quiet

LOCATION을 이전 섹션에서 모델을 만든 리전으로 바꿉니다.

Artifact Registry 저장소 및 그 안의 컨테이너 이미지를 삭제하려면 셸에서 다음 명령어를 실행합니다.
```
gcloud artifacts repositories delete saxml \
  --location=LOCATION \
  --quiet
```
LOCATION을 이전 섹션에서 Artifact Registry 저장소를 만든 리전으로 바꿉니다.

제한사항

Vertex AI Prediction에서 Cloud TPU는 us-west1에서만 지원됩니다. 자세한 내용은 위치를 참조하세요.

다음 단계

Llama2 및 GPT-J와 같은 다른 Saxml 모델을 배포하는 방법을 알아봅니다.

Saxml을 통해 Vertex AI Prediction에서 TPU를 사용하여 Gemma 개방형 모델 제공

배경

Gemma

Saxml

TPU

시작하기 전에

모델 액세스 권한 얻기

라이선스 동의 계약 서명

액세스 토큰 생성

Cloud Shell에 액세스 토큰 업로드

Cloud Storage 버킷 만들기

Cloud Storage 버킷에 모델 복사

Artifact Registry 저장소 만들기

컨테이너 이미지를 Artifact Registry에 푸시

모델 배포

모델 업로드

Gemma 2B-it

Gemma 7B-it

엔드포인트 만들기

엔드포인트에 모델 배포

배포된 모델에서 온라인 예측 가져오기

삭제

제한사항

다음 단계