8xH100 GPU로 3.28 손실 달성, 훈련 시간 단축
NanoGPT 속도 경쟁(Speedrun)은 8개의 NVIDIA H100 GPU를 사용하여 언어 모델 훈련 속도를 경쟁하는 프로젝트임
Andrej Karpathy의 GPT-2 복제 모델을 기반으로, FineWeb 검증 세트에서 3.28의 손실을 목표로 함
현대적인 아키텍처(Modernized Architecture), Muon 옵티마이저(Muon Optimizer) 등 다양한 기술을 통해 훈련 속도 향상
최신 기록(Current Record)은 2.313분으로, 기존 llm.c GPT-2 복제 모델 대비 획기적인 성능 개선
Muon 옵티마이저(Muon Optimizer)의 혁신
본 프로젝트는 Muon 옵티마이저(Muon Optimizer)를 핵심 기술로 사용하며, 이는 메모리 사용량 감소, 샘플 효율성 향상, 그리고 낮은 오버헤드를 특징으로 한다. 특히, Nesterov 모멘텀(Nesterov Momentum)을 활용하고, Newton-Schulz 반복(Newton-Schulz Iteration)을 통해 직교화(Orthogonalization)를 수행한다. 이러한 접근 방식은 기존 Adam 옵티마이저 대비 1.5배의 샘플 효율성을 제공하며, 2% 미만의 런타임 오버헤드를 가진다.
최신 아키텍처 및 최적화 기법
본 레포지토리는 로터리 임베딩(Rotary Embeddings), QK-Norm, ReLU² 등 최신 아키텍처를 채택하여 성능을 향상시켰다. 또한, FP8 matmul, 비대칭 재조정, 소프트캡 로짓(Softcap Logits)을 사용하여 연산 효율성을 높였다. 특히, 임베딩에서 모든 블록으로의 스킵 연결(Skip Connections)과 블록 3에서 6으로의 연결을 통해 정보 흐름을 개선했다. 이러한 기법들은 훈련 속도 향상에 기여한다.
속도 경쟁을 위한 시스템 최적화
본 프로젝트는 훈련 속도를 극대화하기 위해 다양한 시스템 최적화를 적용했다. 플래시 어텐션 3(Flash Attention 3), 긴-짧은 슬라이딩 윈도우 어텐션 패턴(Long-short Sliding Window Attention Pattern)을 활용하여 연산 효율성을 높였다. 또한, 훈련 배치 시작을 EoS(End of Sequence)에 맞추고, 최대 문서 길이를 설정하는 등 데이터 처리 방식을 최적화했다. 이러한 시스템 최적화는 훈련 시간 단축에 기여한다.
Docker를 활용한 환경 구축
README에 따르면, CUDA 또는 NCCL 버전 호환성 문제를 해결하기 위해 Docker를 활용한 환경 구축을 권장한다. Docker를 사용하면 CUDA, NCCL, CUDNN, Python 버전을 표준화하여 종속성 문제를 줄이고 설정을 간소화할 수 있다. 특히, NVIDIA 드라이버가 이미 설치된 시스템에서 유용하다. 구체적인 구현 코드는 문서에 포함되지 않았다.
