100달러 예산으로 GPT-2급 언어 모델을 직접 훈련하고 대화하는 경험
터미널 창에서 방금 훈련한 모델과 이야기를 나누는 장면을 상상해 보세요. 이 저장소는 토큰화부터 추론까지 LLM 학습의 전 과정을 단일 GPU 노드에서 처리하는 최소한의 실험 도구입니다. 8개의 H100 GPU를 약 2시간만 사용하면 GPT-2 수준의 모델을 만들 수 있으며, 비용은 약 48달러(스팟 인스턴스 기준 약 15달러)에 불과합니다.
레이어 수 하나로 결정되는 모델 구조
트랜스포머의 레이어 수인 --depth 값만 지정하면 됩니다. 이 값에 따라 모델의 폭, 헤드 수, 학습률, 학습 기간 등 모든 하이퍼파라미터가 자동으로 계산되어 최적의 상태로 설정됩니다. 사용자는 모델 크기를 조절하는 것만 신경 쓰면 되고, 나머지 세부 설정은 코드 내부에서 처리됩니다. 현재 GPT-2 수준의 성능을 내는 모델은 대략 depth 24에서 26 범위에서 생성됩니다. 이러한 설계는 복잡한 설정 객체나 모델 팩토리를 다룰 필요 없이, 읽기 쉽고 수정하기 쉬운 코드를 기반으로 실험할 수 있게 합니다.
GPT-2 도달 시간 리더보드
개발의 초점은 가장 많은 컴퓨팅 자원을 소모하는 사전 학습 단계의 최적화에 맞춰져 있습니다. 이를 위해 GPT-2 등급의 성능을 달성하는 데 걸리는 실제 시간을 측정하는 리더보드를 운영하고 있습니다. 평가 지표는 DCLM CORE 점수이며, 기준이 되는 GPT-2의 CORE 점수는 0.256525입니다. 2019년 GPT-2 훈련에 약 43,000달러가 들었던 것과 달리, 현재는 8XH100 노드에서 약 1.5시간, 비용으로는 약 48달러(스팟 인스턴스 기준 약 15달러)로 동일한 성능을 달성할 수 있습니다. runs/speedrun.sh 스크립트는 이 기준을 충족하는 모델을 훈련하고 대화할 수 있는 전체 파이프라인을 포함하고 있습니다.
설치 및 실행 방법
의존성 관리는 uv를 사용하며, CUDA 환경에서는 uv sync --extra gpu 명령으로 설치합니다. CPU 또는 Apple Silicon(MPS) 환경에서는 uv sync --extra cpu를 사용합니다. GPU 메모리가 80GB 미만인 경우 --device-batch-size 값을 줄여야 메모리 부족 오류를 피할 수 있습니다. 단일 GPU에서도 실행은 가능하지만, 시간 소요가 8배로 늘어납니다.
연구자와 실험자를 위한 도구
연구 목적이라면 runs/scaling_laws.sh와 runs/miniseries.sh 스크립트가 유용합니다. 빠른 실험을 위해 12층 모델(GPT-1 크기)을 훈련하는 예시도 제공되며, 이는 약 5분 만에 완료됩니다. wandb를 통해 검증 손실(val_bpb), CORE 점수, VRAM 사용률, 모델 FLOPS 활용률(MFU) 등을 모니터링할 수 있습니다. 정밀도(dtype)는 하드웨어에 따라 자동으로 결정되지만, NANOCHAT_DTYPE 환경 변수로 수동으로 지정할 수도 있습니다. CUDA SM 80 이상에서는 bfloat16, 그 미만에서는 float32가 기본값으로 설정됩니다.
실험을 시작하기 전 확인 사항
이 저장소는 $1,000 미만의 예산으로 끝까지 작동하는 마이크로 모델을 개선하는 것을 목표로 합니다. 방대한 설정 옵션을 제공하는 프레임워크가 아니라, 포크하고 수정하기 쉬운 최소한의 코드베이스를 지향합니다. 따라서 특정 기능을 추가하거나 복잡한 파이프라인을 구성하려는 목적보다는, LLM 학습의 기본 원리를 이해하거나 사전 학습 속도를 높이는 연구를 하는 사람에게 적합합니다. GPU 환경이 없다면 CPU나 Apple Silicon에서도 실행은 가능하지만, 모델 크기가 크게 축소되어 의미 있는 결과를 얻기 어렵습니다. 또한, AI를 활용한 코드 기여 시 해당 부분을 명시해야 하는 정책이 있으므로, 기여를 고려한다면 이 점을 유의해야 합니다.
