MaxText 개요 — 고성능·고확장 LLM 라이브러리
MaxText 개요
LLM 훈련에서 샤딩·양자화·체크포인트 같은 최적화를 직접 다루는 건 많은 시간이 걸려요. MaxText는 순수 Python/JAX로 작성된 고성능·고확장 오픈소스 LLM 라이브러리이자 참조 구현으로, Google Cloud TPU와 GPU에서 훈련하는 걸 목표로 해요. JAX와 XLA 컴파일러의 힘 덕분에 단순하고 대부분 "최적화 불필요(optimization-free)"한 채로 높은 MFU와 초당 토큰 수를 냅니다.
특징
- 고성능 (High-performance) — 단일 호스트에서 매우 큰 클러스터까지 높은 Model FLOPs Utilization(MFU)과 tokens/second를 달성
- 사전학습 (Pre-training) — 샤딩·양자화·체크포인트 등 다양한 차원에서 최적 성능을 내는 opinionated 구현 제공
- 포스트트레이닝 (Post-training) — SFT, 지식 증류, GRPO 같은 최신 강화학습(RL) 알고리즘으로 사유·OSS 모델을 파인튜닝하는 확장 가능 프레임워크
JAX AI 스택
MaxText는 Google 내외 연구자들이 유용하다고 여긴 라이브러리 모음인 JAX AI Stack 위에서 동작해요.
- JAX — 핵심 배열 연산과 프로그램 변환
- Flax — 신경망 구축
- Orbax — 체크포인트·지속성 유틸
- Optax — 그래디언트 처리·최적화
- Tunix — 최신 실험 알고리즘·포스트트레이닝 기법의 JAX 라이브러리
- Grain 또는 tf.data — 데이터 로딩