PyTorch에서 JAX로, Olmo 3 7B를 TPU에서 완벽 재현한 대규모 학습 실험기

왜 이 실험이 중요한가

대규모 언어모델(LLM) 학습은 이제 소수 빅테크만의 영역이 아닙니다. 하지만 "우리 모델이 잘 학습되고 있다"는 말은 대개 학습 손실(training loss) 곡선 하나로 증명됩니다. 문제는 이 곡선이 거짓말을 할 수 있다는 점이에요.

Google Cloud TPU 팀은 Allen Institute for AI(Ai2)의 Olmo 3 7B를 MaxText(JAX/XLA 기반 TPU 학습 프레임워크)로 처음부터 재현했습니다. 이 실험이 흥미로운 이유는 세 가지입니다.

• 진짜 오픈 모델: 가중치, 데이터, 학습 레시피, 로그까지 전부 공개된 프론티어급 모델은 극소수입니다.
• 프레임워크 교차 검증: PyTorch+GPU 레시피를 JAX+TPU로 옮겼을 때 수학적으로 동일한 결과가 나오는지 확인할 수 있습니다.
• 실무 함정 공개: 학습 손실이 좋아 보이는데 실제 일반화 성능은 그대로인 데이터 버그를 잡아낸 과정이 고스란히 담겨 있습니다.

이 글은 단순 논문 요약이 아니라, 수 주간 실제로 클러스터를 돌리며 겪은 삽질기를 실무 관점에서 재구성한 것입니다. 근거자료는 Google Developers Blog 원문에서 확인할 수 있어요.

핵심 메시지: "재현했다"고 말하려면 학습 손실이 아니라 홀드아웃(held-out) 지표로 증명해야 합니다.

재현의 기술적 핵심: 무엇을 맞추고 무엇을 단순화했나

Olmo 3 7B는 32레이어, 4096차원 dense transformer입니다. 특이점은 reordered norm 블록, QK-norm, 슬라이딩 윈도우 3:1 + 글로벌 어텐션 조합이에요. MaxText 설정(olmo3-7b-pt.yml)은 이 구조를 정확히 미러링합니다.

시작점: 가중치 변환부터 검증

Ai2의 step-0 PyTorch 체크포인트를 Orbax로 변환한 뒤, 순전파 결과를 HuggingFace 레퍼런스와 비교했습니다. 결과는 KL ≈ 1.5e-3, top-10 토큰 중 9개 일치 — "같은 모델, 다른 프레임워크" 수준의 노이즈 플로어입니다.

# 체크포인트 변환 후 검증 개념도 (의사 코드)
# PyTorch → Orbax 변환 후 KL divergence 측정
kl_div = compute_kl_divergence(maxtext_logits, hf_reference_logits)
# 기대값: ~1.5e-3 (프레임워크 노이즈 플로어)
assert kl_div < 2e-3, "변환 실패: 모델이 다름"

데이터 파이프라인 재현

Olmo-core와 동일하게 구현했습니다.

1. 모든 문서를 토큰화 후 EOS로 연결
2. 8192 토큰 단위로 비중복 슬라이스
3. 고정 시드로 전역 셔플
4. n-gram 반복 필터(32회 초과 반복 인스턴스 마스킹)

MaxText의 dataset_type=olmo_grain(Grain 기반)이 이를 담당합니다.

의도적 단순화 2가지

• LR 스케줄: Ai2는 원래 5T 토큰 계획이었는데 실행 중 5.93T로 연장해서 코사인 곡선이 두 개 이어져 있습니다. MaxText는 전체 구간에 단일 코사인을 적용했습니다.
• 데이터: 공개된 Olmo-3 mix 사용(내부 실행 대비 <0.5% 토큰 누락).

이런 단순화에도 불구하고 홀드아웃 지표 전부에서 매치가 유지됐습니다. 이게 "run-to-run 노이즈 범위 내 재현"이라고 말하는 이유예요.

실무에서 진짜 배울 점: 학습 손실의 배신과 하드웨어 최적화

사례 1: 학습 손실이 거짓말한 순간

약 0.9M 스텝부터 MaxText의 학습 손실이 Ai2보다 낮아지기 시작했고, 1.25M 이후엔 평균 −0.06까지 벌어졌습니다. 곡선만 보면 "우리가 이겼다" 싶죠. 하지만:

• 홀드아웃 C4 손실은 동률(Δ −0.004)
• 다운스트림 정확도는 오히려 Ai2가 근소 우위

원인은 Grain 데이터 로더의 이중 샤딩 버그였습니다. ShardOptions를 샘플러에 넘기는데 샘플러가 이미 내부적으로 샤딩 중이어서, 인덱스 스트림이 32배 빠르게 전진했습니다. 결과적으로 stage-1이 한 번의 깨끗한 epoch이 아니라 Poisson(≈1) 리샘플링이 되어, 코퍼스의 약 37%는 한 번도 안 보이고 26%는 두 번 이상 반복 학습됐어요. 반복된 인스턴스에서 손실이 인위적으로 낮아진 겁니다.

교훈: 학습 손실 곡선만 모니터링하면 데이터 버그를 "성능 향상"으로 착각합니다. 반드시 스텝 정렬된 체크포인트로 홀드아웃 평가를 병행하세요.

사례 2: 헤드 차원 하나 바꿔서 12% 무료 성능 향상

Olmo-3 7B 기본 설정은 32 헤드 × 128 head-dim입니다. num_heads × head_dim = 4096이므로 16 헤드 × 256 head-dim으로 바꿔도 파라미터 수(7.298B)와 FLOPs(1565 TFLOP/step)는 완전히 동일합니다.

Ironwood의 MXU가 256×256 systolic array라서, head-dim 128은 QK matmul에서 배열의 절반을 놀립니다. 256으로 맞추면 유휴 사이클이 사라져서:

