MaxText, 포스트 트레이닝 기능 확장: 단일 호스트 TPU에서 SFT 및 RL 도입
요약
MaxText가 단일 호스트 TPU에서도 지도 파인튜닝(SFT)과 강화학습(RL)을 지원하게 되어, 개발자들이 LLM 모델을 더 쉽고 효율적으로 정교하게 만들 수 있게 되었어요.
인사이트
- MaxText는 이제 단일 호스트 TPU(v5p-8, v6e-8 등)에서 SFT(지도 파인튜닝)와 RL(강화학습)을 지원해서, LLM의 전문화와 성능 개선 과정을 훨씬 간편하게 만들었어요.
- SFT 기능은 Hugging Face 데이터셋 및 체크포인트와 매끄럽게 연동되어, 기존 모델이 특정 지시를 따르거나 특화된 작업을 수행하는 능력을 효율적으로 높일 수 있게 해줘요.
- RL 기능은 GRPO, GSPO 같은 최신 알고리즘을 활용하고 vLLM으로 고성능 추론을 지원해서, 수학이나 코딩처럼 복잡한 추론 능력을 요구하는 작업에서 LLM의 성능을 획기적으로 향상시킬 수 있어요.
왜 중요한가
이 업데이트는 개발자들이 고가의 복잡한 멀티-호스트 TPU 환경 없이도, 단일 TPU만으로도 LLM을 미세 조정하고 강화 학습시킬 수 있는 접근성을 크게 높여줘요. 덕분에 비용 부담과 기술적 진입 장벽이 낮아지면서도, 최첨단 LLM 기술을 활용해 혁신적인 AI 애플리케이션을 더 많은 개발자가 쉽게 만들 수 있게 될 거예요.
2026년 4월 16일
대규모 언어 모델(LLM) 분야는 정말 빠르게 발전하고 있죠. 그런데 모델을 사전 학습(pre-training)하는 건 사실 첫걸음에 불과해요. 기본 모델을 특정 목적의 어시스턴트나 고성능 추론 엔진으로 만들려면, 사전 학습 이후의 과정(post-training)이 정말 중요하거든요. 오늘, MaxText에 이런 과정을 훨씬 간편하게 만들어 줄 새로운 기능들을 발표하게 되어 정말 기뻐요. 바로 단일 호스트 TPU 구성(예: v5p-8, v6e-8)에서 이제 **지도 파인튜닝(SFT)**과 **강화학습(RL)**을 사용할 수 있게 됐다는 소식이에요!
JAX의 강력한 성능과 Tunix 라이브러리의 효율성을 활용해서, MaxText는 개발자들이 최신 포스트 트레이닝 기법으로 모델을 정교하게 만들 수 있는 고성능의 확장 가능한 방법을 제공해요. SFT와 RL에 대한 전체 문서는 각각 여기와 여기에서 확인하실 수 있으니, 지금 바로 TPU에서 포스트 트레이닝 여정을 시작해보는 건 어때요?
Supervised Fine-Tuning (SFT): 정밀한 튜닝을 간편하게
지도 파인튜닝(SFT)은 사전 학습된 모델이 특정 지시를 따르거나 특정 작업에 뛰어난 성능을 보이도록 만드는 데 가장 기본적인 방법이에요. 이번에 새로 추가된 단일 호스트 SFT 지원 덕분에, 이제 사용자들은 기존 MaxText나 Hugging Face 체크포인트를 가져와서 레이블이 지정된 데이터셋으로 아주 간단하게 파인튜닝할 수 있게 됐어요.
주요 특징:
- 매끄러운 통합: Hugging Face 데이터셋(예: ultrachat_200k)을 바로 지원해요.
- 유연한 체크포인트: 기존 MaxText 체크포인트를 사용하거나, Hugging Face 모델(예: Gemma 3)을 MaxText 생태계 내에서 바로 변환해서 쓸 수 있어요.
- 최적화된 실행: 포스트 트레이닝 효율성을 위해 특별히 설계된 JAX 기반 라이브러리인 Tunix의 지원을 받아 최적의 성능을 보여줘요.
Reinforcement Learning (RL): 추론 능력 향상
수학이나 코딩처럼 복잡한 논리 및 추론이 필요한 작업에서는 강화학습(RL)이 정말 중요한 역할을 해요. MaxText는 이제 단일 호스트 TPU에서 몇 가지 최첨단 RL 알고리즘을 지원하는데요, 트레이닝 루프 중에는 고처리량 추론을 위해 vLLM을 활용해요. 예를 들어볼까요?
- Group Relative Policy Optimization (GRPO) GRPO는 PPO(Proximal Policy Optimization)의 메모리 효율적인 변형이에요. 이 방법은 별도의 가치 함수(value function) 모델이 필요 없어요. 대신 하나의 프롬프트당 여러 응답을 생성하고, 이 그룹 내에서 상대적인 이점(relative advantages)을 계산하죠. 덕분에 하드웨어 사용량을 훨씬 줄여서, 단일 TPU 호스트에서도 고급 RL을 할 수 있게 만들어줘요.
- Group Sequence Policy Optimization (GSPO) GSPO는 시퀀스 수준의 중요도 비율(importance ratios)과 클리핑(clipping)에 초점을 맞춰요. 시퀀스 수준에서 모델의 행동에 보상을 줘서 트레이닝의 안정성과 효율성을 높여주고요. 특히 GSM8K 같은 벤치마크에서 성능을 향상시키는 데 아주 효과적이에요.
시작하기
이 새로운 기능들을 사용하려면, 최신 포스트 트레이닝 종속성들을 설치해야 해요:
uv pip install maxtext[tpu-post-train]==0.2.1 --resolution=lowest
install_maxtext_tpu_post_train_extra_deps
Shell
Copied
SFT 실행하기:
train_sft 모듈을 사용해서 SFT를 실행할 수 있어요. 모델, 데이터셋, 그리고 출력 디렉토리를 지정해주면 돼요:
python3 -m maxtext.trainers.post_train.sft.train_sft \
model_name=${MODEL?} \
load_parameters_path=${MAXTEXT_CKPT_PATH?} \
run_name=${RUN_NAME?} \
base_output_directory=${BASE_OUTPUT_DIRECTORY?}
Shell
Copied
RL (GRPO/GSPO) 실행하기:
RL의 경우, train_rl 모듈이 정책(policy) 및 참조(reference) 모델을 로드하고, 트레이닝을 실행하며, 추론 벤치마크에서 자동 평가까지 제공해줘요:
python3 -m maxtext.trainers.post_train.rl.train_rl \
model_name=${MODEL?} \
load_parameters_path=${MAXTEXT_CKPT_PATH?} \
run_name=${RUN_NAME?} \
base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
loss_algo=gspo-token \
chips_per_vm=${CHIPS_PER_VM?}
Shell
Copied
다음은 무엇일까요?
단일 호스트 지원이 많은 개발자에게 강력한 시작점이 되어주지만, MaxText는 기본적으로 스케일 확장을 염두에 두고 만들어졌어요. 지금 말씀드린 이 모든 워크플로우는 나중에 더 큰 모델과 방대한 데이터셋을 학습시킬 때, 멀티 호스트 구성으로도 매끄럽게 전환될 수 있도록 설계되어 있답니다. 앞으로 이 방향으로 더 많은 업데이트 소식을 전해드릴 예정이니 기대해주세요!