JetStream 및 PyTorch와 함께 GKE에서 TPU를 사용하는 LLM 제공

이 가이드에서는 PyTorch를 통해 JetStream과 함께 Google Kubernetes Engine(GKE)에서 TPU(Tensor Processing Unit)를 사용하여 대규모 언어 모델(LLM)을 제공하는 방법을 보여줍니다. 이 가이드에서는 모델 가중치를 Cloud Storage에 다운로드하고 JetStream을 실행하는 컨테이너를 사용하여 GKE Autopilot 또는 Standard 클러스터에 배포합니다.

JetStream에 모델을 배포할 때 Kubernetes 기능에서 제공하는 확장성, 복원력, 비용 효율성이 필요한 경우 이 가이드가 좋은 출발점이 될 수 있습니다.

이 가이드는 PyTorch를 사용하는 생성형 AI 고객, GKE의 신규 또는 기존 사용자, ML 엔지니어, MLOps(DevOps) 엔지니어, LLM 제공을 위해 Kubernetes 컨테이너 조정 기능을 사용하는 데 관심이 있는 플랫폼 관리자를 대상으로 합니다.

배경

JetStream과 함께 GKE에서 TPU를 사용하는 LLM을 제공하면 비용 효율성, 확장성 및 더 높은 가용성을 비롯하여 관리형 Kubernetes의 모든 이점을 갖춘 강력한 프로덕션에 즉시 사용 가능한 서빙 솔루션을 빌드할 수 있습니다. 이 섹션에서는 이 튜토리얼에서 사용되는 주요 기술을 설명합니다.

TPU 정보

TPU는 Google이 TensorFlow, PyTorchJAX와 같은 프레임워크를 활용하여 빌드된 머신러닝 및 AI 모델을 가속화하는 데 사용하기 위해 커스텀 개발한 ASIC(application-specific integrated circuits)입니다.

GKE에서 TPU를 사용하기 전에 다음 학습 과정을 완료하는 것이 좋습니다.

  1. Cloud TPU 시스템 아키텍처를 사용하는 현재 TPU 버전 가용성 알아보기
  2. GKE의 TPU 알아보기

이 튜토리얼에서는 다양한 LLM 모델 서빙을 다룹니다. GKE는 지연 시간이 짧은 프롬프트를 제공하기 위한 모델 요구사항에 따라 TPU 토폴로지가 구성된 단일 호스트 TPUv5e 노드에 모델을 배포합니다.

JetStream 정보

JetStream은 Google이 개발한 오픈소스 추론 서빙 프레임워크입니다. JetStream은 TPU 및 GPU에서 고성능, 높은 처리량, 메모리 최적화된 추론을 가능하게 합니다. JetStream은 연속 일괄 처리, KV 캐시 최적화, 양자화 기법을 비롯한 고급 성능 최적화를 제공하여 LLM 배포를 용이하게 합니다. JetStream은 최적의 성능을 달성하기 위해 PyTorch/XLA 및 JAX TPU 제공을 지원합니다.

연속 일괄 처리

연속 배치 처리는 수신되는 추론 요청을 동적으로 배치로 그룹화하여 지연 시간을 줄이고 처리량을 늘리는 기법입니다.

KV 캐시 양자화

KV 캐시 양자화에는 주의 메커니즘에 사용되는 키-값 캐시를 압축하여 메모리 요구사항을 줄이는 작업이 포함됩니다.

Int8 가중치 양자화

Int8 가중치 양자화는 모델 가중치의 정밀도를 32비트 부동 소수점에서 8비트 정수로 줄여 계산 속도를 높이고 메모리 사용량을 줄입니다.

이러한 최적화에 대한 자세한 내용은 JetStream PyTorchJetStream MaxText 프로젝트 저장소를 참조하세요.

PyTorch 정보

PyTorch는 Meta에서 개발한 오픈소스 머신러닝 프레임워크로, 현재 Linux Foundation에 속해 있습니다. PyTorch는 텐서 계산, 심층 신경망과 같은 고급 기능을 제공합니다.