긴 학습을 시작하기 전에 반드시 확인해야 할 공짜 최적화입니다.

사례 3: 옵티마이저 dtype의 조용한 함정

weight_dtype=bfloat16으로 설정하면 Adam의 m/v 모멘트가 mu_dtype 상속으로 조용히 bf16으로 강등됩니다. bf16의 ~3자리 가수부가 초기 워밍업의 미세 업데이트를 매번 잘라먹어서 1000 스텝 만에 손실 +0.93이 누적됐어요. weight_dtype=float32(기본값)를 유지하면 격차가 30배 줄어듭니다. "왜 재현이 안 되지?"의 최대 원인이었습니다.

사례 4: 토폴로지와 무관한 레시피

JAX/XLA 스택의 최대 장점은 레시피가 칩 개수와 분리된다는 점입니다. 글로벌 배치(512 인스턴스, 4.19M 토큰/step)는 고정, 칩 수는 유동적. 점유된 클러스터에서 남는 용량 아무거나 잡아서 수학은 동일하게 유지할 수 있습니다. 실제로 stage-2는 Ironwood에서 v5p로 세대까지 넘어갔는데, XLA 플래그만 그대로 두고 device type만 바꿔서 57.4% MFU를 찍었습니다(v5p 피크의 57%).

사례 5: 정확한 재개(resume)의 3중 결정론

47,684 스텝 학습 중 16시간 만에 호스트 장애로 step 19,627에서 죽었는데, 19,500 체크포인트에서 재개해 잃은 127 스텝을 재학습한 결과 손실·퍼플렉시티 모두 Δ=0.000. 이게 가능한 이유는 세 계층의 결정론이 쌓여서입니다.

1. XLA 컴파일된 TPU 프로그램
2. Grain 반복자 상태 체크포인트
3. Orbax 복원

특히 stage-2에서 grain_worker_count=4일 때 stateless resume이 깨지는 걸 발견하고, Grain의 GrainCheckpointHandler로 데이터 반복자 위치까지 원자적으로 함께 복원하도록 고쳤습니다.

주의사항과 한계

• 홀드아웃 평가 없이는 재현 주장 금지: 학습 손실 오버레이는 필요조건일 뿐 충분조건이 아닙니다.
• 아키텍처 튜닝은 하드웨어 종속: head-dim 256 최적화는 Ironwood MXU(256×256)에 특화된 것입니다. 다른 TPU 세대나 GPU에서는 재검증이 필요합니다.
• Stage 3와 포스트트레이닝은 미완: 이 글은 stage 1(사전학습)과 stage 2(mid-training anneal)까지만 커버합니다. SFT/RL은 레시피만 작성된 상태입니다.
• 데이터 순서 노이즈는 완전히 제거 불가: 최종 손실 오프셋 ~+0.006 nat는 버그가 아니라 프레임워크 간 데이터 순서 차이입니다. 도메인별 퍼플렉시티를 보면 wikitext103은 −0.0005로 동률, wikipedia만 +0.0135로 최대 — 한 도메인만 튀면 버그, 전부 균등하게 오프셋되면 순서 노이즈입니다.

실무 적용을 위한 정리

이 사례 연구에서 가져갈 핵심은 세 가지입니다.

1. "재현"의 정의를 홀드아웃 지표로 못 박으세요. 학습 손실 곡선만 보면 데이터 버그가 성능 향상으로 위장합니다. 스텝 정렬된 체크포인트에서 홀드아웃 C4 손실, 다도메인 퍼플렉시티, lm-eval 스위트를 최소 3개 랜드마크에서 비교하세요.

2. 하드웨어-소프트웨어 공동 설계는 무료 점심이 있는 영역입니다. head-dim을 MXU 크기에 맞추는 것만으로 파라미터·FLOPs 변화 없이 12% 처리량을 얻습니다. 긴 학습을 커밋하기 전에 반드시 어블레이션을 돌리세요.

3. 재개 정확성은 결정론 3계층의 합입니다. XLA 프로그램, 데이터 반복자 상태, 옵티마이저 상태를 원자적으로 저장하지 않으면, 조용히 리샘플링된 데이터로 학습하게 됩니다.

한국 개발 생태계 맥락에서 덧붙이면, 국내에서도 네이버클라우드·KT클라우드·삼성SDS 등이 TPU 유사 가속기와 관리형 학습 환경을 제공하기 시작했습니다. 이런 환경에서 대규모 학습을 돌린다면 "칩 시간(chip-hours)"을 예산 단위로 삼는 습관을 들이세요. 이 글에서 stage-1은 약 77k Ironwood chip-hours(≈3,200 chip-days)를 소모했는데, 칩 시간은 슬라이스 크기와 무관하게 고정이라 예산 산정이 훨씬 정확해집니다.

다음 단계 학습 방향

• MaxText 공식 문서로 JAX/XLA 학습 스택 입문하기
• Grain 데이터 로더의 체크포인트 핸들러 패턴 익히기 (대규모 학습 재현성의 핵심)
• lm-eval-harness로 홀드아웃 평가 파이프라인 직접 구축해보기
• Olmo 3의 stage 3(장문 컨텍스트 적응) 와 포스트트레이닝(SFT/RL) 레시피는 아직 미실행 상태이므로, 직접 돌려보는 것이 좋은 다음 과제입니다

함께 보면 좋은 글

• "AI 엔지니어, 6개월 만에 된다고? 현실적인 커리어 전환 로드맵" — 대규모 학습 스택을 다루는 커리어로 전환하려는 분께
• "암 진단을 혁신하는 AI, Artera의 AWS 기반 확장 가능한 아키텍처 설계" — 실제 프로덕션 AI 시스템의 인프라 설계 사례