KellerJordan/modded-nanogpt

해결하는 문제

8개의 NVIDIA H100 GPU를 사용하여 특정 성능 목표(FineWeb 검증 세트에서 3.28의 교차 엔트로피 손실)에 도달하는 가장 빠른 알고리즘과 시스템 구성 찾기를 목표로 합니다.

작동 방식

이 프로젝트는 '스피드런' 방식으로, 아키텍처, 최적화, 시스템 수준 기술의 방대한 조합을 반복적으로 적용하고 테스트하여 학습 시간을 단축합니다. 앤드레이 카르파티의 llm.c와 NanoGPT에서 시작하여 45분의 베이스라인에서 3분 미만으로 단축되었습니다. 주요 방법은 다음과 같습니다:

• 아키텍처 변경: Rotary 임베딩, QK-Norm, ReLU², 그리고 임베딩에서 모든 블록에 이어지는 스킵 연결 구현.
• 최적화 기술 혁신: Muon 및 NorMuon 최적화기 사용, 그리고 트윈레일 모멘텀을 갖춘 ANVIL 최적화기 사용.
• 정밀도 및 하드웨어: MLP의 전방/역방향 전파, QKV 프로젝션, 헤드에서 FP8 사용, CUDA 그래프 캡처 및 Flash Attention 3 활용.
• 데이터 및 학습 로직: 배치 시작을 End-of-Sentence (EoS) 토큰과 동기화, 멀티토큰 예측 구현, 배치 크기 스케줄 사용.

대상 사용자

극한의 LLM 학습 효율성, 하드웨어 가속, 최적화 및 아키텍처 설계의 최전선에 관심 있는 연구자 및 엔지니어.

주요 성과

• 8xH100 GPU에서 학습 시간을 45분에서 3분 미만으로 단축.
• 목표 손실에 도달하기 위한 토큰 수를 100억에서 3억 3000만 미만으로 감소.
• Polar Express, 스파스 어텐션 게이트, 해시된 바그램/트리그램 임베딩 테이블 등 다양한 최첨단 기술 구현.