GPT-2 재현 따라하기 — 카파시의 PyTorch·nanoGPT 학습 전 과정

핵심 메시지

• 2019년 공개된 GPT-2의 1억2400만 파라미터 모델을 PyTorch로 바닥부터 재현하는 과정을 단계별로 보여준다.
• 허깅페이스 Transformers로 원본 GPT-2 가중치를 불러와 '목표(target)'로 삼고, 같은 구조를 직접 구현해 따라잡는다.
• 5년 전에는 큰 작업이었지만 이제는 클라우드 GPU로 약 1시간, 10달러 수준이면 같은 수준의 모델을 학습할 수 있다.
• GPT-2 논문은 학습 세부가 부족해 더 구체적인 GPT-3 논문의 하이퍼파라미터와 최적화 설정을 함께 참고한다.
• torch.compile과 플래시 어텐션 같은 최적화로 학습 속도를 크게 끌어올리고, 최종적으로 GPT-2 124M의 검증 손실을 따라잡거나 넘어선다.

쉽게 이해하기

이 영상은 카파시의 'Zero to Hero' 시리즈의 연장으로, 2019년 OpenAI가 블로그·논문·깃허브 코드와 함께 공개한 GPT-2 중 1억2400만 파라미터 버전을 재현한다. GPT-2는 124M부터 1558M까지 여러 크기의 미니시리즈로 공개됐고, 모델 크기를 키울수록 번역·요약·질의응답 같은 다운스트림 지표가 좋아지는 스케일링 법칙을 보여준다.

재현의 출발점은 '목표를 먼저 보는 것'이다. 원본 코드는 텐서플로로 작성돼 있어 다루기 번거롭기 때문에, 허깅페이스 Transformers가 PyTorch로 변환해 둔 GPT-2 구현에서 가중치(state dict)를 불러온다. 토큰 임베딩은 50257 어휘 × 768 차원, 위치 임베딩은 최대 1024 위치 × 768 차원이며, 학습된 위치 임베딩에서 사인·코사인 형태의 구조가 나타나는 것도 시각화로 확인한다.

GPT-2 124M은 트랜스포머 12개 층과 768 차원으로 이루어진다. 카파시는 이전 영상('GPT를 처음부터 만들기')에서 다룬 개념을 전제로 하면서, 학습이 끝나면 검증 데이터에서 다음 토큰을 잘 예측하게 되는 과정을 손실 곡선으로 보여준다. 토큰들은 직전 토큰화 시리즈에서 다룬 GPT-2 토크나이저로 만들어진 것과 정확히 같다.

학습 속도를 높이는 핵심은 GPU 메모리 계층을 이해하는 것이다. torch.compile은 연산을 융합(kernel fusion)해 데이터를 칩 위에 올려둔 채 여러 계산을 한 번에 처리하고 글로벌 메모리 왕복을 줄인다. 한편 torch.compile이 찾지 못하는 최적화로 2022년 스탠퍼드의 플래시 어텐션을 도입하는데, 거대한 어텐션 행렬을 메모리에 실체화하지 않는 알고리즘 재작성으로 더 많은 연산을 하면서도 약 7.6배 빠를 수 있다.

영상은 이 모든 것을 nanoGPT 저장소로 정리하고, 더 나아가 순수 CUDA로 작성한 llm.c 구현과 나란히 실행해 동일한 손실·노름 값을 내면서도 더 빠르게 도는 모습을 비교한다. 사전학습에서 멈추지만, 대화형 데이터로 파인튜닝하면 챗봇 형태로 확장할 수 있다는 점도 짧게 언급한다.

주요 인사이트

• 재현의 정석은 먼저 공개된 가중치를 불러와 '목표'로 고정한 뒤, 같은 구조를 직접 구현해 손실을 맞춰 가는 것이다. OpenAI는 GPT-2 가중치를 공개했기에 이런 비교가 가능하다.
• GPT-2 논문만으로는 학습 세부가 부족하므로, 구조가 크게 다르지 않으면서 하이퍼파라미터가 더 구체적인 GPT-3 논문을 함께 참고하는 것이 실용적이다.
• 학습 성능은 연산량(FLOPs)보다 고대역폭 메모리(HBM)로의 읽기·쓰기 횟수에 좌우된다. 플래시 어텐션이 FLOPs가 더 많은데도 빠른 이유가 여기에 있다.
• 플래시 어텐션의 토대인 '온라인 소프트맥스' 트릭은 2018년 엔비디아 논문에서 이미 제안됐다. 메모리 접근을 줄이면 소프트맥스가 빨라진다는 가설이 4년 뒤 실제로 입증된 셈이다.
• PyTorch 기반 nanoGPT는 가독성 좋은 참조 코드이고, 같은 결과를 내는 llm.c의 CUDA 구현은 시작도 단계당 처리도 더 빠르다. 범용성과 속도의 절충을 보여준다.

자주 묻는 질문

관련 AI 소식