구글, MaxText·TPU로 AI2 OLMo 3 사전학습 재현 성공

ai.euno.news 자동 뉴스 요약

구글이 앨런AI연구소(AI2)의 오픈소스 언어모델 OLMo 3 7B의 사전학습 과정을 자사 TPU용 학습 프레임워크 MaxText로 처음부터 재현하는 데 성공했다고 지난 24일 구글 개발자 블로그를 통해 밝혔다. PyTorch와 GPU로 학습한 원본 모델을 JAX 기반 MaxText와 TPU에서 다시 훈련시켜, 손실 곡선뿐 아니라 별도로 떼어둔 평가 지표에서도 결과가 일치함을 확인했다는 것이다.

구글 팀이 재현 대상으로 OLMo 3를 고른 이유는 세 조건을 동시에 만족하는 모델이 드물기 때문이다. 실제 프로덕션 규모로 학습된 강력한 7B급 모델이면서, AI2가 데이터·코드·설정·체크포인트·로그·평가 결과까지 학습 과정 전체를 거의 그대로 공개했고, 여기에 PyTorch·GPU라는 독립적인 비교 기준까지 갖췄다는 점이다. 이 조합 덕분에 MaxText와 TPU 스택이 단순히 “학습되는 것처럼 보이는” 수준이 아니라 실제로 신뢰할 만한지를 검증할 수 있었다.

재현은 AI2가 공개한 학습 0단계(step-0) PyTorch 가중치를 Orbax 포맷으로 변환해 그 지점부터 시작했다. 변환이 제대로 됐는지부터가 첫 검증 대상이었는데, 변환된 가중치로 순전파를 돌린 결과 허깅페이스 기준 모델과 비교해 KL 발산이 약 1.5e-3, 상위 10개 토큰 중 9개가 일치했다. 구글은 이를 “같은 모델을 다른 프레임워크로 돌렸을 때 나오는 잡음 수준”이라고 설명했다.

이후 약 5.93조 토큰, 141만 스텝에 이르는 1단계 사전학습 전체 구간에서 MaxText 학습 곡선은 AI2가 공개한 손실 곡선을 그대로 따라갔고, 종료 시점에는 두 곡선이 겹쳤다. 2단계인 중간학습(annealing) 애닐링 단계도 함께 재현해 대조에 성공했다. 구글은 이 과정에서 학습률(LR) 스케줄을 AI2가 두 개로 나눠 이어붙인 방식에서 단일 코사인 스케줄로 단순화하고, 데이터 조합도 AI2가 공개한 버전을 그대로 썼는데도 결과가 어긋나지 않았다고 밝혔다.

가중치 상속 여부가 결과를 좌우했는지 확인하기 위한 별도 실험도 진행됐다. MaxText 자체의 무작위 초기화 상태에서 전체 일정의 3.5%에 해당하는 약 5만 스텝만 학습시킨 결과, 손실 곡선이 AI2의 공개 곡선을 근소하게 밑도는 수준으로 따라갔다. 완전한 재현은 아니지만, 일치 여부가 AI2의 초기 가중치를 그대로 물려받았기 때문만은 아니라는 근거로 제시됐다.

데이터 파이프라인은 OLMo-core의 방식을 그대로 따랐다. 문서를 토큰화해 EOS로 이어붙이고, 8192토큰 단위로 겹치지 않게 자른 뒤 고정된 시드로 인덱스를 전역 셔플하고, 32회 넘게 반복되는 n-그램이 있는 인스턴스는 걸러내는 방식이다. MaxText에서는 Grain 기반의 dataset_type=olmo_grain 옵션이 이 과정을 구현한다. 모델 구조 역시 32층, 4096차원 밀집 트랜스포머에 재정렬된 정규화 블록, QK-정규화, 슬라이딩 윈도우와 전역 어텐션을 3:1로 섞은 구성까지 olmo3-7b-pt.yml 설정 파일에 그대로 옮겼다.

다만 구글이 이번에 재현한 것은 OLMo 3의 3단계 커리큘럼 중 1단계 사전학습과 2단계 중간학습까지다. 장문 컨텍스트 적응인 3단계와 Tunix를 활용한 SFT·RL 사후학습은 레시피는 작성했지만 아직 실제로 돌리지는 않았다고 밝혔다.

읽기 목록은 이 브라우저에 저장됩니다.