Andrei Karpathy가 Github에 공개한 nanochat의 구조와 코드를 분석하는 글 입니다.
Andrei Karpathy는 전 OpenAI 창립 멤버이자, 테슬라에서 테슬라 오토파일럿 팀을 이끈 사람으로 AI 업계에서 가장 큰 영향력을 가진 사람들 중 하나 입니다.
nanochat을 소개합니다: 100달러로 만들 수 있는 최고의 ChatGPT
nanochat repo에 올린 첫번째 토론글에서 nanochat에 대한 소개를 하고 있습니다. 아래는 해당 글 중 중요한 부분만 발췌해서 정리해 놓았습니다.
환경 설정
프로젝트를 복제합니다.
git clone git@github.com:karpathy/nanochat.git cd nanochat
100달러로 살 수 있는 최고의 ChatGPT를 "스피드런"이라고 부르는 방식으로 훈련시키고 싶습니다. 빈 상자에서 처음부터 끝까지 바로 실행되도록 설계된 speedrun.sh 스크립트를 참고하세요 . 하지만 이 글에서는 각 섹션에 대해 자세히 설명할 수 있도록 단계별로 살펴보겠습니다. 먼저 new&hot uv 프로젝트 관리자가 설치되어 있는지 확인해야 합니다. uv를 설치하고, 에서 새 가상 환경을 생성하고
.venv, 모든 종속성을 가져온 후, 환경을 활성화합니다. 이렇게 하면 입력 시 python시스템 파이썬이 아닌 가상 파이썬 환경이 사용됩니다.speedrun.sh 스크립트에 모든 실행 코드가 담겨져 있습니다. 위에 언급된대로 H100 GPU 8대가 있는 환경을 쓸 수 있다면, 코드를 받고 환경 설정 후, speedrun.sh 스크립트만 실행하면, 처음부터 끝까지 실행해 볼 수 있습니다.
(중략)
토크나이저를 훈련시키세요
다음으로 1) 토크나이저를 훈련하고 2) 모델을 사전 훈련하기 위해 사전 훈련 데이터가 필요합니다. 사전 훈련 데이터는 많은 웹 페이지의 텍스트일 뿐이며, 이 부분에서는 FineWeb-EDU 데이터 세트를 사용합니다. 일반적으로 huggingface 를 사용하면 되지만
datasets.load_dataset()너무 무겁고 부풀려져 매우 간단한 로직을 가리는 방식이 마음에 들지 않아 전체 데이터 세트를 간단하고 완전히 셔플된 샤드로 다시 패키징하여 원하는 대로 쉽고 효율적으로 액세스할 수 있도록 하고 sample-100B 버전을 karpathy/fineweb-edu-100b-shuffle 로 다시 업로드했습니다 . 이 페이지에서 데이터 세트의 예제 텍스트를 미리 볼 수도 있습니다. 각 샤드는 약 0.25M 문자의 간단한 파켓 파일이며 디스크에서 약 100MB를 차지합니다(gzip 압축). 총 1,822개의 샤드가 있지만, depth=20 모델을 학습시키는 데는 240개만 필요합니다(나중에 자세히 설명하겠습니다). 이제 모든 데이터를 다운로드해 보겠습니다. 여기서는 약 24GB 정도 다운로드해야 하지만, 일반적으로 클라우드 환경에서는 상당히 빠릅니다.python -m nanochat.dataset -n 240
이 모든 것은 기본적으로
~/.cache/nanochat.에 저장됩니다. 다운로드가 완료되면 코드북의 문자열과 기호 시퀀스를 서로 변환하는 토크나이저를 학습해 보겠습니다. 기본적으로 어휘 단위의 2**16 = 65,536토큰(꽤 많은 수)을 학습하고 있으며, 그중 일부는 나중에 채팅 스키마에 사용할 특수 토큰으로 예약되어 있습니다. 학습 세트는 20억 개의 문자로 구성되어 있으며, 약 1분밖에 걸리지 않습니다. 학습 알고리즘은 OpenAI에서 사용하는 알고리즘(정규식 분할, 바이트 수준 BPE)과 동일합니다. 더 자세한 내용은 토큰화 에 대한 제 동영상을 참조하세요. 바로 이어서 토크나이저를 평가할 수 있습니다.
python -m scripts.tok_train --max_chars=2000000000 python -m scripts.tok_eval
평가 결과, 약 4.8의 압축률을 달성하고 있는 것으로 나타났습니다(즉, 원본 텍스트 4.8자가 평균적으로 토큰 1개로 변환됨). GPT-2 및 GPT-4 토크나이저와의 비교도 확인할 수 있습니다. GPT-2(토큰 50,257개)와 비교했을 때, 저희 토크나이저는 수학 연산을 제외하고는 텍스트 압축 측면에서 전반적으로 훨씬 우수합니다.
GPT-4에 비해 크게 앞서 있지는 않지만, GPT-4의 어휘량이 훨씬 많다는 점(100,277개)을 기억해야 합니다. 특히 GPT-4는 다국어 뿐만 아니라 코드와 수학에서도 훨씬 뛰어납니다.
사전 훈련
사전 학습을 시작하기 전에 "eval bundle"이라고 부르는 파일을 하나 더 다운로드해야 합니다. 사전 학습 중에 스크립트는 CORE 지표를 주기적으로 평가합니다 . 자세한 내용은 DCLM 논문에서 확인할 수 있지만, 기본적으로 CORE 지표는 자동 완성 기능을 사용하는 다양한 데이터셋에서 모델의 성능을 평가하는 훌륭하고 정규화된 광범위한 지표입니다. HellaSwag, jeopardy, bigbench QA wikidata, ARC-Easy/Challenge, copa, commonsense qa, piqa, lambada, winograd, boolq 등 총 22개의 데이터셋이 있습니다. eval bundle 디렉터리를 다운로드하고 압축을 풀고 기본 디렉터리에 다음과 같이 넣습니다
~/.cache/nanochat/eval_bundle.curl -L -o eval_bundle.zip https://karpathy-public.s3.us-west-2.amazonaws.com/eval_bundle.zip unzip -q eval_bundle.zip rm eval_bundle.zip mv eval_bundle "$HOME/.cache/nanochat"
제가 추천하는 또 다른 설정은(선택 사항이지만) 훈련 중에 좋은 플롯을 볼 수 있도록 wandb를 설정하는 것입니다 . 위에서 uv가 이미 wandb를 설치했지만, 여전히 계정을 설정하고 다음을 사용하여 로그인해야 합니다.
wandb login
이제 사전 학습을 시작할 수 있습니다! 이 부분은 연산량이 가장 많은 부분으로, 시퀀스의 다음 토큰을 예측하여 인터넷 웹 텍스트를 압축하도록 LLM을 학습시키는 과정입니다. 이 과정에서 LLM은 세상에 대한 많은 지식을 얻게 됩니다.
torchrun --standalone --nproc_per_node=8 -m scripts.base_train -- --depth=20
scripts/base_train.py 스크립트를 통해 8개의 GPU에서 학습을 시작합니다 . 20개의 레이어를 가진 트랜스포머를 학습합니다. 기본적으로 각 GPU는 순방향/역방향 연산당 2,048개의 토큰으로 구성된 32개 행을 처리하며,
32*2048 = 2**19 = 524,288최적화 단계당 총 약 50만 개의 토큰을 처리합니다. wandb가 설정되어 있다면 --run=speedrun(모든 학습 스크립트에서 사용 가능) 실행 이름을 설정하고 로그를 기록하기 위해 append를 사용합니다. 학습을 시작하면 다음과 같은 화면이 표시됩니다(간략화를 위해 여러 부분을 삭제했습니다).Vocab size: 65,536 num_layers: 20 model_dim: 1280 num_heads: 10 num_kv_heads: 10 Tokens / micro-batch / rank: 32 x 2048 = 65,536 Tokens / micro-batch: 524,288 Total batch size 524,288 => gradient accumulation steps: 1 Number of parameters: 560,988,160 Estimated FLOPs per token: 3.491758e+09 Calculated number of iterations from target data:param ratio: 21,400 Total number of training tokens: 11,219,763,200 Tokens : Params ratio: 20.00 Total training FLOPs estimate: 3.917670e+19 Scaling the LR for the AdamW parameters ∝1/√(1280/768) = 0.774597 Muon: Grouping 80 params of shape torch.Size([1280, 1280]), device cuda:0, dtype torch.float32 Muon: Grouping 20 params of shape torch.Size([1280, 5120]), device cuda:0, dtype torch.float32 Muon: Grouping 20 params of shape torch.Size([5120, 1280]), device cuda:0, dtype torch.float32 Step 00000 | Validation bpb: 3.3013 ^[step 00000/21400 (0.00%) | loss: 11.090355 | lrm: 1.00 | dt: 23156.74ms | tok/sec: 22,640 | mfu: 1.00 | total time: 0.00m step 00001/21400 (0.00%) | loss: 10.808654 | lrm: 1.00 | dt: 649.22ms | tok/sec: 807,569 | mfu: 35.64 | total time: 0.00m step 00002/21400 (0.01%) | loss: 10.179083 | lrm: 1.00 | dt: 472.29ms | tok/sec: 1,110,094 | mfu: 48.99 | total time: 0.00m step 00003/21400 (0.01%) | loss: 9.449214 | lrm: 1.00 | dt: 487.47ms | tok/sec: 1,075,523 | mfu: 47.47 | total time: 0.00m step 00004/21400 (0.02%) | loss: 8.903216 | lrm: 1.00 | dt: 487.57ms | tok/sec: 1,075,308 | mfu: 47.46 | total time: 0.00m step 00005/21400 (0.02%) | loss: 8.531662 | lrm: 1.00 | dt: 482.58ms | tok/sec: 1,086,417 | mfu: 47.95 | total time: 0.00m step 00006/21400 (0.03%) | loss: 8.231589 | lrm: 1.00 | dt: 487.21ms | tok/sec: 1,076,113 | mfu: 47.49 | total time: 0.00m step 00007/21400 (0.03%) | loss: 7.993080 | lrm: 1.00 | dt: 484.10ms | tok/sec: 1,083,014 | mfu: 47.80 | total time: 0.00m step 00008/21400 (0.04%) | loss: 7.803373 | lrm: 1.00 | dt: 488.17ms | tok/sec: 1,073,989 | mfu: 47.40 | total time: 0.00m step 00009/21400 (0.04%) | loss: 7.627318 | lrm: 1.00 | dt: 484.78ms | tok/sec: 1,081,486 | mfu: 47.73 | total time: 0.00m step 00010/21400 (0.05%) | loss: 7.491893 | lrm: 1.00 | dt: 487.03ms | tok/sec: 1,076,511 | mfu: 47.51 | total time: 0.00m step 00011/21400 (0.05%) | loss: 7.354157 | lrm: 1.00 | dt: 487.10ms | tok/sec: 1,076,334 | mfu: 47.50 | total time: 0.01m step 00012/21400 (0.06%) | loss: 7.246406 | lrm: 1.00 | dt: 487.99ms | tok/sec: 1,074,383 | mfu: 47.42 | total time: 0.02m step 00013/21400 (0.06%) | loss: 7.159368 | lrm: 1.00 | dt: 486.56ms | tok/sec: 1,077,540 | mfu: 47.55 | total time: 0.02m
Transformer에는 1,280개의 채널과 10개의 Attention 헤드가 있으며, 각 헤드의 dim은 128입니다. 이 모델의 매개변수는 약 5억 6천만 개입니다. Chinchilla 스케일링 법칙 권장 사항을 충족하려면 5억 6천만 X 20개 = 112억 개의 토큰을 사용하여 학습해야 합니다. 최적화의 각 단계는 524,288개의 토큰이므로 11.2B (112억) / 0.5M (50만) = 21,400번의 반복이 필요합니다. 토큰당 추정된 FLOP 수를 총 토큰 수로 곱하면 약
4e19 FLOP의 성능 모델이 됩니다. 더 큰 모델일수록 더 작은 학습률을 선호하므로 학습률은 자동으로 1/sqrt(dim)으로 축소됩니다. 행렬 최적화에는 Muon을 사용하고 임베딩 및 언임베딩 최적화에는 AdamW를 사용합니다. 이 모델에는 다른 학습 가능한 매개변수(바이어스, rmsnorms 매개변수 등)가 없습니다. 학습 과정에서는 주기적으로 "검증 bpb"를 보고합니다. 이는 검증 데이터셋의 바이트당 비트 수를 나타냅니다. 바이트당 비트 수는 일반적인 교차 엔트로피 손실보다 훨씬 더 나은 측정값입니다. 각 토큰의 손실을 해당 토큰의 바이트 수로 정규화하여 토크나이저 메트릭을 불변으로 만들기 때문입니다. 따라서 토크나이저의 어휘 크기가 작든 크든, 이 수치는 원시 교차 엔트로피 손실과는 달리 비교 가능합니다. 각 단계는 약 0.5초가 소요되며, lrm이는 학습률 감소 승수(학습이 끝날 무렵 0으로 선형적으로 감소)입니다. 보고된 MFU(모델 플롭 활용도)는 거의 절반 수준으로 양호한 것으로 보이며, 이는 사용 가능한 bfloat16 컴퓨팅을 많이 활용하고 있음을 의미합니다.이제 4e19 FLOP가 경과할 때까지 약 3시간 동안 기다립니다. wandb 플롯에서 다음과 같은 내용이 표시되어야 합니다.
시간이 지남에 따라 bpb가 감소하는 것은 좋은 현상입니다(모델이 다음 토큰을 더 정확하게 예측하고 있기 때문입니다). 또한, CORE 점수가 상승하고 있습니다. 단순히 근사치 지표를 사용하는 대신, 다음과 같이 모델을 더욱 완벽하게 평가할 수 있습니다.
torchrun --standalone --nproc_per_node=8 -m scripts.base_loss torchrun --standalone --nproc_per_node=8 -m scripts.base_eval
train/val 바이트당 비트 수(bpb)가 약 0.81에 도달하고 CORE 지표가 0.22까지 상승하는 것을 확인할 수 있습니다. 비교를 위해 eval 번들에는 GPT-2 모델의 CORE 점수가 포함되어 있습니다. 특히 CORE 0.22는 GPT-2 large(0.21)보다 약간 높지만, GPT-2 xl(즉, "the" GPT-2, 0.26)보다는 약간 낮습니다. 이 시점의 모델은 자동완성 기능을 갖추고 있으므로, 몇 가지 프롬프트를 실행하여 모델에 저장된 지식을 파악할 수 있습니다. 이 파일은
base_loss.py이러한 프롬프트를 실행합니다. 프롬프트는 다음과 같습니다.prompts = [ "The capital of France is", "The chemical symbol of gold is", "If yesterday was Friday, then tomorrow will be", "The opposite of hot is", "The planets of the solar system are:", "My favorite color is", "If 5*x + 3 = 13, then x is", ]
완성된 텍스트는 다음과 같습니다.
The capital of France is Paris. It is the largest city in France and the second largest city in Europe The chemical symbol of gold is Au. The chemical symbol of silver is Ag. The chemical symbol of copper is If yesterday was Friday, then tomorrow will be Saturday. If yesterday was Monday, then tomorrow will be Monday. If yesterday was The opposite of hot is cold. The opposite of hot is cold. The opposite of hot is cold. The planets of the solar system are: Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, Neptune, My favorite color is red. It is the color of the sun, the color of the sky, If 5*x + 3 = 13, then x is a positive integer.
따라서 이 모델은 파리(프랑스)에 대해 알고, Au는 금이며, 토요일은 금요일 다음이라는 사실, "차가운" 것이 "뜨거운"의 반대라는 사실, 심지어 태양계의 평면까지도 알고 있습니다. 하지만 아직 하늘의 색깔이나 간단한 계산 방법은 잘 모릅니다. 하지만 72달러로 학습한 모델치고는 나쁘지 않습니다. 추론은
Engine효율적인 추론을 위해 KV 캐싱을 사용하는 커스텀 클래스와 두 가지 일반적인 추론 단계인 프리필(prefill)과 디코드(decode)를 간단하게 구현한 클래스를 사용합니다. 엔진(Engine) 클래스는 도구 사용(파이썬 인터프리터)도 지원하며, 이는 GSM8K에서 학습할 때 유용합니다(자세한 내용은 나중에 설명).중간 훈련
기본적으로 nanochat은 H100 GPU 8대를 가지고 4시간 정도 학습을 시키면, GPT2 보다 좋은 ChatGPT 모델을 만들 수 있습니다.