google/tunix

What it solves

Tunix는 대규모 언어 모델(LLMs)의 사후 학습(post-training) 과정을 단순화하고 가속화하기 위해 설계된 경량 라이브러리입니다. 지도 학습 미세 조정 및 강화 학습과 같은 사후 학습 워크플로우의 확장성 복잡성을, 특히 TPU와 같은 고성능 하드웨어에 최적화하여 해결합니다.

How it works

JAX 기반으로 구축된 Tunix는 핵심 JAX 유틸리티(Flax, Optax, Orbax)와 최적화된 모델 구현(MaxText 등) 사이의 중간 계층으로 작동합니다. vLLM 및 SGLang-JAX와 같은 고성능 추론 엔진과 통합되어 효율적인 rollout을 수행하며, Pathways를 사용하여 수천 개의 장치에서 원활한 멀티 호스트 분산 학습을 수행합니다.

Who it’s for

JAX를 사용하여 LLM을 미세 조정하고, TPU 가속을 활용하여 확장 가능하고 안정적이며 효율적인 사후 학습을 수행하고자 하는 AI 연구자 및 개발자들을 위한 것입니다.

Highlights

• Diverse Training Algorithms: SFT, DPO, PPO, 및 GRPO를 지원합니다.
• Agentic RL: 멀티턴 도구 사용, 고처리량 궤적(trajectory) 수집을 위한 비동기 rollout, 그리고 궤적 배칭(batching)에 대한 특화된 지원을 제공합니다.
• TPU Optimization: MaxText와 네이티브 통합되어 고성능 커널(kernels)을 사용하며, vLLM/SGLang-JAX를 통해 빠른 rollout을 수행합니다.
• Scalability: Scalability 지원: Pathways를 통해 멀티 호스트 분산 학습과 결함 허용 체크포인트(checkpointing)를 지원합니다.