sungmin.me

Vertex AI를 이용한 BERT 모델 배포

머신러닝 엔지니어는 확장성과 프로덕션 사용을 위해 클라우드 플랫폼에 커스텀 모델을 배포해야 할 필요성을 자주 접한다. 하지만 Vertex AI에서 사전 빌드된 모델을 사용하는 가이드는 많지만, 완전히 커스텀으로 훈련된 모델을 배포하는 자료는 다소 부족하다.

이것이 Google Cloud Vertex AI에서 커스텀 BERT 기반 텍스트 분류 모델을 훈련, 미세 조정, 배포하는 방법을 보여주는 완전한 엔드투엔드 Colab 노트북을 만들게 된 동기다. 이 가이드는 환경 설정부터 배포된 엔드포인트를 통해 실시간 예측을 실행하는 것까지 모든 필요한 단계를 안내한다.

이 튜토리얼을 마치면 다음에 대해 깊이 이해하게 될 것이다.

  • 훈련된 모델을 Google Cloud Vertex AI에 배포하기:TorchServe와 MAR 파일을 이용해 커스텀 모델 패키징부터 엔드포인트 배포까지의 전 과정.
  • 실시간 예측을 위한 추론 엔드포인트 설정하기:gcloud CLI와 Python SDK를 활용한 엔드포인트 프로비저닝 및 예측 요청 테스트.
  • 텍스트 분류를 위한 PyTorch 기반 BERT 모델 훈련하기:Transformers 라이브러리를 활용한 모델 미세 조정Fine-tuning 및 데이터셋 전처리.
  • ML 엔지니어링 기초 (훈련, 평가, 아카이빙):실험 추적을 넘어선 모델 아티팩트 관리 및 배포 준비를 위한 엔지니어링 실무.

이제 시작해보자.

Vertex AI Colab 코드:

Google Colab

Preview unavailable.

iconhttps://colab.research.google.com/drive/1Q4o4KUjaAb_a3dHHi-KJngUcWsa6wFqv?usp=sharing
preview

잠깐, Vertex AI란 무엇인가?

Vertex AI는 ML 모델을 대규모로 훈련, 배포, 관리하기 위한 엔드투엔드 솔루션을 제공하는 Google Cloud의 통합 머신러닝 플랫폼이다. AutoML, 커스텀 모델 훈련, MLOps와 같은 도구를 단일 환경으로 통합하여 ML 애플리케이션을 더 쉽게 개발, 배포, 모니터링할 수 있게 해준다.

Google Cloud Vertex AI 플랫폼
그림 1: Google Cloud Vertex AI 플랫폼(Source: Google Cloud Blog)

Vertex AI 서비스 유형

Vertex AI는 다양한 ML 요구에 맞는 여러 서비스를 제공한다.

  • Vertex AI Training: 모델 개발을 위한 커스텀 훈련 작업 및 AutoML[1]을 지원한다.
  • Vertex AI Prediction: 훈련된 모델을 온라인 또는 배치 추론을 위한 엔드포인트로 배포할 수 있게 한다.
  • Vertex AI Pipelines: 엔드투엔드 ML 모델 라이프사이클 관리를 위한 워크플로우 자동화를 제공한다.
  • Vertex AI Feature Store: ML 피처를 저장, 제공 및 공유하기 위한 중앙 저장소이다.
  • Vertex AI Workbench: 모델 실험을 위한 관리형 Jupyter 기반 환경이다.
  • Vertex AI Model Monitoring: 프로덕션에서 모델 성능을 추적하고 드리프트를 감지하는 도구이다.

이 글에서는 특히 Vertex AI Prediction에 초점을 맞춘다. 이를 통해 훈련된 모델을 온라인 또는 배치 추론을 위한 엔드포인트로 배포하여 관리형 클라우드 기반 API를 통해 실시간 예측을 제공할 수 있다.

1단계: 환경 설정

모델을 훈련하고 배포하기 전에 Google Cloud 인증을 설정해야 한다. 아직 Google Cloud 서비스 계정을 만들지 않았다면 다음 단계를 따른다.

  1. Google Cloud 콘솔을 연다.
  2. IAM 및 관리자 > 서비스 계정으로 이동한다.
  3. + 서비스 계정 만들기 버튼을 클릭한다.
  4. 계정 이름을 입력하여 서비스 계정을 만들고 다음 역할을 부여한다.

