핵심 요약

  • 안드레이 카파시는 GPT-2의 1억 2,400만 파라미터 모델을 PyTorch로 구현하고, 공개된 가중치를 불러와 구조가 맞는지 확인한 뒤 처음부터 다시 학습했다.
  • 학습에는 FineWeb-Edu의 100억 토큰 표본을 사용했다. 한 차례 학습한 모델은 HellaSwag 평가에서 공개 GPT-2 모델의 성능을 넘어섰고, 네 차례 학습한 모델은 33.24%를 기록했다.
  • 이 비교에는 한계가 있다. 두 모델의 학습 데이터가 다르고, 평가 문항이 학습 데이터에 포함되지 않았는지도 확신할 수 없다고 카파시는 밝혔다.
  • 실행 속도 개선은 연산량뿐 아니라 GPU 메모리 접근, 연산 결합, 입력 크기 조정에서 나왔다. 카파시는 학습 데이터의 순서와 손실 곡선에 남은 문제도 함께 설명했다.

공개 모델을 먼저 재현한 이유

카파시가 대상으로 삼은 것은 GPT-2 계열에서 가장 작은 1억 2,400만 파라미터 모델이다. 이 모델은 트랜스포머 층 12개, 각 층의 표현 차원 768개, 어텐션 헤드 12개를 사용하며 한 번에 다루는 최대 문맥 길이는 1,024토큰이다. 어텐션은 토큰 사이의 정보를 모으는 연산이다.

그는 처음부터 학습하기에 앞서 공개된 GPT-2 가중치를 자신의 구현에 넣었다. 원본 코드는 TensorFlow로 작성됐지만, Hugging Face가 변환한 가중치를 이용하면 PyTorch에서 구조와 생성 결과를 확인할 수 있다. 직접 작성한 모델에서도 공개 가중치를 불러와 문장을 생성할 수 있다는 점은 구현한 구조가 목표 모델과 맞는지 점검하는 수단이 됐다.

구현 과정에서는 입력 토큰을 표현하는 행렬과 출력 토큰을 예측하는 행렬이 같은 가중치를 공유한다는 세부 사항도 반영했다. 두 행렬은 각각 5만 257개 토큰과 768개 차원에 걸쳐 있어, 공유 여부가 파라미터 수에 크게 영향을 준다. 카파시는 초기 가중치 설정과 층을 거듭하며 값을 더하는 경로의 크기 조정도 공개 코드와 논문을 참고해 구현했다.

학습 설정과 GPU 실행 속도

GPT-2 논문에는 학습 설정의 세부 사항이 충분하지 않아 카파시는 구조가 유사한 GPT-3 논문의 설정도 참고했다. 학습률을 처음에는 높이다가 점차 낮추는 방식, 기울기 크기 제한, 가중치 감쇠 등을 적용했다. 큰 배치를 GPU 메모리에 한 번에 올릴 수 없을 때는 작은 배치 여러 개의 기울기를 모아 업데이트했다. 이때 작은 배치의 손실을 누적 횟수로 나눠야 큰 배치로 한 번 계산했을 때와 같은 기울기를 얻는다는 점을 예제로 확인했다.

속도를 높이는 과정에서는 GPU가 계산할 데이터를 메모리에서 읽고 다시 쓰는 비용이 중요했다. torch.compile은 여러 연산을 결합해 메모리 왕복을 줄였다. 어텐션 계산에는 중간의 큰 행렬을 메모리에 만들지 않는 FlashAttention을 적용했다. 카파시의 측정에서 이 변경은 한 단계의 실행 시간을 약 130밀리초에서 96밀리초로 줄였다.

토큰 사전 크기를 원래의 5만 257개에서 5만 304개로 늘린 실험도 했다. 실제 토큰은 추가하지 않고 계산에 쓰는 행렬 크기만 GPU가 처리하기 좋은 값으로 맞춘 것이다. 사용한 PyTorch 개발 버전에서는 실행 시간이 약 4% 줄었다. 이런 개선을 거쳐 단일 GPU에서 측정한 한 단계의 실행 시간은 초기 약 1,000밀리초에서 약 93밀리초로 내려갔다. 이는 전체 학습 시간을 비교한 수치가 아니라 해당 설정의 단계별 측정이다.

100억 토큰 학습 결과

본격적인 학습에는 FineWeb-Edu의 100억 토큰 표본을 사용했다. 카파시는 GPU 8개로 데이터를 나눠 처리하고 각 GPU가 계산한 기울기를 평균 내도록 구성했다. 한 단계에서 처리하는 전체 토큰 수는 52만 4,288개였다. 초기 실행에서 처리량은 초당 약 150만 토큰이었고, 100억 토큰을 한 번 훑는 데 걸리는 시간은 약 1.7시간으로 예상했다.

결과는 두 방식으로 살폈다. 검증 손실은 학습에 쓰지 않은 FineWeb-Edu 데이터에서 다음 토큰을 예측하는 정도를 보여준다. 이 수치는 공개 GPT-2보다 낮아졌지만, 카파시는 원래 GPT-2가 다른 분포의 데이터로 학습했으므로 공정한 직접 비교는 아니라고 선을 그었다.

별도로 사용한 HellaSwag은 주어진 문맥 뒤에 이어질 네 문장 중 자연스러운 것을 고르는 평가다. 카파시의 평가 방식에서 공개 GPT-2 1억 2,400만 파라미터 모델은 약 29.55%를 기록했다. 무작위 선택의 정확도는 25%다. 새 모델은 100억 토큰을 한 번 학습한 뒤 공개 GPT-2의 점수를 넘어섰다.

더 긴 학습과 비교의 한계

카파시는 같은 100억 토큰을 네 차례, 총 약 400억 토큰만큼 학습하는 야간 실행도 했다. 그 결과 HellaSwag 정확도는 33.24%로 올라갔다. GPT-3의 비슷한 크기 모델이 기록한 정확도에는 가까워졌지만 도달하지는 못했다고 설명했다. 원문에서 비교한 학습량은 GPT-2가 약 1,000억 토큰, GPT-3의 해당 모델이 약 3,000억 토큰이다.

카파시는 적은 토큰으로 얻은 성능 차이의 원인을 확정하지 않았다. FineWeb-Edu 표본의 데이터 분포가 원래 GPT-2 학습 데이터와 다를 수 있고, 데이터 품질 차이가 영향을 줬을 수도 있다. 오래된 HellaSwag 평가의 일부가 학습 데이터에 들어갔을 가능성 역시 배제할 수 없다고 했다. 따라서 이 결과만으로 일반적인 학습 효율 향상을 입증할 수는 없다.

학습 손실 곡선에는 반복적인 흔들림도 나타났다. 카파시는 데이터가 충분히 섞이지 않았을 가능성을 의심했지만 원인을 확인하지 못했다. 특히 여러 차례 학습할 때 데이터 로더가 문서와 데이터 조각을 매번 같은 순서로 제공한다는 문제가 있었다. 이번 실행은 공개 모델과 비교할 만한 결과를 얻었지만, 데이터 순서와 평가 조건까지 해결된 완전한 재현으로 제시되지는 않았다.