KellerJordan/modded-nanogpt

해결하는 문제

언어 모델을 가능한 한 빠르게 훈련하는 과제를 다룹니다. 구체적으로, 8개의 NVIDIA H100 GPU에서 특정 성능 목표(FineWeb 검증 세트에서 3.28 cross-entropy loss)를 달성할 수 있는 가장 빠른 알고리즘을 찾는 것을 목표로 합니다.

작동 방식

이 프로젝트는 훈련 파이프라인을 반복적으로 최적화하는 협업 "스피드런"입니다. NanoGPT 및 llm.c 기반의 PyTorch 트레이너에서 시작하여, 이후 다음과 같은 수십 가지의 아키텍처 및 시스템 최적화를 통합했습니다:

• Optimizers: Muon 및 NorMuon 옵티마이저 구현.
• Architecture: Rotary embeddings, QK-Norm, ReLU², 그리고 임베딩에서 모든 블록으로의 skip connections 사용.
• Precision: Head 및 MLP forward pass에는 FP8을, cross-entropy 계산에는 BF16을 활용.
• Attention: Flash Attention 3, long-short sliding window 패턴 및 YaRN 통합.
• Systems: gradient all-reduce/reduce-scatter 최적화 및 계산과 통신의 오버랩.

대상

극한의 훈련 효율성, LLM 최적화, 그리고 H100 GPU의 하드웨어 활용 한계 돌파에 관심이 있는 AI 연구자, 엔지니어 및 애호가를 위한 프로젝트입니다.

주요 특징

• 대규모 가속화: 동일한 목표 손실에 대해 훈련 시간을 45분(베이스라인)에서 2.2분 미만으로 단축.
• 데이터 효율성: 필요 토큰 수를 100억 개에서 4억 개 미만으로 감소.
• 협업적 진화: 훈련 시간을 몇 초라도 줄이기 위해 사용된 기술의 발전 과정을 기록한 상세한 세계 기록 이력.