⚠ 참고: 이 글에서는 예제를 위해 단순화된 형태로 관리자 권한을 부여했다. 그러나 회사 클라우드 계정이나 실제 환경에서는 범위를 적절히 제한하여 세분화된 접근 방식으로 권한을 부여해야 한다.

  • Storage 관리자: Cloud Storage 버킷 및 파일 관리.
  • AI Platform 관리자: AI 모델 및 훈련에 대한 전체 액세스.
  • 서비스 계정 사용자: 서비스 계정으로 활동할 수 있도록 설정.
  • Vertex AI 관리자: Vertex AI 서비스에 대한 전체 제어.
Vertex AI를 위한 Google Cloud IAM 설정
그림 2: 서비스 계정에 필요한 Vertex AI 관련 IAM 권한 설정 화면
  1. 새로 만든 서비스 계정으로 이동하여 키 탭으로 이동한다.
  2. 키 추가 > 새 키 만들기 > JSON을 클릭한다.
  3. JSON 키 파일을 다운로드하고 credentials.json으로 저장한다.
서비스 계정 생성 및 credentials.json 다운로드
그림 3: Google Cloud 콘솔에서 서비스 계정 키를 JSON으로 내려받는 UI

위 화면은 Google Cloud 콘솔에서 서비스 계정 키를 JSON 파일로 내려받는 과정을 보여준다.

이제 Colab이나 로컬 환경에서 credentials.json을 사용할 수 있다.

2단계: 데이터셋 및 모델 준비

AG News 데이터셋을 사용하여 BERT 기반 분류기를 훈련한다.

  • AG News 데이터셋: 주제(예: 세계, 스포츠, 비즈니스, 과학/기술)별로 레이블이 지정된 짧은 텍스트 샘플을 포함한다.
  • BERT: transformers 라이브러리의 사전 훈련된 언어 모델[2]이다.
prepare_dataset.py
# 데이터셋 로드
dataset = load_dataset("ag_news")
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

# 모델 초기화
model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=4)
optimizer = AdamW(model.parameters(), lr=5e-5)

토큰화 및 포맷팅 후 DataLoaders를 생성한다.

preprocess.py
class DatasetItem(TypedDict):
    text: str
    label: str

def preprocess_data(dataset_item: DatasetItem) -> dict[str, torch.Tensor]:
    return tokenizer(dataset_item["text"], truncation=True, padding="max_length", return_tensors="pt")

train_dataset = dataset["train"].select(range(1200)).map(preprocess_data, batched=True)
test_dataset = dataset["test"].select(range(800)).map(preprocess_data, batched=True)

train_dataset.set_format("torch", columns=["input_ids", "attention_mask", "label"])
test_dataset.set_format("torch", columns=["input_ids", "attention_mask", "label"])

train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=8, shuffle=False)

데이터 준비가 완료되면 3 에포크 동안 모델을 훈련한다. T4 GPU를 사용하면 약 5분 안에 훈련을 완료할 수 있다.

train_loop.py
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

num_epochs = 3
losses: list[float] = []

for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for batch in tqdm(train_loader, desc=f"Epoch {epoch + 1}"):
        inputs = {key: batch[key].to(device) for key in batch}
        labels = inputs.pop("label")
        outputs = model(**inputs, labels=labels)
        loss = outputs.loss
        total_loss += loss.item()
        losses.append(loss.item())

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    average_loss = total_loss / len(train_loader)
    print(f"Epoch {epoch + 1}, Average Loss: {average_loss}")
훈련 손실
그림 4: 세대(^Epoch)별 스텝 손실 추이를 시각화한 그래프

위 그래프는 세 번의 에포크 동안 스텝별 손실 추이를 시각화한 것이다.

간단한 NLP 분류 작업을 위해 BERT 모델을 훈련했다. 짧은 훈련 시간에도 불구하고 손실이 수렴하는 것을 볼 수 있다.

분류 작업의 성능을 더 명확하게 평가하려면 Confusion Matrix를 사용하는 것이 가장 좋다.

Confusion Matrix
그림 5: Confusion Matrix heatmap (클래스별 오차 분포)

이 Confusion Matrix에서는 각 클래스별 예측 결과를 한눈에 확인할 수 있다.

위 결과를 보면 레이블 3(과학/기술)이 종종 레이블 2(비즈니스)와 혼동되는 것을 볼 수 있다. 이는 기술 주제가 비즈니스 뉴스에 자주 등장하는 결과로 직관적으로 해석할 수 있다.

3단계: Vertex AI에 모델 배포하기

훈련이 완료되면 모델을 Google Cloud Storage에 업로드하고 Vertex AI에 배포한다.

TorchServe 및 MAR 파일

TorchServe[3]는 PyTorch 모델을 프로덕션 환경에서 서빙하기 위해 특별히 설계된 오픈소스 프레임워크이다. 이를 통해 모델을 가중치, 구성 및 사용자 정의 코드를 포함하는 단일 MARModel Archive[4] 파일로 패키징할 수 있다.

