Moait
홈인기 피드모든 블로그모든 태그
홈인기 피드모든 블로그모든 태그
Reproducing OLMo 3 7B Pre-training in MaxText: case study of large scale training on TPUs 섬네일

Reproducing OLMo 3 7B Pre-training in MaxText: case study of large scale training on TPUs

구글 favicon구글·AI/ML·
PyTorchTPUGoogle CloudJAXXLA
2026년 09월 24일0

AI 요약

이 글은 AI가 요약했어요. 정확한 내용은 꼭 원문을 확인해 주세요!

핵심 요약

구글의 MaxText 팀이 AI2의 OLMo 3 7B를 Google Cloud TPUs에서 JAX/XLA로 재현했고, 프리-트레이닝과 미드-트레이닝 전 과정에서 원본 PyTorch-on-GPU 레퍼런스와 일치하는 성능을 확인했습니다.

구현 방법

  • TPUs에서 JAX/XLA로 재현 수행
  • pre-training 및 mid-training, 모든 held-out 평가에서 일치 검증
  • MFU 57.4% 달성 및 클러스터 리사이즈/TPU 간 이동성 유지
  • 데이터 로더 memorization 버그 발견과 held-out 검증의 중요성 확인

주요 결과

  • MFU 최대 57.4% 달성 및 레퍼런스와의 일치성 확보
  • 다중 TPU/클러스터 환경에서도 재현성 및 이식성 강화
  • 검증 강화의 필요성 재확인

연관 피드

%가 높을수록 이 글과 비슷할 가능성이 높아요!
We terminated a TPU mid-training and it recovered in seconds: Introduction to elastic training with MaxText 섬네일
75%

We terminated a TPU mid-training and it recovered in seconds: Introduction to elastic training with MaxText

구글 favicon구글·2026년 07월 06일
MaxText Expands Post-Training Capabilities: Introducing SFT and RL on Single-Host TPUs 섬네일
73%

MaxText Expands Post-Training Capabilities: Introducing SFT and RL on Single-Host TPUs

구글 favicon구글·2026년 04월 16일
Enterprise-Grade Precision for Long-Context Multimodal Embedding Inference on Cloud TPU 섬네일
72%

Enterprise-Grade Precision for Long-Context Multimodal Embedding Inference on Cloud TPU

구글 favicon구글·2026년 08월 26일