MaxText, 사후 훈련 기능 확장: 단일 호스트 TPU에서 SFT와 RL 도입

Source: Google Developers Blog

APRIL 16, 2026

대규모 언어 모델(LLM)의 급변하는 환경에서 사전 학습은 첫 번째 단계에 불과합니다. 기본 모델을 특화된 어시스턴트나 고성능 추론 엔진으로 전환하려면 사후 학습이 필수적입니다. 오늘 우리는 MaxText에 새로운 기능을 발표합니다: **Supervised Fine‑Tuning (SFT)**와 **Reinforcement Learning (RL)**이 이제 단일 호스트 TPU 구성(예: v5p‑8 및 v6e‑8)에서 사용할 수 있게 되었습니다.

JAX의 강력함과 Tunix 라이브러리의 효율성을 활용하여 MaxText는 최신 사후 학습 기법을 사용해 모델을 정제하려는 개발자를 위한 고성능, 확장 가능한 경로를 제공합니다. 오늘 TPUs에서 사후 학습 여정을 시작하려면 SFT와 RL 전체 문서를 확인해 보세요.

Supervised Fine‑Tuning (SFT): Precision Tuning Made Simple

Supervised Fine‑Tuning은 사전 학습된 모델을 특정 지시를 따르거나 특수 작업에서 뛰어나게 만들기 위한 주요 방법입니다. 새로운 싱글‑호스트 SFT 지원을 통해 사용자는 기존 MaxText 또는 Hugging Face 체크포인트를 가져와 최소한의 설정만으로 라벨이 지정된 데이터셋에 파인‑튜닝할 수 있습니다.

Key Highlights

• Seamless Integration: Native support for Hugging Face datasets (e.g., ultrachat_200k).
• Flexible Checkpoints: Use existing MaxText checkpoints or convert Hugging Face models (like Gemma 3) directly within the ecosystem.
• Optimized Execution: Powered by Tunix, a JAX‑based library specifically designed for post‑training efficiency.

강화 학습 (RL): 추론 능력 향상

복잡한 논리와 추론이 필요한 작업—예를 들어 수학이나 코딩—에 있어 강화 학습은 게임 체인저입니다. MaxText는 이제 vLLM을 활용해 훈련 루프 동안 고처리량 추론을 수행하면서 단일 호스트 TPU에서 여러 최신 RL 알고리즘을 지원합니다.

• 그룹 상대 정책 최적화 (GRPO) – PPO(근접 정책 최적화)의 메모리 효율적인 변형입니다. 별도의 가치 함수 모델이 필요 없으며, 대신 프롬프트당 여러 응답을 생성하고 그룹 내 상대 이점을 계산합니다. 이를 통해 하드웨어 사용량을 줄이고 단일 TPU 호스트에서도 고급 RL을 활용할 수 있습니다.
• 그룹 시퀀스 정책 최적화 (GSPO) – 시퀀스 수준의 중요도 비율과 클리핑에 초점을 맞춥니다. 시퀀스 수준에서 모델 행동에 보상을 부여함으로써 훈련 안정성과 효율성을 향상시키며, GSM8K와 같은 벤치마크에서 특히 효과적임을 입증했습니다.

시작하기

최신 사후 훈련 종속성이 설치되어 있는지 확인하십시오:

uv pip install maxtext[tpu-post-train]==0.2.1 --resolution=lowest
install_maxtext_tpu_post_train_extra_deps

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?}

RL 실행 (GRPO/GSPO)

RL의 경우, train_rl 모듈이 정책 및 레퍼런스 모델 로드를 처리하고, 훈련을 실행하며, 추론 벤치마크에 대한 자동 평가를 제공합니다:

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?}

다음은?

단일 호스트 지원은 많은 개발자에게 강력한 시작점을 제공하지만, MaxText는 규모에 맞게 설계되었습니다. 이러한 워크플로는 더 큰 모델을 학습하고 방대한 데이터셋을 활용하기 위해 멀티‑호스트 구성으로 원활하게 전환됩니다. 향후 업데이트를 기대해 주세요.