Modded-NanoGPT: NanoGPT (124M) 품질, 3.25B 토큰

Modded-NanoGPT: PyTorch GPT-2 트레이너 변형

• Andrej Karpathy의 llm.c 레포에서 파생된 버전
• 효율성 증가: 동일한 검증 손실에 도달하기 위해 10B 대신 3.15B 토큰 사용
• 코드 간소화: 860줄 대신 524줄
• 건축 현대화: 로터리 임베딩, RMSNorm, ReLU^2 구현
• 새로운 옵티마이저: Muon - 뉴턴-슐츠에 의해 직교화된 모멘텀 사용

훈련 명령

• 8xA100 또는 8xH100 노드에서 실행
• 8xH100의 경우 45분 이내에 완료
• 명령어:
  1. pip install -r requirements.txt
  2. python data/cached_fineweb10B.py 35
  3. ./run.sh

훈련 결과

• Fineweb 데이터의 3.15B 토큰을 6000단계 동안 훈련 시키면 약 3.275 검증 손실 도달
• 기본 llm.c PyTorch 트레이너는 10B 토큰 훈련 후 3.28 이상의 검증 손실 도달

제안된 옵티마이저 (Muon):

• Adam 대비 메모리 사용량 절반
• 1.5배 빠른 훈련
• <9% 벽시계 오버헤드 (분산을 통해 추가 감소 가능)
• 뉴턴-슐츠 반복 사용: Orthogonalization 방법으로 사용
• 훈련 중 실험적 발견:
  • 네스테로프 모멘텀 사용
  • 특정 quintic 뉴턴-슐츠 반복 사용
  • bfloat16에서 뉴턴-슐츠 반복 실행

기타 변경 사항:

• 텍스트 생성 기능 제거
• GPT-2 논문의 엄격한 재현에서 벗어나 훈련 속도 향상
• 3배 증가된 학습률
• 트라페조이드 학습률 스케쥴로 변경
• 로터리 임베딩 및 ReLU^2 활성화로 전환
• 모든 선형 층의 특수 초기화를 제거하고, 고정된 스칼라 값으로 출력 크기 조정
• Affine 스케일 및 바이어스 파라미터 제거, RMSNorm으로 전환
• AdamW에서 새로운 옵티마이저로 전환, 학습률 워밍업 제거