AI-Hypercomputer/maxtext

해결하는 문제

MaxText는 대규모 언어 모델(LLM) 훈련을 위한 고성능, 확장 가능한 참조 구현을 제공하도록 설계되었습니다. 단일 호스트에서 매우 큰 TPU 및 GPU 클러스터로 확장할 때 높은 Model FLOPs Utilization(MFU) 및 처리량(초당 토큰 수)을 달성하는 동시에 구현을 단순하고 대체로 "최적화 없이" 유지하는 문제를 해결합니다.

작동 방식

순수 Python과 JAX로 작성된 MaxText는 XLA 컴파일러를 활용하여 성능을 최적화합니다. Flax(신경망), Tunix(사후 훈련), Orbax(체크포인팅), Optax(최적화), Grain(데이터 로딩)을 포함한 JAX AI 라이브러리 제품군과 통합됩니다. 처음부터 사전 훈련과 지도 미세 조정(SFT) 및 GRPO(Group Relative Policy Optimization) 및 GSPO(Group Sequence Policy Optimization)와 같은 강화 학습(RL) 방법을 사용한 확장 가능한 사후 훈련을 모두 지원합니다.

대상 사용자

연구 및 프로덕션 환경에서 야심 찬 LLM 프로젝트를 구축하고 인기 있는 오픈 소스 모델을 사전 훈련하거나 사후 훈련할 수 있는 확장 가능한 프레임워크가 필요한 연구자와 개발자를 대상으로 합니다.

주요 특징

• 광범위한 모델 지원: Gemma, Llama, DeepSeek, Qwen, Mistral을 포함한 다양한 모델을 지원하며, 밀집 및 Mixture-of-Experts(MoE) 아키텍처를 모두 포함합니다.
• 높은 확장성: 수만 개의 칩에서 사전 훈련이 가능합니다.
• 멀티모달 기능: Gemma 3, Gemma 4, Llama 4 VLM의 멀티모달 훈련을 지원합니다.
• 다양한 훈련 기법: SFT 및 RL(GRPO, GSPO) 구현을 제공하고, Multi-Token Prediction(MTP) 및 분산 저통신 훈련을 위한 Streaming DiLoCo와 같은 고급 기능을 통합합니다.
• 하드웨어 대상: Google Cloud TPU 및 GPU에 특별히 최적화되어 있습니다.