MaxText 저장소 — 모델 라이브러리와 훈련 흐름

MaxText 저장소

MaxText는 "simple, performant and scalable Jax LLM"이라는 한 줄로 설명돼요. AI-Hypercomputer org에서 관리하며 Python 3.12를 주 지원 버전으로 권장해요.

출처: https://github.com/AI-Hypercomputer/maxtext

모델 라이브러리

MaxText는 고성능 모델 라이브러리를 제공해요.

  • Google: Gemma 4 (26B MoE, 31B Dense), Gemma 3 (4B, 12B, 27B), Gemma 2 (2B, 9B, 27B), Gemma 1 (2B, 7B)
  • Alibaba: Qwen 3 MoE(30B, 235B), Qwen 3 Dense(0.6B~32B), Qwen 2.5(1.5B, 7B, 14B)
  • DeepSeek: V3.2/V3.1/V3 0324(671B), R1 0528(671B), V2(16B, 236B)
  • Meta: Llama 4 Scout(109B)·Maverick(400B), Llama 3.x(8B405B), Llama 2(7B70B)
  • OpenAI: GPT-OSS(20B, 120B), GPT3(52K, 6B, 22B, 175B)
  • Mistral: Mixtral(8x7B, 8x22B)

사전학습과 포스트트레이닝

사전학습은 새로 모델을 만드는 팀을 위한 참조 구현이에요. Llama 8B 같은 작은 dense 모델부터 DeepSeek-V3 같은 큰 MoE까지, MaxText를 포크·수정해 훈련할 수 있어요.

포스트트레이닝은 Tunix를 이용한 확장 가능 프레임워크예요. RL(GRPO 등)에선 vLLM으로 샘플링을 하고, SFT·RL 튜토리얼을 단일/멀티 호스트 TPU로 제공해요. 목표는 다양한 모델(a)과 기법(b)을 조합해 최적 모델을 훈련할 수 있게 하는 거예요.

사용 사례·연결

  • 텍스트 LLM 외에 Gemma 3/4, Llama 4의 멀티모달 훈련도 지원해요.
  • 디퓨전 모델(Wan 2.1, Flux 등)은 MaxDiffusion을 참고해요.
  • 라이선스는 Apache License 2.0이에요.

더 알아보기