google/tunix

해결하는 문제

Tunix은 대규모 언어 모델(LLMs)의 후학습을 간소화하고 가속화하기 위해 설계된 JAX 기반 라이브러리입니다. 다양한 정렬 및 최적화 기법을 통해 모델을 개선하는 확장 가능한 프레임워크를 제공하며, TPU에서 높은 성능을 발휘하도록 특별히 최적화되어 있습니다.

작동 방식

Tunix은 JAX 생태계 내의 중간 계층으로 작동하며, Flax, Optax, Orbax와 같은 핵심 도구와 통합됩니다. 이러한 도구들을 MaxText와 같은 고성능 모델 구현체 및 vLLM, SGLang-JAX와 같은 추론 엔진과 연결하여 효율적인 롤아웃과 모델 실행을 처리합니다.

대상 사용자

JAX와 TPU 인프라를 사용하여 LLM의 확장 가능한 후학습을 수행해야 하는 연구자 및 개발자를 대상으로 합니다. 특히, 감독된 미세조정 및 강화학습에 종사하는 분들에게 적합합니다.

주요 특징

• 다양한 학습 알고리즘: 감독된 미세조정(SFT), 직접적 선호도 최적화(DPO), PPO, GRPO, GSPO-Token을 포함한 다양한 강화학습(RL) 방법 지원.
• 에이전트 기반 RL: 다단계 도구 사용, 고처리량 트래잭터리 수집을 위한 비동기 롤아웃, 트래잭터리 배치 처리 지원.
• TPU 최적화: vLLM, SGLang-JAX, MaxText와의 네이티브 통합을 통해 TPU 하드웨어에서 최고 수준의 성능 제공.
• 확장성: Pathways를 통해 멀티호스트 분산 학습을 원활하게 지원하며, 수천 개의 장치까지 확장 가능.
• 모델 호환성: Gemma, Llama, Qwen 등 인기 있는 모델 패밀리와 호환됩니다.