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

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

배경

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

Gemma

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

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

Saxml

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

TPU

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

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

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

시작하기 전에

1. 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.

2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

   Go to project selector

3. Make sure that billing is enabled for your Google Cloud project.

4. Enable the Vertex AI API.

   Enable the API

5. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

   Go to project selector

6. Make sure that billing is enabled for your Google Cloud project.

7. Enable the Vertex AI API.

   Enable the API

8. In the Google Cloud console, activate Cloud Shell.

   Activate Cloud Shell

   At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

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

1. Install the Google Cloud CLI.

2. If you're using an external identity provider (IdP), you must first sign in to the gcloud CLI with your federated identity.

3. To initialize the gcloud CLI, run the following command:

   gcloud init

4. Vertex AI용 TPU v5e 칩에 충분한 할당량이 있는지 확인합니다. 기본적으로 이 할당량은 0입니다. 1x1 토폴로지의 경우 1이어야 합니다. 2x2의 경우 4여야 합니다. 두 토폴로지를 모두 실행하려면 5여야 합니다.
5. Kaggle 계정이 없는 경우 만듭니다.

모델 액세스 권한 얻기

Cloud Shell에는 모델 가중치를 다운로드하기에 충분한 리소스가 없을 수 있습니다. 이 경우 해당 작업을 실행할 Vertex AI Workbench 인스턴스를 만들 수 있습니다.

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

라이선스 동의 계약 서명

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

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

액세스 토큰 생성

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

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

1. 브라우저에서 Kaggle 설정으로 이동합니다.

2. API 섹션에서 새 토큰 만들기를 클릭합니다.

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

Cloud Shell에 액세스 토큰 업로드

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

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

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
gcloud storage cp /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/ --recursive

# 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
gcloud storage cp /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/ --recursive

모델 배포

모델 업로드

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

Gemma 2B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-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=us-docker.pkg.dev/vertex-ai/prediction/sax-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 엔드포인트를 통해 모델을 호출하려면 표준 추론 요청 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 리소스를 삭제합니다.

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

   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을 이전 섹션에서 모델을 만든 리전으로 바꿉니다.

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

   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을 이전 섹션에서 모델을 만든 리전으로 바꿉니다.

제한사항

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

다음 단계

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