훈련 후, Torch Model Archiver CLI 도구를 사용하여 MAR 파일을 생성한다.

package_model.sh
torch-model-archiver \
  --model-name bert_model \
  --version 1.0 \
  --serialized-file model.pth \
  --handler ./handler.py \
  --export-path model-store \
  -f

GCS에 업로드 및 배포

bert_model.mar 파일을 Google Cloud Storage 버킷에 업로드한다.

upload_to_gcs.py
bucket = storage_client.bucket("mlops-bucket")
blob = bucket.blob("models/bert_model.mar")
blob.upload_from_filename("model-store/bert_model.mar")

마지막으로 Vertex AI에 모델을 등록하고 배포한다.

deploy_vertex_ai.py
model_path = "gs://mlops-bucket/models"
registry_model = aiplatform.Model.upload(
    display_name="AG News Classification",
    artifact_uri=model_path,
    serving_container_image_uri="gcr.io/cloud-aiplatform/prediction/tf2-cpu.2-3:latest"
)

deployment = registry_model.deploy(
    machine_type="n1-standard-2",
    min_replica_count=1,
    max_replica_count=1,
    accelerator_type="NVIDIA_TESLA_T4",
    accelerator_count=1,
    traffic_percentage=100,
)

⚠ 참고: 배포에는 15-30분이 소요될 수 있다.

4단계: 추론 실행

배포가 성공적으로 완료되면 예측 요청을 보내 추론을 테스트한다.

test_inference.py
endpoint.predict(
    instances=[
        "Google Gemini’s AI coding tool is now free for individual users"
    ]
)
예측 결과
그림 6: Vertex AI 엔드포인트 추론 결과를 출력한 Colab 셀

이 캡처는 Colab에서 Vertex AI 엔드포인트 호출 결과를 보여준다.

5단계: 리소스 정리

의도하지 않은 비용을 피하려면 테스트 후 리소스를 배포 취소하고 삭제해야 한다.

cleanup_resources.py
endpoint.undeploy_all()
endpoint.delete()
registry_model.delete()

마무리하며

이번 프로젝트 과정을 통해 모델 로직과 인프라를 분리하는 배포 파이프라인을 구축해보았다. Flask나 FastAPI로 직접 서빙 서버를 구성할 수 있으나, 이는 모델 코드가 웹 서버 로직과 강하게 결합되는 문제가 있기에 권장하지 않는다.

우리는 TorchServe와 MARModel Archive 형식을 통해 모델의 인터페이스를 표준화할 수 있었다. 덕분에 이 모델은 Vertex AI뿐만 아니라 KServe, AWS SageMaker, 심지어 로컬 쿠버네티스 클러스터 어디든 동일한 방식으로 배포될 수 있다는 장점을 얻은 것이 포인트다.

Vertex AI는 이러한 표준화된 아티팩트를 받아 실행해 주는 안정적인 컨테이너 런타임일 뿐이다. 결국 중요한 것은 특정 클라우드 벤더의 기능이 아니라, 어떤 환경에서도 재현 가능하고 확장 가능한 배포 구조를 설계하는 것이다. 이번 튜토리얼을 토대로 향후 지속 가능한 MLOps 파이프라인을 고민하는 엔지니어들에게 의미 있는 레퍼런스가 되기를 바란다.

각주

  • 1: Automated Machine Learning, 데이터 전처리부터 모델 훈련, 튜닝까지 자동화하는 기술 [↩︎]
  • 2: Bidirectional Encoder Representations from Transformers, 구글이 개발한 자연어 처리 모델 [↩︎]
  • 3: PyTorch 모델을 위한 고성능 모델 서빙 라이브러리 [↩︎]
  • 4: TorchServe에서 사용하는 모델 패키징 포맷 [↩︎]

추천 아티클

GPU는 어떻게 머신러닝 연산을 수행할까?

Python 코드 한 줄이 GPU의 트랜지스터를 움직이기까지의 JAX와 CUDA를 통해 보는 하드웨어 가속의 원리를 탐구해보자.

GPUMachine LearningCUDA
2025-12-05

JAX와 TPU를 이용한 Decoder-based Tiny LLM 사전학습

JAX와 Cloud TPU v6e를 활용하여 Llama 스타일의 Tiny LLM을 바닥부터 구현하고 사전 학습하는 과정을 다룬다. RMSNorm, SwiGLU, RoPE 등 최신 아키텍처와 고성능 I/O 파이프라인 구축 방법을 알아본다.

LLMJAXTPU
2025-12-02