Google Cloud TPUs와 JAX를 사용한 프로덕션 AI 구축

Source: Google Developers Blog

아키텍처적 필수 조건: 모듈성 및 성능

JAX AI Stack은 모듈식이고 느슨하게 결합된 컴포넌트라는 철학 위에 구축되었습니다. 각 라이브러리는 하나의 작업에 최적화되도록 설계되었습니다. 이러한 접근 방식은 사용자가 최적화, 데이터 로딩, 체크포인팅 등 요구 사항에 정확히 맞는 최적의 라이브러리를 선택·조합하여 맞춤형 ML 스택을 만들 수 있게 합니다. 특히, 이 모듈성은 급변하는 AI 분야에서 매우 중요합니다. 새로운 라이브러리와 기법을 개발·통합할 때, 거대한 단일 프레임워크를 수정해야 하는 위험과 비용 없이 빠른 혁신이 가능해집니다.

현대 ML 스택은 추상화 연속성을 제공해야 합니다. 개발 속도를 높이는 자동화된 고수준 최적화와, 마이크로초 단위까지 정밀하게 제어해야 할 때 필요한 세밀한 수동 제어 사이의 연속성을 말합니다. JAX AI Stack은 이러한 연속성을 제공하도록 설계되었습니다.

핵심 “JAX AI Stack”

JAX 생태계의 중심에는 “JAX AI Stack” 이 있습니다. 이는 모델 개발의 기반을 제공하는 네 가지 핵심 라이브러리로 구성되며, 모두 JAX와 XLA의 컴파일러‑우선 설계 위에 구축되었습니다.

• JAX – 가속기‑지향 배열 연산을 위한 기반. 순수 함수형 프로그래밍 모델 덕분에 변환이 조합 가능하고, 워크로드를 다양한 하드웨어와 클러스터 규모에 효과적으로 확장할 수 있습니다.
• Flax – 모델 작성 및 “수술”을 위한 유연하고 직관적인 API. JAX 위에 객체‑지향 레이어를 제공하면서도 성능을 희생하지 않습니다.
• Optax – 조합 가능한 그래디언트 처리 및 최적화 변환 라이브러리. 몇 줄의 코드만으로 Adam 같은 표준 옵티마이저를 그래디언트 클리핑이나 누적과 같은 기법과 체인할 수 있습니다.
• Orbax – “어떤 규모든” 지원하는 체크포인팅 라이브러리. 비동기 분산 체크포인팅을 제공해 하드웨어 장애가 발생해도 훈련 진행 상황을 크게 잃지 않도록 합니다.

jax-ai-stack 메타패키지는 다음과 같이 설치할 수 있습니다:

pip install jax-ai-stack

JAX AI Stack 및 생태계 구성 요소

확장된 JAX AI Stack

이 안정적인 코어 위에, 풍부한 특화 라이브러리 생태계가 전체 ML 라이프사이클에 필요한 엔드‑투‑엔드 기능을 제공합니다.

산업 규모 인프라스트럭처

사용자‑대면 라이브러리 아래에는 JAX가 단일 TPU/GPU에서 수천 개의 GPU/TPU로 원활히 확장될 수 있게 하는 인프라가 자리합니다.

• XLA (Accelerated Linear Algebra) – 도메인‑특화, 하드웨어‑비종속 컴파일러. 전체 프로그램 분석, 연산자 융합, 메모리 레이아웃 최적화를 통해 즉시 강력한 성능을 제공합니다.
• Pathways – 대규모 분산 연산을 위한 통합 런타임. 개발자는 마치 단일 강력한 머신을 사용하는 것처럼 코딩하고, Pathways가 수만 개의 칩에 걸쳐 실행을 조정하며 내장된 내결함성을 제공합니다.

최고 효율을 위한 고급 개발 도구

하드웨어 활용도를 극대화하기 위해, 생태계는 더 깊은 제어와 높은 효율성을 제공하는 특화 도구들을 포함합니다.

• Pallas & Tokamax – TPU와 GPU용 맞춤 커널 작성을 지원하는 확장 기능. 메모리 계층 구조와 병렬성에 대한 정밀한 제어를 가능하게 합니다. Tokamax는 플러그‑인‑플레이 성능을 위한 최신 커널(예: FlashAttention) 라이브러리를 제공합니다.
• Qwix – 비침해적 양자화 라이브러리. JAX 함수를 가로채어 QLoRA 또는 PTQ와 같은 기법을 적용하며, 원본 모델 코드를 최소 혹은 전혀 변경하지 않아도 됩니다.
• Grain – 고성능, 결정론적 데이터 로딩 라이브러리. Orbax와 통합되어 데이터 파이프라인의 정확한 상태를 체크포인팅함으로써 재시작 후 비트‑대‑비트 재현성을 보장합니다.

프로덕션까지의 전체 경로

다른 모듈들은 연구 단계에서 배포 단계까지 연결하는 성숙한 엔드‑투‑엔드 애플리케이션 레이어를 제공하여 JAX AI Stack을 보강합니다.

• MaxText & MaxDiffusion – LLM 및 디퓨전 모델 훈련을 위한 확장 가능한 프레임워크. 좋은 처리량(goodput)과 모델 FLOPs 활용도(MFU)를 기본으로 최적화된 신뢰할 수 있는 시작점을 제공합니다.
• Tunix – 사후 훈련 정렬을 위한 JAX‑네이티브 라이브러리. 감독 미세조정(SFT)과 같은 최신 알고리즘을 제공합니다.