GPT-2(124M)를 처음부터 다시 만들어 보자

영상 설명란 원문

우리는 GPT-2 (124M)를 맨바닥에서부터 재현한다. 이 영상은 그 전 과정을 다룬다. 먼저 GPT-2 네트워크를 짓고, 그다음 학습이 아주 빠르게 돌아가도록 최적화하고, 그다음 GPT-2 논문과 GPT-3 논문 및 그 하이퍼파라미터를 따라 학습 실행을 세팅한 뒤, 실행 버튼을 누르고, 다음 날 아침에 돌아와 결과를 확인하고, 모델이 만들어낸 재미있는 생성물을 즐긴다. 이 영상은 군데군데에서 Zero to Hero 플레이리스트의 앞선 영상들에서 쌓은 지식 위에 서 있다는 점을 염두에 두기 바란다(내 채널 참고). 이 영상을 내 nanoGPT 저장소를 짓는 과정으로 봐도 좋다. 영상이 끝날 무렵의 코드는 그 저장소와 약 90% 같아진다.

링크:

• 이 영상의 모든 변경분을 개별 커밋으로 담은 build-nanogpt GitHub 저장소
• nanoGPT 저장소
• llm.c 저장소
• 내 웹사이트
• 내 트위터
• 우리 Discord 채널

보충 링크:

• Attention is All You Need 논문
• OpenAI GPT-3 논문
• OpenAI GPT-2 논문
• 내가 이 모델을 학습시키고 있는 GPU는 Lambda GPU Cloud의 것이다. 클라우드에서 ssh로 붙을 수 있는 온디맨드 GPU 인스턴스를 띄우는 방법 중 가장 좋고 가장 쉬운 방법이라고 생각한다.

정정 사항: 모든 정오표와 후속 내용은 위에 링크한 build-nanogpt GitHub 저장소에 올리겠다.

SuperThanks: 어제 실험 삼아 내 채널에 활성화했다. 전적으로 선택 사항이고, 여유가 있을 때만 쓰기 바란다. 수익은 전부 AI + 교육 분야에서의 내 작업을 지원하는 데 쓰인다.

도입: GPT-2 (124M)를 재현해보자 00:00:00

안녕하세요 여러분. 오늘은 Zero to Hero 시리즈를 이어간다. 특히 오늘은 GPT-2 모델, 그중에서도 1억 2,400만 파라미터 버전을 재현할 것이다. OpenAI가 GPT-2를 공개한 건 2019년이었고, 블로그 포스트와 함께 내놨다. 거기에 더해 이 논문도 냈고, 또 GitHub에 코드도 공개했다. openai/gpt-2가 그것이다.

그런데 "GPT-2를 재현한다"고 말할 때는 조심해야 한다. 왜냐하면 특히 이 영상에서 우리가 재현할 대상은 1억 2,400만 파라미터짜리 모델이기 때문이다. 여기서 알아둬야 할 점은, 이런 공개가 이루어질 때는 언제나 미니 시리즈라는 게 함께 나온다는 것이다. 그러니까 GPT-2 미니 시리즈라는 게 있고, 그건 서로 다른 크기의 모델들로 이루어져 있으며, 보통 그중에서 가장 큰 모델을 그냥 "GPT-2"라고 부른다. 그런데 왜 이렇게 여러 크기를 함께 내놓느냐 하면, 기본적으로 이런 식의 그래프에서 x축에 모델 크기를 놓고, y축에는 번역이나 요약, 질의응답처럼 우리가 관심을 갖는 여러 다운스트림 지표를 잔뜩 놓을 수 있기 때문이다. 그렇게 해서 이런 스케일링 법칙을 그려낼 수 있다. 그러니까 기본적으로 모델 크기가 커질수록 다운스트림 지표에서 점점 더 잘하게 된다는 것이다.

그래서 GPT-2의 경우, 논문에서 아래로 스크롤해보면 GPT-2 미니 시리즈에는 네 개의 모델이 있다. 1억 2,400만에서 시작해서 15억 5,800만까지 올라간다. 그런데 내가 말하는 방식의 숫자가 이 표와 어긋나는 이유는, 이 표가 틀렸기 때문이다. 실제로 GPT-2 GitHub 저장소에 가보면, 파라미터를 합산하는 방식에 오류가 있었다는 식으로 이야기하고 있다. 아무튼 기본적으로 이것이 1억 2,400만 파라미터 모델이고, 나머지도 그런 식이다.

그래서 1억 2,400만 파라미터 모델은 Transformer에 레이어가 12개 있었고, Transformer 안의 채널이 768개, 즉 768차원이었다. 그리고 나는 이런 용어들이 무슨 뜻인지 여러분이 어느 정도는 알고 있다고 가정하고 진행하려 한다. 왜냐하면 이전 영상 "let's build GPT from scratch"에서 이 모든 걸 다 다뤘기 때문이다. 이 플레이리스트의 바로 그 이전 영상에서 전부 설명했다.

자, 우리가 모든 걸 제대로 해내고 모든 게 잘 풀린다면, 이 영상이 끝날 무렵에는 이런 그림 같은 것을 보게 될 것이다. 여기서 우리가 보고 있는 건 검증 손실(validation loss)인데, 이건 기본적으로 모델이 학습 중에는 보지 못한 어떤 검증 데이터 위에서, 시퀀스의 다음 토큰을 예측하는 일을 우리가 얼마나 잘하고 있는지를 측정하는 값이다. 그리고 우리는 그 작업을 아주 못하는 상태에서 출발한다. 맨바닥에서부터 초기화하기 때문이다. 거기서 시작해서 학습이 끝날 무렵에는 그 작업을 꽤 잘하는 상태까지 가는 걸 보게 된다. 그리고 바라건대 우리는 GPT-2 124M 모델을 이기게 될 것이다.

예전에 그들이 이 작업을 하고 있었을 때는 — 벌써 5년 전 일이다 — 당시로서는 아마 상당히 복잡한 최적화였을 것이고, GPU와 연산 자원도 지금보다 훨씬 작았다. 그런데 오늘날에는 이 모델을 대략 한 시간, 어쩌면 그보다도 짧은 시간 안에 재현할 수 있고, 클라우드에서 하고 싶다면 비용은 10달러 정도다. 클라우드 컴퓨팅, 그러니까 누구나 빌릴 수 있는 종류의 컴퓨터 말이다. 그 컴퓨터에 10달러를 내고 한 시간이 채 안 되는 시간을 기다리면, OpenAI가 공개한 이 모델만큼 좋은 모델을 실제로 손에 넣을 수 있다.

한 가지 더 언급하자면, 다른 많은 모델들과 달리 OpenAI는 GPT-2의 가중치를 공개했다. 그래서 그 가중치들은 전부 이 저장소에서 구할 수 있다. 하지만 GPT-2 논문이 학습의 모든 세부 사항을 잘 알려주는 건 아니다. 그래서 GPT-2 논문에 더해 우리는 GPT-3 논문도 참조할 것이다. GPT-3 논문은 하이퍼파라미터와 최적화 설정 등 많은 부분에서 훨씬 구체적이고, 아키텍처 면에서도 GPT-2 버전과 크게 벗어나지 않는다. 그래서 GPT-2 124M을 재현하려는 과정에서 GPT-2와 GPT-3을 둘 다 참조할 것이다.

참고로 이 영상의 모든 변경분은 커밋 하나하나로 쪼개서 build-nanogpt 저장소에 올려두었고, 이 영상 자체를 내 nanoGPT 저장소를 짓는 과정으로 봐도 된다. 끝날 무렵의 코드는 그 저장소와 약 90% 같아진다. 자, 시작해보자.

GPT-2 (124M) OpenAI 체크포인트 살펴보기 00:03:39

내가 제일 먼저 하고 싶은 일은 사실 끝, 즉 목표 지점에서 출발하는 것이다. 다시 말해 OpenAI가 공개한 그대로의 GPT-2 124M 모델을 불러와서, 한번 굴려보자는 것이다. 토큰을 좀 샘플링해보자.

그런데 거기에 한 가지 문제가 있다. GPT-2의 코드베이스로 들어가서 소스로 들어간 다음 model.py를 클릭해서 열어보면, 이게 실제로는 TensorFlow를 쓰고 있다는 걸 깨닫게 된다. 그러니까 여기 이 원래의 GPT-2 코드는 TensorFlow로 작성된 것이고, TensorFlow는 뭐랄까, 요즘은 그렇게까지 많이 쓰이지는 않는다고 해두자. 그래서 우리는 PyTorch를 쓰고 싶다. 훨씬 더 친근하고 더 쉽고, 개인적으로도 내가 훨씬 더 좋아하기 때문이다. 문제는 원래 코드가 TensorFlow로 되어 있다는 점이다. 그래서 그 대신, 목표 대상을 손에 넣기 위해 우리는 Hugging Face Transformers 코드를 쓸 것이다. 나는 그쪽이 훨씬 더 마음에 든다.

그러니까 Transformers 소스로 들어가서 transformers/models/gpt2/modeling_gpt2.py로 가보면, 그 파일 안에 GPT-2의 Transformer 구현이 들어 있는 것을 볼 수 있다. 읽을 만한 정도는 중간쯤 되고, 완전히 술술 읽히는 수준은 아니다. 하지만 이 코드가 해준 일은, 그 모든 가중치를 TensorFlow에서 PyTorch에 친화적인 형태로 변환하는 작업을 전부 다 해놨다는 것이다. 그래서 불러오고 다루기가 훨씬 쉽다.

특히 여기서 GPT-2 모델을 살펴볼 수 있고, Hugging Face Transformers를 써서 불러올 수 있다. 이쪽으로 넘어와 보면 이런 모양이다. from transformers import GPT2LMHeadModel 그리고 from_pretrained("gpt2"). 여기서 한 가지 어색한 점은, 불러올 모델 이름으로 gpt2를 주면 이게 실제로는 1억 2,400만 파라미터 모델이라는 것이다. 진짜 GPT-2, 즉 15억 파라미터짜리를 원한다면 -xl을 붙여야 한다. 그러니까 이게 우리 목표인 124M이다.

우리가 지금 하고 있는 일은 이렇다. 이걸 실제로 받아오면, 여기 이 클래스에 정의되어 있는 PyTorch nn.Module을 그것으로부터 초기화하게 된다. 그런데 나는 거기서 state dict만 얻고 싶다. state dict는 그냥 날것의 텐서들일 뿐이다. 그러니까 우리는 그 파일에 들어 있는 텐서들만 손에 쥐고 있는 셈이다.

그리고 참고로 말하자면, 여기 이건 Jupyter 노트북인데 VS Code 안에서 돌아가고 있는 Jupyter 노트북이다. 나는 이 모든 걸 하나의 인터페이스 안에서 다루는 걸 좋아하기 때문에 VS Code를 쓴다. 그러니까 이건 VS Code 안에 들어 있는 Jupyter 노트북 확장이다.

state dict를 받아오고 나면 이건 그냥 딕셔너리다. 그래서 키와 값 — 값은 텐서다 — 을 출력해볼 수 있고, 그 shape들을 한번 들여다보자. 그러니까 이것들이 GPT-2 모델 안에 들어 있는 여러 파라미터들과 그것들의 shape이다. 토큰 임베딩의 가중치 wte는 크기가 50257 × 768이다. 이 숫자가 어디서 나오느냐 하면, GPT-2 어휘 안에 50257개의 토큰이 있기 때문이다. 그리고 참고로 이 토큰들은 내 토크나이제이션 시리즈 영상에서 우리가 이야기했던 바로 그 토큰들이다. 이 영상 바로 직전에 올린 영상들에서 토크나이제이션을 아주 자세하게 파고들었다. GPT-2 토크나이저는 마침 이만큼의 토큰을 갖고 있다.

그리고 각 토큰마다 768차원의 임베딩이 있고, 그것이 그 토큰을 대신해 서 있는 분산 표현(distributed representation)이다. 그러니까 각 토큰은 작은 문자열 조각이고, 그다음에 오는 768개의 숫자가 그 토큰을 나타내는 벡터다. 그래서 이건 그냥 토큰에 대한 우리의 룩업 테이블인 셈이다.

그리고 그다음 여기에는 위치에 대한 룩업 테이블이 있다. GPT-2는 최대 시퀀스 길이가 1024이기 때문에, 각 토큰이 과거 쪽으로 주의를 기울일 수 있는 위치가 최대 1024개까지 있다. 그리고 GPT-2에서는 그 위치들 하나하나가 768차원짜리 고정 벡터를 갖고, 이 벡터는 최적화를 통해 학습된다. 그러니까 이쪽은 위치 임베딩이고 저쪽은 토큰 임베딩이다. 그리고 여기 있는 나머지 것들은 전부 이 Transformer의 다른 가중치와 편향, 그리고 그 밖의 모든 것이다.

예를 들어 위치 임베딩만 가져와서 펼친 다음 앞의 20개 원소만 뽑아보면, 이것들이 그냥 파라미터, 가중치, 실수들이라는 걸 볼 수 있다. 그걸 가져와서 그림으로 그릴 수도 있다. 그래서 이게 위치 임베딩이고, 이런 그림이 나온다. 보다시피 구조가 있다. 구조가 있는 이유는, 이 시각화에서 각 행이 서로 다른 위치, 즉 0부터 1024 범위 안의 고정된 절대 위치이기 때문이다. 그리고 여기 각 행은 그 위치의 표현이다.

구조가 나타나는 이유는, 이 위치 임베딩들이 결국 각 위치를 나타내는 사인·코사인 형태를 학습하게 되기 때문이다. 여기 각 행은 그 위치를 대신하고, Transformer가 그걸 처리해서 모든 상대적 위치를 복원하고, 어느 토큰이 어디에 있는지 파악해서, 내용만이 아니라 위치에 따라서도 주의를 기울이게 된다.

이 안의 개별 열 하나를 들여다보면 — 나는 무작위로 세 개의 열을 골랐다 — 예를 들어 여기서 우리는 각각의 채널 하나하나에 집중하고 있고, 그 채널이 위치의 함수로서, 실제로는 0부터 1023까지의 위치에 대해 무엇을 하고 있는지를 보고 있다. 그리고 이 채널들 중 일부는 위치 스펙트럼의 서로 다른 부분에 더 혹은 덜 반응한다는 걸 볼 수 있다. 이 초록색 채널은 200 이후부터 800까지 모든 것에 대해 강하게 발화하는 걸 아주 좋아하지만 그 바깥에서는 훨씬 덜하고, 0 근처에서는 급격히 떨어진다.

이 임베딩들이 무슨 일을 하고 있고 왜 이런 모양인지는 누가 알겠는가. 다만 예를 들어, 이것들이 좀 더 들쭉날쭉하고 다소 노이즈가 있다는 점에서, 이 모델이 완전히 학습되지는 않았다는 걸 알 수 있다. 이 모델이 더 많이 학습되었을수록 이 곡선은 더 매끄러워질 것으로 기대된다. 그러니까 이건 이 모델이 약간 덜 학습된 모델이라는 걸 말해준다.

하지만 원칙적으로는 이 곡선들이 매끄러워야 할 이유조차 없다. 완전히 무작위 노이즈여도 된다. 실제로 최적화 초기에는 완전한 무작위 노이즈다. 이 위치 임베딩 테이블은 완전히 무작위로 초기화되기 때문이다. 그래서 처음에는 들쭉날쭉하고, 결국 뭔가 매끄러운 것이 나온다는 사실 자체가 이미 꽤 인상적이다. 그게 최적화에서 그냥 떨어져 나온다는 것 말이다. 원칙적으로는 여기서 말이 되는 그래프가 하나라도 나오는 게 오히려 이상한 일인데, 실제로는 약간 노이즈가 있긴 해도 대체로 사인파처럼 보이는 것이 나온다.

원래의 Transformer 논문, 즉 Attention is All You Need 논문에서는 내 기억이 맞다면 위치 임베딩이 서로 다른 주파수의 사인과 코사인으로 초기화되고 고정된다. 그게 positional encoding이고 고정값이다. 그런데 GPT-2에서는 이것들이 그냥 파라미터이고, 다른 파라미터와 똑같이 맨바닥에서부터 학습된다. 그리고 그게 대략 비슷하게 잘 작동하는 것 같다. 그래서 최적화 과정에서 이런 사인파 비슷한 특징을 스스로 되찾아내는 셈이다.

여기 있는 다른 행렬들도 볼 수 있다. 나는 Transformer의 첫 번째 레이어를 가져와서 그 가중치 중 하나를, 그것도 앞쪽 300 × 300 블록만 봤다. 여기서도 어떤 구조가 보이긴 하지만, 역시 이게 뭔지는 누가 알겠는가. 기계적 해석가능성(mechanistic interpretability)에 관심이 있다면, 여기서 무슨 일이 벌어지고 있는지, 이 구조가 무엇이고 무슨 의미인지 알아내려는 시도에서 진짜 재미를 느낄 수도 있다. 하지만 이 영상에서는 그걸 하지 않을 것이다. 다만 흥미로운 구조가 있다는 건 분명히 보이고, 그건 꽤 멋진 일이다.

우리가 주로 관심 있는 건, OpenAI가 공개한 이 모델의 가중치를 불러왔다는 사실이고, 이제 Hugging Face Transformers를 써서 날것의 가중치만 얻는 게 아니라 그들이 pipeline이라고 부르는 것도 얻어서 거기서 샘플링할 수 있다는 점이다. 그래서 프리픽스는 "Hello, I'm a language model," 이고, 토큰 30개를 샘플링하고, 시퀀스 다섯 개를 얻는다. 이걸 돌려봤더니 이런 결과가 나왔다. "…language model, but what I'm really doing is making a human readable document. There are other languages, but those are…" 하는 식이다. 원한다면 죽 읽어봐도 된다. 어쨌든 이것들은 GPT-2 124M이 같은 프리픽스에 대해 내놓은 다섯 가지 서로 다른 완성문이다.

그리고 내가 이 예시를 여기서 가져왔는데, 안타깝게도 시드를 고정했는데도 그쪽에서 나온 것과는 다른 생성물이 나온다. 아마 코드가 바뀐 모양이다. 하지만 이 단계에서 중요한 건, 우리가 일관성 있는(coherent) 텍스트를 얻고 있다는 점이다. 즉 모델을 성공적으로 불러왔고, 모든 파라미터를 들여다볼 수 있으며, 키 이름이 이것들이 모델의 어디에서 왔는지 알려준다는 것이다.

그리고 우리는 우리만의 GPT-2 클래스를 직접 쓰고 싶다. 그래야 거기서 무슨 일이 일어나는지 완전히 이해하게 된다. modeling_gpt2.py 같은 것을 가지고 작업하고 싶지는 않다. 너무 복잡하기 때문이다. 우리는 이걸 맨바닥에서부터 직접 쓰고 싶다. 그래서 여기서 병렬로 GPT 모델을 구현할 것이고, 첫 번째 과제로 GPT-2 124M을 우리가 여기서 맨바닥에서 개발할 클래스 안으로 불러올 것이다. 그러면 OpenAI 모델을 불러올 수 있다는 자신감이 생기고, 따라서 정확히 124M 모델에 해당하는 가중치 설정이 존재한다는 걸 확인하게 된다.

하지만 그다음에 우리가 할 일은, 물론 모델을 맨바닥에서부터 초기화하고, 우리가 구해올 문서 더미 위에서 직접 학습시켜서 그 모델을 뛰어넘으려 시도하는 것이다. 그러면 다른 가중치를 얻게 되고 모든 게 달라 보일 것이다. 바라건대 더 좋아지기까지 할 것이다. 하지만 OpenAI 모델을 불러올 수 있기 때문에, 우리가 같은 모델 계열, 같은 모델 클래스 안에 있고 그저 좋은 가중치 설정을 맨바닥에서 다시 발견하기만 하면 된다는 큰 자신감을 갖게 된다. 그러니 이제 GPT-2 모델을 쓰고, 가중치를 불러오고, 일관성 있어 보이는 텍스트도 생성할 수 있는지 확인해보자.

섹션 1: GPT-2 nn.Module 구현하기 00:13:47

자, 이제 모든 것을 시작한 Attention is All You Need 논문으로 넘어가서, 모델 아키텍처 부분, 즉 원래의 Transformer로 스크롤해보자. 기억할 것은 GPT-2가 원래 Transformer에서 약간 수정되었다는 점이다. 특히 우리에게는 인코더가 없다. GPT-2는 우리가 부르는 이름대로 디코더 전용(decoder-only) Transformer다. 그래서 여기 이 인코더 전체가 빠져 있다. 게다가 그 인코더를 사용하던 여기 이 cross-attention도 빠져 있다. 그래서 이 부분 전체를 지운다. 나머지는 거의 똑같지만, 우리가 여기서 살펴볼 몇 가지 차이가 있다.

주요한 차이가 두 가지 있다. GPT-2 논문의 2.3절 Model로 가보면, 첫째로 LayerNorm들의 자리가 재배치되었다는 걸 알아차리게 된다. 그러니까 위치가 바뀐 것이다. 그리고 둘째로, 여기 마지막 self-attention 블록 뒤에 LayerNorm이 하나 더 추가되었다. 그러니까 기본적으로 여기 있는 모든 LayerNorm이 MLP 뒤나 attention 뒤에 놓이는 대신 그 앞쪽으로 옮겨 붙었고, 거기에 더해 마지막 분류기 바로 앞에 LayerNorm이 하나 더 추가된 것이다.

자, 이제 여기 우리 GPT nn.Module 안에 첫 번째 골격에 해당하는 nn.Module들을 좀 구현해보자. 그리고 특히 우리는 Hugging Face Transformers가 쓰고 있는 여기 이 스키마와 맞추려고 시도할 것이다. 왜냐하면 그렇게 해야 이 state dict에서 가중치를 불러오는 일이 훨씬 쉬워지기 때문이다. 그러니까 우리는 여기 이 스키마를 그대로 반영하는 무언가를 원하는 것이다.

그래서 내가 만든 건 이렇다. 보다시피 모든 모듈을 담고 있는 여기 주 컨테이너의 이름이 transformer다. 그래서 나는 그걸 nn.ModuleDict로 반영했다. 이건 기본적으로 딕셔너리처럼 키, 즉 문자열로 하위 모듈에 인덱싱할 수 있게 해주는 모듈이다. 그 안에 토큰 임베딩의 가중치 wte가 있고 이건 nn.Embedding이다. 그리고 위치 임베딩의 가중치도 있는데 이것 역시 nn.Embedding이다.

기억하겠지만 nn.Embedding은 사실 그냥 숫자 하나의 배열, 단일 숫자 블록 하나를 감싼 멋들어진 작은 래퍼 모듈일 뿐이다. 그건 텐서 하나이고, embedding은 그 텐서를 감싸서 행 인덱싱으로 원소에 접근할 수 있게 해주는 겉멋 든 래퍼다.

거기에 더해, 여기 h가 있는데 이건 문자열이 아니라 숫자로 인덱싱된다. h.0, h.1, h.2 이런 식으로 h.11까지 간다. 이 Transformer에 레이어가 12개 있기 때문이다. 그걸 반영하려고 나도 h를 만들었는데, 아마 hidden의 약자일 것이다. 그리고 ModuleDict 대신 이건 ModuleList다. 그래서 여기 보이는 것처럼 0, 1, 2 같은 정수로 인덱싱할 수 있다. 그리고 이 ModuleList는 n_layer개의 블록을 갖는데, 블록 자체는 잠시 뒤에 모듈로 정의할 것이다.

거기에 더해 GPT-2 논문을 따라, 우리에겐 마지막 LayerNorm이 하나 더 필요하고 그걸 거기 넣을 것이다. 그다음에 최종 분류기, 즉 language model head가 있다. 이건 이 GPT의 임베딩 차원 수인 768에서 어휘 크기인 50257까지 투사한다. 그리고 GPT-2는 이 최종 투사에 편향(bias)을 쓰지 않는다.

이게 골격이고, 보다시피 저 스키마를 그대로 반영한다. wte는 토큰 임베딩이다. 그림에서는 output embedding이라고 불리지만 실제로는 토큰 임베딩이다. wpe는 positional encoding이다. 앞서 봤듯이 이 두 정보는 더해진 다음 Transformer로 들어간다. h는 회색으로 칠해진 블록 전부이고, ln_f는 GPT-2 모델이 여기 추가한 새 레이어이며, lm_head는 여기 이 선형 부분이다. 이게 GPT-2의 골격이고, 이제 블록을 구현해야 한다.

자, 이제 블록 자체로 재귀해 들어가자. 블록을 정의하고 싶으니 여기다 놓기 시작하겠다. 나는 블록을 이런 식으로 쓰는 걸 좋아한다. 이것들이 초기화 부분이고, 이건 이 블록이 계산하는 실제 순전파(forward pass)다.

여기서 다시 한번 Transformer로부터의 변화가 있다는 걸 주목하라. GPT-2 논문에서 언급된 그 변화다. 원래 그림에서는 LayerNorm이 attention 또는 feed-forward를 적용한 뒤에 온다. 게다가 정규화가 잔차 스트림(residual stream) 안쪽에 들어가 있다는 점도 주목하라. feed-forward가 적용되고 이 화살표가 정규화를 통과해서 지나가는 걸 보라. 그 말은 잔차 경로 안에 정규화가 들어 있다는 뜻이고, 이건 그다지 좋지도 바람직하지도 않다.

사실 우리는 감독 신호(supervision)가 있는 꼭대기에서부터 입력인 토큰까지 쭉 내려가는, 깨끗한 잔차 스트림 하나를 갖는 편을 선호한다. 이게 아주 바람직하고 좋은 이유는, 여러분이 micrograd에서 기억한다면 덧셈이 역전파 단계에서 그래디언트를 두 갈래 모두에게 똑같이 분배하기 때문이다. 즉 덧셈은 그래디언트에서 분기점이다. 그래서 꼭대기에서 온 그래디언트가 잔차 경로를 통해 입력인 토큰까지 변형되지 않은 채로 곧장 흘러간다. 하지만 거기에 더해 그래디언트는 블록들을 통해서도 흐르고, 블록들은 시간이 지나면서 자기 몫의 기여를 보태며 개입해 최적화를 바꿔 나간다. 어쨌든 최적화 관점에서 깨끗한 잔차 경로는 바람직하다.

그리고 이건 pre-normalization 버전이다. 보다시피 x가 먼저 LayerNorm을 통과하고, 그다음 attention을 거친 뒤 다시 밖으로 나와서 두 번째 LayerNorm과 다층 퍼셉트론(multi-layer perceptron) — 때로는 feed-forward network 혹은 FFN이라고도 불린다 — 으로 간다. 그리고 그게 다시 잔차 스트림으로 들어간다.

한 가지 더 흥미로운 점은, attention이 통신(communication) 연산이라는 걸 떠올리는 것이다. attention은 1024개의 토큰이 시퀀스로 늘어서 있는 곳에서, 토큰들이 서로 소통하고 정보를 교환하는 지점이다. 그러니까 attention은 집계(aggregation) 함수이고, 풀링 함수이고, 가중합 함수이고, reduce 연산이다. 반면 MLP는 각 토큰마다 개별적으로 일어난다. 토큰들 사이에 수집되거나 교환되는 정보가 전혀 없다. 그래서 attention이 reduce이고 MLP가 map이다. 그렇게 보면 결국 Transformer는 map-reduce의 반복 적용일 뿐이다. 여기가 토큰들이 소통하는 곳이고, 여기가 각자 모아온 정보에 대해 개별적으로 생각하는 곳이다. 그리고 이 블록 하나하나가 잔차 스트림에 있는 표현을 반복적으로 정련한다. 이게 우리의 블록이고, 저 그림에서 약간 수정된 형태다.

자, 이제 MLP로 넘어가자. MLP 블록은 다음과 같이 구현했다. 비교적 단순하다. 기본적으로 선형 투사 두 개가 있고, 그 사이에 GELU 비선형성이 끼워져 있다. nn.GELU(approximate='tanh')이다.

PyTorch 문서로 넘어가 보면 nn.GELU가 있고 이런 형식을 갖는다. 두 가지 버전이 있는데, 잠시 뒤에 들여다볼 원래 버전의 GELU와, 'tanh'를 지정해서 요청할 수 있는 근사(approximate) 버전이다. 여기 미리보기에서 보다시피 GELU는 기본적으로 ReLU와 비슷하지만, 정확히 0에서 완전히 평평한 꼬리가 없다. 그 점만 빼면 약간 더 매끄러운 ReLU처럼 보인다.

이건 여기 이 논문, Gaussian Error Linear Units에서 왔다. 이 논문을 따라가 보면 특정 수식에 이르는 해석으로 이끄는 수학적 논증이 좀 있다. 확률적 정규화기(stochastic regularizer)와 adaptive dropout의 변형에 대한 기댓값과 관련이 있다. 원한다면 거기 있는 내용을 다 읽어봐도 좋다.

그리고 왜 GELU의 근사 버전이 존재하는지에 관한 약간의 역사가 있는데, 내가 알기로 그건 여기 이 이슈에서 나온다. 그 이슈에서 Daniel Hendrycks가, 이 비선형성을 개발하던 당시에는 정확한 GELU를 계산하는 데 필요한 erf 함수가 TensorFlow에서 아주 느렸다고 언급한다. 그래서 결국 이 근사를 개발하게 되었고, 그 근사가 BERT와 GPT-2 등에 채택되었다. 하지만 오늘날에는 근사 버전을 쓸 특별히 좋은 이유가 없다. 그냥 정확한 버전을 쓰는 편이 낫다. 내 예상으로는 이제 큰 차이가 없기 때문이고, 이건 일종의 역사적 잔재 같은 것이다. 다만 우리는 GPT-2를 정확히 재현하려는 것이고 GPT-2가 tanh 근사 버전을 썼으므로, 우리도 그걸 그대로 따르는 편을 택한다.

ReLU 대신 GELU를 직관적으로 쓸 또 하나의 이유가 있다. 과거 영상들에서 우리는 죽은 ReLU 뉴런 문제(dead ReLU neuron problem)에 대해 이야기했다. ReLU의 이 꼬리 부분이 정확히 0에서 완전히 평평하기 때문에, 그 영역에 떨어지는 활성값은 정확히 0인 그래디언트를 받는다. 변화도 없고, 적응도 없고, 네트워크의 발전도 없다. 이 활성값들 중 어느 것이든 이 평평한 영역에 들어가면 그렇게 된다. 반면 GELU는 언제나 국소 그래디언트를 기여한다. 그래서 언제나 변화가 있고 언제나 적응이 있다. 그리고 그렇게 매끄럽게 만드는 것이 경험적으로 실제로 더 잘 작동한다. 이 논문에서 입증되었고, BERT 논문과 GPT-2 논문 등에 채택되었다는 사실 자체로도 입증된다. 그런 이유로 우리는 GPT-2 재현에서 이 tanh 근사 GELU 비선형성을 채택한다.

Llama 3 같은 더 현대적인 네트워크에서는 이 비선형성이 다시 SwiGLU를 비롯한 다른 변형들로 바뀐다. 하지만 GPT-2에서는 이 근사 GELU를 썼다.

마지막으로 attention 연산이 있다. 내가 작성한 attention을 붙여넣겠다. 양이 많다는 걸 알기에, 약간 빠르게, 약간 천천히, 그래도 너무 느리지는 않게 훑어가겠다. 이건 이전 영상에서 다뤘으니 그쪽을 참고하기 바란다.

이게 attention 연산이다. 이전 영상에서 기억하겠지만 이건 그냥 attention이 아니라 다중 헤드 attention(multi-headed attention)이다. 이전 영상에서 우리는 이 multi-headed attention 모듈을 갖고 있었고, 그 구현은 이 헤드들이 사실 그렇게 복잡하지 않다는 걸 분명히 보여줬다. 기본적으로 모든 attention 블록 안에는 여러 개의 헤드가 병렬로 있고, 전부 병렬로 작동하며, 그 출력들이 그냥 이어 붙여져(concatenate) multi-headed attention의 출력이 된다. 그러니까 헤드들은 일종의 병렬 스트림 같은 것이고 출력이 concat된다. 아주 단순했고, 그 덕에 헤드의 구현이 꽤 직관적이었다.

여기서 벌어지는 일은, 별개의 모듈 두 개, 사실은 그보다 훨씬 많은 모듈을 만들어 그걸 concat하는 대신, 그 모든 것을 하나의 self-attention 모듈에 몰아넣었다는 것이다. 그리고 그 대신 나는 아주 조심스럽게 transpose와 split을 잔뜩 동원한 텐서 곡예를 부려서 PyTorch에서 아주 효율적으로 돌아가게 만들었다. 하지만 근본적으로, 알고리즘적으로는 앞서 그 저장소에서 본 구현과 다른 게 전혀 없다.

아주 간단히 상기시키자면 — 여기에 너무 많은 시간을 쓰고 싶지는 않다 — 우리에겐 시퀀스로 늘어선 토큰들이 있고 1024개가 있다. attention의 이 단계에서 각 토큰은 세 개의 벡터를 내놓는다. query, key, value다. 그리고 여기서 먼저 일어나는 일은, query와 key가 서로 곱해져서 attention의 양, 즉 서로를 얼마나 흥미롭게 여기는지를 얻는 것이다. 그러니까 이들은 곱셈적으로 상호작용해야 한다.

그래서 우리가 여기서 하는 일은 qkv를 계산하고, 그걸 split하고, 그다음에 앞서 말한 곡예가 잔뜩 들어간다. 작동 방식은 기본적으로 헤드 개수 nh를 배치 차원으로 만드는 것이다. 그래서 그건 B와 마찬가지로 배치 차원이 되고, 이후의 연산들에서 PyTorch는 B와 nh를 배치로 취급해서 배치와 헤드 양쪽 모두에 대해 모든 연산을 병렬로 적용한다.

적용되는 연산은 이렇다. 첫째, query와 key가 상호작용해서 attention을 준다. 이건 자기회귀 마스크(autoregressive mask)로, 토큰이 자기 앞의 토큰에만 주의를 기울이고 미래의 토큰에는 절대 주의를 기울이지 않도록 보장한다. 여기 softmax는 attention을 정규화해서 합이 항상 1이 되게 한다. 그리고 이전 영상에서 기억하겠지만, attention 행렬을 value와 곱하는 것은 기본적으로 각 토큰마다 우리가 흥미롭다고 판단한 토큰들의 value를 가중합하는 방법이다. 그리고 마지막의 transpose, contiguous, view는 그 모든 걸 다시 재조립하는 것이고, 이것이 실제로 concatenation 연산을 수행한다.

원한다면 이 부분을 천천히 짚어봐도 좋다. 하지만 이건 수학적으로 우리의 이전 구현과 동치이고, 단지 PyTorch에서 더 효율적일 뿐이다. 그래서 나는 이 구현을 택했다.

거기에 더해, 나는 변수 이름을 붙이는 방식에도 신경을 썼다. 예를 들어 c_attn은 저쪽의 c_attn과 같은 이름이다. 그래서 우리 키들이 Hugging Face Transformers 코드의 스키마를 기본적으로 정확히 따라가게 된다. 그러면 모든 변수가 같은 이름으로 불리기 때문에, 바로 이 명명 규칙으로부터 모든 가중치를 그대로 옮겨오기가 아주 쉬워진다.

자, 이 시점에서 우리는 GPT-2 구현을 끝냈다. 그 덕분에 우리는 Hugging Face의 이 파일을 쓸 필요가 없어졌다. 그 파일은 꽤 길어서 2,000줄쯤 되는 코드다. 대신 우리에겐 100줄이 채 안 되는 코드가 있고, 이것이 완전한 GPT-2 구현이다. 그러니 이 단계에서 우리는 모든 가중치를 그대로 가져와 세팅하고 생성까지 해낼 수 있어야 한다. 그게 어떤 모습인지 한번 보자.

Hugging Face의 GPT-2 파라미터 불러오기 00:28:08

여기서 나는 GPTConfig도 함께 바꿔 두었다. 이 안에 들어 있는 숫자들, 즉 하이퍼파라미터가 GPT-2 124M 모델의 것과 일치하도록 맞춘 것이다. 그래서 최대 시퀀스 길이 — 여기서 내가 block size라고 부르는 값 — 는 1024이다. 그리고 토큰의 개수는 50257이다. 내 토크나이저 영상을 본 사람이라면 이 숫자가 어디서 나온 것인지 알고 있을 것이다. 50,000개의 BPE 병합(merge), 256개의 바이트 토큰 — 이것들은 BPE 트리의 잎(leaf)에 해당하는 것들이다 — 그리고 서로 다른 문서를 구분하는 동시에 생성을 시작하는 용도로도 쓰이는 특수 토큰 end-of-text 하나, 이 셋을 합친 값이다. 그리고 레이어는 12개이고, 어텐션 안의 헤드도 12개이며, Transformer의 차원은 768이다.

이제 Hugging Face에서 파라미터를 우리 쪽 코드로 불러와서, 그 파라미터로 우리의 GPT 클래스를 초기화하는 방법을 보자. 코드를 한 뭉텅이 복사해서 붙여 넣겠다. 이 코드는 너무 느리게도, 그렇다고 너무 빠르게도 훑지 않을 생각이다. 솔직히 말해 그다지 흥미로운 코드가 아니고, 그다지 신나는 부분도 아니기 때문이다. 우리는 그저 가중치를 불러오고 있을 뿐이라서 좀 건조한 작업이다. 다만 앞에서 언급했듯이, 이 GPT-2 미니시리즈에는 모델이 네 가지가 있다.

지금 보고 있는 이것은 오른쪽 Jupyter 노트북에 있던 코드 일부를 그대로 옮겨 온 것이다. 여기 있는 이 값들이 GPT-2 모델들의 하이퍼파라미터다. 우리는 이것으로 config 객체를 만들고, 우리 자신의 모델을 만든다. 그다음 여기서 벌어지는 일은, 우리 모델과 Hugging Face 모델 양쪽에 대해 state_dict를 만드는 것이다. 그리고 그다음에 우리가 하는 일은, Hugging Face 모델의 키들을 하나씩 순회하면서 그 텐서들을 우리 쪽으로 복사해 오는 것이다. 그 과정에서 버퍼(buffer) 몇 개는 일종의 무시를 하고 지나간다. 그것들은 파라미터가 아니라 버퍼이기 때문이다. 예를 들어 attn.bias 같은 것은 그저 자기회귀(autoregressive) 마스크에 쓰이는 값이다. 그래서 그런 마스크들 중 일부는 무시한다. 그게 전부다.

그리고 여기에 한 가지 성가신 점이 더 있다. 이 가중치들은 원래 TensorFlow 저장소에서 온 것인데, 왜 이런 식인지는 나도 잘 모르겠다. 약간 귀찮은 일이지만, 일부 가중치가 PyTorch가 원하는 형태에서 전치(transpose)되어 있다. 그래서 나는 전치해야 하는 가중치들을 수동으로 하드코딩해 두었고, 해당하는 경우에만 그것들을 전치해 준다. 그러고 나서 이 모델을 반환한다.

그러니까 from_pretrained는 파이썬에서 생성자(constructor) 역할을 하는 클래스 메서드이고, 모델 타입만 넘겨 주면 GPT 객체를 반환해 준다. 우리 경우에 그 타입은 gpt2, 즉 우리가 관심을 두고 있는 가장 작은 모델이다. 이것이 그 코드이고, 이렇게 쓰면 된다.

이제 우리는 VS Code에서 터미널을 열어 볼 수 있고, 거기서 python train_gpt2.py를 실행해 볼 수 있다. 잘 되기를 빈다. 좋다, 크래시가 나지 않았다. 그러니까 우리는 가중치와 편향(bias), 그리고 그 밖의 나머지 모든 것을 우리 nn.Module 안으로 잘 불러올 수 있는 것이다. 하지만 이제 이것이 정말로 제대로 동작하고 있는지에 대해 확신을 좀 더 얻어 보자. 실제로 이 모델로부터 생성(generate)을 해 보도록 하자.

logits를 얻기 위한 forward pass 구현 00:31:00

좋다. 이제 이 모델로 실제로 생성을 하려면, 그 전에 먼저 forward를 할 수 있어야 한다. 그런데 우리는 아직 그 코드를 쓰지 않았다. 그래서 여기 그 forward 함수가 있다. 이 함수가 하는 일을 처음부터 끝까지 따라가 보자.

forward의 입력은 우리의 인덱스, 즉 우리의 토큰들, 토큰 인덱스가 될 것이다. 그리고 그것들은 항상 B × T 형태다. 그러니까 배치 차원 B가 있고, 그다음에 최대 T까지 가는 시간 차원이 있다. 그리고 이 T는 block size보다 클 수 없다. block size가 바로 최대 시퀀스 길이이기 때문이다. 그래서 B × T 인덱스는 일종의 2차원 레이아웃으로 배열되어 있다. 여기서 기억해 둘 것은, 기본적으로 이 배열의 모든 행(row)은 최대 block size 크기이고, 그것이 하나의 시퀀스 안에 들어 있는 T개의 토큰이라는 점이다. 그리고 그런 독립적인 시퀀스가 B개 배치로 쌓여 있어서, 처리가 효율적으로 이루어진다.

이제 여기서 우리는 위치 임베딩(position embedding)과 토큰 임베딩(token embedding)을 forward한다. 이 코드는 이전 강의를 봤다면 아주 익숙하게 느껴질 것이다. 기본적으로 우리는 torch.arange를 쓴다. 이건 PyTorch용으로 만들어진 range의 한 버전 같은 것이다. 그리고 0부터 T까지 순회하면서 이 위치들, 일종의 위치 인덱스를 만든다. 그다음에는 그것들이 idx와 같은 디바이스 위에 있도록 확실히 해 준다. 왜냐하면 우리는 CPU에서만 학습하지는 않을 것이기 때문이다. CPU만 쓰는 것은 너무 비효율적이다. 우리는 GPU에서 학습하고 싶고, 그 이야기는 잠시 뒤에 나온다.

그다음에 위치 임베딩과 토큰 임베딩이 있고, 그 둘을 더하는 덧셈 연산이 있다. 여기서 주목할 점은, 위치 임베딩이 입력의 모든 행에 대해 동일하다는 것이다. 그래서 이 덧셈 안에는 브로드캐스팅(broadcasting)이 숨어 있다. 여기서 차원을 하나 추가로 만들어 줘야 하고, 그렇게 하면 이 둘이 더해진다. 배치로 쌓여 있는 우리 예제들의 모든 행에 똑같은 위치 임베딩이 적용되기 때문이다.

그다음 우리는 Transformer 블록들을 forward로 통과시키고, 마지막으로 최종 LayerNorm과 LM head를 통과시킨다. 그래서 forward 이후에 나오는 것이 바로 logits다. 입력이 B × T 인덱스였다면, 우리는 모든 B × T 위치마다 이 시퀀스에서 다음에 어떤 토큰이 오는지에 대한 logits를 계산하게 된다. 즉 B번째 행의 t+1번째 토큰, 다시 말해 이 토큰의 바로 오른쪽에 있는 토큰이 무엇인가 하는 것이다. 그리고 여기 vocab size는 가능한 토큰의 개수다. 따라서 우리가 얻게 될 텐서는 이런 모양이 된다. 그리고 이 logits는 softmax 한 번만 거치면 곧바로 확률이 된다. 다시 말해 logits와 확률 사이에는 softmax 하나밖에 놓여 있지 않다.

이것이 바로 이 네트워크의 forward pass다. 이제 우리는 이것을 불러올 수 있게 되었고, 그래서 곧 이 모델로부터 생성을 할 수 있게 될 것이다.

샘플링 초기화, prefix 토큰, 토크나이제이션 00:33:31

좋다. 이제 왼쪽에서도 오른쪽 Hugging Face와 똑같은 것을 세팅해 보겠다. 오른쪽에서는 파이프라인으로 샘플링을 했다. "hello I'm a language model"이라는 prefix를 주고, 최대 30개 토큰까지 다섯 번 샘플링을 했다. 그렇게 얻은 완성 결과들이 여기 있다. 이제 우리는 그것을 왼쪽에서 재현해 보려는 것이다.

그래서 num_return_sequences는 5이고 max_length는 30이다. 우선 당연히 우리가 가장 먼저 하는 일은 모델을 초기화하는 것이다. 그다음 그 모델을 평가 모드(evaluation mode)로 바꾼다. 학습을 하지 않고 그냥 쓰기만 할 때는 모델을 eval로 두는 것이 좋은 습관이다. 다만 사실 나는 지금 이것이 무언가를 하고 있는지 잘 모르겠다. 이유는 이렇다. 위쪽에 있는 우리 모델에는 학습 시점과 평가 시점에 실제로 다르게 동작하는 모듈이나 레이어가 전혀 들어 있지 않다. 예를 들어 Dropout이나 BatchNorm, 그리고 그 밖의 여러 레이어들은 그런 종류의 동작 차이를 갖는다. 하지만 우리가 여기서 사용한 이 레이어들은 학습 시점과 평가 시점에 모두 동일해야 한다. 그래서 잠재적으로 model.eval()은 아무 일도 하지 않는다. 그렇지만 나는 정말로 그런지 확신하지는 못하겠고, 어쩌면 PyTorch 내부에서 평가 모드인지 아닌지에 따라 뭔가 영리한 일을 하고 있을지도 모른다. 그래서 어쨌든 습관적으로 이것을 호출해 두는 것이다.

여기서 그다음에 하는 일은 모델 전체를 CUDA로 옮기는 것이다. 즉 이 모든 텐서를 GPU로 옮긴다. 나는 지금 클라우드 박스에 ssh로 접속해 있고, 이 박스에는 GPU가 여러 개 달려 있다. 여기서 나는 모델 전체와 그 모든 멤버, 그 모든 텐서, 그리고 그 밖의 모든 것을 옮긴다. 모든 것이 사실상 완전히 별개의 컴퓨터, 즉 GPU 위에 앉아 있는 컴퓨터로 실려 간다. GPU는 CPU와 연결되어 있어서 서로 통신할 수 있지만, 기본적으로는 자기만의 컴퓨터 아키텍처를 가진 별개의 컴퓨터다. 그리고 신경망을 돌리는 것과 같은 병렬 처리 작업에 아주 잘 맞춰져 있다. 내가 이렇게 하는 이유는 모델이 GPU라는 별개의 컴퓨터 위에 살게 하기 위해서다. 이 모든 것이 GPU에서 훨씬 효율적으로 돌아가기 때문에, 우리 코드는 훨씬 더 효율적이 될 것이다. 여기까지가 모델 자체에 관한 이야기다. 즉 모델은 이제 GPU 위에 올라가 있고, 그 위에서 돌아갈 준비가 되어 있다.

이제 그다음으로 하고 싶은 것은, 생성을 할 때 이 문자열을 prefix로 삼아서 시작하는 것이다. 그러니 이제 그 prefix 토큰들을 실제로 만들어 보도록 하자. 여기 내가 그것을 위해 써 둔 코드가 있다. 우리는 OpenAI의 tiktoken 라이브러리를 import하고, GPT-2 인코딩을 가져온다. 그것이 바로 GPT-2에서 쓰이는 그 토크나이저다. 그다음 이 문자열을 인코딩해서 정수들의 리스트를 얻는다. 그것이 바로 토큰들이다.

그런데 여기 있는 이 정수들은 사실 꽤 직관적으로 확인할 수 있다. 이 문자열을 그대로 복사해서 tiktokenizer에 붙여 넣고 그것이 무엇인지 들여다보면 되기 때문이다. 그냥 붙여 넣으면 이런 토큰들이 나온다. 그러니까 이 정수 리스트가 바로 우리가 토큰으로 기대하는 값들이다. 그리고 내 영상을 봤다면 기억하겠지만, 모든 토큰은 그저 작은 문자열 덩어리일 뿐이다. 그러니 이것은 이 문자열을 GPT-2 토큰으로 쪼갠 결과다.

이 토큰들, 즉 정수 리스트를 얻고 나면 그것으로 torch 텐서를 만들 수 있다. 이 경우에는 8개 토큰이다. 그다음 우리는 이 8개 토큰을 5번 복제해서, 8개 토큰짜리 행 다섯 개를 얻는다. 그리고 그것이 바로 우리의 초기 입력이고, 여기서 내가 x라고 이름 붙여 부르는 것이다. 그리고 이것 역시 마찬가지로 GPU 위에 살고 있다.

이제 x는 forward에 넣을 수 있는 idx가 되었다. 그래서 우리는 이 다섯 행 각각에서 여섯 번째 토큰 — 미안하다, 아홉 번째 토큰 — 으로 무엇이 오는지 알 수 있는 logits를 얻게 된다.

샘플링 루프 00:37:02

좋다. 이제 우리는 생성할 준비가 되었다. 코드 블록을 하나 더 붙여 넣겠다.

이 코드 블록에서 벌어지는 일은 이렇다. 우리에게는 B × T 크기의 x가 있다. 즉 배치 곱하기 시간이다. 그리고 이 루프를 도는 매 반복마다 우리는 이 행들 각각에 새 인덱스로 이루어진 열(column)을 하나씩 추가해 넣을 것이다. 이것들이 새로운 인덱스들이고, 우리는 샘플링을 해 나가면서 그것들을 시퀀스 뒤에 이어 붙인다. 그래서 루프를 한 번 돌 때마다 x에 열이 하나씩 더 붙게 된다. 즉 시간 차원이 한 칸씩 자라나는 것이다.

그리고 이 모든 연산은 torch.no_grad()라는 컨텍스트 매니저 안에서 일어난다. 이것은 그저 PyTorch에게, 우리가 이 중 어떤 것에 대해서도 backward를 부르지 않을 것이라고 알려 주는 것이다. 그래서 PyTorch는 중간 텐서들을 전부 캐시해 둘 필요가 없고, 나중에 있을지도 모를 backward를 위해 어떤 방식으로도 준비를 해 둘 필요가 없다. 그리고 이렇게 하면 공간을 많이 아낄 수 있고, 어쩌면 시간도 조금 아낄 수 있다.

그래서 우리는 logits를 얻는다. 우리는 오직 마지막 위치의 logits만 취하고, 나머지 logits는 전부 버린다. 그것들은 필요 없다. 우리는 마지막 열의 logits에만 관심이 있기 때문이다. 그러니까 이건 낭비적이긴 하다. 하지만 이건 그냥 좀 비효율적인 샘플링 구현일 뿐이다. 정확하기는 하되 비효율적이다.

그래서 마지막 열의 logits를 얻고, 그것을 softmax에 통과시켜 우리의 확률을 얻는다. 그다음 여기서 나는 top-k 50 샘플링을 한다. 내가 그렇게 하는 이유는 이것이 Hugging Face의 기본값이기 때문이다. Hugging Face 파이프라인의 문서를 그냥 한번 들여다보면, Hugging Face에는 온갖 별난 설정들이 들어가 있고, 솔직히 말하면 그 양이 꽤 많다. 하지만 내가 알아챈 중요한 것 하나는, 그들이 기본적으로 top-k를 쓰고 있고 그 값이 50이라는 사실이다.

그래서 그것이 여기서도 똑같이 쓰이고 있다. top-k가 하는 일은 기본적으로 이렇다. 우리 확률을 가져다가 상위 50개 확률만 남기고 싶은 것이다. 그리고 50번째 확률보다 낮은 것은 무엇이든 그냥 0으로 clamp한 다음 다시 정규화(renormalize)한다. 그렇게 하면 우리는 아주 희귀한 토큰을 샘플링하는 일이 절대 없어진다. 우리가 샘플링하게 될 토큰은 항상 가장 가능성 높은 토큰 상위 50개 안에 들게 된다. 그리고 이것은 모델이 어느 정도 궤도를 벗어나지 않도록 붙잡아 주는 데 도움이 된다. 모델이 주절주절 헛소리를 늘어놓지 않고, 길을 잃지 않으며, 그렇게 쉽게 탈선하지 않게 된다. 그리고 어느 정도는 그럴듯한 토큰들의 근처에 훨씬 더 잘 붙어 있게 만들어 준다. 즉 top-k는 확률 분포의 꼬리를 잘라 내서 모델이 이상한 방향으로 새는 것을 막아 주는 장치인 셈이다.

이것이 PyTorch에서 그것을 하는 방법이고, 원한다면 한 줄씩 짚어 가며 따라가 봐도 좋다. 나는 이 부분이 그렇게 대단한 통찰을 주는 부분은 아니라고 생각해서 빠르게 지나가겠다. 대략적으로 말하면, 우리는 이 새로운 토큰 열을 얻어서 그것을 x에 이어 붙이고, 그러면 기본적으로 x의 열들이 계속 자라다가 이 while 루프가 멈추게 되는 지점에 이른다. 그리고 마침내 이 예제의 경우 5 × 30 크기의 x 전체를 얻게 된다. 그러면 우리는 그 개별 행들을 전부 출력하기만 하면 된다. 그래서 나는 모든 행을 가져오고, 샘플링된 모든 토큰을 가져와서, tiktokenizer의 decode 함수를 써서 문자열을 다시 얻어 출력한다.

터미널을 새로 열고 python train_gpt2.py를 실행해 보자. 좋다, 우리가 얻은 생성 결과들이 이렇다. "hello I'm a language model, not a program" 줄바꿈 줄바꿈 등등. 그리고 "hello I'm a language model, and one of the main things that bothers me when they create languages is how easy it becomes to create something that…" 이런 식으로 계속 주절주절 이어진다. 이 모든 경우가 다 그렇다.

그런데 여기서 한 가지 눈치챌 점이 있다. 이 생성 결과들이 여기 있는 Hugging Face의 생성 결과와 같지 않다는 것이다. 그리고 솔직히 말해 나는 그 불일치가 어디서 오는지 찾아내지 못했다. 이 옵션들을 전부 끝까지 파고들지도 않았다. 아마도 top-p 말고 뭔가 다른 것이 추가로 숨어 있을 것이다. 그래서 나는 결과를 서로 맞추지 못했다.

다만 정확성을 위해, 아래쪽 Jupyter 노트북에서 Hugging Face 모델을 직접 써 봤다. 여기 있는 것이 그 Hugging Face 모델인데, 내가 코드를 똑같이 재현해서 이것을 돌려 보면 동일한 결과가 나온다. 그러니까 기본적으로 모델 내부는 틀리지 않았다. 다만 내가 Hugging Face의 파이프라인이 정확히 무엇을 하는지 100% 확신하지 못할 뿐이고, 그래서 우리가 결과를 맞추지 못한 것이다. 그것 말고는 코드가 정확하고, 우리는 모든 텐서를 올바르게 불러왔다. 즉 모델을 올바르게 초기화했고, 여기 있는 모든 것이 동작한다.

길게 말했지만 요약하면 이렇다. 우리는 가중치를 전부 이식했고 GPT-2를 초기화했다. 이것은 바로 그 OpenAI GPT-2이고, 시퀀스를 생성할 수 있으며, 그 결과가 그럴듯해 보인다. 물론 지금 우리는 GPT-2 모델 가중치로 초기화하고 있다. 하지만 이제 우리가 원하는 것은 처음부터, 난수로 초기화해서, 이 결과들만큼 좋거나 그보다 더 나은 품질의 시퀀스를 내놓는 모델을 실제로 학습시키는 것이다. 그래서 우리가 다음으로 향하는 곳이 바로 그 이야기다.

샘플 생성, 그리고 디바이스 자동 감지 00:41:47

알고 보면 랜덤 모델을 쓰는 것은 사실 꽤 간단하다. PyTorch가 이미 기본적으로 우리 모델을 랜덤하게 초기화해 주기 때문이다. 그래서 우리가 GPT 모델과 그 생성자를 만들 때, 이 모든 레이어와 모듈들은 기본으로 들어 있는 랜덤 초기화 방식을 이미 갖고 있다. 그러니까 이 선형(linear) 레이어들 같은 것이 생성될 때, 예를 들어 우리가 예전에 봤던 Xavier 초기화 같은 기본 생성자를 써서 이 레이어들의 가중치를 만들어 준다. 그래서 GPT-2 모델 대신 랜덤 모델을 만드는 것은 사실 꽤 간단하다.

우리는 그냥 여기로 와서, 대신 model = GPT(GPTConfig())처럼 기본 config를 쓰도록 만들면 된다. 그리고 기본 config는 124M 파라미터를 쓴다. 이것이 바로 랜덤 모델 초기화다. 그리고 이것을 실행하면 우리는 결과를 얻을 수 있게 된다. 물론 여기서 나오는 결과는 완전히 쓰레기 같은 횡설수설이다. 이것이 랜덤 모델이기 때문에 그렇다. 우리는 그저 완전히 무작위로 잘린 랜덤 토큰 문자열 조각들을 받고 있을 뿐이다. 지금 이 시점에 우리가 가진 것이 바로 그것이다.

그리고 한 가지 더 짚고 넘어가고 싶은 것이 있다. 혹시 GPU가 없어서 CUDA를 쓸 수 없더라도, 우리가 여기서 하고 있는 것을 어느 정도까지는 따라올 수 있다는 점이다. 아마 맨 끝까지는 아닐 것이다. 마지막에는 GPU를 여러 대 쓰면서 본격적인 학습을 돌릴 것이기 때문이다. 하지만 적어도 지금까지는 꽤 무난하게 따라올 수 있다. 그래서 GPU가 없는 사람들을 위해 여기서 한 가지 요령을 짚고 넘어가겠다.

그래서 내가 PyTorch에서 즐겨 하는 것 하나는, 사용 가능한 디바이스를 자동으로 감지하는 것이다. 구체적으로는 이렇게 할 수 있다. 여기서 우리는 사용 가능한 것들 중에서 연산 능력(compute capability)이 가장 높은 디바이스를 찾아내려고 한다. 그런 식으로 생각하면 된다. 기본값으로 우리는 CPU에서 시작한다. CPU는 당연히 어디에나 있다. 모든 컴퓨터에는 CPU가 있기 때문이다. 하지만 그다음에 우리는 GPU가 있는지 감지해 볼 수 있다. GPU가 있으면 CUDA를 쓴다. 그리고 CUDA가 없다면, 적어도 MPS는 있는가? MPS는 Apple Silicon용 백엔드다. 그래서 비교적 최근에 나온 MacBook을 갖고 있다면 안에 Apple Silicon이 들어 있을 가능성이 높고, 거기에는 어떤 MacBook이냐에 따라 꽤 쓸 만한 GPU가 들어 있다. 그러니 MPS를 쓸 수 있고, 이것은 CPU보다 잠재적으로 더 빠를 수 있다. 그리고 여기서 디바이스를 출력해 볼 수 있다.

이제 디바이스를 정하고 나면, CUDA 자리에 그것을 그대로 넣어서 쓰면 된다. 그냥 바꿔 끼우기만 하면 된다. 여기서 주목할 점은, 우리가 모델에 x를 넣어 호출할 때 이 x가 GPU가 아니라 CPU 위에 있어도 잘 동작한다는 것이다. 왜냐하면 forward 안에서, 즉 위치 인덱스를 만들게 되는 바로 그 지점에서, 우리가 idx의 디바이스를 써서 이 텐서도 만들도록 주의를 기울여 두었기 때문이다. 그래서 한 텐서는 CPU에 있고 다른 텐서는 GPU에 있어서 둘을 결합할 수 없는 상황은 생기지 않는다. 여기서는 이 모델에 들어온 입력이 지시하는 대로 올바른 디바이스 위에서 조심스럽게 초기화를 하고 있다.

그래서 이것은 디바이스를 자동으로 감지해 준다. 나에게는 당연히 GPU가 잡히니까 "using device: cuda"라고 나온다. 하지만 앞에서 말했듯이 다른 디바이스로도 돌릴 수 있고, 그렇게 심하게 느려지지도 않는다. 여기서 디바이스를 강제로 덮어써 보겠다. 앗, 여기서 device = "cpu"로 덮어쓰면, 물론 여전히 "cuda"라고 출력되기는 하지만 실제로는 CPU를 쓰게 된다. 하나, 둘, 셋, 넷, 다섯, 여섯 — 좋다, 약 6초다. 게다가 우리는 아직 torch.compile 같은 것도 쓰고 있지 않은데, 그런 것들이 모든 것을 훨씬 더 빠르게 만들어 줄 것이다. 어쨌든 CPU에서도 어느 정도까지는 따라올 수 있다고 본다. 그 점에 대한 참고 사항이었다.

그리고 나는 언젠가 한 바퀴 돌아와서, PyTorch에서 디바이스가 다르다는 것이 무슨 의미인지, 그리고 module.to(device) 같은 것을 하거나 torch 텐서를 가져다가 .to(device)를 할 때 PyTorch가 뒤에서 정확히 무엇을 해 주는지, 정확히 무슨 일이 일어나고 그것이 어떻게 동작하는지 다루고 싶다. 하지만 지금은 학습으로 넘어가고 싶다.

학습을 시작하자: 데이터 배치 (B,T) → logits (B,T,C) 00:45:50

이제 나는 모델 학습을 시작하고 싶다. 그리고 지금으로서는 그냥 "디바이스는 코드를 빠르게 만들어 준다" 정도로 해 두자. 그러면 실제로 우리가 이 모델을 어떻게 학습시킬 수 있는지로 들어가 보자.

모델을 학습시키려면 우리에게 데이터셋이 필요하다. 그리고 나에게 디버깅용으로 가장 좋고 가장 단순한 데이터셋은 tiny Shakespeare 데이터셋이다. 이것은 이 URL에서 받을 수 있으니 wget으로 내려받아도 되고, 그냥 "tiny Shakespeare dataset"으로 검색해도 된다. 나는 이미 내 파일 시스템에 input.txt로 갖고 있다. ls를 해 보면 그렇게 보인다. 이미 다운로드해 둔 것이다.

여기서 나는 데이터셋을 읽어서 첫 1,000자를 가져오고, 그중 처음 100자를 출력한다. 기억해 둘 것은, GPT-2는 대략적인 압축률 — 정확히는 토크나이저의 압축률 — 이 대략 3 대 1이라는 점이다. 그래서 1,000자는 대략 300개 토큰 정도가 된다. 지금 우리가 잘라 온 이 슬라이스에서 그 정도가 나올 것이다. 그러니까 이것이 처음 몇 글자다.

그리고 이 데이터에 대해 통계를 좀 더 보고 싶다면, input.txt에 wc(word count)를 돌려 보면 된다. 그러면 이 데이터셋이 약 40,000줄이고, 약 200,000단어이며, 이 파일이 약 100만 바이트라는 것을 알 수 있다. 그리고 이 파일이 ASCII 문자만으로 되어 있다는 것을 알고 있으니 — 내가 아는 한 여기에 이상한 유니코드는 없다 — 모든 ASCII 문자는 1바이트로 인코딩된다. 따라서 이것은 대략 같은 숫자, 즉 이 데이터셋 안에 약 100만 개의 문자가 있다는 뜻이다. 그것이 데이터셋의 크기다. 기본적으로 아주 작고 최소한의 데이터셋이며, 우리를 일단 출발시켜 주기 위한 디버깅용 데이터셋이다. 기본값으로 쓰기에 이 정도면 충분히 작고 다루기 쉽다.

이 데이터셋을 토크나이즈하기 위해, 우리는 GPT-2용 tiktoken 인코딩을 가져와서 데이터의 첫 1,000자를 인코딩하고, 그중 처음 24개 토큰만 출력해 보겠다. 그러면 이것들이 정수 리스트 형태의 토큰이다. 그리고 GPT-2 토큰을 읽을 줄 안다면, 여기 있는 198이 개행(newline) 문자라는 것을 알아볼 것이다. 그러니까 그것이 줄바꿈이다. 그리고 예를 들어 여기에는 줄바꿈이 두 번 있어서 198이 두 번 나온다. 그러니 이것은 그냥 처음 24개 토큰에 대한 토크나이제이션 결과일 뿐이다. 토큰 하나하나가 원문의 어떤 조각에 대응하는지 이렇게 눈으로 확인할 수 있다.

이제 우리가 하고 싶은 것은, 이 토큰 시퀀스를 실제로 가공해서 Transformer에 먹이는 것이다. 그리고 구체적으로는, 이 토큰들을 Transformer에 넣어 줄 idx 변수 형태로 재배열하고 싶다. 우리는 아주 긴 1차원 시퀀스 하나를 원하는 것이 아니다. 우리가 원하는 것은 배치 전체다. 그 안에서 각 시퀀스는 최대 T개의 토큰으로 이루어져 있고, 그리고 이 T는 최대 시퀀스 길이보다 클 수 없다. 그렇게 길이 T짜리 토큰 시퀀스들이 있고, 그런 독립적인 예제가 B개 있는 것이다.

그렇다면 이 1차원 시퀀스들로부터, forward에 넣을 수 있는 B × T 텐서를 어떻게 만들 수 있을까? 이걸 해내는 내가 가장 좋아하는 방법이 여기 있다. torch를 가져다가 이 정수 리스트 — 여기서는 앞의 24개 토큰 — 로 텐서 객체를 만든다. 그리고 내가 가장 좋아하는 방식은 기본적으로 .view()를 하는 것이다. 예를 들어 4 × 6으로 보면 서로 곱해서 24가 되니 딱 맞는다. 그러니까 이것은 그냥 이 토큰들을 2차원으로 재배열한 것일 뿐이다.

그리고 보면 알겠지만, 이 1차원 시퀀스를 여기서 2차원 4 × 6으로 볼 때, 여기까지의 첫 6개 토큰이 첫 번째 행이 되고, 그다음 6개 토큰이 두 번째 행이 되는 식이다. 그러니까 기본적으로 이 경우에는 6개 토큰씩을 독립적인 행으로 쌓아 올려서 토큰들의 배치를 만드는 것이다. 예를 들어 Transformer에 이것을 넣어서 이것이 idx가 되었을 때, 우리가 토큰 25라면 이 토큰은 이 세 개의 토큰을 보게 되고, 다음에 198이 온다고 예측하려 시도할 것이다. 이런 식으로 해서 우리는 이 2차원 배치를 만들어 낼 수 있게 된다. 그건 꽤 괜찮은 방식이다.

이제 손실 함수를 계산하기 위해 필요한 타깃, 즉 라벨은 어떻게 얻을까? 사실 forward pass 안에서 코드를 좀 쓸 수도 있다. 시퀀스에서 다음 토큰, 즉 라벨은 바로 우리 오른쪽에 있다는 것을 알고 있기 때문이다. 하지만 눈치챘겠지만, 맨 끝에 있는 이 토큰 13에 대해서는 우리가 다음 정답 토큰을 갖고 있지 않다. 그것을 로드하지 않았기 때문이다. 그래서 여기서는 정보가 충분하지 않다.

그래서 이 배치들을 얻는, 내가 가장 좋아하는 방법을 보여 주겠다. 나는 개인적으로 Transformer에 들어가는 입력 — 내가 x라고 부르기를 좋아하는 것 — 뿐만 아니라, x와 정확히 같은 크기이면서 모든 위치에 타깃이 담겨 있는 라벨 텐서도 함께 만드는 것을 좋아한다.

그리고 내가 그것을 하는 방식은 이렇다. 나는 토큰을 하나 더, 즉 플러스 1 해서 가져오도록 확실히 해 둔다. 마지막 토큰인 13에 대한 정답(ground truth)이 필요하기 때문이다. 그다음 입력을 만들 때는 마지막 토큰을 포함하지 않고 그 앞까지 전부 가져와서 4 × 6으로 view한다. 그리고 타깃을 만들 때는 같은 버퍼를 쓰되 인덱스 0이 아니라 인덱스 1부터 시작한다. 즉 첫 원소를 건너뛰고, 정확히 같은 크기로 view한다.

그렇게 하고 이것을 출력하면 이런 일이 벌어진다. 예를 들어 이 토큰 25의 타깃은 198이었는데, 그것이 이제 타깃 텐서의 정확히 같은 위치에 198로 저장되어 있는 것을 볼 수 있다. 그리고 이 마지막 토큰 13도 이제 자기 라벨을 갖게 되었고, 그것 역시 198이다. 그리고 그것은 그냥 우리가 여기서 플러스 1을 해서 로드해 두었기 때문에 가능한 일이다. 그러니까 기본적으로 이것이 내가 이 일을 하기 좋아하는 방식이다.

그러니까 기본적으로 내가 하는 방식은 이렇다. 긴 시퀀스를 가져와서 2차원 형태로 view해 배치 곱하기 시간을 얻고, 토큰을 하나 더 로드하도록 확실히 해 둔다. 즉 기본적으로 B × T + 1개의 토큰 버퍼를 로드한 다음, 한 칸 어긋나게 offset해서 view한다. 그러면 텐서 두 개가 생긴다. 하나는 Transformer의 입력이고, 다른 하나는 정확히 그에 대응하는 라벨이다. 두 텐서는 크기가 서로 정확히 같고, 위치만 한 칸 어긋나 있다.

이제 이 코드를 다시 정리해서, 기본적으로 이 토큰들을 로드해 Transformer에 먹이고 손실을 계산하는 일을 해 주는, 아주 간단한 데이터 로더 객체를 하나 만들어 보도록 하자.

좋다, 코드를 그에 맞게 다시 섞어 놓았다. 보다시피 여기서 나는 CPU로 돌리기 위해 임시로 디바이스를 덮어쓰고 있고, tiktoken을 import한다. 그리고 이 부분은 전부 익숙해 보일 것이다. 우리는 1,000자를 로드하고, B와 T를 지금은 그냥 4와 32로 설정한다. 지금은 디버깅 중이라서 아주 작은 배치 하나만 있으면 되기 때문이다. 이 모든 것이 이제 익숙해 보일 것이고, 오른쪽에서 우리가 했던 것을 그대로 따라가고 있다.

그다음 여기서 우리는 모델을 만들고 logits를 얻는다. 그래서 보다시피 나는 이미 이것을 실행해 봤고, 실행에는 몇 초밖에 걸리지 않는다. 그런데 우리 배치가 4 × 32이기 때문에, 우리 logits는 이제 4 × 32 × 50257 크기가 된다. 그러니까 이것들이 모든 위치에서 다음에 무엇이 오는지에 대한 logits인 것이다. 그리고 이제 우리에게는 y 안에 저장되어 있는 라벨도 함께 있다.

그러니 이제는 손실을 계산할 차례이고, 그다음에 backward pass를 하고, 그다음에 최적화를 할 차례다. 그러니 먼저 손실부터 계산해 보도록 하자.

크로스 엔트로피 손실 00:52:53

이제 모델 안에 있는 이 nn.Module의 forward 함수를 손볼 것이다. 특히 우리는 이제 logits만 반환하는 것이 아니라 loss도 함께 반환할 것이고, 입력으로 인덱스만 넘기는 것이 아니라 타깃도 y로 함께 넘길 것이다. 그리고 이제는 더 이상 logits.shape를 출력하지 않는다. 실제로 손실 함수 값을 출력할 것이고, 그다음에 sys.exit(0)을 호출해서 샘플링 로직 부분은 건너뛰도록 하겠다.

그러면 이제 거기서 호출되는 그 forward 함수로 올라가 보자. 왜냐하면 이제 이 함수에는 선택적인 타깃 인자가 생겼고, 우리가 타깃을 받았을 때는 loss도 함께 계산할 수 있기 때문이다. 그리고 기억해야 할 것은, 기본적으로 우리가 logits와 loss를 반환하고 싶어 한다는 것이다. 그리고 loss는 기본값이 None이다. 여기에 이런 코드를 넣자. 만약 targets is not None이라면 우리는 loss를 계산하고 싶다. 그리고 코파일럿이 여기서 벌써 신이 나서, 맞아 보이는 loss 계산을 자동으로 채워 넣고 있다. 이것은 크로스 엔트로피 손실을 쓰고 있는데, 여기 문서에 나와 있는 그대로다. 그러니까 이것은 PyTorch의 functional 아래에 있는 함수다.

그런데 여기서 실제로 무슨 일이 일어나고 있는 걸까? 이게 조금 무섭게 생겼기 때문에 짚고 넘어가자. 기본적으로 F.cross_entropy는 다차원 입력을 좋아하지 않는다. 즉 B × T × vocab_size 모양의 텐서를 그대로 받아 주지 못한다. 그래서 여기서 실제로 벌어지는 일은, 이 3차원 텐서를 그냥 2차원으로 평탄화하는 것이다. 첫 번째 차원은 우리가 직접 적지 않고 자동으로 계산되게 두는데, 그 값은 결국 B × T가 될 것이고, 마지막 차원은 vocab_size가 된다. 그러니까 기본적으로 이것이 하는 일은, logits라는 이 3차원 텐서를 그냥 2차원으로 펴는 것이다. 행의 개수는 B × T, 즉 우리가 가진 개별 예시 전부의 개수이고, 각 행의 길이라는 관점에서 보면 마지막 차원은 vocab_size가 된다. 그리고 동시에 타깃도 같이 평탄화한다. 타깃 역시 이 시점에서는 2차원인데, 우리는 이것도 그냥 쭉 펴서 길이가 B × T인 단일 텐서 하나로 만든다. 이렇게 하고 나면 이것을 크로스 엔트로피 함수에 통과시켜 loss를 계산할 수 있고, 우리는 그 loss를 반환한다.

loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))

이 정도면 그렇게 복잡한 게 아니니까 이 시점에서 기본적으로 그냥 돌아갈 것이다. 실행해 보고 loss가 출력되는지 보자. 그리고 여기서 보면 대략 11 정도가 출력되었다. 그리고 주목할 점은, 이것이 원소가 하나뿐인 텐서이고 그 숫자가 바로 이 11이라는 것이다.

그리고 이제 우리는 무작위로 초기화된 이 네트워크에 대해 합리적인 출발점이 어디쯤이어야 하는지도 계산할 수 있으면 좋겠다. 이건 이전 영상들에서 이미 다뤘던 내용인데, 우리의 어휘 크기는 50257이다. 네트워크를 초기화한 그 시점에서는, 모든 어휘 원소가 대략 균일한 확률을 받기를 바라게 될 것이다. 그래야 초기화 시점에 우리가 어떤 토큰도 지나치게 편애하지 않게 되고, 초기화 시점에 자신 있게 틀리는 상태, 즉 confidently wrong인 상태가 아니게 된다. 그러니까 우리가 기대하는 것은, 임의의 토큰 하나의 확률이 대략 1 / 50257이라는 것이다.

이제 이걸로 loss를 상식 검증(sanity check)할 수 있다. 왜냐하면 기억하겠지만 크로스 엔트로피 손실은 기본적으로 음의 로그 가능도(negative log likelihood)이기 때문이다. 그러니까 이제 이 확률을 가져와서 자연로그를 통과시키고, 거기에 음수를 붙이면, 그게 바로 초기화 시점에 우리가 기대하는 loss다. 이것도 이전 영상들에서 다뤘다. 그래서 나는 −ln(1 / 50257) ≈ 10.82 근처의 값을 기대한다. 그리고 지금 우리가 보고 있는 값은 11 근처다. 그러니까 크게 빗나간 것이 아니다. 이 정도면 초기화 시점에 내가 기대하는 확률에 대체로 부합한다. 그건 나에게 이런 것을 말해 준다. 초기화 시점의 확률 분포가 대체로 넓게 퍼져 있다는 것, 즉 좋은 출발점이라는 것이다.

그리고 이제 우리는 최적화를 수행할 수 있다. 네트워크에게 어떤 원소가 어떤 순서로 따라와야 맞는지를 알려줄 수 있다는 말이다. 이 시점에서 우리는 loss.backward()를 호출해서 그래디언트를 계산하고 최적화 스텝을 밟을 수 있다. 그럼 그걸 하러 가보자.

최적화 루프: 단일 배치 오버핏 00:56:42

자, 그럼 이제 최적화를 해보자. 여기 loss가 있다. 이게 우리가 loss를 얻는 방법이다. 그런데 이제 우리에게는 기본적으로 여기에 for 루프가 필요하다. 그러니까 for i in range(50) 정도로 해서 50 스텝쯤 돌려보자. 그리고 PyTorch에서 옵티마이저 객체를 하나 만들자. 그래서 여기서 우리가 쓰는 것은 Adam 옵티마이저인데, 이것은 이전에 우리가 쓰고 있던 확률적 경사 하강법(SGD) 옵티마이저의 대안이다. 그러니까 SGD는 훨씬 더 단순하고, Adam은 조금 더 복잡하다. 그리고 나는 특히 이 AdamW 변형을 좋아한다. 왜냐하면 내 생각에 그건 그냥 버그를 하나 고친 것에 가깝기 때문이다. 그래서 AdamW는 Adam의 버그 픽스다, 라고 나는 말하고 싶다.

AdamW 문서를 열어 보면, 세상에, 하이퍼파라미터가 잔뜩 들어 있고 앞서 우리가 봤던 SGD보다 조금 더 복잡하다는 것을 알 수 있다. 왜냐하면 파라미터를 학습률로 스케일된 그래디언트로 갱신하는 것에 더해서, 이 옵티마이저는 버퍼들을 계속 들고 다니기 때문이다. 그리고 버퍼는 두 개인데, m과 v이고, 논문에서는 이것을 각각 1차 모멘트(first moment)와 2차 모멘트(second moment)라고 부른다. 그러니까 뭔가 모멘텀처럼 보이는 것 하나와, RMSProp처럼 보이는 것 하나가 있는 셈이다. RMSProp에 익숙하다면 그렇다는 뜻이고, 익숙하지 않아도 전혀 상관없다. 이것은 각 그래디언트 원소마다 개별적으로 일어나는 일종의 정규화일 뿐이고, 최적화를 빠르게 해준다. 특히 언어 모델에서 그렇다. 하지만 여기서 그 자세한 내용까지 들어가지는 않겠다. 우리는 이것을 어느 정도 블랙박스로 취급할 것이다. 이것이 SGD보다 목적 함수를 더 빨리 최적화해준다는 것, 그게 이전 강의들에서 우리가 이미 본 것이다. 그러니 우리 경우에도 그냥 블랙박스로 쓰자. 옵티마이저 객체를 만들고, 그다음 최적화 루프를 돈다.

가장 먼저, 항상 확인해야 할 것은 코파일럿이 그래디언트를 0으로 만드는 것을 빠뜨리지 않았는지다. 그러니까 항상 기억하라. 여러분은 반드시 그래디언트를 0으로 놓고 시작해야 한다. 그다음에 loss를 얻고 거기에 .backward()를 호출하는데, 이 .backward()는 그래디언트에 더한다. 그래디언트를 쌓아 놓는다(deposit). 그러니까 기존 그래디언트가 무엇이든 거기에 항상 +=를 한다. 바로 그렇기 때문에 여러분이 그래디언트를 반드시 0으로 설정해야만 하는 것이다. 그러니까 optimizer.zero_grad()로 0을 만든 다음, 이번 loss에서 나온 그래디언트를 여기에 누적하고, 그다음 옵티마이저의 step 함수를 호출해서 파라미터를 갱신하고 loss를 낮춘다. 그리고 그다음 스텝 번호와 loss를 출력한다.

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
for i in range(50):
    optimizer.zero_grad()
    logits, loss = model(x, y)
    loss.backward()
    optimizer.step()
    print(f"step {i}, loss: {loss.item()}")

여기서 loss.item()을 쓰는 이유는, loss가 원소 하나짜리 텐서이기 때문이다. .item()은 그것을 실제로 단일 float으로 변환해주고, 이 float은 CPU 위에 존재하게 된다. 이건 다시 디바이스 내부 동작과 관련된 얘기인데, loss는 원소 하나짜리 텐서이고 나에게는 GPU 위에 올라가 있다. 내가 GPU를 쓰고 있기 때문이다. 그래서 .item()을 호출하면 PyTorch가 뒤에서 그 1차원 텐서를 가져다가 CPU 메모리로 돌려보내고, 우리가 그냥 출력할 수 있는 float으로 변환해준다.

이게 최적화 루프이고, 이건 아마 그냥 잘 돌 것이다. 실제로 무슨 일이 일어나는지 보자. 아, 미안하다. CPU 오버라이드를 쓰는 대신 그 부분을 지우자. 그러면 나에게는 좀 더 빠르고, CUDA 위에서 돌게 된다. 어라, "expected all tensors to be on the same device but found at least two devices, cuda:0 and cpu"라는 에러가 나온다. 여기서 cuda:0은 0번째 GPU를 뜻한다. 나는 이 박스에 실제로 GPU가 8개 있으니까, 내 박스의 0번 GPU라는 뜻이다. 그리고 다른 하나는 CPU다. 모델은 우리가 .to(device)로 옮겼는데, 이 코드를 쓰면서 내가 사실 버그를 하나 집어넣었다. buf를 디바이스로 옮긴 적이 없다는 것이다. 그리고 여기서 조심해야 한다. 왜냐하면 그냥 buf.to(device)라고만 쓰면 안 되기 때문이다. 그건 상태를 바꾸는(stateful) 연산이 아니다. 그것이 buf 자체를 디바이스 위에 있는 것으로 변환해주지 않는다. 대신 디바이스 위에 있는 새 메모리를 가리키는 포인터를 반환한다. 그러니까 model.to(device)처럼 그냥 호출만 하면 되는 방식은 모듈에는 통하지만 텐서에는 적용되지 않는다는 것을 알 수 있다. 텐서에 대해서는 buf = buf.to(device)라고 써야 한다. 그리고 이렇게 하면 동작할 것이다.

자, 그럼 우리가 여기서 기대하는 것은 무엇인가? 우리가 기대하는 것은, 처음에는 합리적인 loss가 나오고, 그다음에 우리가 이 단일 배치 하나만 계속해서 최적화한다는 것이다. 그래서 우리는 이 단일 배치를 오버핏할 수 있는지 보고 싶다. 우리는 이 작은 배치를 완전히 뭉개 버릴 수 있고, 이 작은 배치 위에서만큼은 인덱스를 완벽하게 예측할 수 있기를 바란다. 그리고 실제로 지금 우리가 여기서 보고 있는 게 대략 그 모습이다. 우리는 대략 10.82에서, 이 경우에는 11에서 출발했고, 그다음 새로운 예시를 하나도 로드하지 않은 채 이 단일 배치 위에서 계속 최적화해 나가면서, 우리가 단일 배치를 오버핏할 수 있다는 것을 확인하고 있다. 그리고 우리는 아주 아주 낮은 loss까지 내려가고 있다. 즉 Transformer가 이 개별 배치 하나를 통째로 암기하고 있는 것이다.

그리고 내가 언급하지 않은 게 하나 더 있는데, 여기서 학습률은 3e-4다. 이 값은 아주 초기의 디버깅 단계에서 돌리고 싶은 대부분의 최적화에 대해 꽤 괜찮은 기본값이다. 이것이 우리의 단순한 내부 루프이고, 우리는 단일 배치를 오버핏하고 있으며, 이 정도면 괜찮아 보인다.

그럼 다음에 오는 것은 무엇인가? 우리는 단일 배치 하나만 오버핏하고 싶은 게 아니라 실제로 제대로 된 최적화를 하고 싶다. 그러니까 우리는 실제로 이 x, y 배치들을 하나씩 순회해야 하고, 우리가 항상 새로운 배치를 받아오도록 보장해 주는 작은 데이터 로더를 하나 만들어서, 우리가 실제로 합리적인 목적 함수를 최적화하고 있도록 만들어야 한다. 그럼 다음으로 그걸 해보자.

데이터 로더 라이트 01:02:00

좋다, 내가 만들어 본 것은 이렇다. 나는 작은 데이터 로더 라이트(data loader lite)를 하나 썼다. 그래서 이 데이터 로더가 하는 일은, 위쪽에서 토크나이저를 임포트하고, 이 하나뿐인 input.txt 텍스트 파일 전체를 읽어서 그것을 토크나이즈한 다음, 전체 토큰 수와 이 데이터셋을 한 번 순회하는 한 에폭에 들어 있는 배치의 개수를 그냥 출력하는 것이다. 즉, 우리가 문서의 처음으로 되감아서 그것을 다시 읽기 시작하기 전까지 몇 개의 서로 다른 배치를 내보낼 수 있는가를 출력한다.

그래서 우리는 위치 0에서 시작해서, 그다음 문서를 B × T 크기의 배치 단위로 아주 단순하게 걸어 나간다. 즉 B × T만큼의 덩어리를 가져오고, 그다음 항상 B × T만큼 전진한다. 그리고 여기서 중요하게 짚어둘 점이 하나 있다. 우리는 위치를 정확히 B × T만큼 전진시키지만, 실제로 토큰을 가져올 때는 현재 위치부터 B × T + 1까지 가져온다는 것이다. 그리고 우리에게 이 +1이 필요한 이유는, 기억하겠지만 현재 배치의 마지막 토큰에 대한 타깃 토큰이 필요하기 때문이다. 그래야 우리가 앞에서 했던 것과 정확히 똑같은 방식으로 x와 y를 만들 수 있다. 그리고 만약 데이터가 바닥나게 되면 그냥 0으로 되감으면 된다. 이것이 아주 아주 단순한 데이터 로더를 쓰는 한 가지 방법이다. 그냥 파일을 덩어리 단위로 훑고 지나가는 것이고, 지금 우리 목적에는 이 정도면 충분히 좋다. 이건 나중에 더 복잡하게 만들 것이다.

이제 다시 여기로 돌아와서, 실제로 우리 데이터 로더를 써보고 싶다. import tiktoken은 위로 올라갔고, 사실 이 부분은 이제 전부 쓸모가 없어졌다. 그래서 대신 우리는 학습 데이터용 train_loader만 있으면 된다. 그리고 우리는 앞에서와 같은 하이퍼파라미터를 쓰고 싶다. 배치 크기 B는 4였고, 시간 차원 T는 32였다. 그리고 여기서 우리는 현재 배치에 대한 x, y를 얻어야 한다. 이 정도면 충분히 간단하니 코파일럿이 알아서 맞히는지 한번 보자. 그러니까 우리는 next_batch를 호출하고, 그다음 우리 텐서들을 CPU에서 디바이스로 옮겨야 한다.

여기서, 내가 토큰을 변환할 때 이 토큰들을 실제로 GPU로 옮기지는 않았다는 점을 눈여겨보라. 나는 이것들을 기본값인 CPU에 그냥 두었다. 그리고 그건 단지 이 경우에 내가 GPU 메모리를 너무 낭비하지 않으려고 했기 때문이다. 이건 아주 작은 데이터셋이라 GPU에 올려도 들어가긴 한다. 그러니 지금 우리 목적에서는 그냥 GPU로 보내도 괜찮다. 어쨌든 우리는 다음 배치를 받아오고, 데이터 로더는 단순한 CPU 클래스로 유지하고, 그다음 여기서 실제로 그것을 GPU로 보내서 모든 연산을 한다.

이게 돌아가는지 보자. python train_gpt2.py를 실행한다. 그런데 이게 실제로 실행되기 전에, 우리가 기대하는 것은 무엇인가? 우리가 기대하는 것은 이렇다. 이제는 실제로 매번 다음 배치를 가져오고 있으므로, 우리가 단일 배치를 오버핏하지는 않을 것이라고 예상한다. 그래서 loss가 내려가긴 하지만 그렇게 많이 내려가지는 않을 것이라고 예상한다. 그래도 loss가 내려갈 것이라고 내가 기대하는 이유는, 50257개의 토큰 중에 우리 데이터셋에 한 번도 나타나지 않는 토큰이 아주 많기 때문이다. 그래서 최적화 과정에서 아주 쉽게 얻을 수 있는 이득이 여기에 있다. 예를 들어 한 번도 나오지 않는 모든 logit의 bias를 가져다가 음의 무한대로 몰아버리는 것이다. 기본적으로 그것뿐이다. 온갖 이상한 유니코드나 다른 언어들로 된 그 토큰들은 절대 나오지 않으므로, 그것들의 확률은 아주 낮아야 한다. 그러니까 우리가 보게 될 이득은 기본적으로 한 번도 등장하지 않는 토큰의 사용을 지워버리는 방향의 것이다. 지금 이 규모에서 우리가 보게 될 loss 감소의 대부분은 아마 바로 그것이다. 하지만 0까지 내려가지는 않을 것이다. 왜냐하면 우리는 50번만 반복하고 있고, 지금으로서는 그게 한 에폭을 다 돌기에 충분하지 않다고 나는 보기 때문이다.

그럼 우리가 무엇을 얻었는지 결과를 보자. 토큰이 338,000개 나왔다. 그리고 이것은 우리의 3:1 압축 비율과 잘 맞아떨어진다. 문자가 100만 개쯤 되기 때문이다. 그러니까 지금의 B와 T 설정에서 한 에폭을 도는 데는 2,600 배치가 걸린다. 그런데 우리는 여기서 최적화를 50 배치만 하고 있다. 그래서 우리는 예상대로 익숙한 지점에서 출발하고, 그다음 대략 6.6 정도까지 내려오는 것으로 보인다. 그러니까 기본적으로 지금까지는 우리 기대에 비추어 볼 때 모든 게 잘 돌아가고 있는 것 같다. 좋다.

파라미터 공유 — wte와 lm_head 01:06:14

다음으로 나는 우리 코드에 있는 버그를 하나 고치고 싶다. 큰 버그는 아니지만, GPT-2 학습이 어떻게 이루어져야 하는가 하는 관점에서는 버그다. 버그는 다음과 같다. 우리가 Hugging Face에서 가중치를 로드할 때 충분히 조심하지 않았고, 사실 작은 디테일 하나를 놓쳤다.

여기로 와서 보면, 이 두 텐서의 shape이 같다는 것을 알 수 있다. 여기 이 하나는 Transformer의 맨 아래에 있는 토큰 임베딩이고, 여기 다른 하나는 Transformer의 맨 위에 있는 언어 모델링 헤드다. 그리고 이 둘은 기본적으로 모두 2차원 텐서이고, 그 shape이 서로 동일하다. 그러니까 여기 첫 번째 것은 출력 임베딩, 즉 토큰 임베딩이고, 두 번째 것은 맨 위에 있는 이 선형 레이어, 즉 분류기 레이어다. 그리고 둘 다 50257 × 768 모양이다. 여기 이것은 맨 아래에서 우리에게 토큰 임베딩을 주고, 여기 저것은 Transformer의 768개 채널을 받아서 그것을 50257로 업스케일해 다음 토큰에 대한 logits를 만들려고 한다.

그런데 단지 shape만 같은 게 아니다. 실제로 그 원소들을 서로 비교해 보면 — PyTorch에서 이것은 원소별 동등 비교인데 — 여기에 .all()을 쓰면 모든 원소 하나하나가 동일하다는 것을 볼 수 있다. 그리고 그것보다 더 나아가서, 실제로 data_ptr을 들여다보면, 이것은 PyTorch에서 실제 데이터와 그 스토리지에 대한 포인터를 얻는 방법인데, 그 포인터까지 동일하다는 것을 알 수 있다. 그러니까 이것들은 우연히 같은 shape과 같은 원소를 가진 두 개의 별개 텐서인 것이 아니라, 실제로 하나의 동일한 텐서를 가리키고 있는 것이다.

여기서 벌어지고 있는 일은 흔히 쓰이는 가중치 결합(weight tying) 기법이다. 이것은 사실 원래의 Attention is All You Need 논문에서 왔고, 사실 그보다 더 앞선 참고문헌에서 왔다. 그 논문의 "Embeddings and Softmax" 절로 가보면, 자기들 모델에서는 두 임베딩 레이어와 소프트맥스 직전의 선형 변환 사이에 같은 가중치 행렬을 공유한다고, 그리고 그것이 참고문헌 30번과 비슷하다고 언급한다. 이것은 "이 둘이 공유되고 서로 묶여 있으며 같은 행렬이다"라는 말을 꽤 어색하게 표현한 문장이다. 그리고 그 30번 참고문헌이 바로 이 논문이다. 이 논문은 2017년에 나왔고 전문을 읽어볼 수 있는데, 기본적으로 이 가중치 결합 기법을 옹호하는 논문이다.

직관적으로 왜 이렇게 하고 싶은지에 대한 아이디어는 여기 이 문단에서 나온다고 나는 생각한다. 기본적으로 이렇게 관찰할 수 있다. 우리는 이 두 행렬이 다음과 같은 의미에서 서로 비슷하게 행동하기를 원한다. 만약 두 토큰이 의미적으로 아주 비슷하다면 — 예를 들어 하나는 전부 소문자이고 다른 하나는 전부 대문자이거나, 또는 다른 언어로 된 같은 토큰이거나 하는 식으로 — 즉 두 토큰 사이에 어떤 유사성이 있다면, 짐작건대 우리는 그 둘이 토큰 임베딩 공간에서 서로 가까이 있기를 기대할 것이다. 그런데 정확히 같은 방식으로, 의미적으로 비슷한 두 토큰이 있다면 우리는 Transformer의 출력에서도 그 둘이 비슷한 확률을 받기를 기대할 것이다. 그것들이 의미적으로 비슷하기 때문이다. 그러니까 Transformer의 맨 아래와 맨 위, 이 두 위치 모두 "비슷한 토큰은 비슷한 임베딩 또는 비슷한 가중치를 가져야 한다"는 바로 이 성질을 갖는다. 그리고 이것이 여기서 저자들의 탐구를 동기 부여한 것이다.

논문 전체를 다 훑고 싶지는 않으니 직접 읽어 보면 되지만, 그들이 관찰한 것은 이렇다. 그들은 출력 임베딩 쪽을 들여다보면 그것도 단어 임베딩처럼 행동한다는 것을 관찰한다. 그 가중치를 그냥 단어 임베딩으로 써보려고 하면 그렇게 된다는 것이다. 그렇게 그들은 이 유사성을 관찰하고, 둘을 묶어 보고, 그렇게 했을 때 훨씬 더 좋은 성능을 얻을 수 있다는 것을 관찰한다. 그래서 이 기법이 채택되어 Attention is All You Need 논문에 들어갔고, 그다음 GPT-2에서도 다시 쓰였다.

그런데 나는 이 부분을 Hugging Face transformers 구현에서는 못 찾았다. 그들이 어디서 그 임베딩들을 묶는지는 나도 잘 모르겠다. 하지만 OpenAI가 공개한 원래 GPT-2 코드에서는 그것을 찾을 수 있다. 그래서 OpenAI GPT-2 소스의 model 코드를 보면, 그들이 이 모델을 forward하는 부분에서 — 그리고 이건 TensorFlow로 되어 있지만 괜찮다 — 먼저 wte 토큰 임베딩을 가져오고, 그다음 여기가 그 토큰 임베딩과 위치 임베딩을 인코딩하는 곳이라는 것을 알 수 있다. 그리고 여기 아래쪽에서 그들은 logits를 만들기 위해 wte를 다시 쓴다. 그러니까 그들이 logits를 얻을 때, 그것은 Transformer에서 나온 이 출력과 wte 텐서의 행렬곱이고, 그 wte 텐서가 재사용되는 것이다.

그래서 기본적으로 이 wte 텐서는 두 번 쓰인다. Transformer의 맨 아래에서 한 번, 그리고 Transformer의 맨 위에서 한 번 쓰인다. 그리고 역전파 과정에서 우리는 이 두 갈래 모두로부터 그래디언트 기여를 받게 된다. 그리고 이 그래디언트들은 wte 텐서 위에서 서로 더해진다. 그러니까 우리는 분류기 쪽 logits에서 오는 기여를 하나 받고, 그다음 Transformer의 맨 끝에 가서는 그 맨 아래쪽에서 다시 wte 텐서 안으로 흘러들어오는 기여를 또 하나 받는 것이다.

우리 코드는 현재 wte를 공유하고 있지 않은데, 우리도 그렇게 하고 싶다. 가중치 공유 기법이다. 이걸 구현하는 한 가지 방법이 있는데, 코파일럿이 맞히는지 보자. 오, 맞혔다. 이게 한 가지 방법이고, 기본적으로 상당히 직관적이다. 여기서 우리가 하는 일은 wte.weight를 가져와서, 그것이 lm_head를 가리키도록 단순히 재지정하는 것이다.

self.transformer.wte.weight = self.lm_head.weight

이렇게 하면 기본적으로 데이터 포인터가 복사된다. 즉 참조를 복사하는 것이다. 그러면 이제 wte.weight의 옛 값은 고아가 되어 버리고, PyTorch가 — 정확히는 Python이 — 그것을 알아서 정리해 줄 것이다. 그래서 우리에게는 텐서가 딱 하나만 남게 되고, 그 하나가 forward 패스에서 두 번 쓰이게 된다. 그리고 내가 아는 한, 필요한 것은 이게 전부다. 그러니 우리는 이대로 쓸 수 있을 것이고, 아마 학습도 잘 될 것이다. 우리는 기본적으로 정확히 같은 이 텐서를 두 번 쓰게 된다. 우리가 여기서 가능도를 꼼꼼히 추적하고 있었던 건 아니지만, 논문과 그 결과에 따르면 이렇게 했을 때 실제로 약간 더 나은 결과를 기대할 수 있다.

그리고 거기에 더해서, 이것이 우리에게 아주 아주 좋은 또 하나의 이유가 있다. 이건 어마어마한 양의 파라미터라는 점이다. 여기 이것의 크기가 얼마인가? 768 × 50257이다. 즉 4000만(40M) 파라미터다. 그리고 우리가 다루는 이것은 1억 2400만(124M) 파라미터 모델이다. 그러니까 40을 124로 나눠 보면, 이 가중치 결합 기법을 써서 대략 30%의 파라미터가 절약되고 있는 셈이다. 그래서 이것이, 이 기법이 약간 더 잘 동작하는 이유 중 하나일 수도 있다. 모델을 충분히 오래 학습시키지 않는 경우라면, 가중치 결합 덕분에 학습해야 할 파라미터의 수가 그만큼 줄어들고, 그래서 학습 과정이라는 측면에서 더 효율적이 된다. 파라미터가 더 적기 때문이고, 또 우리가 "이 두 임베딩은 토큰들 사이의 유사성을 공유해야 한다"는 귀납 편향(inductive bias)을 집어넣고 있기 때문이기도 하다. 이것이 바로 가중치 결합 기법이다. 우리는 엄청난 수의 파라미터를 절약했고, 이 기법 덕분에 우리 모델이 약간 더 잘 동작하기를 기대한다.

모델 초기화 — std 0.02와 residual 초기화 01:13:47

다음으로 나는 우리가 초기화에 대해 좀 더 신중해지기를, 그리고 GPT-2가 자기 모델을 초기화한 방식을 따라가 보기를 바란다. 그런데 불행히도 GPT-2 논문과 GPT-3 논문 모두 초기화에 대해 그다지 명시적이지 않다. 그래서 우리는 어느 정도 행간을 읽어야 한다. 그리고 상당히 모호한 논문으로 가는 대신, OpenAI가 공개한 코드에 정보가 조금 들어 있다.

그래서 우리가 그 model.py로 가보면, 그들이 자기들 가중치를 초기화할 때 표준편차 0.02를 쓰고 있다는 것을 알 수 있다. 그러니까 이것은 가중치에 대한 정규분포이고, 그 표준편차는 0.02다. 그리고 bias에 대해서는 그것을 0으로 초기화한다. 그리고 여기서 아래로 스크롤해 보면 — 왜 이게 스크롤이 안 되지 — 토큰 임베딩은 0.02로 초기화하고, 위치 임베딩은 어째서인지 0.01로 초기화한다는 것을 알 수 있다. 그러니까 그것들이 그들의 초기화이고, 우리는 우리 모듈 안에서 이것을 그대로 따라 하고 싶은 것이다.

여기 내가 꽤 빠르게 만들어 본 코드 조각이 있다. 여기서 벌어지는 일은, 우리 GPT 모듈의 생성자 끝에서 nn.Module의 apply 함수를 호출하는 것이다. 그리고 이 apply 함수는 이 모듈의 모든 하위 모듈을 순회하면서 각각에 대해 가중치 초기화 함수를 적용해 준다. 그러니까 여기서 벌어지는 일은, 우리가 이 모든 모듈을 순회하고 있다는 것이다. 그리고 만약 그것이 nn.Linear 모듈이라면, 우리는 그 가중치를 표준편차 0.02의 정규분포를 써서 확실히 초기화하도록 한다. 만약 그 레이어에 bias가 있다면, 그것은 확실히 0으로 초기화하도록 한다. 참고로 bias를 0으로 초기화하는 것은 사실 PyTorch의 기본값이 아니다. 기본적으로 여기 bias는 균등분포로 초기화된다. 그러니 그건 흥미로운 지점이다. 그래서 우리는 확실히 0을 쓴다. 그리고 임베딩에 대해서는 그냥 0.02를 쓰고 그대로 유지할 것이다. 우리는 위치 임베딩을 0.01로 바꾸지는 않을 것인데, 어차피 거의 비슷한 값이기 때문이다.

그리고 우리 모델을 쭉 훑어보면, 초기화가 필요하면서 파라미터를 가지고 있는 다른 레이어는 LayerNorm뿐이라는 것을 알 수 있다. 그런데 PyTorch의 기본 초기화는 LayerNorm의 스케일을 1로, 그리고 LayerNorm의 오프셋을 0으로 설정한다. 그리고 그게 정확히 우리가 원하는 것이다. 그래서 우리는 그것을 그냥 그대로 두겠다. 그러니까 이것이, 우리가 OpenAI가 공개한 그 GPT-2 소스 코드를 따라갈 때의 기본 초기화다.

그런데 여기서 한 가지 짚어두고 싶은 것이 있다. 보통 이 초기화에서의 표준편차는, 만약 여러분이 Xavier 초기화를 따른다면 이 레이어로 들어오는 입력 피처 수의 제곱근의 역수, 즉 1 / √d가 될 것이다. 그런데 잘 보면 0.02라는 값은 사실 기본적으로 그것과 일관된다. 왜냐하면 GPT-2에서 이 Transformer들 내부의 모델 크기는 대략 768, 1600 같은 값들이기 때문이다. 그래서 예를 들어 1 / √768을 계산하면 우리에게 0.03이 나온다. 만약 1600을 대입하면 0.02가 나온다. 만약 그것의 세 배를 대입하면 0.014 같은 값이 나온다. 그러니까 기본적으로 0.02는 이런 초기화들에 대해 합리적인 값들의 근방에 있는 셈이다. 그래서 여기에 0.02를 하드코딩하는 것이 완전히 미친 짓은 아니다. 다만 보통은 모델 크기에 따라 같이 커지거나 작아지는 값을 쓰고 싶어 할 것이다. 하지만 우리는 이것을 이대로 유지할 것이다. 소스 코드상으로 확인되는 GPT-2의 초기화가 바로 그것이기 때문이다.

그런데 초기화가 아직 완전히 끝난 것은 아니다. 여기에 단서 조항이 하나 더 있기 때문이다. GPT-2 논문은 이렇게 말한다. 모델 깊이에 따른 residual 경로 위의 누적을 고려한 수정된 초기화를 사용한다고, 그래서 초기화 시점에 residual 레이어의 가중치를 1 / √N이라는 인자만큼 스케일한다고. 여기서 N은 residual 레이어의 개수다. 이것이 GPT-2 논문이 말하는 바다. 우리는 아직 이것을 구현하지 않았고, 이제 구현할 수 있다.

그런데 그들이 여기서 무슨 말을 하는 건지 조금 동기 부여를 해보고 싶다. 대략 이런 이야기라고 생각한다. residual 스트림을 0에서 시작한다고 하자. 기억하겠지만 각 residual 스트림은 x = x + (무언가) 형태다. 어떤 종류의 기여가 계속 더해지는 것이다. 그러니까 residual 네트워크의 모든 블록이 각각 어떤 양만큼 기여하고 그것이 더해진다. 그래서 결국 벌어지는 일은, residual 스트림에 있는 활성값들의 분산이 커진다는 것이다.

여기 작은 예시가 하나 있다. 0에서 시작해서 100번 반복한다고 하자. 즉 768개의 0으로 이루어진 일종의 residual 스트림이 있고, 거기에 100번에 걸쳐 무작위 값을 더한다. 이 무작위 값은 평균이 0이고 표준편차가 1인 정규분포에서 나온 값이다. 이것을 계속 더해 나가면, 끝에 가서는 이 residual 스트림의 표준편차가 10까지 커져 있다. 그리고 그건 그냥 우리가 이 숫자들을 계속해서 더하고 있기 때문이다. 그리고 그들이 여기서 쓰는 이 스케일링 인자가 바로 그 증가를 정확히 상쇄해 준다. 그러니까 우리가 N을 가져와서, residual 스트림으로 들어가는 이 기여들 하나하나를 1 / √N만큼 줄여 주면 된다. 즉 1 / √N은 N^−0.5다. N의 0.5제곱이 곧 제곱근이고, 그 제곱근의 역수가 N의 −0.5제곱이기 때문이다. 그리고 이런 식으로 스케일해 주면, 우리는 실제로 표준편차가 1이 되는 것을 볼 수 있다. 그러니까 이것은 forward 패스에서 residual 스트림 내부의 활성값이 커지는 것을 제어하는 하나의 방법이다.

그래서 우리도 같은 방식으로 초기화하고 싶다. 각 블록의 끝에 있는 이 가중치들, 즉 이 c_proj 레이어의 가중치를 GPT 논문은 residual 레이어 개수의 제곱근의 역수만큼 줄이라고 제안한다.

그래서 이것을 구현하는 조잡한 방법 하나는 다음과 같다. 이게 PyTorch가 공인한 방식인지는 모르겠지만, 나에게는 잘 동작한다. 우리는 초기화 부분에서 이렇게 한다. 해당 모듈에 NANOGPT_SCALE_INIT = 1을 설정한다. 그러니까 우리는 그 모듈에 대해 일종의 플래그를 달아 놓는 것이다. PyTorch에 이보다 더 나은 방법이 분명히 있을 텐데, 나는 그걸 모르겠다. 어쨌든 우리는 기본적으로 이 플래그를 붙이고, 그것이 기존의 어떤 것과도 충돌하지 않도록 확실히 이름을 정한다.

그다음에 아래로 내려오면, 여기서 std는 기본적으로 0.02여야 하지만, 만약 hasattr로 이 모듈에 그 플래그가 달려 있는지 확인해 보고 달려 있다면 std에 곱셈을 해준다. 코파일럿이 여기서는 제대로 못 맞히고 있는데, 우리가 원하는 것은 레이어 개수의 제곱근의 역수다. 그러니까 여기서 residual 레이어의 개수는 2 * self.config.n_layer이고, 여기에 −0.5제곱을 해준다. 즉 우리는 그 표준편차를 그만큼 줄이고 싶은 것이다. 그리고 이렇게 하면 맞을 것이고, 그것이 이 부분의 구현이다.

참고로 분명히 해두자면, 여기서 레이어 수에 2를 곱하는 이유는, Transformer의 모든 레이어 하나하나가 실제로 residual 경로에 더해지는 블록을 두 개씩 갖고 있기 때문이다. 어텐션이 있고 그다음 MLP가 있다. 그래서 2배가 나오는 것이다.

그리고 언급해 둘 것이 하나 더 있다. 조금 어색하긴 하지만 우리가 고치지는 않을 부분인데, 우리가 wte와 lm_head를 가중치 공유하고 있기 때문에, 하위 모듈들을 순회하는 이 과정에서 사실 우리가 그 텐서를 두 번 만나게 된다는 점이다. 그러니까 우리는 먼저 그것을 임베딩으로서 0.02로 초기화하고, 그다음 다시 한 바퀴 돌아와서 이번에는 선형 레이어로서 또 0.02로 초기화하게 된다. 그리고 두 번째도 0.02가 되는 이유는, lm_head는 당연히 스케일 대상이 아니기 때문이다. 그러니까 그것은 스케일 분기로는 들어가지 않고, 기본적으로 동일한 초기화를 두 번 적용받게 되는 것뿐이다. 하지만 그래도 괜찮다.

그리고 여기로 스크롤해 보면, 내가 재현성을 확보하기 위해 시드를 설정하는 코드를 여기에 조금 추가해 두었다. 그래서 이제 우리는 python train_gpt2.py를 실행하고 이게 그대로 돌아가게 둘 수 있다. 그리고 내가 아는 한, 지금 우리가 구현해 놓은 이 방식이 바로 GPT-2의 초기화다. 그러니까 나에게는 이 정도면 합리적으로 보인다.

SECTION 2: 빠르게 만들자 — GPU, 혼합 정밀도, 1000ms 01:22:18

이 시점에서 우리는 GPT-2 모델을 손에 쥐고 있다. 그것이 올바르게 구현되었다는 것에 대해 어느 정도의 확신도 갖고 있고, 초기화도 제대로 해 두었으며, 데이터 배치를 하나씩 순회하는 데이터 로더도 갖췄다. 그래서 학습을 돌릴 수 있는 상태다. 이제 재미있는 부분이 온다. 나는 우리가 학습 속도를 아주 많이 끌어올렸으면 한다. 여기서 우리가 쓰고 있는 하드웨어에 대해 지불한 값어치를 제대로 뽑아내자는 것이다. 그리고 우리는 실제로 학습 속도를 꽤 많이 끌어올릴 것이다.

이런 일은 언제나 내가 가진 하드웨어가 무엇인가, 그것이 무엇을 제공하는가, 그리고 나는 그것을 온전히 활용하고 있는가라는 질문에서 출발해야 한다. 내 경우에는 nvidia-smi를 실행해 보면, GPU가 여덟 장 있고 그 GPU 각각이 A100 SXM 80GB라는 것을 알 수 있다. 그러니까 이것이 이 박스 안에서 내가 쓸 수 있는 GPU다.

그런데 이런 종류의 박스를 띄울 때 내가 제일 좋아하는 곳은, 참고로 말하자면, Lambda GPU Cloud다. 이들은 내 개발과 내 프로젝트들을 후원해 주기도 한다. 하지만 그와 별개로 이곳이 내가 제일 즐겨 찾는 곳이고, 여기서 이런 머신 중 하나를 띄울 수 있다. 시간 단위로 요금을 내면 되고, 아주아주 간단하다. 나는 머신을 하나 띄운 다음 거기에 VS Code를 붙이는 것을 좋아하고, 그런 식으로 개발한다.

자, 여기 올라와 있는 A100들을 보면, A100 80GB SXM이 바로 내가 지금 여기서 쓰고 있는 GPU다. 그리고 이 GPU에서 어느 정도의 계산을 기대할 수 있는지에 대한 숫자들이 여기 잔뜩 적혀 있다. 이제 코드 쪽으로 돌아와서, 여기 바로 다음 지점에 브레이크를 걸어 보자. 그러니까 python train_gpt2.py를 인터랙티브 모드로 돌려서 로짓과 손실을 계산한 바로 다음에 멈추는 것이다. 그리고 여기서 여러분이 주목했으면 하는 흥미로운 점은, logits.dtype을 찍어 보면 이게 torch.float32를 출력한다는 것이다.

그러니까 PyTorch에서 텐서를 만들면 기본값이 이렇다는 뜻이다. 그리고 이건 모든 활성값에 대해서도 그렇고, 신경망의 파라미터에 대해서도 그렇고, 그 밖의 모든 것에 대해서도 그렇다. 기본적으로 모든 것이 float32다. 즉 모든 숫자 하나하나가, 그게 활성값이든 가중치든 그 밖의 무엇이든, 32비트를 가진 부동소수점 표현을 쓰고 있다는 뜻이다. 그리고 그건 사실 꽤 많은 메모리다. 그런데 경험적으로 드러난 사실은, 딥러닝이라는 계산 워크로드에 대해서는 이게 너무 과하다는 것이다. 딥러닝, 그리고 이런 신경망들의 학습은 훨씬 더 낮은 정밀도를 충분히 견뎌낼 수 있다.

모든 계산 워크플로가 작은 정밀도를 견딜 수 있는 것은 아니다. 예를 들어 다시 데이터시트로 돌아가 보면, 이 GPU들이 실제로 FP64까지 지원한다는 것을 볼 수 있다. 내가 이해하기로 이건 많은 과학 계산 응용에서 꽤 유용하고, 그쪽에서는 정말로 그 정밀도가 필요하다. 하지만 우리는 딥러닝 학습을 위해 그만큼의 정밀도가 필요하지 않다.

그래서 현재 우리는 여기 FP32에 있고, 지금 이 상태 그대로의 코드로는 기껏해야 19.5 테라플롭스의 성능을 기대할 수 있다. 그 말은 우리가 초당 19.5조 번의 연산, 즉 부동소수점 연산을 하고 있다는 뜻이다. 여기서 말하는 부동소수점 연산은 십중팔구 부동소수점 곱셈-덧셈(multiply-add)일 것이다. 그러니까 이것들이 바로 부동소수점 연산이다.

이제 주목할 점은, 우리가 정밀도를 낮출 의향이 있다면 — 그러니까 TF32라는 더 낮은 정밀도 포맷이 있고 이건 잠시 뒤에 볼 것이다 — 여기서 실제로 8× 개선을 얻을 수 있다는 것이다. 그리고 float16이나 bfloat16까지 내려갈 의향이 있다면 실제로 16× 성능, 즉 312 테라플롭스까지 쭉 갈 수 있다.

여기서 NVIDIA가 별표가 붙은 숫자를 인용하기 좋아한다는 것이 보인다. 이 별표는 "희소성(sparsity) 적용 시"라는 뜻이다. 하지만 우리 코드에서 희소성을 쓸 일은 없고, 나는 지금 이게 업계에서 아주 널리 쓰이고 있다고도 생각하지 않는다. 그래서 대부분의 사람들은 여기 이 숫자, 즉 희소성 없는 쪽 숫자를 본다. 그리고 여기서 더 큰 숫자도 얻을 수 있었다는 것을 눈치챘겠지만, 그건 INT8이고 INT8은 학습이 아니라 추론에 쓰인다. 왜냐하면 INT8은 기본적으로 간격이 균일하기 때문이다. 그리고 우리는 사실 부동소수점이 필요하다. 신경망 학습 중에 나타나는 정규분포에 더 잘 들어맞아야 하기 때문이다. 학습 중에는 활성값도 가중치도 정규분포를 따라 분포하고, 그래서 그 표현에 들어맞으려면 부동소수점이 정말로 정말로 중요하다. 그러니 우리는 학습에는 보통 INT8을 쓰지 않는다. 하지만 추론에는 쓴다.

그리고 정밀도를 내리면 GPU에 들어 있는 텐서 코어에서 훨씬 더 많은 테라플롭스를 뽑아낼 수 있다. 이 이야기는 잠시 뒤에 하겠다. 그런데 그것에 더해서, 이 숫자들이 전부 더 적은 비트로 표현된다면 그것들을 이리저리 옮기는 일 자체도 훨씬 쉬워진다. 그리고 바로 거기서 우리는 메모리 대역폭, 그리고 모델의 메모리라는 이야기로 들어가게 된다.

우리 GPU가 저장할 수 있는 비트 수에 유한한 용량이 있다는 것만이 문제가 아니다. 그에 더해서 이 메모리에 접근할 수 있는 속도라는 것이 있고, 우리는 일정한 메모리 대역폭을 갖는다. 이건 아주 귀중한 자원이고, 실제로 학습을 위한 딥러닝 워크로드 상당수가 메모리 바운드다. 메모리 바운드라는 말이 뜻하는 바는, 이 극도로 빠른 곱셈을 전부 수행하는 텐서 코어들이 대부분의 시간 동안 그저 기다리며 놀고 있다는 것이다. 우리가 그들에게 데이터를 충분히 빨리 먹여 줄 수가 없기 때문이다. 우리가 메모리에서 데이터를 충분히 빨리 불러올 수가 없는 것이다.

그래서 여러분의 하드웨어의 전형적인 활용률이라는 것이 어떠냐 하면, 여러분이 60% 활용률을 얻고 있다면 여러분은 사실 극도로 잘하고 있는 것이다. 다시 말해 잘 튜닝된 애플리케이션에서조차 절반의 시간 동안은 텐서 코어가 곱셈을 하고 있지 않다. 데이터가 준비되어 있지 않기 때문이다. 그러니 여기서 메모리 대역폭도 마찬가지로 극도로 중요하다. 그리고 우리가 모든 float에 대해, 모든 숫자와 가중치와 활성값에 대해 정밀도를 내리면, 갑자기 그것들이 더 적은 메모리를 요구하게 된다. 그래서 우리는 더 많이 저장할 수 있고 그것에 더 빨리 접근할 수 있다. 그래서 모든 것이 빨라지고, 그건 놀라운 일이다.

그러니 이제 그것의 이득을 거둬들여 보자. 그리고 우선 첫 번째로 TensorFloat-32 포맷부터 한번 보도록 하자.

텐서 코어, 코드 시간 측정, TF32 정밀도, 333ms 01:28:14

좋다, 우선 텐서 코어란 무엇인가? 텐서 코어는 그냥 A100 아키텍처에 들어 있는 하나의 명령어다. 그러니까 이 명령어가 하는 일은 기본적으로 작은 4×4 행렬 곱셈이다. 그러니까 이건 그냥 4×4 행렬들의 행렬 곱셈이다. 그리고 이 행렬들 각각이 어떤 정밀도인지, 내부 누산(accumulate)이 어떤 정밀도로 일어나는지, 그리고 출력 정밀도는 무엇이고 입력 정밀도는 무엇인지 등에 대해 여러 가지 구성이 있다. 그래서 스위치가 몇 개 있기는 하지만, 기본적으로는 4×4 곱셈이다.

그리고 행렬 곱셈이 필요한 연산이 있을 때마다 그것들은 전부 이 작은 4×4 곱셈이라는 명령어 단위로 쪼개진다. 그러니까 모든 것이 이 명령어로 분해된다. 그게 행렬을 곱하는 가장 빠른 방법이기 때문이다. 그런데 알고 보면 우리가 위에서 하고 있는 계산 작업의 대부분이, 정말로 그 전부가 행렬 곱셈이다. 계산의 대부분은 선형 레이어에서 일어난다. linear, linear, 계속 이런 식이다. 그 사이사이에 끼어 있는 것들이 몇 개 있기는 하다. 잔차 연결에 덧셈이 몇 개 있고, GELU 비선형성이 몇 개 있고, LayerNorm 같은 것들이 있다. 하지만 그것들의 시간을 그냥 재 보면, 이런 것들은 아무것도 아니라는 걸 알게 될 것이다. 기본적으로 Transformer는 정말로 그냥 행렬 곱셈 덩어리다.

그리고 특히 이 정도로 작은 규모, 그러니까 1억 2400만 파라미터 모델에서는 사실 단연코 가장 큰 행렬 곱셈이 맨 위에 있는 분류기 레이어다. 768에서 50257로 가는 어마어마한 행렬 곱이고, 대략적으로 말해서 그 행렬 곱이 그 신경망에서 일어나는 다른 어떤 것보다도 압도적이다. 그러니까 훨씬 빨라지게 되는 것은 행렬 곱셈이고, 그것들은 우리의 선형 레이어 안에 숨어 있으며, 텐서 코어를 통해 가속된다.

이제 텐서 코어에 대한 최고의 레퍼런스는, 내가 말하자면, 그냥 A100 아키텍처 백서로 가는 것이다. 꽤 상세하긴 하지만, 무슨 일이 일어나고 있는지 절반쯤 이해하고 있다면 대체로 상대적으로 읽을 만하다고 생각한다. 거기 그림 9가 TensorFloat-32다. 그러니까 이게 기본적으로 TF32에 대한 설명이고 여기서 무슨 일이 일어나는지를 보여 준다. 그리고 여기에도 쓸 수 있는 구성 옵션이 많다는 것이 보인다. 입력 피연산자가 무엇이고 그것들이 어떤 정밀도에 있는지, 누산기는 어떤지, 그리고 이 행렬 곱셈의 누산을 수행할 때 명령어 내부의 표현은 기본적으로 어떤지 같은 것들이다. 그러니까 여기 이 작은 벡터 곱셈들의 중간 +=, 그 중간 누산은 전부 FP32에서 일어난다. 그리고 내가 앞서 말했듯이, 이것이 우리가 얻는 연산량에 대해 8× 개선이다.

TF32를 구체적으로 보면, 우리는 여기 이 행을 보고 있는 것이고, 이게 작동하는 방식은 이렇다. 보통 FP32는 32비트를 갖는다. TF32는 정확히 같은 비트 구성을 쓴다. 부호 비트가 하나 있고 지수 비트가 여덟 개 있다. 다만 float에서 가수부(mantissa) 비트가 잘려 나간다. 그래서 기본적으로 우리는 32비트 대신 그냥 19비트를 갖게 된다. 마지막 13비트가 절삭되어 버려지기 때문이다.

그리고 이 모든 것이 그 명령어의 내부에서 일어난다. 그래서 그중 어느 것도 우리 PyTorch 코드 안의 그 무엇에게도 보이지 않는다. 우리 PyTorch 코드는 하나도 바뀌지 않을 것이고, 모든 숫자는 똑같아 보일 것이다. 그저 하드웨어에서 내부적으로 텐서 코어 명령어를 호출할 때 이 13비트를 잘라낼 뿐이고, 그것이 이 작은 행렬 곱셈을 훨씬 더 빠르게, 8× 더 빠르게 계산할 수 있게 해 준다.

물론 이 속도 향상에는 대가가 따르고, 그 대가는 우리가 정밀도를 줄이고 있다는 것이다. 우리 누산은 여전히 FP32이고, 우리 출력도 FP32이며, 우리 입력도 FP32다. 하지만 내부적으로는 연산을 더 빨리 수행하기 위해 피연산자 안에서 값이 절삭된다. 그래서 우리 결과는 약간 더 근사적이 되기 시작한다. 그런데 경험적으로, 이걸 가지고 실제로 학습을 돌려 보면 기본적으로 차이를 알아챌 수 없다.

내가 TF32를 좋아하는 이유가 바로 그것이다. 정밀도를 약간 눙치는 것을 견딜 수 있다면, 이건 공짜다. 우리 코드 중 어느 것도 이걸 보지 못하고, 완전히 그 연산 내부에 국한되며, 우리 입장에서 그 연산은 그냥 8× 더 빨라진다. 그리고 조금 더 근사적일 뿐이다. 그래서 나는 이게 최적화에서 꽤 달콤한 지점이라고 말하겠다.

그게 어떤 모습인지 우선 보자. 나는 우리 코드를 고쳐서 반복(iteration) 시간을 재도록 해 두었다. import time을 했고, 하이퍼파라미터도 바꿔서 우리가 돌리고 싶은 워크로드를 좀 더 반영하는 무언가를 갖도록 했다. 이 강의 끝에서는 꽤 큰 학습을 한번 돌려 보고 싶기 때문이다. 그러니 배치 크기 16을 쓰고, 그리고 이제는 실제 GPT-2의 최대 시퀀스 길이인 1,024 토큰을 쓰자. 이게 설정이다. 그리고 그다음 50회 반복 동안, 나는 여기서 아주 게으른 방식으로 time.time()을 불러 현재 시각을 얻고 있다. 그리고 이게 최적화 루프이고, 이제 나는 이게 얼마나 걸리는지 재고 싶다.

그런데 GPU를 가지고 작업할 때 한 가지 문제가 있다. CPU가 돌아갈 때, CPU는 그저 GPU에 작업을 스케줄링하고 있을 뿐이다. 작업을 주문하고 있는 것이다. 요청을 보내고 나서 계속 자기 일을 진행한다. 그래서 실제로는 우리가 이 지점을 휙 지나쳐 버려서, GPU에서 돌릴 커널을 잔뜩 큐에 쌓아 놓은 다음 CPU가 여기까지 와서 시각을 재는 일이 일어날 수 있다. 하지만 실제로는 GPU가 아직 돌아가고 있다. 스케줄된 작업을 실제로 처리하는 데 시간이 걸리기 때문이다. 그래서 우리는 그냥 GPU를 위한 큐를 쌓아 올리고 있는 셈이 된다.

그러니 필요하다면 여러분은 torch.cuda.synchronize()를 넣어서 기다리고 싶을 것이다. 그리고 이건 GPU가 위에서 스케줄된 모든 작업을 끝낼 때까지 기다려 준다. 그런 다음에야 우리가 실제로 시각을 잴 수 있다. 그러니까 기본적으로 우리는 GPU가 멈추기를 기다리고, 이 반복의 시각을 재고, 그다음 그냥 그걸 출력하는 것이다.

자, 여기서 나는 학습 루프를 돌릴 것이고 오른쪽에서는 nvidia-smi를 지켜보고 있다. 우리는 0에서 시작한다. GPU를 쓰고 있지 않다. 그리고 기본적으로 PyTorch는 GPU 0을 쓸 것이므로, 그게 채워지는 것이 보이고 우리는 가용한 80GB 중 35GB를 쓰고 있다. 그리고 여기 왼쪽을 보면, 우리가 배치 크기를 키웠기 때문에 이제 우리의 타이니 셰익스피어 데이터에서 한 에폭을 도는 데 배치 20개밖에 되지 않는다. 그리고 여기 반복당 대략 1,000밀리초가 나오는 것이 보인다.

첫 번째 반복은 때때로 더 느린데, 그건 PyTorch가 아주 첫 번째 반복에서 초기화 작업을 많이 하고 있을 수 있기 때문이다. 그래서 아마 그래디언트를 전부 담을 이 텐서들과 버퍼들을 초기화하고 있을 것이다. 여기서 정확히 무슨 작업이 다 일어나는지 나도 100% 확신하지는 못하지만, 어쨌든 이건 더 느린 반복일 수 있다. 자기 로직의 시간을 잴 때는 항상 그 점을 조심하고 싶을 것이다. 하지만 기본적으로 우리는 반복당 1,000밀리초를 보고 있다. 그래서 지금 이 상태 그대로라면 대략 50초 동안 돌아갈 것이다. 그게 float32에서의 우리 기준선이다.

하나 더 언급하고 싶은 것은, 이게 여러분의 GPU에 들어가지 않아서 out-of-memory 에러가 난다면 배치 크기를 줄여 나가면서 맞추라는 것이다. 그러니까 16 대신 8을 시도해 보거나 4를 시도해 보거나, 배치가 여러분의 GPU에 들어가는 데 필요한 값이면 무엇이든 좋다. 그리고 더 큰 GPU를 가지고 있다면 32 같은 값으로도 실제로 해낼 수 있을 것이다. 기본적으로 여러분은 여러분의 GPU에 들어가는 한도까지 배치 크기를 최대로, 최대한으로 키우고 싶고, 그러면서도 그것을 좋은 숫자로 유지하고 싶다. 그러니 2의 거듭제곱 인수를 많이 가진 숫자를 써라. 16은 좋은 숫자이고, 8, 24, 32, 48도 좋은 숫자다. 하지만 17 같은 것은 쓰지 마라. 그건 GPU에서 아주 비효율적으로 돌아갈 것이다. 이것도 조금 뒤에 다시 보게 될 것이다. 그러니 지금으로서는 그냥 16과 1024를 그대로 고수하도록 하자.

그리고 내가 여기 하나 더 추가해서 다시 돌린 것이 있는데, 학습 중에 초당 토큰 처리량을 계산하는 것이다. 왜냐하면 우리가 시간이 지나면서 배치 크기를 이리저리 바꾸게 될 수도 있기 때문이다. 하지만 초당 토큰 수야말로 우리가 실제로 정말 신경 쓰는 객관적인 척도다. 우리가 얼마나 많은 데이터 토큰으로 학습하고 있는가, 그리고 우리 최적화에서 얻고 있는 토큰 처리량은 얼마인가. 지금 우리는 초당 대략 163,000 토큰을 처리하며 학습하고 있고, 그게 좀 더 객관적인 지표다.

좋다, 이제 TF32를 켜 보자. 다행히 PyTorch가 이걸 우리에게 꽤 쉽게 해 준다. TF32를 켜려면 딱 한 줄만 쓰면 되고, 그게 이것이다.

torch.set_float32_matmul_precision('high')

그리고 PyTorch 문서에서 이 함수를 보면, 기본적으로 이건 PyTorch에게 어떤 종류의 커널을 돌릴지 알려 준다. 그리고 기본값은 행렬곱에 대해 highest, 즉 최고 정밀도라고 알고 있다. 그 말은 이전과 똑같이 모든 것이 float32에서 일어난다는 뜻이다. 하지만 지금 우리가 하는 것처럼 high로 설정하면, 행렬 곱셈은 가능한 경우 TensorFloat-32를 쓰게 된다. 내 GPU는 A100이므로 Ampere 계열이고, 따라서 TF32를 쓸 수 있다. 더 오래된 GPU를 가지고 있다면 이게 여러분에게는 불가능할 수도 있지만, 내 GPU에서는 가능하다.

그래서 내가 PyTorch에게 기대하는 것은, nn.Linear가 보이는 모든 지점마다 그 안에 행렬 곱셈이 있으니, 그 행렬 곱셈이 이제 텐서 코어 위에서 TF32를 활용해 돌아가는 것이다. 그러니까 이것이 내가 믿기로 필요한 유일한 한 줄의 변경이고, 이제 이걸 다시 돌려 보자.

우리가 약속받은 처리량으로 보자면 우리는 대략 8×를 받아야 한다. 그러니 무슨 일이 일어나는지 보자. 그리고 그 8×는 여기서 나온 것이고, 또 이쪽에서 19.5 대신 156 테라플롭스라는 것을 본 데서 나온 것이다.

자, 실제로 무슨 일이 일어났는가. 우리가 보고 있는 것은 처리량이 대략 3배 향상이지 8×가 아니라는 사실이다. 우리는 1,000밀리초에서 300밀리초로 내려가고 있고, 우리 처리량은 이제 초당 약 50,000 토큰이다. 그러니까 8× 대신 대략 3배를 얻은 것이다.

무슨 일이 일어난 걸까? 그리고 기본적으로 여기서 벌어지고 있는 일은 또다시, 이런 워크로드 상당수가 메모리 바운드라는 것이다. 그래서 TF32가 원리적으로는 훨씬 더 빠른 처리량을 제공하더라도, 도처에 있는 이 숫자들은 여전히 전부 float32다. 그리고 메모리 시스템을 통해 여기저기로 실려 다니는 것은 float32 숫자들이고, 이 데이터를 전부 이리저리 실어 나르는 데 우리에게 너무나 많은 시간이 든다. 그래서 우리가 곱셈 자체는 훨씬 더 빠르게 만들었지만 우리는 메모리 바운드이고, 이 냅킨 계산에서 나올 법한 온전한 이득을 실제로 보고 있지는 못하다.

그렇다고는 해도 우리는 3배 빠른 처리량을 얻고 있고, 이건 공짜다. PyTorch에서 한 줄이고, 여러분의 모든 변수는 여전히 도처에서 float32이며, 그저 더 빨리 돌아갈 뿐이다. 그리고 약간 더 근사적이지만 기본적으로 우리가 그걸 알아차리지는 못할 것이다. 그러니까 그게 바로 TF32다.

float16, GradScaler, bfloat16, 300ms 01:39:38

좋다, 이제 계속 가 보자. 우리는 이 행을 소진했고, 연산 자체의 내부에서 정밀도의 일부를 잘라낼 수 있다는 것을 보았다. 하지만 우리는 여전히 메모리 바운드라는 것도 보았다. 다른 부분에서는 여전히 이 float들을 전부 이리저리 옮기고 있고, 그것 때문에 우리는 대가를 치르고 있다.

그러니 이제 우리가 옮겨 다닐 것의 양 자체를 줄여 보자. 그리고 우리는 그걸 bfloat16으로 내려감으로써 할 것이다. 그러니까 우리는 float당 16비트만 유지할 것이고 bfloat16을 쓸 것이다. 그리고 FP16과의 차이는 잠시 뒤에 설명하겠다. 그리고 우리는 이 행에 있게 될 것이다.

그래서 A100 문서로 다시 돌아가 보면, 여기 쓸 수 있는 정밀도들이 보인다. 그리고 이게 원래의 FP32이고, TF32는 정밀도를 잘라낸다. 그리고 여기 BF16을 보면, 이게 TF32와 아주 비슷하지만 이 float의 가수부를 잘라내는 데 있어 훨씬 더 공격적이라는 것을 볼 수 있다.

그래서 bfloat16에서 중요한 점은 지수 비트와, 물론 부호 비트가 바뀌지 않은 채로 남아 있다는 것이다. 여러분이 부동소수점 숫자에 익숙하다면 — 그리고 나는 이게 아마 그 자체로 하나의 영상이 되어야 한다고 생각한다 — 지수는 여러분이 표현할 수 있는 숫자의 범위를 정하고, 정밀도는 여러분의 숫자들에 대해 얼마나 세밀한 해상도를 갖는지를 정한다. 그러니까 숫자의 범위는 동일하지만, 우리가 가수부를 절삭하고 있기 때문에 그 범위 안에서 가능한 값의 가짓수가 줄어든다. 그 범위 안에서 정밀도가 덜하다는 뜻이다.

그게 뜻하는 바는, 상황이 사실 꽤 괜찮다는 것이다. 왜냐하면 우리는 float으로 표현 가능한 원래의 수 범위를 그대로 갖되 그것에 대한 정밀도만 덜 갖기 때문이다. 그리고 FP16과의 차이는, FP16은 실제로 범위를 건드려서 바꾼다는 것이다. 그래서 FP16은 FP32의 전체 범위를 표현할 수 없다. FP16은 축소된 범위를 갖고, 바로 거기서 여러분이 실제로 문제에 부딪히기 시작한다. 이제 GradScaler 같은 것이 필요해지기 때문이다. 그리고 나는 이 영상에서 그 세부까지 들어가지는 않겠다. 그것도 그 자체로 하나의 영상거리이기 때문이다.

그런데 사실 역사적으로는 FP16이 먼저 나왔다. Ampere 이전의 Volta 계열에서 쓸 수 있었고, 그래서 FP16이 먼저였으며 모두가 FP16으로 학습하기 시작했다. 하지만 모두가 이 그래디언트 스케일링 연산들을 전부 써야만 했고, 그건 좀 짜증나는 일이며, 추가적인 상태와 복잡성의 원천이다. 그리고 그렇게 된 이유는 FP16에서 지수 범위가 축소되었기 때문이다. 그러니까 그게 IEEE FP16 규격이다. 그러다 그들이 BF16과 Ampere를 내놓았고, 훨씬 더 단순하게 만들어 놓았다. 왜냐하면 우리는 그냥 가수부를 절삭할 뿐이고, 우리는 정확히 동일한 범위를 가지며, 따라서 GradScaler가 필요 없다. 모든 것이 훨씬훨씬 더 단순하다.

다만 우리가 BF16을 실제로 쓸 때는, 우리 PyTorch 코드에서 보게 될 숫자들에 영향을 주게 된다. 이 변화는 TF32와 달리 연산 자체에만 국한되지 않는다. 그러니 그게 어떻게 동작하는지 보자.

여기 문서가 좀 있는데, 나는 이 페이지가 PyTorch에서 혼합 정밀도를 어떻게 쓰는지 설명하는 데 아마 가장 좋은 페이지라고 생각한다. 왜냐하면 다른 튜토리얼이 많고, PyTorch 문서 안에서조차 훨씬 더 혼란스러운 것들이 있기 때문이다. 그래서 나는 특별히 이 페이지를 추천한다. 내가 추천하지 않을 다섯 벌의 다른 사본이 있기 때문이다. 그리고 여기로 오면 GradScaler에 관한 모든 것을 무시하고, torch.autocast만 보면 된다. 그리고 기본적으로 이것도 마지막에 딱 한 줄의 코드로 귀결된다. 그러니까 이게 우리가 원하는 컨텍스트 매니저이고, 우리는 그걸 우리 신경망에서 쓰고 싶은 것이다.

torch.autocast, 그러니까 autocasting 문서로 들어가 보면 여러분을 위한 지침이 조금 더 나와 있다. 여러분의 텐서 어느 것에도 bfloat16을 직접 호출하지 말고 그냥 autocast를 쓰라고, 그리고 모델의 순전파와 손실 계산만 감싸라고 말한다. 그 두 가지만이 여러분이 감싸야 할 대상이고, 역전파와 옵티마이저 스텝은 건드리지 말고 놔두라고 한다. 그러니까 그게 PyTorch 팀에서 나온 지침이다.

그래서 우리는 그 지침을 따를 것이다. 그리고 우리 경우에는 손실 계산이 모델의 순전파 안에 들어 있으므로, 우리는 이렇게 할 것이다. 그리고 우리는 torch.float16을 쓰고 싶지 않다. 그렇게 하면 GradScaler도 같이 쓰기 시작해야 하기 때문이다. 그래서 우리는 bfloat16을 쓸 것이다.

with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
    logits, loss = model(x, y)

이건 Ampere에서만 가능한 일이지만, 어쨌든 이 말은 변경 사항이 극도로 적다는 뜻이다. 기본적으로 그냥 이 한 줄의 코드가 전부다.

이걸 실제로 돌리기 전에 우선 여기서 한번 브레이크를 걸어 보겠다. 그러니까 로짓 바로 다음 지점이다. 여기서 여러분에게 보여 주고 싶은 것은, 우리가 본 TF32와 달리 이번에는 실제로 우리 텐서에 영향이 간다는 것이다. 그래서 이 로짓 텐서를 이제 보고 dtype을 보면, 우리는 갑자기 이게 bfloat16이 되어 있는 것을 보게 된다. 더 이상 float32가 아니다. 그러니까 우리의 활성값이 바뀐 것이다. 활성값 텐서가 이제 bfloat16이다.

하지만 모든 것이 바뀐 것은 아니다. model.transformer.wte, 이건 가중치, 그러니까 토큰 임베딩 테이블인데, 그 안에 가중치가 있고 이 가중치라는 파라미터의 dtype은 여전히 torch.float32다. 그러니까 우리 파라미터는 여전히 float32에 있는 것으로 보이지만 우리 활성값, 즉 로짓은 이제 bfloat16에 있다. 그래서 명백히 이것이 우리가 혼합 정밀도를 얻는 이유다. 어떤 것들은 PyTorch가 float32로 유지하고, 어떤 것들은 PyTorch가 더 낮은 정밀도로 변환한다.

그리고 어느 시점에 무엇이 변환되는지는 그리 명확하지 않다. 내가 기억하기로 아래로 스크롤을 내리면… 여기인가, 못 찾겠다. 여기인 줄 알았는데. 자, 여기 있다. 여러분이 이 autocast를 쓸 때 무엇이 bfloat16으로 변환되고 언제 그렇게 되는지에 대한 문서가 몇 개 있다. 그래서 예를 들어 이런 행렬 곱셈류 연산만 저정밀도로 변환되고, 많은 연산은 float32로 남는다.

특히 LayerNorm 같은 정규화 계열의 많은 연산들은, 그런 레이어가 전부 변환되지는 않을 수도 있다. 그러니까 일부 레이어만 선택적으로 bfloat16으로 돌아갈 것이다. 하지만 softmax 같은 것들, LayerNorm 같은 것들, log softmax 같은 것들, 그러니까 손실 함수 계산 같은 많은 것들은 float32로 남을 수 있다. 왜냐하면 그것들이 정밀도 변화에 더 민감하기 때문이다. 행렬 곱셈은 정밀도 변화에 꽤 견고한 편이다. 그래서 신경망의 부분마다 정밀도 변화에 더 영향을 받거나 덜 영향을 받는다. 그러니까 기본적으로 모델의 일부만이 감소된 정밀도로 돌아가는 것이다.

한번 돌려 보고 여기서 우리가 어떤 개선을 달성했는지 실제로 보자. 자, 우리는 예전에 333밀리초였는데 이제 300이다. 그리고 예전에는 초당 50,000 토큰 언저리였는데 이제 55,000이다. 그러니까 우리는 확실히 더 빨리 돌아가고는 있지만, 어쩌면 아주 많이 빨라지지는 않았다. 그리고 그건 우리 GPT-2에 여전히 아주아주 많은 병목이 남아 있기 때문이다. 우리는 이제 막 시작했을 뿐이다.

하지만 우리는 현재 내 GPU, 즉 A100으로 갈 수 있는 데까지 정밀도를 낮추기는 했다. 우리는 PyTorch autocast를 쓰고 있다. 아쉽게도 나도 PyTorch autocast가 정확히 무엇을 하는지는 사실 정확히 알지 못한다. 무엇이 bfloat16에 있고 무엇이 float32에 있는지 정확히는 모른다. 들어가서 그걸 하나하나 뜯어보기 시작할 수도 있겠지만, 이런 것들은 PyTorch가 내부적으로 갖고 있는 규칙이고 아쉽게도 그들은 그걸 아주 잘 문서화해 놓지 않았다. 그래서 우리는 거기까지 너무 자세히 들어가지는 않겠다.

어쨌든 지금 우리는 bfloat16으로 학습하고 있고 GradScaler는 필요하지 않다. 그리고 여러 가지가 더 빨리 돌아가고 있는 이유는 우리가 bfloat16에서 텐서 코어를 돌릴 수 있게 되었기 때문이다. 즉 그 말은 우리가 이제 이 행에 있다는 뜻이다. 다만 우리는 이것에 대해 정밀도로 대가를 치르고 있다. 그래서 우리는 원래의 FP32에 비해 약간 덜 정확한 결과를 예상한다. 하지만 경험적으로 많은 경우 이건 그만한 가치가 있는 맞바꿈이다. 왜냐하면 그게 여러분을 더 빨리 돌릴 수 있게 해 주고, 그러면 예를 들어 여러분은 더 오래 학습해서 그 정밀도 감소를 메울 수도 있기 때문이다. 그러니까 여기까지가 바로 그 bfloat16에 대한 이야기다.

torch.compile, Python 오버헤드, 커널 퓨전, 130ms 01:48:15

자, 지금 보다시피 우리는 현재 한 iteration당 약 300밀리초 수준에 와 있다. 이제부터는 PyTorch 무기고에서 정말 무거운 무기들을 꺼내 들 차례다. 그중에서도 제일 먼저 꺼낼 것이 torch.compile이다. torch.compile은 PyTorch 팀이 만든 정말 굉장한 인프라다. 기본적으로 신경망을 위한 컴파일러라고 보면 된다. C/C++ 코드에 GCC가 있는 것처럼, 신경망에 대한 GCC가 바로 이것이다. 나온 지는 꽤 됐고, 쓰는 법은 극도로 간단하다.

torch.compile을 쓰는 방법은 이게 전부다. 코드 한 줄로 모델을 컴파일해서 돌려받는다. 이 한 줄은 컴파일 시간이라는 비용을 치르게 만든다. 하지만 짐작할 수 있듯이, 그 대가로 코드가 훨씬 빨라진다. 그러니 실제로 돌려보자. 실행되는 데 시간이 좀 걸릴 것이다. 지금 우리가 300밀리초에 있다는 걸 기억해 두고, 무슨 일이 일어나는지 보자.

이게 돌아가는 동안 torch.compile이 내부에서 무슨 일을 하는지 조금 설명하겠다. PyTorch의 이 페이지를 마음껏 읽어봐도 좋다. 다만 기본적으로, 여러분의 PyTorch 코드에서 torch.compile을 쓰지 않을 이유는 사실상 없다. 나는 거의 기본값으로 쓰는 게 맞다고 본다. 디버깅을 하는 중이고 코드가 정말 빨리 돌 필요가 없는 경우가 아니라면 말이다. 그리고 torch.compile 문서에 이게 왜 빠른지를 정확히 짚어 주는 한 줄이 있는 걸 발견했다. 속도 향상은 주로 Python 오버헤드를 줄이고 GPU의 읽기·쓰기를 줄이는 데서 나온다는 것이다. 이 문장을 좀 풀어서 설명하겠다.

자, 결과가 나왔다. 300밀리초에서 시작해서 지금은 129밀리초로 돌고 있다. 300에서 129면 약 2.3배 개선이다. PyTorch에서 코드 한 줄로 얻은 결과다. 정말 대단하다.

그러면 도대체 무슨 일이 벌어진 걸까. 내부에서 어떤 일이 일어나는지 보자. 여러분이 모델을 torch.compile에 넘길 때, 우리가 여기 nn.Module 안에 가지고 있는 것은 사실 우리 네트워크에서 어떤 일이 일어나기를 바라는지에 대한 알고리즘적 서술일 뿐이다. 그 이상도 이하도 아니다. torch.compile은 이 전체를 통째로 분석해서, 여러분이 어떤 연산들을 쓰려고 하는지 들여다본다. 그리고 앞으로 무슨 일이 일어날지 정확히 알고 있다는 이점 덕분에, 이른바 eager 모드로 돌 필요가 없어진다. Python 인터프리터가 보통 하듯 레이어를 하나씩 순서대로 밟아 가며 실행할 필요가 없다는 뜻이다.

Python 인터프리터는 forward에서 출발해서 이렇게 움직인다. "좋아, 이 연산을 하자." 그다음 "이제 저 연산을 하자." 이런 식으로 진행하면서 모든 연산을 그때그때 물질화한다. 즉 이 계산들은 그 순서대로 디스패치되어 실행되는데, Python 인터프리터와 이 코드는 나중에 어떤 종류의 연산이 올지 전혀 모른다. 반면 torch.compile은 여러분의 코드 전체를 동시에 본다. 그래서 어떤 연산들을 돌리려고 하는지 알 수 있고, 그 과정을 최적화할 수 있다.

가장 먼저 하는 일은 forward pass에서 Python 인터프리터를 통째로 걷어내는 것이다. 그리고 이 신경망 전체를 Python 인터프리터가 전혀 개입하지 않는 단일 객체로 컴파일한다. 그러니까 무엇이 실행될지 정확히 알고, 그냥 그것을 실행하며, 전부 효율적인 코드로 돌아간다.

두 번째로 일어나는 일이 문서에서 아주 짧게 언급한 그 읽기·쓰기 문제다. 이걸 설명하기에 좋은 예가 우리가 계속 보고 있던 GELU 비선형성이라고 생각한다. 여기서 우리는 nn.GELU를 쓰고 있다. 그런데 지금 이 코드는 내가 tanh 근사 GELU를 그 공식 그대로 풀어 쓴 것이다. 기억하겠지만 GELU에는 이런 형태의 공식이 있다. 그러니까 여기 이것은 nn.GELU 안에서 알고리즘적으로 일어나고 있는 것과 동등한 구현이다. 완전히 동일하다.

이제 기본 상태에서, 우리가 nn.GELU 대신 이 풀어 쓴 코드를 쓴다면 torch.compile 없이는 무슨 일이 벌어질까. Python 인터프리터가 여기까지 와서 이렇게 판단한다. "좋아, 입력이 있군. 먼저 이 입력을 3제곱하자." 그리고 입력을 받아 3제곱하는 커널을 디스패치한다. 그 커널이 실행되는데, 이 커널이 실행될 때 실제로 벌어지는 일은 이렇다. 이 입력은 GPU의 메모리에 저장되어 있다.

여기 무슨 일이 일어나는지에 대한 배치도가 도움이 될 것이다. 먼저 CPU가 있다. 모든 컴퓨터에 다 있는 것이고, 그 안에 코어가 몇 개 들어 있다. 그리고 RAM, 즉 메모리가 있고, CPU는 이 메모리와 대화할 수 있다. 여기까지는 다들 아는 얘기다. 그런데 이제 GPU가 추가됐다. GPU는 물론 CPU와 통신할 수 있지만 구조가 조금 다르다. CPU보다 코어가 훨씬 많고, 그 코어들은 하나하나로 보면 훨씬 단순하다. 그리고 GPU도 메모리를 가지고 있다. 바로 이 고대역폭 메모리, HBM이다. 이것은 기본적으로 컴퓨터의 RAM에 대응하는 물건이라고 보면 된다.

그래서 지금 벌어지는 일은 이렇다. 입력은 그 메모리 안에 살고 있다. 그런데 여러분이 입력을 3제곱하면, 이 데이터는 GPU로, 코어들로, 그리고 실제 GPU 칩 위의 온갖 캐시와 레지스터까지 이동해야 한다. 그리고 모든 원소를 3제곱으로 계산한 다음, 그 결과를 다시 메모리에 저장한다. 바로 이 이동 시간이 실제로 많은 문제를 일으킨다.

여기서 이 메모리 대역폭을 다시 보자. 우리는 초당 약 2테라바이트로 통신할 수 있다. 큰 값이긴 하다. 하지만 어쨌든 이 링크를 건너야 하고, 그건 매우 느리다. GPU 안, 즉 온칩에서는 모든 것이 엄청나게 빠르다. 그런데 메모리로 나가는 것은 극도로 비싸고 극도로 오랜 시간이 걸린다. 그래서 우리는 입력을 불러오고, 계산하고, 출력을 다시 실어 보낸다. 이 왕복 한 번에 아주 많은 시간이 든다.

그런데 그걸 하자마자 우리는 이 상수를 곱한다. 그러면 무슨 일이 일어나나. 또 다른 커널이 디스패치되고, 결과가 다시 이동해서 모든 원소가 상수와 곱해지고, 그 결과가 다시 메모리로 돌아간다. 그다음 우리는 그 결과에 입력을 다시 더한다. 그러면 이 전체가 또 GPU로 이동해서 입력을 더하고 다시 써 나간다. 즉 우리는 메모리와 실제 계산이 일어나는 곳 사이를 계속 왕복하고 있는 것이다. 텐서 코어든 ALU든 그 모든 것이 GPU 칩 위에 있기 때문이다. 우리는 엄청난 횟수의 왕복을 하고 있다.

그리고 torch.compile을 쓰지 않으면 PyTorch는 이걸 최적화할 줄 모른다. 나중에 어떤 종류의 연산을 돌릴 건지 모르기 때문이다. 여러분은 그냥 "3제곱해라, 그다음 이걸 해라, 그다음 저걸 해라"라고 말하고 있을 뿐이고, PyTorch는 그 순서대로 그냥 할 뿐이다.

그런데 torch.compile은 코드 전체를 본다. 여기까지 와서 이렇게 깨닫는다. "잠깐, 이것들은 전부 원소별(element-wise) 연산이잖아. 그러면 내가 할 일은 입력을 GPU로 한 번만 실어 보내고, 그 메모리가 GPU 위에 올라와 있는 동안 — 정확히는 그 덩어리 단위로 — 모든 원소에 대해 이 연산들을 전부 수행한 다음, 딱 한 번만 다시 써 내보내는 것이다." 그러면 이 왕복들이 사라진다. 이것이 이른바 커널 퓨전(kernel fusion)의 한 예이고, 모든 것이 빨라지는 주요 경로다. 그러니까 기본적으로 전체를 미리 본다는 이점이 있고 무엇을 계산할지 정확히 알고 있다면, 메모리로의 왕복을 최적화할 수 있고 메모리 대역폭 비용을 치르지 않아도 된다. 이것이 근본적으로 이런 연산들 중 일부를 훨씬 빠르게 만드는 것이고, 문서에서 말하는 읽기·쓰기가 바로 이 얘기다.

자, 이 풀어 쓴 코드는 쓰지 않으니 지우겠다. 그렇다. 우리는 torch.compile을 써야 하고, 우리 코드는 이제 상당히 빨라져서 초당 약 125,000 토큰을 처리하고 있다. 그래도 아직 갈 길이 멀다.

다음으로 넘어가기 전에, 이 논의를 그림 몇 장으로 좀 더 보충하고 싶다. 꽤 복잡한 주제지만, 여기서 무슨 일이 일어나고 있는지를 큰 틀에서 이해해 둘 가치가 있다. 사실 이 주제 하나만으로도 두 시간짜리 영상을 통째로 하나 더 만들 수 있을 것이다. 그래도 그 맛보기 정도는 지금 해 두자.

여기 이 칩이 GPU다. 이 칩이 대부분의 계산이 실제로 일어나는 곳이다. 이 칩도 안에 메모리를 어느 정도 가지고 있긴 하다. 하지만 메모리의 압도적 대부분은 여기 고대역폭 메모리, 즉 HBM 쪽에 있다. 둘은 서로 연결되어 있지만, 기본적으로는 별개의 두 칩이다.

이제 이것이 GPU 개념도를 확대한 그림이다. 여기서 우리가 보고 있는 것은, 첫째로 이 HBM이다. 여러분 화면에서는 아주 작게 보이겠지만 양옆에 HBM이라고 적혀 있다. 그게 HBM으로 가는 링크다. 다시 말하지만 HBM은 칩 바깥에 있다. 그리고 칩 위에는 스트리밍 멀티프로세서가 아주 많이 있다. 이것 하나하나가 SM이고, 전부 합쳐 120개가 있다. 여기서 많은 계산이 일어난다.

그리고 이것이 SM 하나를 확대한 그림이다. 네 개의 사분면으로 나뉘어 있는 것이 보인다. 예를 들어 텐서 코어가 있는데, 여기서 행렬 곱 관련 계산이 많이 일어난다. 그 외에도 fp64, fp32, 정수 등 온갖 종류의 계산을 담당하는 다른 유닛들이 다 들어 있다.

그래서 계산을 수행하는 이 모든 로직이 있는데, 그에 더해 칩 위에는 메모리가 곳곳에 뿌려져 있다. L2 캐시는 칩 위에 사는 어느 정도 크기의 메모리다. 그리고 SM 자체에는 L1 캐시가 있다. 이것도 여러분 화면에서는 아주 작게 보이겠지만, 이 파란 막대가 L1이다. 그리고 레지스터도 있다. 그러니까 여기에도 메모리가 저장되어 있다. 다만 이 메모리가 저장되는 방식은 HBM에 메모리가 저장되는 방식과 아주 다르다. 실리콘이 어떻게 생겼는지의 수준에서 보면 완전히 다른 구현이다. 저쪽은 트랜지스터와 커패시터를 쓰는 구현이고, 여기는 SRAM이라는 아주 다른 구현을 쓴다. 길게 말할 것 없이 핵심은 이것이다. 칩 안에도 메모리가 있긴 하다. 하지만 그 양이 많지 않다. 이게 결정적인 지점이다.

이건 약간 다른 GPU의 예시 다이어그램인데, 여기서도 같은 얘기가 보인다. 예를 들어 CPU 쪽의 DRAM 메모리, 여기 이것은 전형적으로 1테라바이트쯤 가지고 있을 수도 있다. 하지만 접근 비용이 극도로 비싸다. 특히 GPU 입장에서는 여기 CPU를 거쳐 가야 하므로 더욱 그렇다. 그다음이 HBM이다. 전형적인 GPU에 수십 기가바이트의 HBM 메모리가 달려 있다. 하지만 말했듯이 접근 비용이 매우 비싸다. 그리고 칩 자체 위에서는 모든 것이 극도로 빠르지만, 칩 전체를 통틀어 우리가 가진 메모리는 수십 메가바이트 정도밖에 안 된다. 그래서 공간이 애초에 충분하지 않다. 칩 위의 메모리는 아주 비싸기 때문에 양이 적을 수밖에 없다. 하지만 상대적으로 보면 접근 속도가 번개처럼 빠르다.

그래서 우리가 이런 커널들을 돌릴 때 실제로 일어나는 일의 더 정확한 그림은 이렇다. 기본적으로 우리는 글로벌 메모리에 기본적으로 살고 있는 이 입력들을 가져다가 어떤 계산을 수행해야 한다. 그래서 우리는 글로벌 메모리에서 칩으로 데이터를 스트리밍하기 시작하고, 칩 위에서 계산을 수행한 다음, 그것을 다시 스트리밍해서 글로벌 메모리에 되돌려 저장한다. torch.compile이 없다면 우리는 데이터를 칩으로 흘려보내고, 계산하고, 메모리에 저장하는 이 왕복을 몇 번이고 몇 번이고 반복하게 된다.

그런데 torch.compile을 쓰면, 처음에 메모리를 스트리밍하기 시작하는 것까지는 같다. 하지만 칩 위에 올라와 있는 동안, 우리는 처리하려는 데이터의 한 덩어리를 손에 쥐고 있다. 그 덩어리가 지금 칩 위에 살아 있는 것이다. 칩 위에 있는 동안에는 연산이 극도로 빠르다. 그래서 커널 퓨전이 있으면 그 자리에서 모든 연산을 원소별로 전부 처리해 버릴 수 있고, 그 연산들은 아주 싸다. 그런 다음 글로벌 메모리로 딱 한 번만 왕복한다. 즉 연산자 퓨전은 여러분의 데이터 덩어리를 칩 위에 붙들어 두고, 되써 내보내기 전에 그 위에서 많은 계산을 해치울 수 있게 해 준다. 이것이 엄청난 절약을 낳고, 그래서 torch.compile이 훨씬 빨라지는 것이다. 적어도 주된 이유 중 하나다. 여기까지가 메모리 계층 구조에 대한 아주 짧은 소개와, torch.compile이 대략 무엇을 해 주는지에 대한 설명이다.

FlashAttention, 96ms 02:00:18

torch.compile은 놀랍다. 하지만 torch.compile이 찾아내지 못하는 연산들이 있다. 그 대표적인 예가 FlashAttention이고, 이제 그쪽으로 넘어간다.

FlashAttention은 2022년 스탠퍼드에서 나온 논문에서 출발한다. 어텐션을 수행하는, 그리고 훨씬 빠르게 돌리는 놀라운 알고리즘이다. FlashAttention은 여기 이 네 줄을 들어내고, 그 네 줄을 정말 정말 빠르게 구현해 낸다.

어떻게 그렇게 하는가. FlashAttention은 커널 퓨전 연산이다. 논문의 이 그림을 보면 PyTorch 쪽에 네 개의 연산이 나열되어 있다. 거기엔 Dropout도 포함되어 있는데, 우리는 여기서 Dropout을 쓰지 않으므로 우리 코드에는 이 네 줄만 있다. 그리고 그 네 줄 대신 우리는 그것들을 FlashAttention이라는 단일 퓨즈드 커널 하나로 융합해 버린다. 그러니까 이건 커널 퓨전 알고리즘이다. 다만 torch.compile이 찾아낼 수 없는 종류의 커널 퓨전이다. 찾아낼 수 없는 이유는, 이 경우 어텐션이 실제로 어떻게 구현되는지에 대한 알고리즘 수준의 재작성이 필요하기 때문이다. 단순히 연산 순서를 재배치하는 문제가 아니라는 뜻이다.

여기서 정말 놀라운 점이 있다. FLOPS 수만 세어 보면 FlashAttention은 여기 이 어텐션보다 더 많은 FLOPS를 쓴다. 그런데도 FlashAttention이 훨씬 빠르다. 논문에서는 잠재적으로 7.6배까지 빠르다고 인용한다. 그 이유는 방금 설명한 메모리 계층 구조를 아주 세심하게 의식하기 때문이다. 무엇이 고대역폭 메모리에 있고 무엇이 공유 메모리에 있는지를 아주 신경 쓰고, 고대역폭 메모리에 대한 읽기와 쓰기가 더 적어지도록 계산 순서를 아주 조심스럽게 짜 맞춘다. 그래서 FLOPS를 더 쓰는데도 더 빠르다. 비싼 부분은 HBM으로의 로드와 스토어이고, 그들이 피하는 것이 바로 그것이기 때문이다.

특히 그들은 이 N × N 어텐션 행렬, 즉 우리 코드의 att를 절대로 물질화하지 않는다. FlashAttention은 이 행렬이 어느 시점에도 물질화되지 않도록, 그리고 HBM에 읽히거나 쓰이지 않도록 설계되어 있다. 그리고 이건 아주 큰 행렬이다. 여기가 바로 모든 쿼리와 키가 상호작용하는 지점이기 때문이다. 헤드마다, 배치 원소마다 T × T 크기의 어텐션 행렬이 나오고, 단일 헤드·단일 배치 인덱스 하나만 놓고 봐도 원소가 백만 개다. 그러니까 기본적으로 이건 어마어마한 메모리다. 그런데 그게 결코 물질화되지 않는다.

이것이 어떻게 달성되는가. 기본적으로 여기서의 근본적인 알고리즘 재작성은 이전에 제안된 online softmax 트릭에 의존한다. 그 논문은 잠시 뒤에 보여 주겠다. 이전 논문에서 온 이 online softmax 트릭은, 정규화를 하기 위해 softmax의 모든 입력을 다 실체화하지 않고도 softmax를 점진적으로 평가할 수 있는 방법을 보여 준다. 그 방법은 m과 l이라는 중간 변수를 두고, 이 변수들에 대한 갱신 규칙을 통해 softmax를 온라인 방식으로 평가하는 것이다.

그리고 최근에는 FlashAttention 2도 나왔다. 그 논문도 여기 띄워 뒀다. FlashAttention을 계산하는 방식에 추가적인 이득을 더한 것이다. 한편 이 모든 것의 바탕이 된 원래 논문은 기본적으로 softmax를 위한 online normalizer 계산에 관한 것이다. 놀랍게도 이 논문은 엔비디아에서 나왔고, 그것도 아주 이른 2018년에 나왔다. FlashAttention보다 4년 앞선 것이다. 이 논문은 더 적은 메모리 접근으로 고전적인 softmax를 계산하는 방법을 제안하며, 이 메모리 접근 감소가 실제 하드웨어에서 softmax 성능을 개선할 것이라고 가설을 세운다. 결과적으로 그들의 가설은 극도로 옳았다. 그런데 나로서는 이게 정말 흥미롭다. 이들은 엔비디아 소속이었고 이 통찰을 가지고 있었는데, 정작 그것을 실제 FlashAttention까지 밀고 가지는 않았다. 그건 4년 뒤 스탠퍼드에서 나와야 했다. 역사적으로 왜 이렇게 됐는지는 나도 완전히 이해하지 못한다. 어쨌든 그들은 기본적으로 softmax에 대한 이 온라인 갱신을 제안했고, FlashAttention은 근본적으로 그것을 재사용해서 softmax를 스트리밍 방식으로 계산한다. 그리고 다른 모든 연산들을 이 online softmax 계산과 함께 단일 퓨즈드 커널로 융합할 수 있다는 걸 깨달았다. 그게 FlashAttention이고, 우리가 지금 쓰려는 것이다.

이건 몇 가지를 아주 잘 보여 주는 훌륭한 사례라고 생각한다. 메모리 계층 구조를 의식하는 것의 중요성, FLOPS가 중요한 게 아니라 전체 메모리 접근 패턴이 중요하다는 사실, 그리고 torch.compile이 놀랍긴 하지만 그것이 찾아내지 못할 최적화가 여전히 많이 남아 있다는 점이다. 언젠가는 찾아낼 수 있을지도 모르지만, 지금 시점에서는 요구가 너무 큰 것 같다.

그래서 우리가 할 일은 이렇다. FlashAttention을 쓸 것이고, PyTorch에서 그렇게 하는 방법은 기본적으로 이 네 줄을 주석 처리하고 한 줄로 대체하는 것이다. 여기서 우리는 PyTorch의 복합 연산인 F.scaled_dot_product_attention을 호출한다. 이런 식으로 쓰면 PyTorch가 FlashAttention을 호출해 준다. 왜 torch.compile이 이 네 줄이 바로 이 형태의 FlashAttention을 호출하면 된다는 걸 알아채지 못하는지 나도 100% 확신은 없다. 우리가 다시 한번 손으로 알려 줘야 한다는 게 내 생각엔 좀 이상하다. 어쨌든 이렇게 이 복합 연산을 써야 한다.

torch.compile이 처리할 때까지 잠깐 기다려 보자. 그리고 우리가 앞서 6.05661이라는 손실값을 얻었다는 걸 기억해 두자. 여기 적어 뒀다. 그게 우리가 보게 될 것으로 기대하는 손실이고, 이 변경 전에는 130밀리초가 걸렸다. 그러니 iteration 49에서 정확히 같은 결과를 볼 것으로 기대한다. 다만 실행 시간은 더 빠를 것으로 기대한다. FlashAttention은 순전히 알고리즘적 재작성이고 더 빠른 커널일 뿐, 계산 자체를 바꾸지는 않기 때문이다. 최적화 결과는 완전히 동일해야 한다.

좋다. 훨씬 빨라졌다. 약 95밀리초에 와 있고, 손실은 앞서와 사실상 같은 값이 나왔다. 부동소수점 오차 수준의 차이일 뿐이다. 즉 계산은 동일한데 속도는 상당히 빨라졌다. 130에서 대략 90에서 96 사이로 내려왔다. 96을 130 정도로 나누면, 대략 27% 정도 개선이다. 정말 흥미롭다. 이것이 FlashAttention이다.

좋은 수와 나쁜 수. vocab size 50257 → 50304, 93ms 02:06:54

이제 내가 가장 좋아하는 최적화 중 하나로 넘어간다. 이건 동시에 가장 멍청하면서 가장 기발한 최적화이고, 나에게는 볼 때마다 조금씩 놀랍다.

몇 분 전에 어떤 수는 좋은 수고 어떤 수는 나쁜 수라고 말한 적이 있다. 64는 아름답고 좋은 수다. 128은 더 좋다. 256도 아름답다. 이 수들을 아름답게 만드는 건 그 안에 2의 거듭제곱이 많이 들어 있다는 점이다. 2로 여러 번 나눌 수 있다. 반면 나쁜 수의 예는 13이나 17 같은 것들이다. 소수, 짝수가 아닌 수 따위다.

그래서 신경망이나 CUDA를 다루는 코드에서는 사실상 항상 좋은 수를 쓰고 싶어진다. CUDA의 모든 것이 2의 거듭제곱 단위로 돌아가기 때문이다. 수많은 커널이 2의 거듭제곱을 기준으로 작성되어 있고, 16이나 64 같은 크기의 블록이 잔뜩 있다. 모든 게 그런 항으로 쓰여 있고, 입력이 좋은 수로 이루어지지 않았을 때를 위한 온갖 특수 케이스 처리가 항상 따라붙는다.

그게 어떤 모습인지 보자. 기본 요령은 코드를 훑으면서 나쁜 수를 찾는 것이다. 3배는 약간 나쁘다. 100% 확신하진 못하겠고 개선의 여지가 있을지도 모르지만, 어쨌든 이건 나쁘고 이상적이지 않다. 4배는 좋다. 그건 괜찮다. 1024는 아주 좋다. 2의 거듭제곱이다. 12는 약간 의심스럽다. 2의 거듭제곱이 그리 많지 않다. 768은 훌륭하다. 그리고 50257은 정말 정말 나쁜 수다. 우선 홀수다. 그리고 안에 2의 거듭제곱이 별로 없다. 그러니 아주 나쁜 수이고 대단히 의심스럽다.

아래로 내려가 보면 나머지 수들은 다 좋은데, 25가 하나 있다. GPT-2 XL 구성에서 헤드 수가 25다. 정말 나쁜 수다. 홀수다. 실제로 이 수는 최근에 우리가 이걸 빠르게 돌리려고 커널 몇 개를 최적화할 때 큰 골칫거리였고, 특수 케이스 처리를 잔뜩 요구했다.

그러니까 우리에겐 나쁜 수가 몇 개 있고, 그중 어떤 건 고치기 쉽고 어떤 건 어렵다. 특히 vocab size가 50257이라는 건 아주 나쁘고 아주 의심스러우니 고치고 싶다.

이런 걸 고치는 쉬운 방법 중 하나는, 여러분 마음에 드는 가장 가까운 2의 거듭제곱 정렬에 닿을 때까지 그 수를 키우는 것이다. 그래서 여기 훨씬 좋은 수가 있다. 50304다. 왜 좋은가. 50304는 8로도, 16으로도, 32로도, 64로도 나누어떨어진다. 심지어 128로도 나누어떨어진다. 실제로 50304 = 128 × 393이다. 아주 좋은 수다.

그래서 우리가 할 일은 GPTConfig에서, 지금 vocab_size가 50257로 초기화되어 있는데 바로 그 원소만 50304로 덮어쓰는 것이다. 나머지는 전부 그대로다. 우리는 단지 어휘 크기를 키우는 것이다. 말하자면 가짜 토큰을 추가하는 셈이다. 그래서 vocab size가 2의 거듭제곱 인수를 갖도록 만든다.

그런데 여기서 내가 지금 하고 있는 일은 사실 우리 네트워크가 수행할 계산량을 늘리는 것이다. FLOPS를 세어 보면, 우리가 하는 연산량 계산을 해 보면, 우리는 더 많은 FLOPS를 수행하게 된다. 그리고 이것이 아무것도 망가뜨리지 않는지도 따져 봐야 한다. 일단 그냥 돌려 보고 뭐가 나오는지 보자. 현재 스텝당 96.5밀리초 정도로 돌고 있었다. 눈대중이긴 하다. 어떤 결과가 나오는지 보자.

컴파일되는 동안, 우리 코드가 실제로 제대로 동작하는지 생각해 보자. 이렇게 vocab size를 키웠을 때 vocab size가 실제로 어디에 쓰이는지 보자. 위로 올라가 __init__을 보면, 당연히 임베딩 테이블 안에서 쓰인다. 즉 Transformer의 맨 아래다. 그리고 Transformer의 맨 위, 분류기 레이어에서도 쓰인다. 이렇게 두 곳이다.

그리고 결과를 보면 93밀리초로 돌고 있다. 96.5밀리초 대신 93밀리초다. 대략 4% 개선이다. 계산을 더 많이 해서 얻은 개선이다. 그 이유는 우리가 나쁜 수를 좋은 수로 바꿨기 때문이다. 이 설명은 잠시 뒤에 다시 하겠다. 일단은 우리가 이렇게 해도 아무것도 망가뜨리지 않는다는 걸 스스로 납득해 보자.

첫째, 우리는 토큰용 임베딩 테이블인 wte를 더 크게 만들었다. 말하자면 맨 아래에 토큰을 더 집어넣은 셈이다. 그런데 이 토큰들은 결코 쓰이지 않는다. GPT 토크나이저는 50256까지의 토큰만 가지고 있기 때문에, 우리가 추가한 행들을 인덱싱할 일이 절대 없다. 그러니 우리는 결코 접근되지도, 사용되지도 않을 메모리를 만들어서 공간을 약간 낭비하고 있는 것이다.

그런데 이게 완전히 맞는 말은 아니다. 이 wte 가중치는 결국 공유되어 맨 끝의 분류기에서도 쓰이기 때문이다. 그러면 여기 이 분류기에는 무슨 일이 일어나는가. 이제 분류기에서 우리는 추가된 차원들까지 예측하게 된다. 즉 학습 데이터에는 결코 등장하지 않을 토큰들에 대한 확률을 예측하게 된다. 따라서 네트워크는 그 확률들이 0으로 가야 한다는 것을 학습해야 한다. 다시 말해 네트워크가 만들어 내는 logit이 출력의 그 차원들을 음의 무한대(−∞)로 밀어내야 한다.

하지만 이건 이미 우리 데이터셋에 없는 다른 모든 토큰들과 전혀 다르지 않다. 셰익스피어 데이터는 50257개 토큰 중 아마 천 개 정도만 쓸 것이다. 그러니 대부분의 토큰은 이미 최적화에 의해 확률 0으로 밀려나고 있다. 우리는 단지 같은 방식으로 결코 쓰이지 않고 확률 0으로 밀려날 토큰을 몇 개 더 추가했을 뿐이다. 그러므로 기능적으로는 아무것도 깨지지 않는다. 메모리를 약간 더 쓰긴 하지만, 그 외에는 내가 보기에 무해한 연산이다.

그리고 우리는 계산을 더 추가했는데도 더 빨리 돌고 있다. 왜 더 빠른가. 앞서 말했듯이 CUDA에서는 정말 많은 커널이 블록 타일(block tile)을 쓰고, 이 블록 타일은 보통 좋은 수, 즉 2의 거듭제곱이다. 그래서 계산이 64 단위나 32 단위 덩어리로 이루어진다. 그런데 여러분이 원하는 계산이 그 블록 타일에 깔끔하게 들어맞지 않으면, 마지막 남은 부분을 처리하기 위해 온갖 경계 처리 커널들이 끼어든다. 기본적으로 많은 커널에서 입력을 덩어리로 쪼개 좋은 부분을 먼저 처리한 다음, 완전히 별개의 두 번째 단계로 돌아와서 남은 부분을 처리한다. 그런데 그걸 담당하는 커널들은 매우 비효율적일 수 있다. 결국 이 모든 추가 연산을 띄우게 되고 극도로 비효율적이다. 그럴 바에는 차라리 입력을 패딩해서 깔끔하게 들어맞게 만드는 게 낫고, 경험적으로 보통 그쪽이 실제로 더 빠르게 돌아간다.

그래서 이것도 우리가 얻어낸 4% 개선의 한 예다. 그리고 이것 역시 torch.compile이 우리 대신 찾아 주지 못한 것이다. 언젠가는 torch.compile이 이런 최적화를 알아낼 수 있게 되기를 바라지만, 지금으로서는 이 상태다.

한 가지 더 짚어야 할 게 있다. 우리는 PyTorch nightly를 쓰고 있고, 그래서 4%밖에 못 보고 있는 것이다. PyTorch 2.3.1이나 그 이전 버전을 쓴다면, 50257을 50304로 바꾸는 이 변경 하나만으로 30% 정도의 개선을 실제로 보게 될 것이다. 그러니 이것도 내부가 어떻게 돌아가는지를 이해해야 하는 이유, 그리고 코드 성능을 밀어붙이기 위해 어떤 것들을 만지작거려야 하는지를 알아야 하는 이유를 보여 주는 내가 가장 좋아하는 사례 중 하나다.

SECTION 3: 하이퍼파라미터, AdamW, gradient clipping 02:14:55

여기까지 오면서 우리는 성능을 대략 11배 개선했다. 그렇다, 11배다. 처음 시작할 때는 스텝당 대략 1,000밀리초였는데 지금은 93밀리초 정도까지 내려왔으니 말이다. 그러니 이건 꽤 괜찮은 결과다. 그리고 우리는 GPU 자원을 훨씬 더 잘 활용하는 상태가 되었다. 이제부터는 방향을 바꿔서, 최적화 그 자체에 대한 좀 더 알고리즘적인 변경과 개선으로 넘어가려 한다. 우리가 하고 싶은 일은 GPT-2 논문, 아니 정확히는 GPT-3 논문에 언급된 하이퍼파라미터를 그대로 따라가는 것이다.

그런데 안타깝게도 GPT-2 논문은 이 부분에 대해 사실 별로 말해주는 게 없다. 그들이 모델 가중치와 코드를 공개해준 것은 정말 고마운 일이다. 하지만 논문 자체는 최적화 세부사항에 관한 한 극도로 모호하다. 그리고 그들이 함께 공개한 코드도 마찬가지다. 우리가 지금까지 계속 들여다보던 그 코드는 그냥 추론(inference) 코드일 뿐이다. 학습 코드는 거기에 들어 있지 않고, 하이퍼파라미터도 극히 일부만 나와 있다. 그러니까 그 코드 역시 우리에게 많은 것을 알려주지는 못한다.

그래서 그 부분을 채우려면 우리는 GPT-3 논문으로 눈을 돌려야 한다. GPT-3 논문의 부록으로 들어가 보면, 우리가 가져다 쓸 수 있는 하이퍼파라미터가 훨씬 더 많이 들어 있다. 그리고 GPT-3 논문은 전반적으로 훨씬 더 상세하다. 모델 학습에 들어가는 온갖 자잘한 세부사항까지 다 적어두었다. 다만 GPT-3 모델들은 끝내 한 번도 공개되지 않았다. 그러니 상황을 정리하면 이렇다. GPT-2는 가중치는 있는데 세부사항이 없고, GPT-3는 세부사항은 많은데 가중치가 없다.

그래도 대략적으로 말해서 GPT-2와 GPT-3의 아키텍처는 아주 아주 유사하다. 기본적으로 바뀐 것이 거의 없다. 컨텍스트 길이가 1024에서 2048로 확장된 것이 일종의 주요 변경이고, Transformer 주변의 하이퍼파라미터 몇 개가 바뀐 정도다. 그 외에는 사실상 거의 같은 모델이다. 차이는 GPT-3가 훨씬 더 큰 데이터셋 위에서 훨씬 더 오래 학습되었다는 것, 그리고 훨씬 더 철저한 평가가 붙어 있다는 것이다. 또 GPT-3 모델은 1,750억 파라미터인 반면 GPT-2는 16억 파라미터다. 긴 이야기를 짧게 줄이자면, 우리는 하이퍼파라미터 몇 가지를 따라가기 위해 GPT-3 논문으로 간다.

논문에는 이렇게 적혀 있다. "GPT-3의 모든 버전을 학습시키기 위해 우리는 β₁ = 0.9, β₂ = 0.95인 Adam을 사용한다." 그러니 코드로 돌아가서 확인해보자. 여기 betas 인자가 보이는데 기본값이 0.9와 0.999로 되어 있다. 이걸 실제로 0.9와 0.95로 설정해주자. 그 다음 eps 인자를 보면 기본값이 10⁻⁸이다. 그리고 논문에서 말하는 값도 마찬가지로 ε = 10⁻⁸이다. 같은 값이긴 하지만 그냥 명시적으로 적어 넣자. 그러면 이 부분은 그대로 반영된다.

다음으로 논문은 이렇게 말한다. "우리는 그래디언트의 global norm을 1.0에서 클리핑한다." 이게 무슨 이야기냐면 이런 뜻이다. 우리가 그래디언트를 계산하고 나면, 그러니까 loss.backward() 바로 다음이다, 기본적으로 모든 파라미터 텐서에 그래디언트가 붙어 있는 상태가 된다. 그리고 사람들이 여기서 흔히 하고 싶어 하는 일은, 그 그래디언트들이 어떤 최대 norm을 갖도록 잘라내는 것이다. PyTorch에서는 이게 꽤 쉬운 편이다. 그래디언트를 계산한 직후에 코드 한 줄만 여기 끼워 넣으면 된다.

norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

그리고 이 유틸리티 함수가 하고 있는 일은 파라미터들의 global norm을 계산하는 것이다. 즉 모든 파라미터 하나하나에 붙어 있는 모든 그래디언트를 가져와서, 그것을 제곱하고, 그 값을 전부 더하고, 그 합에 큰 제곱근을 하나 씌운다. 그리고 그게 바로 그 파라미터 벡터의 norm이다. 그러니까 그건, 여러분이 그런 식으로 바라보고 싶다면, 그 벡터의 길이다. 그리고 우리가 기본적으로 하고 있는 일은 그 길이가 1.0보다 크지 않도록 확실히 해두는 것이다. 그리고 만약 크다면 우리는 그것을 잘라낼 것이다. 그게 클리핑이다.

사람들이 이걸 쓰기 좋아하는 이유는 이렇다. 최적화 도중에 가끔 운이 나쁠 수가 있다. 어쩌면 나쁜 데이터 배치가 걸릴 수도 있고, 뭐 그 비슷한 일이 생길 수도 있다. 어떤 배치에서 아주 운이 나쁘면 정말 높은 손실이 나올 수 있고, 정말 높은 손실은 정말 높은 그래디언트로 이어질 수 있다. 그리고 이건 기본적으로 여러분의 모델에 충격을 주고 최적화에 충격을 줄 수 있다. 그래서 사람들은 gradient norm clipping을 쓴다. 모델이 그래디언트 크기 측면에서 너무 큰 충격을 받지 않도록 막고, 이런 식으로 상한을 씌우는 것이다. 사실 이건 좀 해킹에 가까운 해법이다. 더 깊은 문제 위에 패치를 하나 얹는 것에 가깝다. 하지만 그래도 사람들은 여전히 꽤 자주 쓴다.

그리고 clip_grad_norm_은 그래디언트의 norm을 반환한다. 나는 이 값을 항상 시각화해두는 걸 좋아한다. 유용한 정보이기 때문이다. 그래디언트의 norm을 들여다봤을 때 그게 잘 behave하고 있으면 상황이 좋은 것이다. 반대로 계속 올라가고 있으면 상황이 나쁜 것이고, 학습 중에 불안정해지고 있다는 뜻이다. 가끔은 norm에 스파이크가 튀기도 하는데, 그건 뭔가 문제가 있거나 불안정성이 있다는 의미다. 그러니 여기서 반환값을 norm이라는 이름으로 받자. 그리고 출력 포맷은 .4f 같은 걸로 하자. 이건 그냥 float일 거라고 생각하니, 그대로 출력할 수 있을 것이다. 자, 이렇게 해서 global gradient clipping이 붙었다.

그 다음 논문은 학습률 스케줄러의 세부사항으로 들어간다. 그들은 지금 우리가 여기서 하듯이 3 × 10⁻⁴ 같은 고정 학습률만 쓰는 게 아니다. 실제로는 기본적으로 코사인 감쇠(cosine decay) 학습률 스케줄을 쓴다. warmup이 붙어 있고, 어떤 구간(horizon)에 걸쳐 10%까지 코사인 감쇠하는 형태다. 이건 잠시 뒤에 구현할 예정이다. 그 전에 일단 norm이 출력되는 걸 좀 보고 싶다.

좋다, 나왔다. 여기서 무슨 일이 일어났는지 보자. norm이 초반에 정말 높다. 30 정도 된다. 그리고 학습을 계속 진행하면서 1 아래 값들로 안정화되는 게 보인다. 사실 아주 초기 단계에서 norm이 높은 것 자체는 그렇게까지 드문 일이 아니다. 여기서 벌어지고 있는 일은 기본적으로 이렇다. 모델이 완전히 랜덤한 상태이기 때문에 네트워크 아주 초반에 엄청나게 많은 학습이 일어난다. 그런데 그 학습이라는 게 대부분 출력 토큰들의 bias를 배우는 일이다. 그래서 좀 불안정한 시기다. 하지만 네트워크는 보통 아주 적은 반복 안에 안정화된다. 그래서 이 그림은 나에게는 상대적으로 꽤 합리적으로 보인다. 다만 보통 내가 기대하는 것과 비교하면 살짝 이상해 보이는 부분이 있다. 28에서 6으로 갔다가 2로 내려가더니 다시 10으로 튀는 게 좀 묘하다. 완전히 정신 나간 값은 아니지만 그냥 조금 웃긴 정도다. 좋다, 이제 학습률 스케줄러로 넘어가자.

학습률 스케줄러: warmup + 코사인 감쇠 02:21:06

여기 GPT-3에서 사용되는 학습률 스케줄은 warmup이 붙은 코사인 감쇠 학습률 스케줄이라고 불리는 것이다. 모양은 이렇게 생겼다. 학습률이 기본적으로 0 근처에서 시작해서 일정 시간 동안 선형으로 올라간다. 그리고 그 다음에는 코사인 같은 형태를 그리며 내려와서, 어떤 최소 학습률에 도달한다. 그 최소값이 얼마인지는 여러분이 정하기 나름이다. 여기 이 그림에서는 최소 학습률이 0이다. 하지만 GPT-3 논문에서는 첫 2,600억 토큰에 걸쳐 학습률을 원래 값의 10%까지 코사인 감쇠시킨다고 말한다. 그리고 그 이후로는 10% 상태로 학습을 계속 이어간다. 또 첫 3억 7,500만 토큰에 걸쳐 선형 warmup이 들어간다. 학습률에 관한 이야기는 대략 그렇다.

이제 이걸 구현해보자. 사실 나는 이미 여기에 구현해두었다. 어떻게 작동하는지 보기 위해 먼저 아래로 스크롤해보자. 학습 루프를 조금 바꿨다. 원래 for i in range(max_steps)였던 것을 step이라는 이름으로 바꿨다. for 루프 안에서 한 스텝이 하나의 최적화 스텝이라는 개념을 갖게 하기 위해서다. 그리고 여기서 get_lr이라는 새 함수를 써서 이번 최적화 스텝의 학습률을 가져온다. 그리고 PyTorch에서 학습률을 설정하는 방법은, 내가 알기로는 이런 식이다. 좀 지저분한 편이다. 옵티마이저 안에 서로 다른 파라미터 그룹이라는 개념이 존재할 수 있기 때문에, 지금 우리는 파라미터 그룹이 하나뿐인데도 그것들을 순회해야 한다. 그리고 이 for 루프 스타일로 lr을 설정해줘야 한다. 적어도 지금 내가 받은 인상으로는 그렇다.

lr = get_lr(step)
for param_group in optimizer.param_groups:
    param_group['lr'] = lr

그래서 이렇게 학습률을 가져오는 부분이 있고, 학습률을 설정해주고, 그리고 아래쪽에서 그 값도 함께 출력한다. 이 루프에 내가 가한 변경은 그게 전부다. 그리고 물론 get_lr이 내가 만든 스케줄러다.

여기서 짚고 넘어갈 만한 게 있다. PyTorch에도 실제로 학습률 스케줄러들이 있고, 그걸 그냥 써도 된다. 그리고 PyTorch 안에 코사인 학습률 스케줄도 있는 걸로 알고 있다. 다만 나는 그 코드를 쓰는 걸 별로 좋아하지 않는다. 솔직히 말해서 이건 다섯 줄짜리 코드이고, 그 줄들 안에서 무슨 일이 벌어지는지 나는 완전히 이해하고 있기 때문이다. 나는 속을 들여다볼 수 없는 추상화를 써놓고 그게 뭘 하는지 모르는 상태가 되는 걸 좋아하지 않는다. 이건 개인적인 스타일이다.

그래서 여기 최대 학습률은 일단 3 × 10⁻⁴라고 해두자. 그런데 GPT-3 논문 여기를 보면 모델 크기별로 최대 학습률이 얼마인지 적힌 표가 있다. 우리 경우를 보면, 기본적으로 12 레이어에 768 차원인 GPT-3, 그러니까 GPT-3 small은 대략 GPT-2 124M에 해당한다. 그리고 여기서 보면 그들은 6 × 10⁻⁴의 학습률을 쓴다. 그러니까 우리도 사실 더 올려도 된다. 실제로 그걸 따라가보는 것도 좋겠다. 여기 최대 학습률을 그냥 6으로, 즉 6 × 10⁻⁴로 설정하자. 그게 최대 학습률이다. 그리고 최소 학습률은 논문 설명대로 그것의 10%다. 그리고 warmup을 진행할 스텝 수가 있고, 최적화의 최대 스텝 수가 있는데, 이 값은 이제 아래쪽 for 루프에서도 함께 쓴다.

그리고 원한다면 이 코드를 한번 훑어봐도 된다. 딱히 통찰이 넘치도록 흥미로운 부분은 아니다. 나는 그냥 반복 번호를 기준으로, 학습률이 어떤 값이어야 하는지를 조절하고 있을 뿐이다. 그러니까 여기 이 부분이 warmup 구간이고, 여기 이 부분은 최적화가 다 끝난 뒤의 구간이며, 그리고 여기 이 부분이 말하자면 그 둘 사이에 놓인 구간이다. 그리고 바로 여기가 내가 코사인 학습률 스케줄을 계산하는 자리다. 원한다면 여기를 한 줄씩 짚어가며 자세히 들여다볼 수 있다. 하지만 기본적으로 이건 아까 봤던 그 곡선을 구현하고 있는 것이다.

이미 실행해봤는데 결과는 이렇게 생겼다. 실행해보면 아주 낮은 숫자에서 시작한다. 여기서 주의할 점은, 우리가 정확히 0에서 시작하지는 않는다는 것이다. 학습률 0으로 업데이트를 하는 건 아무 쓸모가 없기 때문이다. 그래서 it + 1이 들어 있는 것이다. 0번째 반복에서 정확히 0을 쓰지 않고 아주 아주 낮은 값을 쓰도록 하기 위해서다. 그 다음 최대 학습률까지 선형으로 warmup한다. 내가 돌렸을 때 그 값은 3 × 10⁻⁴였지만, 이제는 6 × 10⁻⁴가 될 것이다. 그리고 그때부터 감쇠하기 시작해서 3 × 10⁻⁵까지 쭉 내려간다. 그건 당시 기준으로 원래 학습률의 10%였다.

한 가지 우리가 정확히 따르지 않는 부분이 있다. 다시 찾아볼 수 있는지 보자. 우리가 그들이 한 것을 정확히 따르지 않는 이유는 이렇다. 그들은 자기들의 학습 지평(training horizon)이 3,000억 토큰이라고 말한다. 그런데 2,600억 토큰 지점에서 초기 학습률의 10%까지 내려오고, 2,600억 이후로는 10%로 학습한다. 그러니까 기본적으로 그들의 감쇠 시간은 최대 스텝 시간보다 짧다. 반면 우리 쪽에서는 둘이 정확히 같다. 그래서 완전히 충실한 재현은 아니다. 하지만 이건 우리에게, 그리고 지금 우리 목적에는 괜찮은 수준이다. 우리는 그냥 이대로 쓸 것이다. 솔직히 이게 그렇게까지 큰 차이를 만들 것 같지는 않다.

그리고 짚어두어야 할 게 있다. 어떤 학습률 스케줄을 쓸지는 전적으로 여러분에게 달린 문제다. 종류가 아주 많다. 코사인 학습률은 GPT-2와 GPT-3 덕분에 많이 대중화되었다. 하지만 사람들은 온갖 종류의 다른 학습률 스케줄을 고안해왔다. 그리고 이건 어떤 것이 이런 네트워크를 학습시키는 데 가장 효과적인가에 대해 지금도 활발히 연구가 진행 중인 영역이다. 좋다, 다음으로 논문이 이야기하는 내용으로 넘어가자.

배치 크기 스케줄, weight decay, fused AdamW, 90ms 02:26:21

논문은 점진적인 배치 크기 증가에 대해 이야기한다. 그러니까 배치 크기에 선형 램프가 걸려 있어서, 아주 작은 배치 크기로 시작해 시간이 지나면서 큰 배치 크기까지 끌어올리는 것이다. 우리는 사실 이건 건너뛸 것이고, 이걸 가지고 작업하지는 않을 것이다. 내가 이걸 쓰는 걸 별로 좋아하지 않는 이유는, 산술을 많이 복잡하게 만들기 때문이다. 최적화의 매 스텝마다 여러분이 처리하는 토큰 수가 바뀌게 되는데, 나는 그 계산을 아주 아주 단순하게 유지하는 걸 좋아한다.

또 내가 이해하기로는 이게 그렇게 대단한 개선도 아니다. 그리고 역시 내가 이해하기로는, 이건 알고리즘적인 최적화 개선이 아니라 시스템과 속도 측면의 개선에 더 가깝다. 대략적으로 말하자면 이유는 이렇다. 최적화의 초기 단계에서 모델은 다시 말하지만 아주 비전형적인 상태에 놓여 있다. 그리고 이때 여러분이 주로 배우고 있는 것은, 학습 데이터에 자주 등장하지 않는 토큰들을 무시하는 법이다. 아주 단순한 bias 같은 것들, 뭐 그런 종류의 것을 배우고 있는 것이다. 그래서 여러분이 네트워크에 집어넣는 모든 예제 하나하나가 기본적으로 "이 토큰들은 써라, 저 토큰들은 쓰지 마라"라고 말해주는 것에 지나지 않는다. 따라서 예제 하나하나에서 나오는 그래디언트가 실제로 극도로 높은 상관관계를 갖는다. 최적화의 원래 초기 구간에서는 전부 대충 같아 보인다. 전부 그냥 이 토큰들은 나타나지 않고 저 토큰들은 나타난다는 이야기만 하고 있기 때문이다.

그래디언트가 전부 아주 비슷하고 높은 상관관계를 가진다면, 도대체 왜 수백만짜리 배치 크기를 쓰는가? 3만 2천 정도의 배치 크기만 써도 학습 초기에는 기본적으로 정확히 같은 그래디언트를 얻는데 말이다. 그러다 최적화 후반으로 가서 단순한 것들을 다 배우고 나면, 그때부터 진짜 일이 시작된다. 그리고 바로 그때부터 예제별 그래디언트가 서로 탈상관되기 시작하고, 어떤 의미에서 실제로 통계적 힘을 제공하기 시작한다. 어쨌든 우리는 이게 일을 복잡하게 만들기 때문에 그냥 건너뛴다.

그리고 다음으로 넘어가면, 논문은 "학습 동안 데이터는 복원 없이(without replacement) 샘플링된다"고 말한다. 에포크 경계에 도달할 때까지는 그렇다는 뜻이다. 복원 없이 샘플링한다는 것은 이런 의미다. 그들은 고정된 풀에서 뽑아서 시퀀스를 하나 가져다 학습한 뒤, 그 시퀀스를 다시 풀에 돌려놓는 식이 아니다. 그들은 풀을 소진시킨다. 그래서 시퀀스를 하나 뽑으면 다음 학습 에포크가 올 때까지 그건 사라진 상태다. 우리는 이미 그렇게 하고 있다. 우리 데이터 로더는 데이터 청크들을 순회하기 때문에 복원이 없다. 다음 에포크가 되기 전까지는 다시 뽑힐 자격이 생기지 않는다. 그러니까 우리는 기본적으로 이미 그걸 하고 있는 셈이다.

그 다음 항목은 이렇다. "모든 모델은 약간의 정규화를 제공하기 위해 0.1의 weight decay를 사용한다." 그러니 weight decay를 구현해보자. 그리고 보면 알겠지만 나는 이미 어느 정도 변경을 해두었다. 특히 옵티마이저를 바로 여기서 만드는 대신, 모델 안에 configure_optimizers라는 새 함수를 만들고 하이퍼파라미터 일부를 인자로 넘기고 있다. 그러니 옵티마이저 객체를 반환하게 되어 있는 configure_optimizers를 한번 보자.

복잡해 보이지만 사실은 정말 단순하다. 우리는 그냥 아주 조심스럽게 처리하고 있을 뿐이고, 짚고 넘어갈 설정이 몇 개 있을 뿐이다. 이 줄과 관련해 가장 중요한 건, 여기 weight_decay 파라미터가 있고 내가 그걸 optim_groups라는 것에 넘기고 있으며, 그게 결국 AdamW 옵티마이저 안으로 들어간다는 점이다. 그리고 여기 AdamW에서 기본으로 쓰이는 weight decay는 0.01이다. 즉 GPT-3 논문에서 쓰는 값보다 10배 낮다. 그래서 weight decay는 optimizer group을 통해 AdamW까지 전달되는 구조다.

이 함수에서 또 무슨 일이 벌어지고 있나? 여기서 중요한 건 두 가지다. 첫째, 나는 파라미터를 weight decay를 적용해야 할 것과 적용하지 말아야 할 것으로 나누고 있다. 특히 bias와 그 밖의 1차원 텐서들에는 weight decay를 적용하지 않는 것이 일반적이다. 그래서 1차원 텐서들은 nodecay_params에 들어간다. 그리고 여기에는 LayerNorm의 스케일과 bias 같은 것들도 포함된다. 그런 것들에 weight decay를 거는 건 별로 말이 되지 않는다. 여러분이 주로 weight decay를 걸고 싶은 대상은 행렬 곱셈에 참여하는 가중치들이다. 그리고 임베딩에도 걸 만하다.

그리고 가중치를 decay시키는 것이 왜 말이 되는지는 이전 영상에서 이미 다뤘다. 여러분은 이것을 일종의 정규화로 볼 수 있다. 왜냐하면 모든 가중치를 아래로 끌어내리고 있으면, 여러분은 최적화가 더 많은 가중치를 사용하도록 강제하게 되기 때문이다. 그리고 어느 한 가중치가 혼자서만 지나치게 커지는 것을 허용하지 않게 된다. 여러분은 네트워크가 그 일을 더 많은 채널에 걸쳐 분산하도록 강제하고 있는 것이다. 가중치 자체에 일종의 중력 같은 당김이 걸려 있기 때문이다. 그래서 우리는 여기서 이런 방식으로 그것들을 분리한다. 우리는 임베딩과 행렬 곱에 참여하는 가중치들만 decay시킨다. 그리고 우리는 decay하고 있는 파라미터의 개수와 그렇지 않은 파라미터의 개수를 출력한다. 그리고 대부분의 파라미터는 decay될 것이다.

여기서 하는 일이 하나 더 있다. 또 다른 최적화를 하나 넣었다. 예전 AdamW에는 이 옵션이 없었지만, 이후 버전의 PyTorch에서 도입되었다. 그래서 내가 inspect.signature로 가드를 걸어둔 것이다. 이건 기본적으로 fused라는 인자가 AdamW 안에 존재하는지 확인한다. 그리고 존재하면 그걸 쓰고 여기로 넘긴다. 일부 옛 버전에는 fused=가 없기 때문이다. 그러니까 여기 AdamW의 fused=는 원래 존재하지 않았고 나중에 추가되었다. 그리고 여기 무슨 일이 일어나는지 설명하는 문서가 있다. 기본적으로 그들은 fused가 비교적 새로운 기능이라 충분한 숙성 시간을 주고 싶기 때문에 기본값으로는 쓰지 않는다고 말한다. 그러니까 기본은 fused를 쓰지 않는 것이다. 하지만 fused는 사용 가능한 상황이고 CUDA 위에서 돌고 있다면 훨씬 빠르다.

fused가 하는 일은 이렇다. 원래는 모든 파라미터 텐서를 for 루프로 순회하면서 그것들을 업데이트하는데, 그렇게 하면 커널이 아주 많이 런치된다. 그렇지 않은가. 그래서 fused는 그 모든 커널이 하나의 커널로 융합된다는 뜻이다. 오버헤드가 아주 많이 사라지고, 모든 파라미터에 대해 단 한 번 커널을 호출해서 그것들을 업데이트한다. 그러니까 기본적으로 텐서들을 하나씩 순회하는 대신 AdamW 업데이트에 대해 커널 퓨전을 하는 것이다. 이게 내가 즐겨 쓰는 configure_optimizers 함수다.

다시 실행해보면 이전에 봤던 것과 큰 차이는 보이지 않을 것이다. 하지만 여기서 나오는 출력 몇 개는 보게 될 것이다. 그러니 그게 어떻게 생겼는지 한번 보자. decay되는 텐서의 수는 50개이고, 이게 대부분의 파라미터다. 그리고 decay되지 않는 텐서의 수는 98개인데, 이건 주로 bias와 LayerNorm 파라미터들이다. 그리고 그건 10만 개밖에 되지 않는다. 그러니까 대부분이 decay된다. 그리고 우리는 AdamW의 fused 구현을 쓰고 있는데, 이게 훨씬 빠를 것이다. 그러니 가능하다면 쓰라고 권하고 싶다. 사실 나도 그들이 왜 이걸 기본값으로 하지 않는지 100% 확신하지는 못하겠다. 꽤 무해하고 해롭지 않아 보이는데 말이다.

그리고 우리가 fused 구현을 쓰고 있기 때문에, 내 생각에는 바로 이것 덕분에 시간이 줄어든 것 같다. 보다시피 실행 시간이 원래 스텝당 93밀리초였는데, 지금은 fused AdamW 옵티마이저를 쓰기 때문에 스텝당 90밀리초로 내려왔다. 그러니까 이 하나의 커밋 안에서 우리는 fused Adam을 도입해 시간에서 개선을 얻었고, 동시에 weight decay를 추가하고 바꿨다. 다만 우리는 2차원 파라미터, 즉 임베딩과 linear에 참여하는 행렬들에만 weight decay를 건다. 그게 이 부분이고, 이건 이제 지워도 된다. 그래, 이 항목은 여기까지다. 계속 진행하기 전에 짧은 메모 하나만 더 하겠다.

gradient accumulation 02:34:09

계속 진행하기 전에 여기서 한 가지만 짚고 싶다. weight decay, 학습률, 배치 크기, 그리고 Adam의 파라미터인 β₁, β₂, ε 등등 사이의 관계는 최적화 문헌에서 아주 복잡한 수학적 관계다. 그리고 대체로 이 영상에서 나는 그냥 OpenAI가 사용한 설정을 복사해 붙여넣으려 하고 있을 뿐이다. 하지만 이건 복잡한 주제이고 꽤 깊은 주제다. 그래서 이 영상에서는 그냥 파라미터를 복사하고 싶다. 그걸 제대로 다루면서 고수준 직관만 늘어놓는 게 아니라 정당하게 대접하려면, 그건 완전히 별개의 영상이 되어야 하기 때문이다.

이제 다음으로 넘어가고 싶은 건 이거다. 참고로 여기 이 문단은 나중에 데이터 로더를 개선할 때 다시 돌아올 것이다. 지금은 이 표로 다시 돌아가고 싶다. 표를 보면, 당연히 모델마다 Transformer 네트워크의 크기를 결정하는 하이퍼파라미터가 다르다. 그리고 학습률도 다르다. 여기서 우리는 큰 네트워크일수록 약간 더 낮은 학습률로 학습된다는 패턴을 보게 된다. 그리고 배치 크기도 보이는데, 작은 네트워크에서는 더 작은 배치 크기를 쓰고 큰 네트워크에서는 더 큰 배치 크기를 쓴다.

그런데 우리 입장에서 문제는, 50만짜리 배치 크기를 그냥 쓸 수가 없다는 것이다. 여기 들어와서 이 B를 설정하려고 하면 — 내 B가 어디 있더라, 데이터 로더를 어디서 부르지, 그래 여기 B = 16이 있다 — 여기서 조심해야 한다. 50만이 그대로 들어가는 게 아니다. 왜냐하면 이건 토큰 수로 센 배치 크기이기 때문이다. 우리 행 하나하나가 1024 토큰이다. 그래서 0.5 × 10⁶, 즉 100만의 절반을 1024로 나누면 대략 488 정도의 배치 크기가 필요하다는 계산이 나온다. 그런데 문제는 내가 여기에 488이라고 적어 넣을 수 없다는 것이다. 내 GPU가 터질 테니까. 이건 확실히 메모리에 들어가지 않는다.

그런데도 우리는 여전히 이 배치 크기를 쓰고 싶다. 앞서 말했듯이 배치 크기는 다른 모든 최적화 하이퍼파라미터, 그리고 학습률 같은 것들과 상관되어 있기 때문이다. 그래서 우리는 하이퍼파라미터 전체를 충실하게 재현하고 싶고, 따라서 대략 50만 토큰짜리 배치 크기를 써야 한다. 하지만 질문은 이것이다. 작은 GPU밖에 없는데 어떻게 50만을 쓰는가? 바로 그걸 위해 필요한 것이 이른바 gradient accumulation이다. 그래서 다음으로 그걸 다루려 한다. gradient accumulation은 우리가 설정한 임의의 배치 크기를 직렬(serial) 방식으로 시뮬레이션할 수 있게 해준다. 그래서 우리는 50만짜리 배치 크기를 쓸 수 있다. 다만 더 오래 돌려야 하고, 여러 시퀀스를 처리하면서 기본적으로 거기서 나온 그래디언트를 전부 더해서 50만 배치 크기를 시뮬레이션해야 한다. 그러면 그걸 다음으로 해보자.

구현은 바로 여기, 이 몇 줄을 추가하는 것으로 시작했다. 내가 한 일은 기본적으로 이렇다. 먼저 우리가 원하는 총 배치 크기를 설정한다. 이건 정확히 50만이다. 그리고 나는 2의 거듭제곱이라는 예쁜 숫자를 썼다. 2¹⁹ = 524288이니까 대략 50만이고, 좋은 숫자다.

total_batch_size = 524288  # 2**19, ~0.5M tokens
B = 16   # micro batch size
T = 1024 # sequence length
grad_accum_steps = total_batch_size // (B * T)

그리고 이제 우리가 마이크로 배치 크기라고 부르는 값은 16이다. 그러니까 이건 이런 뜻이다. 여전히 B × T짜리 시퀀스들이 Transformer로 들어가서 forward와 backward를 한다. 하지만 우리는 업데이트를 하지 않는다. 그렇다, 우리는 forward-backward를 여러 번 할 것이고, 그 그래디언트들은 전부 파라미터 그래디언트 위에 +=로 더해질 것이다. 전부 합산되는 것이다. 그래서 우리는 grad_accum_steps 횟수만큼 forward-backward를 하고, 그게 전부 누적되고 나면 그때 단 한 번 업데이트를 한다.

특히 우리의 마이크로 배치 크기는 이제 한 번의 forward-backward에서 몇 개의 토큰을, 몇 개의 행을 처리하느냐만 제어한다. 그래서 여기서 우리는 16 × 1024를 하고 있다. 즉 한 번의 forward-backward당 16,384 토큰을 처리하고 있는 것이다. 그리고 우리가 해야 하는 총합은 2¹⁹이다. 잠깐, 내가 지금 뭘 하고 있지. 총합이 2¹⁹이니까 gradient accumulation은 32가 된다. 그러니까 여기서 grad_accum_steps는 32로 계산된다. 그리고 우리는 32번의 forward-backward를 하고, 그 다음에 단 한 번의 업데이트를 해야 한다. 그리고 우리는 한 번의 forward-backward에 대략 100밀리초가 걸린다는 걸 알고 있다. 그러니 그걸 32번 하면 매 스텝이 대략 3초쯤 걸리게 될 것이다. 그냥 냅킨 계산이다.

자, 이게 grad_accum_steps인데 이제 실제로 그걸 구현해야 한다. 그러니 학습 루프로 가보자. 이제 여기 이 부분과 여기 이 부분, 즉 forward와 backward를, 그 뒤에 따라오는 다른 모든 것을 하기 전에 32번 반복해야 한다. 그러면 그걸 어떻게 구현할 수 있는지 보자. 여기로 오자. 그리고 사실 우리는 매번 새 배치를 로드해야 한다. 그러니 그 줄을 이쪽으로 옮기자. 그리고 이제 여기가 안쪽 루프가 있는 자리다. for micro_step in range(grad_accum_steps): 안에서 이 일을 하는 것이다.

그리고 기억하라. loss.backward()는 항상 그래디언트를 쌓아 넣는다(deposit). 즉 loss.backward() 내부에는 그래디언트에 대한 +=가 항상 들어 있다. 그래서 매번 loss.backward()를 호출할 때마다 그래디언트 텐서 위에 그래디언트가 더해진다. 그러니 loss.backward()를 하고, 그러면 우리는 그 모든 그래디언트를 저기 얻게 되고, 그 다음에 정규화를 하고 나머지는 그대로 따라오면 된다. 우리는 아주 가까이 왔다. 그런데 사실 여기에는 미묘하고 깊은 문제가 하나 있다. 그리고 지금 이 코드는 틀렸다. 그러니 왜 이것만으로는 아직 충분하지 않은지 한번 생각해보기를 권한다. 그럼 이제 고쳐보겠다.

이해를 돕기 위해 Jupyter 노트북을 다시 가져왔다. 단순한 장난감 상황 안에서 무슨 일이 벌어지는지 조심스럽게 생각해보자. 아주 간단한 신경망을 하나 만들자. 16개 숫자로 이루어진 벡터를 받아서 숫자 하나를 반환한다. 그리고 여기서 랜덤한 예제 x와 타깃 y를 몇 개 만든다. 그리고 여기서 mean squared loss를 써서 손실을 계산한다. 그러니까 기본적으로 이건 네 개의 개별 예제이고, 우리는 그 네 예제에 대해 mean squared loss로 단순 회귀를 하고 있는 것이다.

이제 우리가 손실을 계산하고 loss.backward()를 한 다음 그래디언트를 들여다보면, 이게 바로 우리가 얻어내는 그래디언트다. 그리고 이제 여기 손실 목적함수를 보자. MSELoss에서 손실 함수의 기본값은 reduction='mean'이라는 점에 주목하라. 그러니까 우리는 평균 손실을 계산하고 있는 것이다. 여기 네 개의 예제에 대한 평균 손실 말이다. 그래서 이게 정확한 손실 목적함수이고, 이건 평균이다. 즉 4분의 1이 곱해진 것이다. 왜냐하면 여기에는 독립적인 예제가 네 개 있기 때문이다. 그리고 그 다음 우리는 네 개의 예제를 갖고 있고 각각에 대한 mean squared error, 즉 제곱 오차를 갖고 있다. 그리고 그것이 이것을 mean squared error로 만들어준다. 그러니까 따라서 우리는 제곱 오차를 계산한 다음, 그것을 예제들에 대한 평균이 되도록 정규화하는 것이다. 그리고 여기에는 예제가 네 개 있다.

이제 gradient accumulation 버전으로 와보자. 여기 이게 gradient accumulation 버전인데, grad_accum_steps가 4이고 나는 그래디언트를 초기화한 상태다. grad_accum_steps가 4이고, 이제 나는 예제들을 하나씩 개별적으로 평가하면서 그것들에 대해 loss.backward()를 여러 번 호출한다. 그리고 그렇게 해서 얻은 그래디언트를 들여다본다. 그러니까 기본적으로 이제 우리는 함수를 forward하고, 정확히 같은 손실을 계산하고, backward를 한다. 그리고 그걸 네 번 한다. 그리고 그래디언트를 보면, 두 그래디언트가 서로 일치하지 않는다는 걸 알게 될 것이다. 여기서 우리는 크기 4짜리 배치 하나를 처리했고, 저기서는 배치 크기 1짜리 gradient accumulation 스텝을 네 번 했다. 그런데 그래디언트가 같지 않다.

그리고 기본적으로 그것들이 같지 않은 이유는 정확히, 이 mean squared error에서 손실 안의 그 4분의 1이 사라져버리기 때문이다. 여기서 무슨 일이 벌어지느냐면, 각 루프에서의 손실 목적함수는 그냥 mean squared error인데, 이 경우에는 예제가 단 하나뿐이므로 그냥 여기 이 항 하나가 된다. 그게 0번째 반복의 손실이었고, 첫 번째도 같고, 세 번째도 같고, 계속 그런 식이다. 그리고 loss.backward()를 하면 우리는 그래디언트를 누적하게 되는데, 여기서 벌어지는 일은 그래디언트에서의 그 누적이 기본적으로 손실에서의 합(sum)과 동치라는 것이다. 그래서 우리 손실은 사실상 바깥에 4분의 1이라는 인자가 없는 이 형태다. 즉 우리는 정규화 인자를 놓치고 있다. 그래서 우리 그래디언트가 어긋나는 것이다.

그리고 이것을 고치는 방법, 혹은 그중 한 가지는 이렇다. 기본적으로 우리는 여기로 와서 loss = loss / 4라고 말해줄 수 있다. 그러면 지금 무슨 일이 일어나느냐면, 우리가 손실을 스케일링하고 있는 것이고, 이 모든 자리 앞에 4분의 1이라는 인자를 도입하고 있는 것이다. 그래서 개별 손실들이 이제 전부 4분의 1로 스케일된다. 그리고 그 다음 우리가 backward를 하면 이것들이 전부 합으로 누적되는데, 다만 이제는 이 구성요소 하나하나 안에 4분의 1이 들어 있다. 그러면 이제 우리 손실들이 동치가 될 것이다. 그래서 이걸 실행해보면, 그래디언트가 이제 완전히 동일해진 것을 볼 수 있다.

긴 이야기를 짧게 줄이면 이렇다. 이 단순한 예제를 하나하나 짚어보면, 기본적으로 아까 그게 왜 맞지 않았는지 알 수 있다. 여기 MSELoss에서와 똑같은 방식으로, 모델에서 우리가 계산하는 손실도 reduction이 mean이기 때문이다. 손실이 어디 있더라, cross entropy 다음에 있다. 그리고 기본적으로 여기 cross entropy의 reduction도 — 왜 표시를 안 해주는지 모르겠지만 — 그것도 mean이다. 즉 B × T개 원소 전부에 대한 평균 손실이다. 그렇다. 그러니까 거기 안에 mean에 의한 reduction이 들어 있는데, 우리가 여기서 그냥 gradient accumulation만 하면 그걸 놓치게 되는 것이다.

그래서 이것을 고치는 방법은 단순히 gradient accumulation 스텝의 수만큼 보상해주는 것이다. 그리고 우리는 똑같은 방식으로 이 손실을 나눠줄 수 있다. 특히 여기서 우리가 하고 있는 스텝의 수가 있으니, loss = loss / grad_accum_steps가 된다. Copilot조차 이 수정을 알아맞힌다. 하지만 정확히 같은 방식으로 우리는 손실을 축소하고 있는 것이다. 그래서 우리가 loss.backward()를 할 때, 그건 기본적으로 목적함수에서의 합에 대응하는데, 우리는 이미 정규화되어 있는 손실을 합산하고 있게 된다. 그리고 따라서 grad_accum_steps로 나눈 손실들을 합산할 때, 우리는 빠져 있던 그 추가 정규화 인자를 되찾는 것이다. 그래서 이제 이 둘은 동치가 될 것이다. 이제 이것은 원래의 최적화와 사실상 같아질 것이다. 왜냐하면 그래디언트가 똑같이 나올 것이기 때문이다.

좋다, 그리고 몇 가지 손질을 더 해서 여기서 최적화를 실행했다. 특히 우리가 하고 싶은 것 중 하나는 출력을 예쁘게 하는 것인데, 그러려면 우선 손실에 대한 누적기(accumulator) 같은 걸 만들어야 한다. 그냥 loss를 출력할 수는 없다. 그러면 마지막 마이크로 스텝에서의 마지막 손실만 출력하게 되기 때문이다. 그래서 그 대신 loss_accum이라는 걸 두고, 나는 그걸 0으로 초기화한 다음 거기에 손실을 누적한다. 그리고 detach를 쓴다. 텐서를 그래프에서 떼어내서 값만 추적하려고 하는 것이다. 그래서 더할 때 그것들을 리프 노드로 만드는 셈이다. 그게 loss_accum이고, 그 다음 loss 대신 그걸 여기서 출력한다.

그리고 그에 더해서, 처리한 토큰 수 안에서도 grad_accum_steps를 계산에 넣어줘야 했다. 이제 스텝당 처리되는 토큰 수는 B × T × grad_accum_steps이기 때문이다.

긴 이야기를 짧게 줄이면, 여기 우리 최적화가 있고 꽤 합리적으로 보인다. 그렇지 않은가. 우리는 좋은 지점에서 시작하고 있고, grad_accum_steps는 32로 계산되었고, 여기서 대략 3초가 나온다. 그래서 이 정도면 꽤 좋아 보인다.

이제 만약 여러분 쪽의 최적화와 여기 구현이 맞는지 확인하고 싶고 자기 쪽에서 작업하고 있다면, 이렇게 보면 된다. 이제 우리에게 총 배치 크기와 gradient accumulation 스텝이 있기 때문에, B를 얼마로 두느냐는 순전히 성능 최적화를 위한 설정일 뿐이다. 그래서 큰 GPU가 있다면 이걸 32로 올려도 되고, 아마 조금 더 빨라질 것이다. 아주 작은 GPU가 있다면 8이나 4를 시도해볼 수 있다. 하지만 어느 경우든 여러분은 정확히 같은 최적화와 같은 답을 얻어야 한다. 부동소수점 오차 수준의 차이만 빼면 그렇다. gradient accumulation이 개입해서 모든 것을 직렬로 처리해줄 수 있기 때문이다. 자, gradient accumulation은 여기까지인 것 같다. 좋다, 이제 다음으로 넘어가자.

분산 데이터 병렬 처리(DDP) 02:46:52

이제 중화기를 꺼낼 시간이다. 여러분도 눈치챘겠지만, 지금까지 우리는 학습에 GPU를 딱 한 장만 써 왔다. 그런데 나는 여기서 GPU 한 장이 아니라 여덟 장에 대한 비용을 내고 있다. 그러니 그 여덟 장을 놀려 두지 말고 전부 일하게 만들어야 마땅하다. 특히 이 여덟 장은 앞으로 서로 협업하게 될 것이다. 같은 시각에 각자 토큰을 놓고 최적화를 수행하고, 서로 통신을 주고받고, 그래서 결국에는 하나의 최적화 과정을 여덟 장이 다 같이 굴리는 형태가 된다. 한 장이 여덟 배 빨라지는 게 아니라, 여덟 장이 한 덩어리처럼 움직이는 것이다.

이걸 하기 위해서 우리는 PyTorch에 들어 있는 DistributedDataParallel, 줄여서 DDP를 쓸 것이다. PyTorch에는 예전부터 있던 DataParallel이라는 것도 있는데, 그건 쓰지 말라고 권하고 싶다. 그건 말하자면 옛날 방식, 레거시에 해당하는 물건이다. 반면 DistributedDataParallel은 아주 단순한 방식으로 동작한다. GPU가 여덟 장 있으니 프로세스를 여덟 개 띄운다. 그리고 각 프로세스에 GPU를 한 장씩 배정한다. 각 프로세스 입장에서 보면, 학습 루프를 비롯해서 지금까지 우리가 작업해 온 모든 것이 거의 그대로다. 각 GPU는 자기가 아는 한, 우리가 지금까지 만들어 온 바로 그 코드를 아무 변화 없이 그대로 돌리고 있을 뿐이다. 다만 몰래, 그런 GPU가 여덟 장 있고, 각자가 데이터의 서로 조금씩 다른 부분을 처리하고 있는 것이다. 그리고 여기에 한 단계가 더 붙는다. 여덟 프로세스가 각자 자기 그래디언트를 계산하고 나면, 그 그래디언트들의 평균을 내는 단계다. 바로 이것이 이 여덟 장이 계산 부하를 두고 협업하는 방식이다. 다시 말해 데이터는 여덟 갈래로 쪼개지지만 모델의 갱신은 하나로 합쳐진다. 그래서 이름이 분산 데이터 병렬, 즉 데이터를 병렬로 나누는 방식인 것이다.

torchrun과 랭크 환경변수

여덟 장을 전부 쓰려면, 이제부터는 스크립트를 python train_gpt2.py 같은 식으로 띄우지 않는다. 대신 PyTorch에 들어 있는 torchrun이라는 특수한 명령으로 돌린다. 잠시 뒤에 그걸 직접 보게 될 것이다. torchrun은 우리 파이썬 스크립트를 실행할 때, 그 스크립트를 반드시 여덟 개 병렬로 띄워 준다. 그리고 각각의 프로세스가 자기가 여덟 개 중 어느 것인지를 직접 조회해 볼 수 있도록, 몇 가지 환경변수를 만들어 준다. 예를 들어 torchrun은 RANK, LOCAL_RANK, WORLD_SIZE라는 환경변수를 설정해 준다. 그러니 이 변수들이 설정되어 있는지 아닌지가, DDP가 지금 돌고 있는지를 알아내는 방법이 된다. 조금 조잡한 판별법이기는 하지만 실제로는 잘 통한다. 즉 우리가 torchrun으로 돌리고 있다면, DDP가 돌고 있는 것이고, 그렇다면 CUDA가 사용 가능한지도 확인해 줘야 한다. 이걸 이제 와서 CPU에서 돌릴 수 있는지도 모르겠고, 설령 돌린다 한들 그게 말이 되는지도 모르겠기 때문이다.

여기 세팅 코드가 조금 들어가 있는데, 그중 중요한 부분은 이렇다. 먼저 월드 사이즈라는 것이 있고, 우리에게 이 값은 8이 된다. 이건 지금 돌고 있는 프로세스의 총 개수를 뜻한다. 그다음 랭크라는 것이 있다. 각 프로세스는 기본적으로 정확히 같은 코드를 거의 같은 시각에 실행한다. 이 프로세스들 사이의 유일한 차이는 각자 DDP 랭크가 다르다는 것뿐이다. 그러니까 0번 GPU는 DDP 랭크가 0이고, 1번 GPU는 랭크가 1이고, 그런 식으로 쭉 간다. 그것 말고는 여덟 개 전부가 완전히 동일한 스크립트를 돌리고 있다. 단지 DDP 랭크만 서로 조금씩 다른 정수일 뿐이다. 그리고 바로 그것이, 예를 들어 프로세스들이 같은 데이터 위에서 돌지 않도록 우리가 조율할 수 있는 유일한 수단이 된다. 우리는 이 프로세스들이 데이터셋의 서로 다른 부분에서 돌기를 원하기 때문이다. 그러니 이 랭크라는 정수 하나가, 여덟 개의 동일한 복사본을 서로 구별해 주는 유일한 표식인 셈이다.

이제 LOCAL_RANK는 멀티 노드 환경에서만 쓰이는 것이다. 우리에게는 GPU 여덟 장이 달린 노드가 딱 하나뿐이다. LOCAL_RANK는 단일 노드 안에서의 GPU 랭크를 뜻한다. 예를 들면 0부터 7까지의 값이다. 하지만 우리는 대부분 박스 한 대 위에서 돌릴 것이므로, 우리가 실제로 신경 쓸 것은 RANK와 WORLD_SIZE 두 가지다. 월드 사이즈는 8이고, 랭크는 이 스크립트의 특정 인스턴스가 어느 GPU 위에서 돌고 있느냐에 따라 달라지는 값이 된다.

그리고 여기서는 로컬 랭크에 따라 디바이스를 cuda: 뒤에 콜론을 붙이고 번호를 적은 형태로 설정한다. 콜론 뒤의 번호는 GPU가 여러 장 있을 때 그중 어느 GPU를 쓸지를 가리킨다. 그러니 이 프로세스의 로컬 랭크가 무엇이냐에 따라, 딱 그에 맞는 적절한 GPU 하나만 쓰게 된다. 그 덕분에 어느 GPU를 어느 프로세스가 쓰느냐를 두고 충돌이 생기는 일이 없다.

그리고 마지막으로, 내가 즐겨 만드는 불리언 변수가 하나 있다. DDP 랭크가 0인지 아닌지를 담는 변수다. 마스터 프로세스는 임의로 0번 프로세스로 정한다. 그리고 이 마스터 프로세스가 출력이라든가 로깅, 체크포인트 저장 같은 일을 많이 담당하게 된다. 나머지 프로세스들은 대체로 옆에서 거들어 주는 계산 프로세스로 생각하면 된다. 그러니까 마스터 프로세스인 0번은 추가로 할 일이 좀 있고, 나머지 프로세스들은 대부분 그냥 포워드와 백워드만 수행한다. 그리고 만약 우리가 DDP를 쓰지 않는 상황이라서 이 환경변수들이 하나도 설정되어 있지 않다면, 코드는 원래의 단일 GPU 학습으로 되돌아간다. 즉 랭크는 0 하나뿐이고, 월드 사이즈는 그냥 1이며, 우리가 곧 마스터 프로세스이고, 디바이스는 자동으로 감지하려고 시도한다. 원래 하던 그대로다.

여덟 개의 인터프리터를 동시에 머릿속에 떠올리기

여기까지 우리가 한 일은 init_process_group을 호출해서 DDP를 초기화한 것뿐이다. 그리고 잠시 뒤에 보게 될 torchrun으로 실행하는 경우, 복사본 여덟 개가 병렬로 돌게 되고, 그 각각이 서로 다른 랭크를 갖게 된다. 그러니 이제 그 뒤에 오는 모든 것이 올바르게 동작하도록 확실히 만들어 줘야 한다. 초기화만 해 놓고 나머지 코드를 손대지 않은 채 그대로 두면, 여덟 개의 프로세스가 전부 똑같은 일을 중복해서 하게 되기 때문이다. 그러면 GPU를 여덟 장 쓰는 의미가 없어진다.

여러 프로세스를 돌릴 때 까다로운 점은, 언제나 프로세스 여덟 개가 병렬로 돌고 있다고 상상해야 한다는 것이다. 그러니까 이제부터 코드를 읽을 때는, 파이썬 인터프리터 여덟 개가 이 코드 줄들을 따라 나란히 내려가고 있다고 상상해야 한다. 그들 사이의 유일한 차이는 DDP 랭크가 서로 다르다는 것뿐이다. 여덟 개가 전부 여기로 와서, 전부 정확히 같은 시드를 고르고, 이 모든 계산을 다른 복사본들의 존재는 전혀 모르는 채로 수행한다. 대략적으로 말하면 그렇다. 여덟 개 전부가 정확히 같은 계산을 한다. 그래서 이제 우리는, 실제로는 어떤 월드 사이즈와 어떤 랭크라는 게 존재한다는 사실을 반영하도록 이 계산들을 조정해 줘야 한다.

특히 마이크로 배치 크기와 시퀀스 길이, 이 두 가지는 전부 GPU 한 장당 값이다. 그런데 이제는 이런 게 프로세스 개수만큼 병렬로 돌아간다. 그러니 이 부분을 조정해야 한다. 그래디언트 누적 스텝 수는 이제 총 배치 크기를 B × T × DDP 월드 사이즈로 나눈 값이 된다. 각 프로세스가 B × T만큼을 처리하고, 그런 프로세스가 이만큼 있기 때문이다. 여기에 더해서, 이 값이 총 배치 크기에 깔끔하게 나누어떨어지는지도 확인하고 싶다. 우리 경우에는 다행히 떨어진다. 16 × 1024 × 8, 즉 GPU가 여덟 장이니 이건 131K가 된다. 그리고 총 배치 크기가 524,288이므로, 현재 설정에서 그래디언트 누적 스텝 수는 4가 된다. 그러니까 각 GPU에서 16 × 1024가 돌고, 그런 GPU가 여덟 장이니, 여덟 장 전체로 한 번의 포워드-백워드에 131,000 토큰을 처리하게 되는 것이다. 그러니 이 값이 딱 떨어지는지 확인해서, 거기서 깔끔한 그래디언트 누적 스텝 수를 유도할 수 있게 해야 한다. 주석도 여기에 × DDP 월드 사이즈를 넣어서 같이 고쳐 두자. 정리하면, 마이크로 배치 크기와 시퀀스 길이는 GPU 한 장 기준의 값이고, 거기에 프로세스 개수를 곱한 것이 한 번의 포워드-백워드로 실제로 소비되는 토큰의 양이며, 총 배치 크기를 그 값으로 나눈 것이 그래디언트 누적 스텝 수가 된다.

좋다. 이제 각 GPU가 이 계산을 수행한다. 그런데 바로 여기서부터 문제가 생기기 시작한다. 각 프로세스가 print를 만나면 전부 출력을 한다. 그러니 이 출력이 여덟 벌이나 나오게 된다. 이걸 처리하는 한 가지 방법이 바로 우리가 만들어 둔 마스터 프로세스 변수다. if master_process:로 이 부분을 감싸 주면, 딱 한 번만 출력된다. 어차피 모든 프로세스가 정확히 같은 변수들을 계산했을 테니, 똑같은 내용을 여덟 번씩 출력할 이유가 전혀 없다. 이런 식으로 앞으로도 출력이나 로깅에 해당하는 부분은 전부 마스터 프로세스 안쪽으로 몰아넣게 된다.

첫 실행 — 5번 프로세스가 먼저 도착하다

데이터 로더로 들어가기 전에, 그리고 데이터 로더는 당연히 리팩터링해야 하는데, 어쩌면 이 시점에서 출력을 좀 찍어 보고 한번 돌려 보는 게 좋겠다. 그리고 여기서 바로 종료해 버리자. import sys를 하고 sys.exit를 넣고, "나는 GPU 몇 번이다"라고 DDP 랭크를 찍고, 그다음 "bye"를 출력한다.

이제 이걸 돌려서 어떻게 동작하는지 한번 보자. 어떤 모습인지 볼 수 있게 그냥 한번 굴려 보는 것이다. 평소에는 python train_gpt2.py로 띄웠지만, 이제는 torchrun으로 돌릴 것이고, 명령은 이렇게 생겼다.

torchrun --standalone --nproc_per_node=8 train_gpt2.py

프로세스 개수는 우리 경우 GPU가 여덟 장이니 8이고, 그다음에 train_gpt2.py가 온다. 이게 명령의 모양이다. 그리고 torchrun은 다시 말하지만 이걸 여덟 개 돌려 준다. 무슨 일이 벌어지는지 보자.

일단 화면이 조금 바빠진다. 여기 벌어지는 일이 꽤 많다. 프로세스가 여덟 개나 동시에 올라오기 때문에, 한 화면에 여덟 개 몫의 출력이 뒤섞여 쏟아진다. 우선 distributed 쪽에서 경고가 몇 개 나온다. 솔직히 이게 무슨 의미인지는 나도 잘 모르겠다. 내 생각엔 그냥 코드가 세팅되는 중이고 프로세스들이 올라오는 중이고, 프로세스들이 올라오는 동안 뭔가 수집에 잠깐 실패하는 걸 보고 있는 것 같다. 100% 확신하지는 못하겠다. 그러다가 실제 출력이 나오기 시작한다. 프로세스들이 전부 내려가고 나서, 첫 번째 출력이 나오는데 그건 우연히도 5번 프로세스에서 나온 것이다. 그러니까 5번 프로세스가 여기 제일 먼저 도착해서, "나는 GPU 5번에서 돌고 있다, bye"라고 찍은 것이다. 그리고 그다음에 이어지는 출력들은 마스터 프로세스에서 나온 것이다. 5번 프로세스가 무슨 이유에서든 그냥 제일 먼저 끝난 것뿐이고, 거기에 특별한 의미는 없다. 그건 운영체제가 프로세스 실행을 어떻게 스케줄했느냐에 달린 문제다. 그다음엔 0번 GPU가 끝났고, 그다음에 3번, 그다음에 2번이 끝났다.

그리고 아마 5번 프로세스쯤이 먼저 종료되었는데, DDP가 이걸 굉장히 싫어한다. 우리가 멀티 GPU 설정을 제대로 정리해 주지 않았기 때문이다. 그래서 "프로세스 그룹이 소멸되기 전에 파괴되지 않았다"는 식의 불평이 나온다. 그러니까 DDP가 이걸 정말 싫어하는 것이다. 실제 애플리케이션이라면 destroy_process_group을 호출해서 DDP를 깔끔하게 정리해 줘야 한다. DDP는 이걸 별로 좋아하지 않는다. 그다음 나머지 GPU들이 끝나고, 그걸로 끝이다. 결국 정리하자면, 이 프로세스들이 정확히 언제 어떤 순서로 도는지는 우리가 보장할 수 없다. 그 순서는 완전히 제멋대로다. 하지만 여덟 개가 실제로 병렬로 돌고 있다는 것만큼은 이 출력으로 확인된 셈이다. 그리고 우리는 이것들이 이런 식으로 출력을 남발하는 것을 원하지 않는다. 그래서 출력은 마스터 프로세스에만 맡기는 것이다.

데이터 로더를 프로세스별로 엇갈리게 하기

이건 지우고 다음으로 넘어가자. 다음으로 할 일은, DataLoaderLite를 만들 때 이 멀티 프로세스 설정을 인지하게 만드는 것이다. 모든 프로세스가 정확히 같은 데이터를 로드하는 것을 우리는 원하지 않기 때문이다. 우리는 모든 프로세스가 각자 자기 몫의 데이터 덩어리를 받기를 원한다. 당연히 데이터셋의 서로 다른 부분에서 작업하도록 말이다. 그러니 그 부분을 조정하자.

이걸 하는 특히 단순하고 순진한 방법이 하나 있다. 우선 데이터 로더에 랭크와 프로세스 개수를 넘겨줘야 한다. 그리고 위쪽으로 올라오면, 이제 랭크와 프로세스 개수를 받아서 저장한다. 그리고 이제 현재 위치는 0이 아니게 된다. 우리가 원하는 것은 모든 프로세스를 서로 엇갈리게 펼쳐 놓는 것이기 때문이다. 이걸 하는 한 가지 방법은 self.B × self.T를 계산한 다음, 거기에 프로세스 랭크를 곱하는 것이다. 그러면 프로세스 랭크 0은 0에서 시작하고, 프로세스 랭크 1은 이제 B × T에서 시작하고, 프로세스 랭크 2는 2 × B × T에서 시작한다. 그런 식으로 쭉 간다. 이것이 초기화 부분이다.

이제 배치를 꺼내오는 방식 자체는 여전히 동일하게 하지만, 위치를 전진시킬 때는 B × T만큼 전진하지 않는다. B × T × 프로세스 개수만큼 전진한다. 기본적으로 우리가 소비하는 토큰의 총량이 B × T × 프로세스 개수이고, 그것들이 각자 서로 다른 랭크로 흩어져 있으니, 위치는 그 덩어리 전체만큼 전진해야 하는 것이다. 그리고 여기서, B × T × self.num_processes + 1이 토큰 개수를 넘어서게 되면 루프를 돌아야 한다. 그리고 루프를 돌 때는 당연히 정확히 같은 방식으로 돌기를 원하므로, 처음 상태로 리셋해 준다.

이것이 아주 단순한 분산 DataLoaderLite를 위해 내가 찾을 수 있는 가장 단순한 변경이다. 그러니까 여덟 개의 랭크가 각자 다른 출발점에서 시작하고, 매번 여덟 개 몫만큼 한꺼번에 건너뛰기 때문에, 어느 두 랭크도 같은 토큰을 보지 않게 되는 것이다. 각 랭크가 데이터 위를 일정한 간격으로 엇갈리게 훑어 나가는 셈이고, 여덟 개를 합치면 결국 데이터 전체를 빠짐없이 한 번씩 덮게 된다. 그리고 여기서 눈여겨볼 점은, 프로세스 랭크가 0이고 프로세스 개수가 1이면 전체가 이전과 완전히 동일해진다는 것이다. 하지만 이제 우리는 실제로 여러 프로세스를 돌릴 수 있게 되었고, 이건 잘 동작해야 한다. 자, 이게 데이터 로더다.

모델을 DDP 컨테이너로 감싸기

다음으로, 여덟 프로세스가 전부 데이터 로더를 초기화하고 나면 이리로 와서 전부 GPT 모델을 만든다. 그래서 프로세스 여덟 개 위에 GPT 모델이 여덟 개 만들어진다. 하지만 여기서 시드가 고정되어 있기 때문에, 여덟 개가 전부 완전히 동일한 초기 가중치를 가진 모델을 만들게 된다. 전부 자기 랭크에 해당하는 디바이스로 모델을 옮기고, 전부 모델을 컴파일한다. 그리고 모델들이 서로 동일하므로, 완전히 똑같은 컴파일이 여덟 개 병렬로 일어나게 된다. 낭비처럼 보이지만 그건 괜찮다.

이 중 어떤 것도 바뀌지 않는다. 그건 스텝 단위의 이야기이고, 우리는 지금 스텝 안쪽에서 작업하고 있기 때문이다. 우리가 지금 하고 있는 변경은 전부 스텝 내부의 변경이다. 그런데 여기서 중요한 것은, 모델을 구성할 때 실제로 할 일이 좀 있다는 점이다. get_logits는 이제 더 이상 쓰지 않는다. 모델을 만들고 나서, 그 모델을 분산 데이터 병렬 컨테이너 안으로 감싸 줘야 한다.

이것이 모델을 DDP 컨테이너로 감싸는 방법이다. 그리고 이게 DDP의 문서인데, 문서가 상당히 방대하고, 주의사항이라든가 조심해야 할 것이 아주 많다. 여러 프로세스가 끼어들면 모든 게 열 배로 복잡해지기 때문이다. 대략적으로 말하면, 이 device_ids라는 것은 이제 넘겨줘야 하는 값인 것으로 알고 있다. 그런데 안타깝게도 device_ids가 무엇인지에 대한 문서가 극도로 불명확하다. 실제로 여기까지 들어와서 보면, device_ids가 무엇인지 설명하는 이 주석이 대체로 말이 되지 않는다. 하지만 나는 이게 DDP 로컬 랭크여야 한다고 꽤 확신한다. DDP 랭크가 아니라 로컬 랭크다. 그러니 여기에 그걸 넘겨준다.

이렇게 하면 모델이 감싸진다. 그리고 특히 DDP가 우리를 위해 해 주는 일은 이렇다. 포워드 패스에서는 사실상 동일하게 동작한다. 내가 이해하기로는 포워드 패스에서 바뀌는 것은 아무것도 없다. 하지만 백워드 패스에서는 이야기가 다르다. 단순한 상황에서 백워드 패스가 끝나고 나면, 각각의 독립적인 GPU는 모든 파라미터에 대한 그래디언트를 각자 갖고 있게 된다. 이때 DDP가 해 주는 일은, 백워드 패스가 끝나는 순간 이른바 all-reduce라는 것을 호출하는 것이다. 기본적으로 모든 랭크의 그래디언트에 대해 평균을 내고, 그 평균을 다시 모든 랭크 위에 하나씩 내려놓는다. 그래서 모든 랭크가 결국 그 평균값을 갖게 된다. 그러니까 이게 바로 통신이다. 그저 그래디언트를 동기화하고 평균 내는 것이다. 그것이 DDP가 여러분에게 제공하는 것이다. 다시 말해, 여덟 장의 GPU가 각자 서로 다른 데이터 조각으로 계산한 그래디언트가, 이 all-reduce 한 번으로 모두 같은 하나의 평균 그래디언트가 되고, 그래서 여덟 장이 마치 하나의 거대한 배치를 처리한 것처럼 동일한 파라미터 갱신을 하게 되는 것이다.

사실 DDP는 그보다 조금 더 복잡하다. 트랜스포머의 레이어들을 거슬러 올라가며 백워드 패스를 진행하는 도중에도, 백워드 패스가 아직 진행 중인 상태에서 이미 그래디언트에 대한 통신을 내보낼 수 있기 때문이다. 그러니 그래디언트 통신 및 그 동기화와, 백워드 패스 자체가 시간적으로 서로 겹쳐서 진행된다. 백워드가 다 끝날 때까지 기다렸다가 통신을 시작하는 게 아니라, 위쪽 레이어의 그래디언트가 준비되는 대로 그것부터 먼저 주고받기 시작하는 것이다. 그리고 그렇게 하는 편이 훨씬 더 효율적이다. 그것이 DDP가 여러분을 위해 해 주는 일이다. 포워드는 그대로이고 백워드도 대부분 그대로이며, 우리는 거기에 평균 내는 동작을 하나 얹는 셈이다. 그건 잠시 뒤에 보게 될 것이다.

그래디언트 누적과 DDP의 충돌 — require_backward_grad_sync

좋다. 이제 최적화 쪽으로 가 보자. 여기서 바뀌는 건 없다. 최적화의 안쪽 루프로 가서, DDP에서 이 그래디언트들의 동기화를 어떻게 처리할지 차근차근 생각해 보자. 기본적으로 기본 동작은, 내가 앞서 말했듯이 loss.backward()를 하면 백워드 패스를 하고 그다음 그래디언트를 동기화한다는 것이다. 그런데 여기서 문제는, 그래디언트 누적 스텝 루프가 있기 때문에 매번 loss.backward()를 할 때마다 동기화를 하고 싶지는 않다는 점이다. 우리는 그냥 그래디언트를 쌓아 올리고 있을 뿐이고, 그것도 순차적으로 하고 있으며, 그저 더해지기만 하면 된다. 그러니 매번 동기화하고 싶지는 않다. 그건 극도로 낭비다. 어차피 중간 단계의 부분 합계를 여덟 장이 서로 맞춰 봐야 아무 쓸모가 없는데, 그걸 위해 GPU 간 통신 비용을 네 번 다 치르는 꼴이기 때문이다. 그래서 기본적으로 우리가 원하는 것은, 그래디언트를 계속 더해 나가다가 맨 마지막 스텝에서, 즉 micro_step이 grad_accum_steps - 1이 되는 바로 그 마지막 스텝에서만 실제로 all-reduce를 해서 그래디언트를 평균 내는 것이다.

그걸 하려고 여기로 온다. 참고로 공식적으로 승인된 방법은 no_sync 컨텍스트 매니저를 쓰는 것이다. PyTorch는 이것이 "DDP 프로세스들 사이의 그래디언트 동기화를 비활성화하는 컨텍스트 매니저"라고 설명한다. 이 컨텍스트 안에서는 그래디언트가 누적되고, 기본적으로 no_sync를 하면 통신이 전혀 일어나지 않는다. 그래서 그들은 with ddp.no_sync(): 안에서 그래디언트 누적을 하고, 그다음 그 바깥에서 또 다른 입력으로 DDP를 한 번 더 호출하고 백워드를 하라고 안내한다. 나는 이게 정말이지 마음에 들지 않는다. 정말 좋아하지 않는다. 코드를 여기에 복사해서 붙여넣어야 하고, 컨텍스트 매니저를 써야 하고, 그냥 극도로 보기 흉하다.

그래서 나는 여기 소스 코드를 뒤져 봤다. 여기를 보면, 컨텍스트에 진입할 때 그냥 이 변수 하나를 토글한다는 걸 알 수 있다. 바로 require_backward_grad_sync다. 이 변수가 여기저기서 왔다 갔다 토글되고 값이 바뀐다. 그리고 실제로 한 줄씩 따라가 보면, 이 변수가 그래디언트를 동기화할지 말지를 결정하는 데 쓰인다는 걸 알 수 있다. 그래서 나는 그냥 그걸 직접 쓰는 쪽을 좋아한다.

그러니까 대신 내가 하고 싶은 건 이렇다. 바로 여기, loss.backward() 바로 앞에서, 우리가 DDP를 쓰고 있다면 기본적으로 이 변수가 마지막 이터레이션일 때만 True가 되게 하고, 마이크로 스텝 안쪽의 다른 모든 이터레이션에서는 False가 되게 한다. 그래서 그냥 이런 식으로 토글한다. require_backward_grad_sync는 micro_step이 마지막 스텝일 때만 켜져야 한다. 그렇게 나는 이 변수를 직접 토글하고, 그것이 마지막 loss.backward()에 영향을 미치기를 바라는 것이다.

이건 좀 못된 짓이다. 왜냐하면 알다시피 그들이 DDP를 바꾸면 이 변수가 사라져 버릴 수도 있기 때문이다. 공식 API가 아니라 내부 구현 세부를 직접 건드리는 것이기 때문이다. 하지만 지금으로서는 이게 동작한다고 믿고 있고, 이 덕분에 컨텍스트 매니저를 쓰는 것과 코드 중복을 둘 다 피할 수 있다. 나는 그냥 변수 하나를 토글할 뿐이고, 그러면 loss.backward()가 대부분의 스텝에서는 동기화하지 않다가 맨 마지막 스텝에서만 동기화한다. 그래서 이게 끝나고 우리가 빠져나오면, 모든 랭크가 갑자기 마법처럼 모든 랭크에 저장되어 있던 그래디언트들의 평균을 갖게 된다. 그러니까 마이크로 스텝 네 번 중 앞의 세 번은 통신이 전혀 일어나지 않고 그냥 로컬에서 그래디언트가 더해지기만 하다가, 네 번째 마지막 마이크로 스텝에서 딱 한 번 all-reduce가 돌면서 여덟 장의 누적 그래디언트가 한꺼번에 평균되는 것이다.

loss 누적값도 all-reduce로 평균 내기

이제 그것이 정말 우리가 원하는 것인지, 그리고 이것으로 충분한지, 그리고 이게 loss와는 어떻게 맞물리는지, 그리고 loss_accum이란 게 무엇인지를 생각해 봐야 한다. 그러니 그걸 지금 따져 보자. 내가 짚으려는 문제는 이렇다. 우리는 그래디언트를 평균 냈고, 그건 아주 좋다. 하지만 loss_accum은 아직 아무런 영향도 받지 않았다. 그리고 이건 DDP 컨테이너 바깥에 있으므로 평균이 내려지지 않는다. 그래서 여기서 loss_accum을 출력할 때, 아마 우리는 마스터 프로세스, 즉 랭크 0에서만 출력할 텐데, 그건 그저 자기 프로세스에서 본 손실만 출력하게 된다. 하지만 우리가 원하는 것은 모든 프로세스에 걸친 손실, 그리고 그 손실의 평균을 출력하는 것이다. 그래디언트의 평균을 냈으니 손실의 평균도 원하는 것이다.

그래서 그냥 여기, 이 뒤에 넣는다. 이건 내가 예전부터 쓰던 코드다. 그리고 loss가 아니라 loss_accum에 대해 하고 싶다. 그러니까 다시 DDP라면, 이건 PyTorch distributed다. 이걸 어디서 임포트하지? 아이고, 이 파일이 슬슬 통제 불능이 되어 가는구나. 어쨌든 import torch.distributed as dist를 넣는다. 그리고 dist.all_reduce로 loss_accum에 대해 평균을 낸다. 이 loss_accum 텐서는 모든 랭크 위에 존재한다. 평균 연산으로 all-reduce를 호출하면 그 숫자들의 평균이 만들어지고, 그 평균이 모든 랭크 위에 놓인다. 그래서 이 호출이 끝난 뒤에는 모든 랭크가 평균된 loss_accum을 담고 있게 된다. 그러니 여기 마스터 프로세스에서 출력할 때, 그 loss_accum은 다른 모든 랭크에서도 동일한 값이다. 그러니까 랭크 0이 출력하는 손실은 이제 랭크 0 혼자 본 손실이 아니라 여덟 개 랭크 전체에 걸친 평균 손실이 되고, 그래야 우리가 화면에서 보고 있는 손실이 실제로 그래디언트가 하고 있는 일과 일치하게 된다. 그래서 여기 if master_process: 안에서, 어이쿠, 이런 식으로 출력해 준다.

그리고 마지막으로 조심할 게 하나 더 있다. 우리는 이제 훨씬 더 많은 토큰을 처리하고 있으므로, 처리한 토큰 수에 × DDP 월드 사이즈를 곱해 줘야 한다. 위쪽에서 처리한 토큰 수 말이다. 나머지는 다 괜찮을 것이다. 조심할 게 하나 더 있다면, 앞서 말했듯이 destroy_process_group을 불러서 프로세스 그룹을 파괴해 줘야 한다는 것이다. 시작할 때 init_process_group으로 열었으니, 끝낼 때는 짝을 맞춰 닫아 주는 셈이다. NCCL에게 예의를 지키기 위해서다. 그래야 우리가 여기서 종료할 때 DDP가 우리한테 불평을 늘어놓지 않는다. 앞서 첫 실행에서 봤던 그 경고가 바로 이걸 빼먹었을 때 나오는 것이다.

8×A100에서 초당 150만 토큰

이 정도면 됐을 것이다. 그럼 한번 굴려 보자. 좋다, 스크립트를 띄웠고, 곧 여기에 출력이 나오기 시작할 것이다. 이제 우리는 GPU 여덟 장으로 동시에 학습하고 있다. 그래서 그래디언트 누적 스텝 수는 이제 32가 아니라 8로 나눈 값, 즉 그냥 4다. 그것 말고는 최적화가 이런 모습이고, 와, 정말 빠르다. 이제 초당 150만 토큰을 처리하고 있다. 이건 진지한 숫자다. 그리고 tiny Shakespeare 데이터셋은 너무나 작아서, 아마 그 위를 수도 없이 많은 에폭으로 돌고 있을 것이다. 여덟 장이 초당 150만 토큰을 먹어 치우는데 데이터는 그것밖에 없으니, 같은 텍스트를 계속 다시 보고 있는 셈이다. 하지만 어쨌든 대략 이런 모습이다.

참고로 내가 하나 고쳐야 했던 게 있다. model.configure_optimizers였는데, 이제 model이 DDP 모델이 되어 버려서 이게 동작하지 않는다. 그래서 이건 raw_model.configure_optimizers가 되어야 한다. 여기서 raw_model은 내가 여기서 따로 만들어 주는 것이다. 그러니까 모델을 DDP로 감싼 직후에 raw_model을 만들어 줘야 한다. DDP의 경우 그건 model.module이다. 거기에 원본이 저장되고, 그 원본이 바로 우리가 갖고 있는 GPT-2 모듈이며, 우리가 호출하고 싶은 configure_optimizers 함수를 담고 있다. 그러니까 DDP로 감싸고 나면 우리 클래스에 직접 정의한 메서드에는 더 이상 바로 접근할 수 없고, .module을 통해 안쪽 원본 모델을 꺼내 와야 하는 것이다. 이게 내가 고쳐야 했던 한 가지다. 그것 말고는 잘 도는 것 같다.

단일 GPU와 숫자가 딱 맞지 않는 이유

한 가지 눈에 띄는 게 있다. 이 실행 결과의 숫자를 그냥 단일 GPU로 돌린 것, 그러니까 그래디언트 누적 32로 단일 GPU를 돌린 것과 비교해 보면, 숫자가 정확히 맞아떨어지지 않는다는 점이다. 그리고 그 이유는 꽤나 시시하다. 이유는, 데이터 로더에서 우리가 배치를 순회하는 방식이 살짝 달라졌기 때문이다. 이제는 데이터를 한 페이지 통째로 찾아보게 되고, 모든 GPU를 합친 그 덩어리가 토큰 수를 넘어가면 그냥 루프를 돈다. 그래서 실제로 단일 GPU 프로세스와 8 GPU 프로세스가 약간 다른 방식으로 리셋되게 되고, 그래서 배치가 조금 달라지고, 그래서 숫자가 조금 달라지는 것이다.

하지만 이게 괜찮다는 걸 스스로 납득하는 한 가지 방법은, 총 배치 크기를 훨씬 작게 만들고 B와 T도 작게 만들어 보는 것이다. 나는 4 × 1024 × 8을 썼던 것 같다. 그러니까 총 배치 크기로 32,768을 쓴 것이다. 그렇게 해서 단일 GPU는 그래디언트 누적 스텝을 여덟 번 하게 만들고, 멀티 GPU도 같이 돌려 본다. 그러면 데이터 로더의 경계 효과가 줄어들고, 숫자가 딱 맞아떨어지는 걸 보게 된다. 즉 숫자가 어긋난 원인은 최적화나 그래디언트 평균 쪽에 있는 것이 아니라, 그저 데이터 로더가 데이터셋 끝에서 되감기는 지점이 두 설정에서 서로 다르기 때문이라는 것이 이렇게 확인되는 셈이다. DDP 자체가 뭔가를 잘못하고 있는 게 아니라는 뜻이다.

길게 말했지만 요점은, 이제 우리가 정말 정말 빠르게 돌고 있고, 최적화가 GPT-2 및 GPT-3의 하이퍼파라미터와 대체로 일치한다는 것이다. 그리고 우리는 이제 tiny Shakespeare 파일을 졸업했으니, 그걸 업그레이드하고 싶다. 여덟 장이 이 속도로 도는데 데이터가 그렇게 작으면 계산 자원이 아깝기 때문이다. 그러니 다음 주제로 넘어가자.

GPT-2, GPT-3에 쓰인 데이터셋과 FineWeb(EDU) 03:10:21

이제 GPT-2와 GPT-3가 어떤 데이터셋을 썼는지 살펴보자. GPT-2 논문은 WebText라는 데이터셋을 썼는데, 이건 끝내 한 번도 공개되지 않았다. 이걸 재현해 보려는 시도로 OpenWebText라는 것이 있다. 논문에서 그들이 말하는 바를 대략 정리해 보면 이렇다. 그들은 Reddit에서 밖으로 나가는 아웃바운드 링크를 전부 긁어모았고, 그것도 카르마가 최소 3 이상인 것만 골랐다. 그게 일종의 출발점이었다. 그리고 그 링크가 가리키는 웹페이지들을 전부 수집해서, 그 안에 들어 있는 텍스트를 전부 모았다. 그래서 링크가 4,500만 개였고, 최종적으로 40GB의 텍스트가 되었다. 대략 이것이 GPT-2가 자기 데이터셋에 대해 말하는 전부다. 그러니까 기본적으로 Reddit에서 밖으로 나가는 아웃바운드 링크인 것이다. 카르마 3 이상이라는 조건 하나로, 사람이 어느 정도 가치 있다고 판단한 링크만 남기는 일종의 인간 큐레이션 필터를 대신한 셈이다.

GPT-3의 데이터 혼합

이제 GPT-3 논문으로 넘어가 보면 학습 데이터셋 섹션이 있고, 바로 거기서 Common Crawl 이야기가 시작된다. Common Crawl은 실제로 훨씬 더 많이 쓰인다. 사실 GPT-2도 Common Crawl 이야기를 했던 것 같다. 하지만 기본적으로 Common Crawl은 그 자체로는 품질이 아주 높은 데이터셋이 아니다. 극도로 노이즈가 많기 때문이다. 이건 인터넷의 완전히 무작위한 부분집합이고, 여러분이 생각하는 것보다 훨씬 더 나쁘다. 그래서 사람들은 Common Crawl을 걸러내는 데 엄청난 공을 들인다. 그 안에 좋은 것도 분명히 들어 있기는 하지만, 대부분은 광고 스팸이라든가 무작위한 표와 숫자, 주식 시세 같은 것들이고, 그냥 완전히 엉망진창이다. 그래서 필터링을 얼마나 잘하느냐가 곧 데이터셋의 품질을 결정한다. 바로 그래서 사람들은 자기가 직접 큐레이션하고 신경 써서 만든 데이터 혼합물 위에서 학습시키기를 좋아하는 것이다.

이런 데이터 혼합물의 큰 덩어리는 보통 Common Crawl 같은 것이다. 예를 들어 토큰의 50%가 Common Crawl인 식이다. 그런데 GPT-3에서는 여기에 더해 앞서 나온 WebText2도 쓰고 있다. 그러니까 GPT-2에서 쓰던 그 Reddit 아웃바운드 링크다. 여기에 더해 예를 들면 책 말뭉치도 넣고, Wikipedia도 넣는다. 그 밖에도 넣기로 마음먹을 수 있는 것은 얼마든지 많다. 그래서 GPT-3의 표는 Common Crawl, WebText2, Books1, Books2, Wikipedia 다섯 갈래로 정리되어 있다.

여기서 핵심은, 각 소스마다 그 안에 들어 있는 토큰 수와, 학습 혼합물에서 차지하는 샘플링 비중이 따로따로 잡혀 있다는 점이다. 즉 어떤 소스는 그 안에 실제로 들어 있는 토큰 수에 비해 훨씬 더 자주 뽑히게 되고, 그래서 학습 도중에 같은 데이터를 여러 에폭에 걸쳐 다시 보게 된다. 그런 식으로 고품질 소스는 여러 번 돌게 하고, Common Crawl처럼 양만 많고 품질은 떨어지는 소스는 한 바퀴도 다 못 돌게 배치하는 것이다. 그러니까 어떤 데이터를 몇 번 보여 줄지까지가 설계의 일부인 셈이다.

오늘날 쓸 만한 공개 데이터셋들

GPT-3를 위한 이 데이터셋 역시 끝내 공개되지 않았다. 그래서 오늘날 내가 알고 있는 데이터셋 중에서 상당히 좋고, 이 계열을 대표할 만한 것들을 몇 개 꼽아 보겠다. 첫째는 RedPajama 데이터셋이다. 더 구체적으로는, 예를 들어 RedPajama의 부분집합인 SlimPajama가 있는데, 이건 RedPajama를 정제하고 중복을 제거한 버전이다. 다시 감을 잡게 말해 보면, 그 안에는 Common Crawl이 잔뜩 들어 있고, C4도 있다. C4도 내가 알기로는 역시 Common Crawl인데 처리 방식이 다르다. 그리고 GitHub, 책, arXiv, Wikipedia, Stack Exchange가 들어 있다. 이런 것들이 바로 이런 데이터 혼합물에 들어가는 종류의 데이터셋들이다. 그 밖에도 The Pile처럼 비슷한 성격으로 여러 출처를 한데 묶어 만든 공개 말뭉치들이 같은 계열에 속한다. 결국 어느 것이든 큰 그림은 비슷하다. 웹 크롤 데이터를 바탕에 깔고, 그 위에 코드·책·백과사전·포럼 같은 상대적으로 품질이 높은 소스를 얹는 구조다.

그중에서도 최근에 나온 것으로 내가 특히 좋아하는 것이 FineWeb 데이터셋이다. 이건 기본적으로 정말 고품질의 Common Crawl 데이터를 모으고, 이 경우 15조 토큰으로 필터링해 낸 시도다. 그리고 여기에 더해, 더 최근에는 Hugging Face가 FineWeb-EDU라는 부분집합을 따로 공개했다. 이건 교육적 콘텐츠 1.3조 토큰과 고도로 교육적인 콘텐츠 5.4조 토큰으로 이루어져 있다. 교육적 가치를 판정하는 분류기를 한 번 더 태워서 남긴 것들이다. 그러니까 기본적으로 Common Crawl을 아주 높은 품질의 교육적 부분집합으로 걸러내려는 시도인 것이다. 그리고 우리가 쓸 것이 바로 이 FineWeb-EDU다. 원래의 FineWeb이 15조 토큰 규모라면, 거기서 교육적 가치를 기준으로 한 번 더 체를 친 것이 EDU 쪽이라고 보면 된다. 그러니까 같은 Common Crawl을 두 단계에 걸쳐 거른 결과물인 셈이다.

FineWeb에 대해서는 긴 웹페이지가 하나 있는데, 데이터를 어떻게 처리했는지에 대해 엄청나게 자세히 설명해 준다. 참고로 읽어 보면 정말 매혹적이다. 데이터 혼합물이라든가, 이런 규모에서 데이터가 실제로 어떻게 처리되는지에 관심이 있다면, 이 페이지를 꼭 한번 보라고 강력히 권하고 싶다. 그리고 더 구체적으로 우리는 FineWeb-EDU를 쓰게 될 것 같다. 이건 기본적으로 인터넷에서 가져온 교육적 콘텐츠다. 그들은 자기들이 쓰는 지표에서, 이렇게 교육적 콘텐츠로 학습시키는 것이 정말 정말 잘 동작한다는 걸 보여 준다. 같은 토큰 수를 쓰더라도 더 나은 모델이 나온다는 이야기다.

그리고 우리는 그중에서 sample-10BT라는 100억 토큰짜리 부분 샘플을 쓸 것이다. 우리가 수조 개의 토큰으로 학습시킬 건 아니기 때문이다. 우리는 그냥 FineWeb-EDU의 100억 토큰 샘플로만 학습시킬 것이다. 경험적으로, 내가 이전에 해 본 몇 번의 실험에서 이 정도면 실제로 GPT-2 성능에 꽤 근접하기에 충분했기 때문이다. 게다가 다루기에도 충분히 단순하다. 15조 토큰을 내려받아 굴릴 필요 없이, 100억 토큰짜리 부분 샘플 하나만 있으면 우리 목적에는 충분한 것이다. 그러니 이 sample-10BT를 가지고 작업해 보자. 우리 목표는 이걸 내려받고, 전처리하고, 우리 데이터 로더가 이걸 그대로 다룰 수 있게 만드는 것이다. 그럼 그걸 하러 가자.

다운로드·토크나이즈·샤드 저장 스크립트

좋다. 여기 파일을 하나 더 도입했다. 이 파일은 기본적으로 Hugging Face 데이터셋에서 FineWeb-EDU를 내려받고, 데이터를 전부 전처리하고 미리 토크나이즈한 다음, 로컬 디스크의 어느 폴더에 데이터 샤드로 저장한다. 그리고 이게 돌아가는 동안 잠깐 언급하고 싶은 게 하나 있다. 여기 있는 데이터셋 뷰어를 훑어보면서, 이 데이터셋 안에 실제로 무엇이 들어 있는지 감을 잡아 볼 수 있다는 것이다. 꽤 흥미롭다. 보면 기본적으로 꽤 잘 동작하고 있는 것 같다. 프랑스의 원자력 에너지 이야기가 있고, 멕시코계 미국 이야기가 있고, 맥 파자마 같은 것도 있다. 그러니까 실제로 이들이 걸어 둔 필터가 꽤 잘 동작하고 있는 것처럼 보인다. 쓰레기가 아니라 읽을 만한 설명문들이 남아 있는 것이다.

참고로 여기 필터는 Llama 3 70B를 써서 자동으로 적용되었다고 알고 있다. 그러니까 기본적으로 LLM들이 어떤 콘텐츠가 교육적인지를 판단하는 심판 노릇을 하고, 그렇게 교육적이라고 판단된 것만 결국 필터를 통과하게 되는 것이다. 사람이 일일이 규칙을 적어 거르는 대신 모델에게 교육적 가치를 평가시킨 셈이다. 꽤 멋지다.

이제 스크립트 자체에 대해서는 전체를 훑지는 않겠다. 그렇게 흥미롭지도 않고 LLM 중심적이지도 않기 때문이다. 하지만 이걸 돌리면, 첫째로 데이터셋을 로드한다. 이건 전부 Hugging Face 코드다. 이걸 돌리려면 pip install datasets를 해 줘야 한다. 그렇게 데이터셋을 내려받고, 그다음에는 이 데이터셋 안에 들어 있는 모든 문서를 토크나이즈한다.

문서를 토크나이즈할 때 눈여겨볼 점이 있다. 문서 하나를 토크나이즈할 때 우리는 먼저 토큰 목록을 end-of-text 토큰으로 시작한다. 이건 여러분도 알다시피 GPT-2 토크나이저의 특수 토큰이다. 그러니까 50256이 end-of-text의 ID이고, 이름은 end-of-text지만 실제로는 이것이 문서를 시작하는 첫 번째 토큰이다. 그다음 그 문서의 모든 토큰을 이어 붙인다. 그다음 그것으로 NumPy 배열을 만들고, 모든 토큰이 범위 안에 들어 있는지 확인한다. 즉 모든 문서는 end-of-text 토큰 하나로 시작해서 그 문서의 토큰들이 뒤따르는 형태가 되고, 그런 문서들이 죽 이어 붙여져 하나의 긴 토큰 스트림이 된다. 아, 잠깐, 이걸 좀 디버깅하자.

좋다, 미안하다. 이건 그냥 내가 파이썬에서 부동소수점 나눗셈을 쓴 게 문제였다. 정수 나눗셈이어야 한다. 그래야 이게 정수가 되고 모든 게 괜찮아진다. 어쨌든 여기 토크나이제이션은 비교적 단순하다. 토큰을 np.uint16으로 반환한다. 공간을 조금이나마 아끼려고 16비트 부호 없는 정수를 쓰는 것이다. 2¹⁶ − 1이 65,000 정도이고, GPT-2의 최대 토큰 ID는 그보다 한참 아래에 있기 때문이다. 토큰 ID가 그 범위 안에 전부 들어가니, 굳이 32비트를 쓸 이유 없이 절반의 용량으로 같은 정보를 담을 수 있다.

그리고 여기에는 멀티프로세싱 코드가 잔뜩 들어 있는데, 솔직히 그렇게 흥미롭지 않아서 한 줄씩 짚어 가지는 않겠다. 어쨌든 우리는 데이터셋을 로드하고, 토크나이즈하고, 전부 샤드로 저장한다. 그리고 그 샤드는 NumPy 파일 형태다. 그냥 NumPy 배열 하나를 통째로 저장하는 것이고, 이 배열은 토치 텐서와 아주 아주 비슷하다. 그리고 첫 번째 샤드인 000000은 검증용 샤드로 쓰고, 나머지 샤드는 전부 학습용 샤드로 쓴다. 그리고 앞서 말했듯이 샤드 하나에 전부 정확히 1억 토큰씩 들어 있다. 그렇게 크기를 맞춰 두면 작업하기가 훨씬 더 쉬워진다. 파일을 샤드로 쪼개는 이유는, 거대한 파일 하나만 덜렁 있으면 디스크에서 다루기가 힘들 때가 있기 때문이다. 그래서 1억 토큰 단위로 샤딩해 두는 편이 그런 관점에서 그냥 더 낫다. 필요한 샤드만 하나씩 메모리에 올려 가며 순회할 수 있고, 검증용과 학습용을 파일 단위로 깔끔하게 갈라 둘 수도 있다.

그래서 이건 그냥 돌게 두겠다. 아마 30분쯤 걸릴 것이다. 그리고 그다음에 실제로 이 데이터 위에서 학습하러 돌아오겠다. 이번에는 진짜배기 사전학습을 하게 된다. 이건 좋은 데이터셋이고, 초당 처리하는 토큰 수도 충분히 많고, GPU가 여덟 장 있고, 코드도 이미 준비되어 있다. 그러니 이번에는 정말로 진지한 학습 실행을 하게 될 것이다. 잠깐 뒤에 다시 돌아오자.

100개의 샤드, 그리고 스텝 수 계산

좋다, 돌아왔다. edu_fineweb을 ls해 보면 이제 그 안에 샤드가 100개 들어 있다. 각 샤드가 1억 토큰이니 100개면 총 100억 토큰이다. 말이 된다.

이제 메인 파일로 넘어가서, 데이터 로더를 다시 좀 손봤다. 더 이상 Shakespeare로 돌지 않고 FineWeb 샤드를 쓰고 싶기 때문이다. 그래서 여기 이 샤드들을 추가로 로드할 수 있는 코드가 좀 보일 것이다. uint16 NumPy 파일을 로드하고, 그걸 torch.long 텐서로 변환한다. 위쪽 레이어들 상당수가 기본적으로 그 타입을 기대하기 때문이다. 그다음 여기서 모든 샤드를 열거한다. 그리고 DataLoaderLite에 split도 하나 추가해서, train 스플릿도 로드할 수 있고 val 스플릿, 즉 0번 스플릿도 로드할 수 있게 했다. 그러면 그 샤드들을 로드할 수 있다.

그리고 여기서는 이제 현재 위치뿐 아니라 현재 샤드도 갖게 된다. 그러니까 샤드 하나 안에서의 위치가 있고, 한 샤드 안의 토큰이 떨어지면 먼저 샤드를 전진시키고, 필요하면 루프를 돌고, 그다음 토큰을 가져오고 위치를 다시 맞춘다. 그래서 이제 이 데이터 로더는 샤드 하나를 다 쓰면 다음 샤드로 넘어가면서, 100개의 샤드 전체를 차례로 순회하게 된다. 그걸 그렇게 바꿨다.

그리고 데이터가 처리되는 동안 내가 한 다른 일은, 이제 train_loader의 스플릿을 당연히 train으로 두는 것이다. 그리고 아래쪽에 숫자를 몇 개 세팅해 두었다. 우리는 스텝당 2¹⁹ 토큰을 처리하고 있고, 대략 100억 토큰을 돌고 싶다. 우리가 가진 고유 토큰이 그만큼이기 때문이다. 그러니 100억 토큰을 2¹⁹로 나누면 19,073 스텝이 나온다. 그게 바로 그 숫자가 나온 출처다. 즉 스텝 수는 우리가 임의로 고른 값이 아니라, 우리가 가진 고유 토큰 전체를 정확히 한 바퀴 도는 데 필요한 스텝 수로 자연스럽게 정해진 것이다. 그래서 같은 데이터를 두 번 보지 않고 딱 1에폭을 돌게 된다.

그리고 GPT-3 논문은 학습률을 3억 7,500만 토큰에 걸쳐 워밍업한다고 말한다. 그래서 여기 와서 375 × 10⁶ 토큰을 2¹⁹로 나누면 715 스텝이다. 그래서 워밍업 스텝을 715로 잡아 두었다. 이렇게 하면 GPT-3가 쓴 워밍업 스케줄과 정확히 일치한다. 참고로 내 생각에 715는 아주 온건한 값이고, 훨씬 더 공격적으로 잡아도 될 것 같다. 아마 100 정도만 해도 충분할 것이다. 하지만 괜찮다. GPT-3의 하이퍼파라미터를 정확히 그대로 가져가기 위해 지금은 그냥 715로 두자. 그걸 고쳤고, 그러면 사실상 준비는 끝이다. 이제 실제로 돌릴 수 있다.

마이크로 배치 64 — 그래디언트 누적 없이

스크립트가 여기 있으니 띄울 수 있다. 아, 미안하다, 하나만 더 하자. 내 GPU라면 배치 크기를 더 키울 수 있다. 마이크로 배치 크기로 64까지 넣을 수 있다고 믿는다. 그러니 한번 해 보자. 내 기억이 틀렸을 수도 있다. 하지만 그게 맞다면 GPU당 64 × 1024이고 GPU가 여덟 장이니, 이게 그대로 곱해져서 총 배치 크기 전체가 되어 버린다. 즉 그래디언트 누적을 아예 하지 않게 된다. 이게 메모리에 들어간다면 꽤 빠르게 돌 것이다. 그러니까 마이크로 배치를 16에서 64로 네 배 키우면, 앞서 4였던 그래디언트 누적 스텝 수가 1이 되어 누적 루프 자체가 사라지는 셈이다. 가자, 가자.

이게 동작한다면 이건 기본적으로 진지한 사전학습 실행이다. 우리는 아직 로깅도 하고 있지 않고, 검증 스플릿으로 평가도 하지 않고 있고, 어떤 평가도 돌리고 있지 않다. 그러니 아직 T에 가로줄을 긋고 I에 점을 찍은 건 아니다. 마무리 손질이 아직 남아 있다는 뜻이다. 하지만 이걸 한동안 그냥 돌게 내버려 두면 실제로 꽤 좋은 모델을 얻게 될 것이고, 그 모델은 GPT-2 124M과 대등하거나 어쩌면 그보다 더 나을 수도 있다.

좋다, 전부 잘 굴러가고 있는 것 같다. 초당 150만 토큰을 처리하고 있고, 여기 보이는 것들이 전부 좋아 보인다. 이터레이션당 330밀리초가 걸린다. 그리고 우리가 해야 할 총 스텝 수는, 이걸 어디에 출력하고 있더라, 19,073이다. 그러니 19073 × 0.33은 이만큼의 초이고, 이만큼의 분이다. 즉 이건 1.7시간 동안 돌게 된다. 대략 한 시간 반짜리 실행인 셈이다. 그리고 우리는 그래디언트 누적을 쓸 필요조차 없다. 좋은 일이다. 여러분의 GPU에서는 그런 호사를 누리지 못할 수도 있는데, 그런 경우에는 들어맞을 때까지 배치 크기를 계속 줄여 나가면 된다. 다만 총 배치 크기가 딱 나누어떨어지도록 깔끔한 숫자로 유지하라.

꽤 신난다. 지금은 학습률을 워밍업하는 중이라 아직 아주 낮은 1×10⁻⁴다. 이게 앞으로 다음 몇 스텝에 걸쳐 여기 6×10⁻⁴까지 쭉 올라가게 된다. 아주 좋다.

이제 내가 하고 싶은 것은 T에 가로줄을 긋고 I에 점을 찍는 일이다. 검증 스플릿에서 평가를 해 보고, 평가를 어떻게 돌릴지, 로깅을 어떻게 할지, 손실을 어떻게 시각화할지, 그 밖의 좋은 것들을 어떻게 갖출지 알아내 보자. 실제 본 실행에 들어가기 전에 그것부터 먼저 하자.

검증 데이터 분할, validation loss, 샘플링 되살리기 03:23:10

코드를 손봐서 이제 validation split 위에서도 평가가 돌아가도록 만들었다. 방법은 아주 단순하다. 데이터 로더를 하나 더 만들되 split="val"을 넘겨주면, 그것이 곧 validation shard만을 대상으로 하는 데이터 로더가 된다. 그러니까 train 로더와 완전히 같은 클래스인데 어떤 shard를 읽느냐만 다른 것이다. 앞에서 우리가 shard를 만들 때 train용과 validation용을 따로 떼어 두었기 때문에, 여기서 split 인자 하나만 바꿔 주면 검증 전용 로더가 그대로 나온다. 학습에 쓰이는 shard와 검증에 쓰이는 shard가 애초에 분리되어 있다는 게 여기서 전부를 가능하게 하는 전제다.

그 다음에 내가 한 일은 데이터 로더 안에 reset이라는 새 함수를 도입한 것이다. 이 함수는 초기화 시점, 즉 __init__에서 호출되고, 하는 일은 기본적으로 데이터 로더를 처음 상태로 되돌리는 것이다. 다시 말해 초기화 때 하던 세팅을 그대로 한 번 더 수행하는 함수다. 이게 왜 그렇게 유용한가 하면, 이제 메인 학습 루프로 내려가 보면 그 이유가 드러난다. 검증을 한 번 돌 때마다 항상 validation 데이터의 같은 지점에서 시작하고 싶기 때문이다. 그래야 매번 재는 validation loss가 서로 비교 가능한 값이 된다. 만약 로더를 되돌리지 않고 이어서 읽는다면, 매 검증마다 서로 다른 문서 조각을 보게 되어 값이 들쭉날쭉해지고 비교가 어려워질 것이다.

이것이 내가 추가한 코드다. 기본적으로 100번째 반복마다, 그리고 여기에는 0번째 반복도 포함되는데, 모델을 평가 모드로 전환한다. 즉 model.eval()을 부른다. 그 다음 validation 로더를 reset한다. 그리고 그 다음에는 그래디언트를 전혀 쓰지 않는다. 즉 torch.no_grad() 컨텍스트 안에서 돈다. 그 안에서 대략 20 스텝에 걸쳐 loss를 누적하고, 전부 평균 내서 validation loss를 출력한다.

그러니까 이 검증 루프는 대체로 학습 루프와 정확히 똑같은 논리인데, 다만 loss.backward()가 없다. 순수하게 추론만 하는 것이다. 우리는 그냥 loss를 재고 그걸 더해 갈 뿐이다. 그 외의 모든 것 — 자동 혼합정밀도, 마이크로 배치를 여러 번 돌려서 누적하는 방식 — 은 앞에서 우리가 본 그대로 적용되고 아무것도 달라지지 않는다. 그러니까 여기서 달라지는 것은 딱 두 가지뿐이다. 역전파를 하지 않는다는 것, 그리고 읽는 데이터가 학습 shard가 아니라 validation shard라는 것이다.

그래서 이 코드는 100번째 반복마다, 그리고 맨 첫 반복에도 validation loss를 출력한다. 그건 꽤 좋은 일이다. 우리가 얼마나 과적합되고 있는지에 대해 어느 정도, 아주 약간이나마 알려 주기 때문이다. 다만 솔직히 말하면 우리에겐 사실상 무한에 가까운 데이터가 있어서, train loss와 validation loss가 거의 같게 나올 것으로 기대하고 있다.

그런데 내가 이걸 궁금해하는 다른 이유가 하나 더 있다. OpenAI가 공개한 그대로의 GPT-2 124M을 가져다가 거기서 초기화한 다음, 그 모델이 이 validation loss에서 어떤 값을 내는지를 볼 수 있기 때문이다. 그러면 그 모델이 우리가 쓰는 데이터, 정확히는 FineWeb-Edu의 validation split에 얼마나 일반화되는지에 대한 감이 잡힌다.

물론 이건 GPT-2에 대해 완전히 공정한 비교는 아니다. GPT-2는 아주 다른 데이터 분포 위에서 학습되었기 때문이다. 그래도 여전히 흥미로운 데이터 포인트이긴 하다. 그리고 어떤 경우가 됐든, 이런 식의 학습 실행에는 validation split을 반드시 두고 싶어진다. 자신이 과적합하고 있지 않다는 것을 확인할 수 있어야 하기 때문이다.

이건 특히 우리가 학습 데이터를 여러 에폭 돌게 되면 훨씬 더 큰 관심사가 된다. 예를 들어 지금 우리는 딱 한 에폭만 돌고 있다. 하지만 만약 10 에폭 같은 식으로 학습하는 지점까지 가게 된다면, 그리고 모델이 충분히 크다면, 그 데이터를 너무 외워 버리고 있는 것은 아닌지를 정말로 조심해야 한다. 그리고 validation split은 그런 일이 실제로 벌어지고 있는지를 알아낼 수 있는 하나의 방법이다. train loss는 계속 내려가는데 validation loss가 더 이상 내려가지 않거나 오히려 올라간다면, 그것이 바로 모델이 학습 데이터를 외우기 시작했다는 신호다.

그리고 그것에 더해서, 기억할지 모르겠지만 우리 스크립트 맨 아래쪽에는 한참 전에 써 두고 그대로 방치해 둔 샘플링 코드가 고아처럼 남아 있었다. 그래서 그 코드를 거기서 삭제하고 위쪽 여기로 옮겨 왔다. 그러니까 이제 구조는 이렇다. 가끔 한 번씩 validation을 한다. 그리고 가끔 한 번씩 샘플링을 해서 생성 결과를 뽑아 본다. 그리고 그 두 가지는 100 스텝마다만 하고, 학습은 매 스텝마다 한다. 지금 내가 잡아 둔 구조가 바로 그것이다.

나는 이걸 10,000번 반복 돌려 왔다. 그래서 여기 반복 1,000 시점의 샘플들이 있다. 모델은 자기가 언어 모델이라고 말한 뒤에 더 창의적으로는 못 가겠다는 식으로 이어 가고, 또 다른 샘플에서는 언어 모델이라고 말한 뒤에 컴퓨터의 시작이 어쩌고 하는 식으로 문장이 흘러간다. 보다시피 이건 아직 죄다 뒤죽박죽이다. 하지만 우리는 겨우 반복 1,000에 와 있고, 워밍업을 거쳐 이제 막 최대 학습률에 간신히 도달한 참이다. 그러니 아직 배우는 중인 것이다. 곧 또 다른 샘플들이 나올 텐데, 보다시피 모델은 여전히 갓난아기 상태다.

기본적으로 내가 여기에 넣어 둔 이 샘플링 코드 전부는 여러분에게 이미 익숙할 것이다. 앞에서 그대로 가져온 것이기 때문이다. 내가 한 유일한 변경은 PyTorch에서 generator 객체를 하나 만든 것이다. 그렇게 해서 난수 샘플링을 직접 통제할 수 있게 했다. 왜 그랬냐면, 학습에 쓰이는 전역 난수 생성기의 RNG 상태를 이 샘플링이 건드리는 걸 원하지 않기 때문이다. 나는 이 부분이 학습 루프 바깥에 완전히 따로 떨어져 있기를 바란다.

그래서 샘플링 전용 RNG를 따로 쓰고, 시드를 줄 때는 모든 랭크가 서로 다른 시드를 갖도록 확실히 해 둔다. 즉 랭크 번호를 시드에 반영해서, DDP로 띄운 프로세스들이 각각 다른 난수열을 쓰게 만든다. 그렇게 하지 않으면 DDP의 모든 프로세스가 완전히 똑같은 샘플을 뱉게 될 것이고, 여러 장의 GPU로 여러 개의 생성 결과를 보는 의미가 사라진다. 그 다음, 난수를 실제로 소비하는 지점, 즉 샘플링이 실제로 일어나는 torch.multinomial 호출에 그 generator 객체를 인자로 넘겨 준다. 그 외에는 이전과 완전히 동일하다.

또 한 가지는, 지금 실행 속도가 조금 느려진 걸 알아챘을 것이다. 그건 샘플링이 되게 하려고 내가 실제로 torch.compile을 꺼야 했기 때문이다. 그래서 조금 느리게 돌고 있다. 무슨 이유에서인지 torch.compile 없이는 잘 돌아가는데, 모델을 컴파일하면 PyTorch에서 정말 무시무시한 에러가 튀어나오고, 나는 지금 그걸 어떻게 해결해야 할지 전혀 모르겠다.

그러니 아마도 여러분이 이 코드가 공개된 걸 보게 될 즈음이나 그 비슷한 시점에는 고쳐져 있을지도 모른다. 하지만 지금은 그냥 False로 두겠다. 그리고 torch.compile을 다시 켤 텐데, 그러면 샘플은 못 보게 된다. 이건 나중에 고치면 되겠다고 생각하고 있다.

그건 그렇고, 나는 이 코드 전부를 공개할 것이다. 사실 나는 우리가 뭔가를 추가할 때마다 git 커밋을 남기는 일에 아주 신경을 써 왔다. 그래서 완전히 맨바닥에서 시작해서 지금까지, 그리고 이 이후까지 전부를 담은 저장소를 통째로 공개할 생각이다. 그러면 모든 것이 git 커밋 히스토리 안에 정확히 문서화되어 있는 셈이 된다. 그래서 나는 그게 꽤 괜찮을 것이라고 생각한다. 그러니 여러분이 build-nanogpt GitHub repo에 갈 때쯤에는 이 부분이 제거되어 있고 제대로 동작하기를, 그리고 내가 그 버그를 고쳐 뒀기를 바란다.

평가: HellaSwag, 학습 실행 시작 03:28:23

여기서 최적화가 돌아가고 있고, 한 스텝씩 밟아 나가고 있고, 우리는 지금 스텝 6,000 언저리쯤에 와 있다. 그러니까 지금 학습의 약 30% 지점을 지나고 있는 셈이다. 이게 학습되고 있는 동안, 나는 우리가 validation set을 보완하기 위해 쓰게 될 평가 하나를 소개하고 싶다. 바로 HellaSwag eval이다.

HellaSwag란 무엇인가

HellaSwag는 2019년에 나온 이 논문에서 왔다. 그러니까 이제 5년 묵은 eval이다. HellaSwag가 동작하는 방식은 이렇다. 기본적으로 문장 완성 데이터셋이다. 즉 객관식 형식이다. 이 문항들 하나하나에 대해 우리는 모든 선택지가 공유하는 맥락 문장을 갖는다. 예를 들면 한 여자가 양동이와 개를 데리고 밖에 나와 있고, 개는 목욕을 피하려고 이리저리 뛰어다니고 있다, 그리고 "그녀는"에서 문장이 끊긴다. 그리고 선택지가 넷이다. A는 양동이를 비누로 헹구고 개의 머리를 드라이기로 말린다. B는 호스를 써서 개가 비누투성이가 되지 않게 한다. C는 개를 물에 적시고 개는 다시 도망간다. D는 개와 함께 욕조에 들어간다.

그러니까 기본 아이디어는 이렇다. 이 객관식 선택지들은 그중 딱 하나만이 그 문장의 자연스러운 이어짐이 되고 나머지는 그렇지 않도록 구성된다. 그리고 나머지 것들은 아예 말이 안 될 수도 있다. 호스를 써서 비누투성이가 되지 않게 한다는 건 말이 되지 않는다. 그래서 무슨 일이 벌어지느냐면, 제대로 학습되지 않은 모델들은 이것들을 서로 구분해 내지 못한다. 반면 세계 지식을 많이 갖고 있고 세상에 대해 많은 것을 말할 수 있는 모델들은 이런 완성을 만들어 낼 수 있다.

그리고 이 문장들은 ActivityNet과 WikiHow에서 가져온 것이다. 논문 아래쪽에는 WikiHow의 어떤 도메인들에서 문장이 왔는지를 보여 주는 꽤 근사한 차트가 있다. 컴퓨터와 전자기기 쪽에서 온 문장이 많고, 주택과 정원 쪽에서 온 것도 많다. 그리고 가장 그럴듯한 완성이 무엇인지, 또 그 완성의 정체가 무엇인지를 찾아내려면 알아야 하는 세상 지식의 종류를 꽤 폭넓게 커버하고 있다.

adversarial filtering과 당시의 수치

HellaSwag에서 또 하나 흥미로운 점은 이게 만들어진 방식이다. 틀린 선택지들이 의도적으로 적대적(adversarial)인 방식으로 확보되었다는 것이다. 그러니까 그것들은 그냥 아무렇게나 가져온 문장이 아니다. 실제로는 언어 모델이 생성한 문장들이고, 게다가 언어 모델은 그것들을 어려워하지만 사람은 쉬워하도록 그런 식으로 생성되었다.

그래서 논문에서 그들이 언급하기를, 사람은 이 집합에서 95% 정확도를 내는 반면 당시 최고 성능 언어 모델들은 겨우 48%였다고 한다. 사람에게는 쉽고 모델에게는 어렵다는 것, 그 간격이 바로 이 eval의 설계 의도였다. 그러니 그 시점에는 이게 좋은 벤치마크였던 것이다.

이 논문의 세부는 직접 읽어 보면 더 많은 걸 알 수 있다. 다만 여기서 짚어 둘 점은 이게 5년 전 일이라는 것이고, 그 이후로 HellaSwag에 벌어진 일은 그냥 완전히 풀려 버렸다는 것이다. 이제 여기 언어 모델들은 96%를 낸다. 그러니까 기본적으로 남은 4%, 그 마지막 4%는 아마 데이터셋 자체의 오류이거나 정말로 정말로 어려운 문항일 것이다. 그래서 이 데이터셋은 언어 모델에 관한 한 사실상 박살난 셈이다. 하지만 그때 당시에는 최고의 언어 모델조차 50% 정도에 불과했다. 5년 사이에 일이 여기까지 온 것이다.

그래도 여전히 사람들이 HellaSwag를 좋아하는 이유가 있다. 그건 그렇고 이건 GPT-2에서는 쓰이지 않았다. 하지만 GPT-3에서는 HellaSwag eval이 쓰이고, 아주 많은 사람들이 HellaSwag를 쓴다. 그래서 GPT-3에 대해서는 여기 인용된 결과들이 있고, GPT-3가 이 서로 다른 모델 체크포인트들 각각에서 HellaSwag eval에 대해 몇 퍼센트의 정확도를 달성하는지를 우리가 알고 있다.

왜 작은 모델에 적합한 eval인가 — smooth eval과 early signal

사람들이 이걸 좋아하는 이유는 HellaSwag가 매끄러운(smooth) eval이고, 이른바 "초기 신호(early signal)"를 제공하는 eval이기 때문이다. 초기 신호라는 건, 작은 언어 모델들도 무작위 확률인 25%에서 출발하지만 거기서부터 천천히 나아진다는 뜻이다. 25, 26, 27 하는 식으로 야금야금 올라가는 게 눈에 보인다. 그래서 모델이 아주 작고 학습이 아주 초기 단계일 때조차 그 느린 개선을 관측할 수 있다. 이것이 우리 같은 작은 모델에 이 eval이 잘 맞는 이유다. 정답 하나를 골라내느냐 마느냐로 0이냐 1이냐가 갈리는 과제였다면 작은 모델은 오랫동안 아무 신호도 주지 못했을 것이다. 하지만 후보마다 토큰 확률의 평균을 재서 비교하는 방식이기 때문에, 모델이 아주 조금 나아져도 그 차이가 지표에 매끄럽게 드러난다. 그러니까 매끄럽고, 초기 신호가 있고, 게다가 오래 전부터 있었던 eval이다. 그래서 사람들이 이 eval을 꽤 좋아하는 것이다.

이제 우리가 이걸 평가하는 방식은 다음과 같다. 앞서 말했듯 우리에겐 공유 맥락이 있고 이건 일종의 객관식 과제다. 그런데 모델에게 객관식 문제를 주고 A냐 B냐 C냐 D냐를 물어보는 대신 다른 방식을 써야 한다. 우리는 그렇게 물어볼 수가 없기 때문이다. 지금 여기서 우리가 보고 있는 것처럼 모델이 이렇게 작을 때, 모델들은 실제로 객관식을 풀지 못한다. 라벨 하나를 객관식 선택지 중 하나에 대응시킨다는 개념 자체를 이해하지 못한다. 그걸 이해하지 못하기 때문에, 우리는 모델에게 그들에게 네이티브한 형태로 문제를 줘야 한다. 그리고 그 네이티브한 형태란 바로 토큰 완성이다.

배치 구성과 마스킹

그래서 우리가 하는 건 이렇다. 네 개의 행과 T개의 토큰 — 여기서 T가 얼마가 되든 상관없다 — 으로 이루어진 배치를 하나 구성한다. 그러면 공유 맥락, 즉 네 선택지 모두에 대해 공통인 그 문맥에 해당하는 토큰들은 네 행 전부에 걸쳐 똑같이 들어간다. 그리고 그 뒤에 네 개의 선택지를 각 행마다 하나씩 이어 붙인다. 말하자면 그런 식으로 네 줄로 나란히 펼쳐 놓는 것이다. 그리고 그중 딱 하나만이 정답이다. 이 경우에는 라벨 3, 즉 세 번째 선택지가 정답이다. 그러니까 이게 맞는 선택지이고, 1번과 2번과 4번은 틀린 선택지다. 정리하면, 네 행은 앞부분이 전부 동일하고 뒷부분만 서로 다른 구조가 된다.

그런데 이 선택지들은 길이가 서로 다를 수 있다. 그래서 우리가 하는 일은 가장 긴 길이를 취해서 그것을 배치의 크기 B × T로 삼는 것이다. 그러면 여기 일부 자리는 패딩된 차원이 되어 쓰이지 않은 채로 남는다. 그래서 우리에게 필요한 것은 세 가지다. 토큰들이 필요하고, 정답 라벨이 필요하고, 그리고 어떤 토큰이 활성인지를 알려 주는 마스크가 필요하다. 그리고 마스크는 그 패딩된 영역에서는 0이 된다. 우리가 이 배치들을 구성하는 방식이 바로 이것이다. 마스크가 필요한 이유는 분명하다. 나중에 평균을 낼 때 공유 맥락 부분과 패딩 부분은 빼고, 오직 선택지에 해당하는 완성 토큰들만 가지고 계산해야 하기 때문이다.

그 다음, 언어 모델이 A, B, C, D를 예측하게 만들기 위해 하는 일은 기본적으로 이렇다. 우리는 그냥 그 토큰들과 그들의 확률을 들여다보고, 그 토큰들에 대해 가장 낮은, 혹은 가장 높은 평균 확률을 얻는 선택지를 고른다. 왜냐하면 그것이 그 언어 모델이 보기에 가장 그럴듯한 완성이기 때문이다. 그러니까 우리는 여기 확률들을 보고, 선택지별로 그것들을 평균 내고, 대략적으로 말해 가장 높은 확률을 가진 것을 고른다. 우리가 HellaSwag를 하는 방식이 이것이다.

그리고 내가 믿기로는 GPT-3도 이렇게 했다. 내가 아는 한 GPT-3가 한 방식이 이것이다. 다만 유의할 점이 있다. 여러분이 HellaSwag를 보게 될 다른 eval 모음들 중 일부는 이런 식으로 하지 않을 수도 있다. 그들은 객관식 형식으로 할 수도 있다. 즉 맥락을 한 번만 주고 그 뒤에 네 개의 완성을 전부 보여 줘서, 모델이 최선의 선택지를 고르기 전에 네 선택지를 모두 볼 수 있게 하는 방식이다.

그리고 그건 모델 입장에서는 실제로 더 쉬운 과제다. 자기 선택을 고르는 동안 다른 선택지들을 함께 볼 수 있어서 서로 비교가 가능하기 때문이다. 하지만 안타깝게도 우리 크기의 모델들은 그걸 하지 못한다. 더 큰 크기의 모델들만 그게 가능하다. 그래서 우리 모델은 이 점에서 약간 불리한 조건에 놓인다. 다른 선택지를 보지 못하고 한 번에 하나씩만 보게 되며, 그저 확률을 매길 뿐이고, 이 지표 안에서 정답 선택지가 알아서 이겨 내야만 한다.

hellaswag.py 구현

좋다, 이제 이걸 아주 간단히 구현해서 우리 스크립트에 통합해 보자. 내가 여기서 한 일은 hellaswag.py라는 새 파일을 도입한 것이다. 한번 들여다봐도 되지만, 나는 그 전부를 한 줄씩 훑지는 않겠다. 이건 딱히 깊은 코드라고 할 만한 게 아니기 때문이다. 솔직히 말하면 오히려 좀 지루한 쪽에 가깝다. 거기서 벌어지는 일은 GitHub에서 HellaSwag를 내려받아 그 예제들을 전부 렌더링하는 것이다. 총 10,000개의 예제가 있고, 나는 그것들을 우리가 방금 이야기한 그 형식으로 렌더링하고 있다.

그래서 이 render_example 함수 끝부분을 보면 내가 무엇을 반환하는지가 보인다. 토큰들, 즉 4 × T 크기의 토큰 배열을 반환한다. 그리고 어느 부분이 선택지에 해당하는지를 알려 주고 그 밖의 모든 자리는 0인 마스크를 반환한다. 그리고 정답 라벨을 반환한다. 그러면 그걸로 예제들을 순회하면서 하나씩 렌더링할 수 있게 된다. 그리고 여기에 evaluate라는 함수가 있는데, 이건 Hugging Face에서 GPT-2를 불러온 다음 그 모델로 eval을 돌린다. 하는 일은 기본적으로 내가 방금 설명한 그대로다. 가장 낮은, 혹은 가장 높은 확률을 갖는 선택지를 예측한다.

그걸 실제로 하는 방법은 이렇다. 우리는 기본적으로 cross-entropy loss를 평가하면 된다. 그러니까 시퀀스 안에서 다음 토큰을 예측하는 loss를 재는 것이다. 그런 다음 평균 loss가 가장 낮은 행을 찾고, 그 행에 해당하는 선택지를 우리의 예측으로 고른다. 확률을 평균 내서 가장 높은 것을 고르는 것과 loss를 평균 내서 가장 낮은 것을 고르는 것은 결국 같은 이야기다. 그리고 그 뒤에 몇 가지 통계를 내고 출력하고 하는 것들을 한다. 이것이 HellaSwag를 평가하는 하나의 방법이다. 요점은, 모델에게 글자 A나 B를 뱉으라고 요구하지 않고 오직 다음 토큰 확률만 물어봐도 객관식 문제를 풀게 만들 수 있다는 것이다.

GPT-2의 HellaSwag 수치

이제 여기 위쪽을 보면, GPT-2 124M에 대해 이 스크립트를 돌리면 HellaSwag가 29.5%를 낸다는 것을 보여 주고 있다. 그러니까 여기서 우리가 얻는 성능이 바로 그것이다. 무작위 확률이 25%였다는 걸 기억하면, 우리는 거기서 그렇게까지 멀리 가지 못한 셈이다. 그리고 GPT-2 XL, 즉 GPT-2 중에서 가장 큰 모델은 대략 49%까지 올라간다. 오늘날의 최고 성능이 95%에 더 가깝다는 걸 생각하면 이건 확실히 꽤 낮은 값들이다. 그러니 이들은 이제 확실히 옛날 모델이 되었다.

그리고 한 가지 더, Eleuther harness라는 게 있다. 언어 모델용 eval을 돌리기 위한 아주 중요한 인프라 조각인데, 거기서는 약간 다른 숫자들이 나온다. 나는 그 불일치가 정확히 무엇 때문인지 100% 확신하지는 못한다. 그들이 단순 완성 대신 실제로 객관식 방식으로 평가하는 것일 수도 있고, 그게 그 차이의 원인일 수도 있다. 하지만 그것에 대해서도 100% 확신하지는 못한다. 그건 직접 들여다봐야 알 것이다. 어쨌든 지금으로서는 우리 스크립트가 29.55%를 보고한다. 그러니까 우리가 GPT-2 124M을 맨바닥에서 우리 손으로 직접 학습시킨다면, 그것이 우리가 넘어서고 싶은 숫자가 된다.

학습 스크립트에 eval 통합하기

이제 이 eval을 우리 메인 학습 스크립트에 실제로 통합하는 쪽으로 넘어가겠다. 그리고 그렇게 하려는 이유는 기본적으로 우리가 이걸 한 번만 재고 마는 게 아니라 주기적인 방식으로 반복해서 평가하고 싶기 때문이다. 그래야 HellaSwag 정확도가 시간에 따라 어떻게 변해 가는지를 계속 추적할 수 있고, 우리가 저 29.55% 구간을 언제 넘는지, 혹은 애초에 넘기는 하는지를 눈으로 볼 수 있다. 그러니 이제 train_gpt2.py에 가한 변경 사항들을 하나씩 짚어 보자.

여기서 내가 처음 한 일은 사실상 컴파일 사용을 선택 가능하게 만든 것이고, 기본값으로는 꺼 두었다. 그 문제는 이렇다. 컴파일의 문제는, 안타깝게도 코드를 더 빠르게 만들어 주기는 하지만 평가 코드와 샘플링 코드를 망가뜨린다는 것이다. 아주 고약한 메시지를 뱉는데 나는 그 이유를 모르겠다. 그러니 여러분이 내가 GitHub에 올려 둔 코드베이스에 도달할 때쯤에는 그게 고쳐져 있기를 바란다. 하지만 지금은 torch.compile 없이 돌리고 있고, 그래서 이게 조금 느리게 보이는 것이다. 우리는 torch.compile 없이 돌고 있다.

나는 또 log라는 로그 디렉터리를 만들었고 거기에 log.txt를 놓을 수 있게 했다. 이 파일은 train loss, validation loss, 그리고 HellaSwag 정확도를 기록한다. 아주 단순한 텍스트 파일이다. 우리는 이걸 쓰기 모드로 열어서 처음에는 비어 있는 상태로 시작하게 만들고, 그 이후로는 학습이 도는 내내 거기에 한 줄씩 계속 덧붙여 나간다. 그리고 우리가 지금 마지막 스텝에 와 있는지를 알려 주는 데 도움이 되는 단순한 변수도 하나 만들었다. 마지막 스텝에서는 검증과 샘플링을 반드시 한 번 더 돌려야 하기 때문에 이 플래그가 필요하다.

그 다음 기본적으로 이 루프 안에서 주기적으로 일이 벌어진다. 250번째 반복마다, 혹은 마지막 스텝에서 validation loss를 평가한다. 그리고 역시 250번째 반복마다 HellaSwag를 평가하는데, 단 우리가 컴파일을 쓰고 있지 않을 때만 그렇게 한다. 컴파일이 그걸 망가뜨리기 때문이다. HellaSwag를 평가하는 이 코드에는 잠시 뒤에 다시 돌아오겠다. 그리고 또한 250번째 반복마다 모델에서 샘플링도 한다. 이건 우리가 영상을 시작하던 무렵에 썼던 아주 오래된 코드라는 걸 알아볼 것이다. 우리는 그냥 모델에서 샘플을 뽑고 있는 것이다.

마지막으로 여기, validation을 하고 샘플링을 하고 HellaSwag를 평가한 다음에는 실제로 학습 스텝을 한 번 밟는다. 이게 학습의 한 스텝이고, 이것이 무엇을 하는지는 이제 여러분에게 충분히 익숙할 것이다. 여기서는 아무것도 바뀌지 않았다. 그리고 여기 끝부분에서 train loss를 얻고 나면 그 값을 파일에 기록한다.

여러 GPU가 협력하는 HellaSwag 평가

그러니까 정말로 바뀐 것, 내가 실제로 새로 추가한 것은 HellaSwag eval을 위한 이 섹션 전체다. 이게 동작하는 방식은 이렇다. 나는 모든 GPU가 HellaSwag 위에서 서로 협력하게 만들려고 하고 있다. 그래서 모든 예제를 순회하되, 각 프로세스는 자기에게 배정된 예제만 골라 간다. 그러니까 인덱스 i를 world_size로 나눈 나머지를 취하고, 그 값이 rank와 같아야만 한다. 같지 않으면 continue로 그냥 넘어간다.

그런 다음 예제를 렌더링해서 GPU 위에 올리고 그 다음 logits를 얻는다. 그리고 여기에는 가장 낮은 loss를 갖는 선택지를 예측하도록 도와주는 헬퍼 함수를 하나 만들어 뒀다. 그래서 여기서 예측이 나오고, 그게 정답이면 맞힌 개수를 세어 둔다. 이렇게 하면 각 GPU가 자기 몫의 예제만 처리하게 되므로 전체 eval이 그만큼 빨리 끝난다.

그리고 여러 프로세스가 이 모든 일에 대해 협력하고 있다면 그들의 통계를 동기화해야 한다. 그렇게 하는 한 가지 방법은 여기서 우리 통계를 텐서로 포장한 다음, 그 텐서에 대해 dist.all_reduce를 합 연산으로 부르는 것이다. 그러고 나서 여기서 다시 텐서 바깥으로 풀어내 그냥 정수 값으로 되돌린다. 그렇게 모든 프로세스의 통계를 합쳐야 비로소 10,000개 예제 전체에 대한 정확도가 나온다. 그리고 여기서 마스터 프로세스가 그 HellaSwag 정확도를 출력하고 로그 파일에 남긴다. 대체로 이게 전부이고, 이것이 지금 내가 여기서 돌리고 있는 것이다.

학습 실행의 중간 지점

여기 최적화가 돌아가는 게 보이고, 방금 생성을 한 번 했고, 지금은 스텝 10,000이다. 전체가 약 20,000 스텝이니까 우리는 절반을 온 셈이다. 그리고 이 단계에서 우리가 얻고 있는 샘플들이 이런 것들이다. 한번 보자. 어떤 샘플에서는 모델이 자기가 언어 모델이라고 말한 뒤에 그걸 이용해 어떤 종류의 출력을 만들어 내고 싶다고 이어 간다. 또 다른 샘플에서는 자기가 언어 모델이고 여러 회사의 개발자라는 식으로 이어 간다. 재미있는 걸 하나 찾을 수 있는지 보자.

잘 모르겠다. 이건 여러분이 직접 훑어봐도 된다. 다만 확실한 건 예측이 점점 덜 무작위해지고 있다는 것이다. 모델이 조금 더 자기 인식이 생긴 것처럼 보이고, 자기가 언어 모델이라는 사실에 좀 더 특화된 언어를 쓰고 있는 것 같다. 어떤 샘플은 언어가 소통에 쓰이는 방식에 대해 말하고, 또 어떤 샘플은 자기가 영어와 독일어로 말하게 될 것이라고 이어 간다.

좋다, 잘 모르겠다. 그러니 이 최적화가 끝날 때까지 그냥 기다려 보자. 그리고 그때 우리가 어떤 종류의 샘플을 얻게 되는지 보겠다. 또한 그때 train loss와 validation loss, 그리고 HellaSwag 정확도를 함께 들여다보면서, 우리가 앞서 정해 둔 그 기준점 대비 어떻게 하고 있는지를 확인할 것이다.

SECTION 4: 아침에 확인한 결과 — GPT-2와 GPT-3 재현 03:43:05

GPT-2. 좋다, 좋은 아침이다. 잠시 동안 오른쪽에 있는 이 Jupyter 노트북에 집중해 보자. 나는 여기에 새로운 셀을 하나 만들었다. 그리고 이 셀이 기본적으로 우리에게 해 주는 일은, 학습 손실(train), 검증 손실(validation), 그리고 HellaSwag 점수를 우리가 시각화할 수 있게 해 주는 것이다. 그리고 이 셀은 한 줄씩 따라가며 확인해 볼 수 있는데, 이것이 하는 일이라고는 사실상 우리가 지금까지 계속 써 온 그 로그 파일을 파싱하는 것이 전부다. 그리고 이 코드의 나머지 부분은 대부분 그냥 지루한 matplotlib 코드일 뿐이다. 하지만 어쨌든, 기본적으로 이것이 바로 우리의 최적화가 실제로 어떤 모습이었는지를 보여 주는 그림이다.

우리는 총 19,073스텝을 돌렸고, 이것은 100억 토큰, 그러니까 10B 토큰에 해당한다. 그리고 그 100억 토큰은 FineWeb sample-10B의 정확히 1에폭에 해당하는 양이다. 그러니까 우리는 이 데이터셋을 딱 한 바퀴만 돌린 셈이다. 왼쪽 그래프에는 손실이 그려져 있다. 파란색이 학습 손실이고, 주황색이 검증 손실이다. 그리고 빨간색 수평선으로 그어 놓은 것은, OpenAI의 GPT-2 124M 모델 체크포인트를 그저 바로 이 FineWeb-EDU의 검증 세트 위에서 그대로 평가했을 때 나오는 값이다. 그래서 여기서 여러분이 볼 수 있는 것은, 우리가 그 선을 넘어서고 있다는 사실이다. 이 주황색 선이 빨간색 선 아래에 놓여 있다. 그러니까 우리는 이 데이터셋의 검증 세트를 기준으로 삼았을 때 GPT-2를 능가하고 있는 것이다.

다만 내가 앞에서 언급했듯이, 이 데이터셋의 분포는 GPT-2가 학습했던 데이터의 분포와 매우 다르다. 그러므로 이것은 정확하게 공정한 비교라고 할 수 없다. 그래도 하나의 교차 점검(cross check)으로 삼아 들여다보기에는 꽤 괜찮은 지표다. 이제 이상적으로 우리가 보고 싶은 것은 따로 떼어 둔(withheld) 데이터, 서로 비교가 가능한 데이터, 그리고 어느 정도 표준으로 통용되는 무언가다. 그리고 우리에게 그 역할을 해 주는 것이 바로 HellaSwag라는 평가다.

그래서 여기 오른쪽에서 우리는 HellaSwag의 진행 상황을 볼 수 있다. 우리는 25%에서 출발해서 여기까지 쭉 올라왔다. 즉 이 곡선이 우리가 밤새 밟아 온 궤적이다. 그리고 여기서 빨간색으로 표시된 것이 바로 OpenAI의 GPT-2 124M 모델이다. 그 모델이 달성하는 HellaSwag 점수가 바로 이 선이다. 그리고 3천억 토큰으로 학습된 GPT-3 124M 모델이 달성하는 점수는 초록색으로 표시되어 있고, 그건 저 위쪽에 있다. 그래서 여러분도 보다시피, 우리는 기본적으로 바로 이 지점에서 GPT-2 124M 모델을 넘어섰다. 그리고 이건 정말로 기분이 좋은 일이다.

그런데 여기서 흥미로운 점이 하나 있다. 우리는 겨우 100억 토큰만 학습하고도 그렇게 할 수 있었다는 것이다. GPT-2는 1,000억 토큰, 그러니까 100B 토큰으로 학습되었는데 말이다. 그러니까 어떤 이유에서인지, 우리는 학습에 훨씬 더 적은 토큰만 쓰고도 그냥 넘어갈 수 있었던 셈이다. 겨우 100억 토큰을 학습한 것만으로 그 정확도를 맞추거나 넘어설 수 있었던 이유에는 여러 가지 가능성이 있다. 그러니 그 가능성들을 하나씩 짚어 보자.

첫 번째로, OpenAI의 GPT-2가 훨씬 더 넓은 데이터 분포 위에서 학습되었을 수 있다는 것이다. 특히 FineWeb-EDU는 전부 영어로 되어 있다. 그것은 다국어가 아니고, 수학과 코드도 그렇게까지 많이 들어 있지 않다. 그래서 수학과 코드, 그리고 다국어 데이터가 원래 GPT-2 모델의 용량(capacity)을 갉아먹고 있었을 가능성이 있다. 그리고 기본적으로 그것이, 이 비교가 딱 맞아떨어지지 않는 이유의 일부일 수 있다는 뜻이다.

그리고 이것 말고도 다른 이유들이 많이 있다. 예를 들어 HellaSwag 평가 자체가 꽤나 오래된 것이라는 점이다. 아마 5년쯤, 또는 그 정도 되었을 것이다. 그러므로 HellaSwag의 어떤 측면들이 어떤 식으로든, 심지어는 완전히 똑같은 형태로 FineWeb의 학습 세트 안에 그대로 들어가 버렸을 가능성도 있다. 우리로서는 그것을 확실하게 알 수가 없다. 하지만 만약 실제로 그런 경우라면, 우리는 사실상 검증 곡선을 보고 있는 것이 아니라 기본적으로 학습 곡선을 보고 있는 셈이 된다. 그러니 길게 말할 것 없이 요약하자면, 이것은 완벽한 평가가 아니고 여기에는 몇 가지 유보 조건이 붙어 있다는 뜻이다.

그래도 적어도 우리는 우리가 완전히 잘못된 무언가를 하고 있지는 않다는 정도의 확신은 얻을 수 있다. 그리고 아마 사람들이 이런 데이터셋을 만들 때는, 아주 흔하게 쓰이는 테스트 세트가 학습 세트 안에 들어가 버리지 않도록 확실히 신경을 쓰는 것이 보통일 것이다. 예를 들어 Hugging Face가 FineWeb-EDU를 만들었을 때, 그들은 HellaSwag를 평가 지표로 사용했다. 그러니 나는 그들이 중복 제거(dedup)를 확실하게 해서, 그 학습 세트 안에 HellaSwag가 섞여 들어가 있지 않도록 만들어 두었기를 바란다. 하지만 우리로서는 그것을 확신할 수 없다.

그리고 내가 짧게 짚고 넘어가고 싶은 또 다른 것은, 여기 이 손실 곡선을 한번 보라는 것이다. 이건 정말로 이상해 보인다. 여기 이 부분은 정말로 잘못된 것처럼 보인다. 솔직히 말해서 나도 이게 정확히 무엇인지 100% 알지는 못한다. 내가 의심하는 것은, FineWeb-EDU의 100억 토큰 샘플이 제대로 셔플되지 않았기 때문이라는 것이다. 여기 데이터 쪽에 뭔가 문제가 있는데 나도 아직 완전히 이해하지 못했고, 거기에는 뭔가 이상한 주기성(periodicity)이 있다.

그리고 우리는 일종의 아주 게으른 방식으로 모든 토큰을 그냥 직렬화해 두고, 그것들을 처음부터 끝까지 순서대로 훑고 있을 뿐이다. 우리 쪽에서 어떤 순열(permutation)도 적용하지 않고, 어떤 무작위 샘플링(random sampling)도 하지 않는다. 그래서 내 생각에 우리는 그들이 그 데이터셋 안에 이미 넣어 둔 순서를 일부 그대로 물려받고 있는 것 같다. 그러니 이것은 이상적인 상태가 아니다. 다만 여러분이 이 저장소를 실제로 보게 될 즈음에는, 그건 그렇고, 이런 것들 중 일부가 이미 고쳐져 있기를 바란다. 나는 이 build-nanogpt 저장소를 공개할 예정인데, 지금 당장 그것은 조금 지저분하고 예비적인 상태로 보인다. 그러니 여러분이 여기까지 도달할 즈음에는 그것이 더 깔끔해져 있기를 바란다.

그리고 여기 아래쪽에서 나는 정오표(errata)를 보여 줄 것이고, 이 영상이 나간 뒤에 일어난 일들 중 몇 가지에 대해서도 이야기할 것이다. 그리고 나는 우리가 그때쯤이면 이 작은 문제를 이미 수정해 두었을 것이라고 예상한다. 하지만 지금으로서 이 그림이 기본적으로 보여 주는 것은, 우리의 학습이 완전히 틀리지는 않았다는 것, 그리고 우리가 10분의 1에 해당하는 토큰 예산만으로 그 정확도를 넘어설 수 있었다는 것이다.

그리고 어쩌면 데이터셋 자체가 그동안 더 좋아졌을 가능성도 있다. 원래 GPT-2의 데이터셋은 WebText였다. 그 데이터셋에는 그렇게까지 많은 주의와 정성이 들어가지 않았을 가능성이 있다. 그때는 LLM의 아주 초창기였으니까 말이다. 반면 지금은 중복 제거, 필터링, 품질 필터링 같은 좋은 관행을 두고 훨씬 더 많은 검증과 주목이 이루어지고 있다. 그래서 우리가 학습하고 있는 데이터가 그냥 토큰당 품질이 더 높을 가능성이 있고, 그것이 우리에게 추가적인 이득을 주고 있을 수도 있다. 그러니 우리가 생각해 보아야 할 유보 조건이 여러 개 있는 셈이다. 하지만 지금으로서는, 우리는 이 결과에 대해 꽤 만족하고 있다.

자, 그리고 그다음으로 내가 관심을 가졌던 것은 바로 이것이다. 여러분도 보다시피 지금은 아침이다. 그러니까 그사이에 하룻밤이 지나간 것이다. 그래서 나는 기본적으로 이 결과를 내가 얼마나 더 멀리까지 밀어붙일 수 있는지를 보고 싶었다. 밤새 돌리는 실행을 하기 위해 나는 1에폭 대신 — 1에폭은 대략 두 시간이 걸렸다 — 그냥 거기에 4배를 곱했다. 그러면 내가 자는 동안 여덟 시간이 걸릴 테니까 말이다. 그래서 우리는 4에폭, 그러니까 대략 400억 토큰, 40B 토큰만큼의 학습을 했고, 나는 우리가 얼마나 멀리까지 갈 수 있는지 보려고 했던 것이다.

그러니까 이것이 내가 준 유일한 변경 사항이었고, 그런 다음 나는 스크립트를 다시 돌렸다. 그리고 400억 토큰 지점에서 로그 파일을 가리켜 읽어 보았더니 곡선이 이런 모습이었다. 좋다, 이 그림을 하나씩 설명해 보자. 첫 번째로, 우리는 여기서 서로 다른 에폭들을 가로지르는 이 주기성 문제를 여전히 보고 있다. FineWeb-EDU 데이터셋에 뭔가 정말로 이상한 일이 벌어지고 있는 것인데, 그것은 아직 규명되어야 할 무언가로 남아 있다.

하지만 그것 말고는, 우리는 HellaSwag 점수가 실제로 아주 많이 올라간 것을 볼 수 있다. 우리는 거의, 정말 거의 GPT-3 124M의 정확도까지 도달했다. 바로 저 위에 있는 그 선까지 말이다. 다만 완전히 거기에 닿지는 못했다. 그러니 내가 조금만 더 오래 자지 않은 것이 아쉬울 뿐이다. 그리고 내 생각에는, 만약 이것이 5에폭짜리 실행이었다면 우리가 저기까지 도달했을 수도 있다.

여기서 짚고 넘어갈 것이 하나 있다. 만약 여러분이 다중 에폭(multi-epoch) 실행을 하고 있다면, 우리의 데이터 로더(data loader)는 사실 그렇게까지 신중하게 만들어져 있지 않다는 것이다. 이 데이터 로더는 데이터를 정확히 같은 형식으로, 그리고 정확히 같은 순서로 훑고 지나간다. 이것은 일종의 차선책이고, 최적의 상태와는 거리가 있다. 그래서 여러분은 실제로 데이터를 무작위로 뒤섞어 주는 쪽의 확장을 한번 들여다보고 싶을 것이다. 즉 매번 새로운 에폭이 돌 때마다 각각의 샤드(shard) 안에 들어 있는 문서들을 무작위로 뒤섞고, 심지어는 그 샤드들 자체의 순서까지 뒤섞어 주는 것이다. 그렇게 하면 주기성을 줄이는 데 큰 도움이 될 것이고, 최적화 자체에도 더 좋다. 똑같은 것을 똑같은 형식으로 반복해서 보지 않게 되고, 문서들이 서로 어떤 순서로 이어지는지에 어느 정도의 무작위성이 들어가기 때문이다.

왜냐하면 여러분이 기억해야 할 것은, 모든 각각의 행(row) 안에서 이 문서들이 서로 이어 붙어 있다는 사실이기 때문이다. 문서 하나가 끝나면 그다음에 end-of-text 토큰이 오고, 그러고 나서 그다음 문서가 온다. 그러니까 현재 이 문서들은 정확히 똑같은, 동일한 방식으로 서로 붙어 있는 것이다. 하지만 우리는 사실 그 문서들을 쪼개서 여기저기 뒤섞고 싶다. 문서의 순서라는 것은 중요하지 않아야 하고, 기본적으로 우리는 그 의존성을 끊어 버리고 싶기 때문이다. 그것은 일종의 허위 상관(spurious correlation)이다. 그런데 우리 데이터 로더는 현재 그것을 하고 있지 않다. 그리고 그것이 여러분이 만들어 볼 만한 개선점 중 하나다.

그리고 또 하나 짚고 넘어갈 것은, 우리가 겨우 400억 토큰만으로 GPT-3의 정확도에 거의 근접했다는 점이다. GPT-3는 3천억 토큰, 그러니까 300B 토큰으로 학습되었다. 그러니 여기서도 우리는 학습 효율 측면에서 대략 10배 정도의 개선을 보고 있는 셈이다. 그리고 이것을 정확히 무엇에 귀속시켜야 하는지는, 앞선 실행에 대해 내가 이미 언급했던 그 이유들 말고는 나도 사실 정확히 알지 못한다.

그리고 내가 짧게 언급하고 싶은 또 다른 것은, 여기 있는 최대 학습률(max learning rate)이다. 나는 이전에 이것과 관련된 저장소에서 사람들이 이 값을 가지고 조금 실험해 보는 것을 이미 보았다. 그리고 결과적으로 이 값을 거의 3배까지 올릴 수 있다는 것이 드러났다. 그러니까 최대 학습률이 훨씬 더 높아도 될 가능성이 있고, 어떤 이유에서인지 우리가 물려받은 GPT-3의 하이퍼파라미터들이 사실 극도로 보수적이라는 뜻이다. 더 높은 학습률로도 충분히 굴러가고, 그러면 학습이 더 빨라질 것이다. 그러니 이런 하이퍼파라미터들 중 상당수는 꽤 조정이 가능한 것들이다. 여러분 마음껏 만져 보아도 좋다. 이 값들은 아마 정확하게 맞춰진 값이 아닐 것이고, 기본적으로 여러분이 이렇게 밀어붙여도 괜찮을 가능성이 있다.

그리고 만약 여러분이 GPT-3에 정확하게 충실하고 싶다면, 다음과 같은 차이도 함께 반영하고 싶을 것이다. 여기로 와서 보면, GPT-3의 시퀀스 길이(sequence length)는 우리 것의 2배다. 즉 1,024가 아니라 2,048이다. 그러니 여기로 와서 T 값을 2,048로 바꾸면 된다. 그리고 만약 반복(iteration)당, 또는 스텝당 토큰 수를 정확히 똑같은 50만 개 남짓으로 유지하고 싶다면, 이 값을 32로 줄이면 된다. 그러면 둘을 곱했을 때 여전히 50만 대가 나온다 — 32 × 2048 × 8 = 524,288. 그렇게 하면 여러분의 모델 시퀀스 길이가 GPT-3의 것과 같아진다. 그리고 그 경우에는, 기본적으로 두 모델이 내가 아는 한 거의 동일해진다. 다시 말하지만 GPT-2와 GPT-3는 아주아주 비슷한 모델들이기 때문이다.

이제 밤새 학습된 그 모델에서 나온 샘플 몇 개도 볼 수 있다. 이것이 그 최적화의 결과다. 그리고 보다시피 우리는 76,292스텝까지, 그러니까 대략 그 정도까지 쭉 나아갔다. 이것은 앞의 19,073스텝에 4에폭을 곱한 값이다. 그리고 우리가 거기서 달성한 HellaSwag 점수는 33.24였다. 그리고 여기 이것들이 바로 그 모델에서 나온 샘플 몇 개다. 여러분이 이것들을 죽 읽어 보고 영상을 잠깐 멈춰 본다면, 이 샘플들이 훨씬 더 일관성 있다는 것을 알 수 있을 것이다. 그리고 이 샘플들은 실제로 자기가 언어 모델이라는 사실을 거의 직접 다루고 있다.

hello I'm a language model and I try to be as accurate as possible

I'm a language model, not a programming language

I know how to communicate. I use Python

한국어로 옮기면 이렇다. "안녕, 나는 언어 모델이고 가능한 한 정확하려고 노력한다." "나는 언어 모델이지 프로그래밍 언어가 아니다." "나는 소통하는 법을 안다. 나는 파이썬을 쓴다." 잘 모르겠다면, 여기서 영상을 한번 멈추고 이 샘플들을 직접 들여다본 다음, 100억 토큰만 학습했던 그 모델의 샘플과 비교해 보라. 그러면 여러분은 이쪽이 훨씬 더 일관성이 있다는 것을 알게 될 것이다. 그리고 이것은 여러분이 직접 가지고 놀아 볼 수도 있는 것이다.

그건 그렇고 내가 코드에 추가한 것이 하나 더 있는데, 바로 여기 있는 이 코드 덩어리다. 그러니까 기본적으로 우리가 검증 손실(validation loss)을 평가한 직후에, 만약 우리가 마스터 프로세스라면, 5,000스텝마다 그 검증 손실을 로그로 남기는 것에 더해서 체크포인트(checkpoint)도 함께 저장하게 된다. 그리고 이 체크포인트라는 것은 사실상 그냥 모델의 state dictionary일 뿐이다. 그리고 체크포인팅은 좋은 일이다. 단지 여러분이 모델을 저장해 둘 수 있고, 나중에 어떤 식으로든 그것을 다시 쓸 수 있기 때문이다.

그리고 만약 여러분이 최적화를 중단한 지점에서 이어서 재개하고 싶다면, 모델을 저장하는 것에 더해서 옵티마이저의 state dict도 함께 저장해야 한다. 기억하겠지만 옵티마이저는 Adam 때문에 몇 개의 추가적인 버퍼를 가지고 있다. m과 v가 있고, 여러분은 옵티마이저도 제대로 복원해야 한다. 그리고 RNG 시드, 즉 난수 생성기 같은 것들도 조심해야 한다. 그러니 최적화를 정확히 그대로 재개할 수 있기를 원한다면, 학습 프로세스의 상태 전체를 처음부터 끝까지 꼼꼼히 따져 보아야 한다. 하지만 여러분이 그냥 모델만 저장하고 싶은 것이라면, 바로 이렇게 하면 된다.

여러분이 이 체크포인팅을 하고 싶어 할 만한 좋은 이유가 하나 있는데, 그것은 모델을 훨씬 더 신중하게 평가해 보고 싶을 수 있기 때문이다. 여기서 우리는 HellaSwag 평가를 다소 즉흥적으로 때우고 있다. 하지만 여러분은 예를 들어 EleutherAI의 evaluation harness 같은 더 좋은 도구를 쓰고 싶을 수 있다. 이것 역시 언어 모델을 평가하는 하나의 방법이다. 그러니까 기본적으로 다른 인프라를 써서 모델을 여러 평가 항목에 걸쳐 훨씬 더 철저하게 평가하고, 예를 들어 수학이나 코드나 다른 언어들이 관련된 다양한 과제에서 그것을 OpenAI GPT-2 모델과 직접 비교해 보고 싶을 수 있다는 뜻이다. 그래서 이것도 갖춰 두면 좋은 기능 중 하나다.

그리고 내가 언급하고 싶은 또 다른 것은, 우리가 여기서 만들어 온 이 모든 것이 오직 사전학습(pre-training) 단계일 뿐이라는 점이다. 그러니까 여기 있는 이 GPT는 문서를 꿈꾸는 물건이다. 그것은 그저 다음 토큰을 예측할 뿐이다. 여러분이 ChatGPT에게 말을 거는 것처럼 이 모델에게 말을 걸 수는 없다. 그러니 만약 여러분이 이 모델과 대화하고 싶다면, 우리는 그것을 채팅 형식으로 파인튜닝해 주어야 한다.

그리고 사실 그것은, 여러분이 지도 파인튜닝(supervised fine-tuning), 즉 SFT를 들여다보고 있는 것이라면, 그렇게까지 복잡한 일이 아니다. 그것이 진짜로 의미하는 바는 그냥 데이터셋을 훨씬 더 대화체인 데이터셋으로 바꿔 끼우는 것이다. 사용자-어시스턴트, 사용자-어시스턴트 같은 구조가 들어 있는 데이터셋으로 말이다. 그리고 우리는 그 위에서 파인튜닝을 한다. 그러고 나면 우리는 기본적으로 사용자 토큰을 채워 넣고 어시스턴트 토큰을 샘플링한다. 그것보다 훨씬 더 깊이 들어가는 무언가는 아니다. 하지만 기본적으로 우리가 하는 일은 데이터셋을 갈아 끼우고 학습을 계속 이어 가는 것이다. 다만 지금으로서는 우리는 사전학습 단계에서 멈추기로 하겠다. 그리고 하나만 더 보여 주겠다.

llm.c 소개 — 같은 일을 순수 C/CUDA로, 더 빠르게 03:56:21

내가 짧게 하나 더 보여 주고 싶은 것이 있다. 물론 우리가 오늘 쌓아 올린 것은 nanoGPT repo를 향해 가는 작업이었다. 앞서 보았던 바로 그 저장소 말이다. 그런데 사실 nanoGPT 구현이 하나 더 있다. 그리고 그것은 내가 최근에 작업해 온 조금 더 새로운 프로젝트 안에 숨어 있다. 그 프로젝트가 바로 llm.c repo다.

그리고 llm.c는 GPT-2, 또는 GPT-3의 학습을 순수 CUDA로 구현한 것이다. 그것은 그냥 CUDA를 직접 사용하고, 그 자체가 CUDA로 작성되어 있다. 여기서 nanoGPT는 그 C 구현에 대한 PyTorch 레퍼런스 코드 역할을 한다. 그러니까 우리는 이 둘을 정확하게 일치시키려고 애쓰고 있지만, 동시에 C/CUDA 쪽이 더 빠르기를 기대하고 있다. 그리고 물론 현재로서는 실제로 그런 것으로 보인다. 왜냐하면 그것은 직접적으로 최적화를 해 놓은 구현이기 때문이다.

llm.c 안에 있는 train_gpt2.py는 기본적으로 nanoGPT다. 그리고 여러분이 이 파일을 쭉 훑어 내려가 보면, 우리가 이 강의에서 쌓아 올린 것과 아주 많이 닮은 것들을 잔뜩 발견하게 될 것이다. 그리고 train_gpt2.cu를 들여다보면, 이것이 바로 그 C/CUDA 구현이다. 그래서 여기에는 MPI, NCCL, GPU, CUDA, C/C++ 같은 것들이 잔뜩 들어 있고, 여러분은 그런 것들에 익숙해져 있어야 한다. 어쨌든 이것이 다 빌드되고 나면, 우리는 실제로 이 둘을 나란히 놓고 돌려 볼 수 있고, 그러면 둘은 정확히 같은 결과를 내놓는다. 다만 llm.c가 실제로는 더 빠르게 돈다. 그러면 그것을 한번 보자.

그래서 왼쪽에는 PyTorch로 된, nanoGPT처럼 생긴 것이 있다. 그리고 오른쪽에는 llm.c 호출이 있다. 그리고 여기서 나는 이 둘을 동시에 실행할 것이다. 이 둘은 모두 단일 GPU 위에서 돌아갈 것이다. 여기서 나는 llm.c를 GPU 1에 올리고, 이쪽은 기본값으로 GPU 0을 잡을 것이다. 그러면 여기서 우리가 볼 수 있듯이 llm.c가 컴파일되고, 그다음에 공간을 할당하고, 그러고 나서 스텝을 밟기 시작한다. 그러니까 기본적으로 그동안 PyTorch는 아직 컴파일 중이다. 왜냐하면 여기서는 torch.compile이 llm.c의 nvcc CUDA 컴파일보다 조금 느리기 때문이다. 그래서 이쪽 프로그램은 이미 돌기 시작했는데, 우리는 여기서 아직 torch compile을 기다리고 있는 것이다.

물론 이것은 GPT-2와 GPT-3에 아주 특화된 구현이고, 반면 PyTorch는 아주 범용적인 신경망 프레임워크(neural network framework)다. 그러므로 이 둘이 정확하게 비교 가능한 대상은 아니다. 하지만 만약 여러분이 오직 GPT-2와 GPT-3를 학습시키는 데에만 관심이 있는 것이라면, llm.c는 매우 빠르다. 그것은 공간도 더 적게 차지하고, 시작하는 것도 더 빠르고, 스텝당 속도도 더 빠르다.

그래서 이제 PyTorch도 여기서 스텝을 밟기 시작했다. 그리고 보다시피 우리는 이쪽에서 초당 약 223,000 토큰으로 돌고 있고, 저쪽에서는 초당 약 185,000 토큰으로 돌고 있다. 그러니까 한쪽이 꽤 느린 셈이다. 다만 나는 PyTorch 구현에서 즙을 남김없이 다 짜냈다고 완전히 확신하지는 못한다. 하지만 여기서 중요한 것은, 만약 내가 스텝들을 나란히 맞춰 놓고 본다면 이 둘 사이에 출력되는 손실(loss)과 노름(norm)이 동일하다는 점을 알아차리게 된다는 것이다. 그래서 왼쪽에 있는 것이 PyTorch이고 오른쪽에 있는 것이 이 C 구현인데, 둘은 서로 같다. 다만 이쪽이 더 빠르게 돌 뿐이다.

그래서 나는 여러분에게 llm.c도 짧게나마 보여 주고 싶었다. 그리고 이것은 병렬 구현이고, 여러분이 한번 가지고 놀아 보거나 들여다볼 만한 것이기도 하다. 그리고 그건 꽤 흥미로운 물건이다.

정리 — build-nanogpt GitHub 저장소 03:59:39

좋다, 이 시점에서 나는 아마 이 영상을 슬슬 마무리하기 시작해야 할 것 같다. 왜냐하면 내가 예상했던 것보다 훨씬 더 길어지고 있다고 생각하기 때문이다. 하지만 우리는 아주 많은 땅을 밟았고, 그 모든 것을 밑바닥부터 직접 만들었다.

그래서 아주 짧게 요약하자면 이렇다. 우리는 GPT-2 논문과 GPT-3 논문을 들여다보았고, 이런 학습 실행을 어떻게 설정하는지, 그리고 거기에 관련된 모든 고려사항들이 무엇인지를 살펴보았다. 그리고 우리는 그 모든 것을 밑바닥부터 직접 작성했다. 그리고 두 시간짜리 학습 실행이든 밤새 돌린 실행이든, 그 시간 동안 우리가 실제로 GPT-2와 GPT-3의 1억 2,400만 파라미터 체크포인트, 즉 124M 체크포인트를 아주 상당한 수준까지 따라잡을 수 있다는 것을 보았다.

원칙적으로, 우리가 작성한 이 코드는 여러분에게 인내심이나 컴퓨팅 자원만 있다면 훨씬 더 큰 모델도 학습시킬 수 있을 것이다. 그러니 여러분은 잠재적으로 그보다 더 큰 체크포인트들을 학습시키는 것에 대해서도 생각해 볼 수 있다.

그리고 아직 해결해야 할 문제가 몇 가지 남아 있다. 여기 이 손실 곡선에서 대체 무슨 일이 벌어지고 있는 것인가 — 나는 이것이 FineWeb-EDU 데이터 샘플링과 관련이 있다고 의심하고 있다. 왜 우리는 torch.compile을 켤 수 없는가 — 현재 그것은 생성(generation)과 HellaSwag를 망가뜨린다. 그건 대체 무슨 일인가. 그리고 데이터 로더에서는, 우리가 경계에 도달했을 때 아마도 데이터를 뒤섞어 주어야 할 것이다. 그러니 이와 같은 종류의 문제가 몇 가지 더 남아 있다. 그리고 나는 시간이 지나면서 그중 일부를 바로 여기, 이 build-nanogpt GitHub repo에 문서로 남겨 둘 생각이다. 그리고 이 저장소는 내가 이 영상과 함께 공개할 예정이다.

만약 여러분에게 질문이 있거나, 우리가 다룬 어떤 것에 대해서든 이야기를 나누고 싶다면 Discussions 탭으로 가 달라. 그래서 우리가 거기서 함께 이야기할 수 있도록 말이다. 아니면 여러분이 무엇을 기여하고 싶은지에 따라서 Issues나 Pull Request 쪽으로 가 달라. 또는 our Discord channel, 즉 Zero to Hero Discord도 한번 들여다봐 달라. 그리고 나는 거기 nanoGPT 쪽에서 얼쩡거리고 있을 것이다.

그 외에는, 지금으로서 나는 우리가 도달한 이 지점에 대해 꽤 만족하고 있다. 그리고 여러분이 이 영상을 즐겼기를 바란다. 그럼 나중에 또 보자.