-
llm.c를 사용하여 GPT-2 (124M) 모델을 90분 안에 $20로 재현하는 방법 설명
- GPT-2 (124M)은 OpenAI가 2019년에 발표한 가장 작은 모델
- Lambda에서 8X A100 80GB SXM 노드를 사용하면 시간당 약 $14, 총 비용은 약 $20
- 단일 GPU로도 훈련 가능하지만 시간이 더 오래 걸림 (4-24시간)
결과 비교
- FineWeb 검증 데이터셋에서 OpenAI가 공개한 체크포인트보다 좋은 성능을 보임
- 단, GPT-2는 WebText로 학습했기에 완전히 공정한 비교는 아님
- HellaSwag 정확도도 측정했는데, GPT-3 Small(124M)의 33.7에 근접한 29.9 달성
- GPT-2(124M)의 29.4는 이미 능가함
- 단, 여기서는 10B 토큰으로 학습했고, GPT-3는 300B 토큰으로 학습함
최소 환경 설정
- GPU 필요 (Lambda labs 추천)
- Linux x86 64bit Ubuntu 22.04 with CUDA 12 기준 가이드
- miniconda 설치 후 PyTorch nightly 버전 설치 (선택사항)
- tokenizer를 위해 필요한 패키지 설치
- 속도를 위해 cuDNN 설치 (선택사항)
- 여러 GPU 사용 시 MPI 설치 (선택사항)
- FineWeb 10B 토큰 데이터셋 전처리 (~1시간 소요)
- llm.c 컴파일 (혼합 정밀도, cuDNN FlashAttention 사용)
학습 실행
- 단일 GPU 사용 시 예시 커맨드
- 멀티 GPU (8개) 사용 시 mpirun으로 실행
- 주요 인자 설명
- -i, -j : 학습/검증 데이터 경로
- -o : 로그, 체크포인트 저장 경로
- -e : 모델 초기화 (depth 12 GPT-2)
- -b : 마이크로배치 크기 (메모리 부족 시 줄이기)
- -t : 최대 시퀀스 길이
- -d : 총 배치 크기 (GPT-3 논문 참조)
- -r : Recompute 설정 (메모리 절약)
- -z : ZeRO-1 (옵티마이저 상태 sharding)
- 기타 weight decay, 학습률, 체크포인트 주기 등 설정
학습 과정
- 10B 학습 토큰, 0.5M 배치 크기 기준 약 20K 스텝 예상
- A100 40GB PCIe GPU 기준 스텝별 소요 시간, MFU, 토큰 처리량 출력됨
- 학습 초기에 gradient exploding 현상 있으나 clipping으로 해결
시각화
- 로그 파일 파싱하여 학습 곡선 시각화하는 Jupyter 노트북 제공
Tokenizer
- 정수 토큰을 문자열로 변환하기 위해 필요
- PyTorch 스크립트로 생성 가능
Sampling
- 현재는 추론에 최적화되어 있지 않음
- 약간의 코드 수정으로 unconditional/conditional sampling 가능
코드 구조
- train_gpt2.cu 파일에 대부분의 구현 포함
- 초반 500 줄은 MPI, NCCL, cuDNN, cuBLAS 등 설정
- 이후 1500 줄은 Transformer의 forward/backward
- 다음 1000 줄은 GPT-2 모델 구현
- 마지막 1000 줄은 학습 루프, 인자 파싱 등
350M 모델
- 10B 토큰으로는 부족하고 30B 토큰 사용
- 8X A100 80GB로 14시간 소요, 비용 약 $200
FAQ
-
샘플링 가능 여부: 가능하지만 비효율적임.
-
채팅 가능 여부: 현재는 사전 훈련만 가능, 채팅 미세 조정은 불가.
-
다중 노드 분산 훈련: 가능하지만 아직 테스트되지 않음.
-
비트 단위 결정론적 여부: 거의 결정론적이지만 일부 커널 패치 필요.
-
FP8 훈련 가능 여부: 현재는 BF16으로 훈련 중, FP8은 곧 지원 예정.
-
비 NVIDIA GPU 지원 여부: 현재는 C/CUDA만 지원.
GN⁺의 의견
- GPT-2는 현대적 LLM의 시초로서 매우 중요한 모델임. 이후 GPT-3나 다른 LLM들도 GPT-2와 크게 다르지 않음.
- 이 프로젝트는 GPT-2 수준의 모델을 누구나 합리적인 비용으로 직접 학습시켜볼 수 있게 해줌. LLM에 대한 이해도를 높이는데 큰 도움이 될 것으로 보임.
- 다만 아직 추론에 최적화되어 있지 않아서 실제 서비스에 활용하기에는 제한적임. 대화형 모델로 파인튜닝하는 것도 지원되지 않음.
- 현재는 NVIDIA GPU만 지원하는데, 향후 AMD나 애플 실리콘 등 다양한 플랫폼 지원이 기대됨.
- 유사한 목적의 오픈소스 프로젝트로는 Megatron-LM, DeepSpeed, FairSeq 등이 있음. 각자의 장단점이 있으니 용도에 맞게 선택하면 좋겠음.
- LLM 개발 생태계 활성화 차원에서 매우 고무적인 프로젝트임. 앞으로가 더 기대됨.