MaxText — JAX로 TPU·GPU에서 훈련하는 고성능 오픈소스 LLM
MaxText
LLM을 수천 개 칩에서 훈련하려면 병렬 샤딩부터 양자화·체크포인트까지 손댈 게 많은데, JAX의 힘으로 그 복잡도를 크게 줄일 수 있어요. MaxText는 순수 Python/JAX로 작성된 고성능·고확장 오픈소스 LLM 라이브러리이자 참조 구현으로, Google Cloud TPU와 GPU에서 훈련을 겨냥해요. Model FLOPs Utilization(MFU)이 높고, 단일 호스트에서 매우 큰 클러스터까지 사전학습(수만 칩)과 확장 가능한 포스트트레이닝(SFT, GRPO 같은 RL)을 지원해요. Gemma·Llama·DeepSeek·Qwen·Mistral 같은 모델 라이브러리를 갖췄고 Apache 2.0 라이선스예요.
이 카테고리의 문서
- 개요: ReadTheDocs 문서와 JAX AI 스택
- 저장소: 모델 라이브러리·사전/포스트트레이닝
- 시작하기: 설치·사전학습·포스트트레이닝·추론