JetStream 및 Pathways를 사용하여 GKE에서 멀티 호스트 TPU를 사용하여 LLM 서빙

이 가이드에서는 여러 노드에서 텐서 프로세싱 단위(TPU)를 사용하여 Google Kubernetes Engine(GKE)에서 Llama 3.1 405B와 같은 최신 대규모 언어 모델(LLM)을 서빙하는 방법을 보여줍니다.

이 가이드에서는 이동식 오픈소스 기술(Kubernetes, JetStream, Pathways on Cloud, LeaderWorkerSet(LWS) API)을 사용하여 GKE의 세밀한 제어, 확장성, 복원력, 이동성, 비용 효율성을 활용하여 GKE에서 AI/ML 워크로드를 배포하고 서빙하는 방법을 설명합니다.

배경

대규모 언어 모델의 크기가 커져 더 이상 단일 호스트 TPU 슬라이스에 수용할 수 없게 되었습니다. ML 추론의 경우 Cloud의 Pathways를 사용하여 상호 연결된 여러 TPU 노드에서 GKE에 대규모 멀티 호스트 추론을 실행할 수 있습니다. 이 가이드에서는 멀티 호스트 TPU 슬라이스로 GKE 클러스터를 프로비저닝하고, Cloud의 Pathways 바이너리를 사용하고, MaxText 프레임워크로 JetStream 서버를 실행하고, 멀티 호스트 추론 요청을 만드는 방법을 살펴봅니다.

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

TPU 정보

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

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

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

이 튜토리얼에서는 Llama 3.1-405B 모델 서빙을 설명합니다. GKE는 지연 시간이 짧은 프롬프트를 서빙하기 위한 모델 요구사항에 따라 TPU 토폴로지가 구성된 멀티 호스트 TPU v6e 노드에 모델을 배포합니다.

Pathways on Cloud

Pathways는 액셀러레이터를 위한 대규모 오케스트레이션 레이어입니다. Pathways는 현재 모델의 최첨단 성능을 유지하면서 새로운 시스템과 ML 연구 아이디어를 탐색할 수 있도록 명시적으로 설계되었습니다. Pathways는 단일 JAX 클라이언트 프로세스가 하나 이상의 대규모 TPU 슬라이스에서 계산을 조정할 수 있도록 지원하여 수백 또는 수천 개의 TPU 칩에 걸쳐 있는 ML 계산을 간소화합니다.

JetStream

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

MaxText

MaxText는 Flax, Orbax 및 Optax와 같은 오픈소스 JAX 라이브러리를 기반으로 구축된 성능이 뛰어나고 확장 가능하며 적응력이 뛰어난 JAX LLM 구현입니다. MaxText의 디코더 전용 LLM 구현은 Python으로 작성됩니다. XLA 컴파일러를 많이 활용하여 커스텀 커널을 빌드할 필요 없이 고성능을 실현합니다.

MaxText에서 지원하는 최신 모델 및 매개변수 크기에 관한 자세한 내용은 MaxText 프로젝트 저장소를 참고하세요.

Llama 3.1 405B

Llama 3.1 405B는 텍스트 생성, 번역, 질의 응답을 포함한 다양한 자연어 처리 태스크를 위해 설계된 Meta의 대규모 언어 모델입니다. GKE는 이러한 규모의 모델의 분산 학습 및 서빙 니즈를 지원하는 데 필요한 인프라를 제공합니다.

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

아키텍처

이 섹션에서는 이 튜토리얼에서 사용되는 GKE 아키텍처를 설명합니다. 이 아키텍처는 TPU를 프로비저닝하고 모델을 배포하고 서빙하기 위해 JetStream 및 Pathways 구성요소를 호스팅하는 GKE Standard 클러스터를 포함합니다.

다음은 이 아키텍처의 구성요소를 보여주는 다이어그램입니다.

JetStream 및 Pathways 구성요소가 포함된 멀티 호스트 TPU 노드 풀이 있는 GKE 클러스터의 아키텍처

이 아키텍처에는 다음 구성요소가 포함됩니다.

  • GKE Standard 리전 클러스터
  • JetStream 배포 및 Pathways 구성요소를 호스팅하는 멀티 호스트 TPU 슬라이스 노드 풀
  • Pathways resource manager는 가속기 리소스를 관리하고 사용자 작업의 가속기 할당을 조정합니다.
  • Pathways client는 Pathways resource manager와 협력하여 컴파일된 프로그램이 실행을 위해 배치되는 위치를 결정합니다.
  • Pathways worker는 액셀러레이터 머신에서 실행되고 계산을 수행하며 IFRT 프록시 서버를 통해 데이터를 워크로드로 다시 전송합니다.
  • IFRT proxy client는 OSS Interim Framework Runtime (IFRT) API를 구현하고 워크로드와 Pathways 구성요소 간의 통신 브리지 역할을 합니다.
  • IFRT proxy server는 IFRT proxy client에서 요청을 수신하고 Pathways client에 전달하여 작업을 분산합니다.
  • JetStream-Pathways 컨테이너는 추론 요청을 수신하고 실행 프로세스를 Pathways workers에 위임하는 JAX 기반 추론 서버를 제공합니다.
  • 서비스 구성요소는 수신 트래픽을 모든 JetStream HTTP 복제본으로 전파합니다.
  • JetStream HTTP는 요청을 JetStream의 필수 형식에 관한 래퍼로 수락하고 JetStream의 GRPC 클라이언트로 전송하는 HTTP 서버입니다.