Google, MaxText와 TPU로 Olmo 3 7B 사전 학습 재현 사례 공개

재현 프로젝트 상세 해설: Ai2 OLMo 2 7B를 MaxText/TPU에서 재현

재현에서는 Ai2의 step-0 PyTorch 체크포인트를 Orbax로 변환한 뒤 같은 가중치에서 학습을 시작했다. 변환 후 로짓은 HuggingFace 참조와 KL ≈ 1.5e-3으로 일치했다. stage-1은 약 5.93T 토큰, 1.41M 스텝에 걸쳐 진행됐으며, 홀드아웃 C4 손실 차이는 +0.003, 8개 태스크 정확도 차이는 +0.0002였다. stage-2는 100B 토큰이며, 손실 차이는 +0.006, 정확도 차이는 −0.0023으로 둘 다 오차 범위 내로 간주된다고 한다.

검증의 핵심은 홀드아웃 평가다. 약 0.9M 스텝 이후 MaxText의 학습 손실이 계속 Ai2보다 낮아 우위가 있어 보였지만, 홀드아웃 C4 손실은 동등한 수준을 유지했다. 원인은 Grain 데이터로더의 이중 샤딩이었으며, 데이터 커서가 32배 빠르게 진행되어 코퍼스의 약 37%는 사용되지 않았고 26%는 두 번 이상 사용됐다. 수정은 `grain.sharding.NoSharding()` 한 줄이다. 재개 스텝 감지에도 off-by-one 버그가 있었으며, 두 가지를 모두 수정한 뒤 체크포인트에서 재개했을 때 99개 스텝 전부 Δ = 0.000이 나왔다.

운영 측면에서는 약 1.05M 스텝에서 용량의 4분의 3을 잃었고, 레시피를 바꾸지 않고 4분의 1 크기의 슬라이스로 재개했다. 디바이스당 처리량은 1% 이내로 유지됐다. stage-2는 같은 런처를 Ironwood에서 TPU v5p로 지정하는 것만으로 동작했다. SparseCore로의 집합 통신 오프로드, rematerialization 조정, 최적의 샤딩을 적용해 Ironwood의 MFU는 44.5%에 도달했으며, 이를 통해 계산 예산의 약 3분의 1을 되찾았다고 밝혔다.

하드웨어와의 협업 설계로, 어텐션을 32헤드×head-dim 128에서 16×256으로 바꾸는 실험도 진행했다. 파라미터 수와 FLOPs는 동일했으며, Ironwood의 MXU가 256×256이기 때문에 12.4% 빨라졌고, 손실 곡선은 120B 토큰까지 원래와 일치했다. 다만 이는 부차적인 어블레이션이며, 재현 본체는 원래의 아키텍처를 사용한다.

향후에는 YaRN을 이용한 65k 컨텍스트로의 stage 3와 Tunix를 활용한 SFT 및 GRPO 기반 포스트트레이닝이 예정되어 있다. 레시피, 런처, 변환 도구는 MaxText에 공개되어 있다.