MaxText 시작하기 — 설치·훈련·추론

MaxText 시작하기

단일 호스트에서든 멀티호스트로 Cloud TPU·NVIDIA GPU까지 확장에서든, 첫 MaxText 워크로드를 돌리는 시작점을 알려드릴게요. 설치 → 첫 모델 훈련 → 추론 순으로 따라가면 돼요.

출처: https://maxtext.readthedocs.io/en/latest/getting_started.html

사전 준비

로그·체크포인트를 저장할 Cloud Storage 버킷을 프로젝트에 만들고, TPU·GPU VM이 버킷에 읽기/쓰기 권한이 있게 해요. 권한은 STORAGE ADMIN 같은 서비스 계정 역할로 부여돼요.

설정

MaxText는 yaml 파일로 설정을 읽어요. configs/base.yml에 약 1B 파라미터의 decoder-only 모델이 있고, 명령줄에서 옵션을 덮어쓸 수 있어요. 예를 들어 steps나 log_period를 configs/base.yml 수정 또는 train.py 호출에 인자로 넘겨 변경할 수 있어요. base_output_directory는 만든 버킷의 폴더로 설정해요.

HuggingFace 체크포인트를 쓰려면 먼저 MaxText/Orbax 형식으로 변환해야 해요. 자세한 건 Checkpoint Conversion Guide를 봐요.

단일 호스트에서 실행

  1. 설치 — TPU VM에선 maxtext[tpu](사전학습) 또는 maxtext[tpu-post-train](포스트트레이닝)을, GPU VM에선 maxtext[cuda12]를 설치해요.
  2. 사전학습 — 첫 모델 훈련은 Pre-training Tutorial 참고
  3. 포스트트레이닝 — SFT·RL 같은 기법을 적용하려면 Post-training Tutorial 참고
  4. 추론(디코딩) — 오프라인·온라인 추론과 vLLM 통합은 Inference Tutorial 참고

멀티호스트·노트북

여러 호스트로 확장할 땐 Google Kubernetes Engine (GKE) 가 권장 방식이에요. Cluster Toolkit 또는 XPK로 실행할 수 있어요. 노트북에서는 Jupyter, Google Colab, VS Code에서 인터랙티브하게 돌릴 수 있어요.

다음 단계

워크로드가 돌아가기 시작하면 성능을 높이는 사전 최적화(preflight optimizations)를 적용할 수 있어요. 자세한 내용은 Optimization Tips를 봐요.

더 알아보기