FareedKhan-dev/train-llm-from-scratch: 밑바닥부터 LLM 학습시키기

저는 "Attention is All You Need" 논문을 바탕으로 PyTorch를 사용하여 트랜스포머 (Transformer) 모델을 밑바닥부터 구현했습니다. 여러분은 제 스크립트를 사용하여 단일 GPU로 자신만의 수십억 (billion) 또는 수백만 (million) 파라미터 규모의 LLM을 학습시킬 수 있습니다.

다음은 학습된 1,300만 (13 million) 파라미터 LLM의 출력 결과입니다:

In ***1978, The park was returned to the factory-plate that
the public share to the lower of the electronic fence that
follow from the Station's cities. The Canal of ancient Western
...

• 학습 데이터 정보 (Training Data Info)
• 사전 요구 사항 및 학습 시간 (Prerequisites and Training Time)
• 코드 구조 (Code Structure)
• 사용법 (Usage)
• 단계별 코드 설명 (Step by Step Code Explanation)
• 다음 단계 (What’s Next)

학습 데이터는 언어 모델 학습을 위한 다양하고 오픈 소스이며 대규모인 데이터셋인 Pile 데이터셋에서 가져왔습니다. Pile 데이터셋은 책, 기사, 웹사이트 등의 텍스트를 포함한 22개의 다양한 데이터셋 모음입니다. Pile 데이터셋의 총 크기는 825GB이며, 아래는 학습 데이터의 샘플입니다:

Line: 0
{
"text": "Effect of sleep quality ... epilepsy.",
...

코드를 이해하기 위해서는 객체 지향 프로그래밍 (OOP), 신경망 (NN) 및 PyTorch에 대한 기본적인 이해가 있어야 합니다. 시작하는 데 도움이 될 몇 가지 리소스는 다음과 같습니다:

주제	비디오 링크
OOP	OOP 비디오
...

모델을 학습시키려면 GPU가 필요합니다. Colab 또는 Kaggle T4는 1,300만(13+ million) 파라미터 이상의 모델을 학습시키는 데 적합하지만, 수십억(billion) 파라미터 학습에는 실패할 것입니다. 다음 비교를 확인해 보세요:

GPU 이름	메모리	데이터 크기	2B LLM 학습	13M LLM 학습	최대 실용적 LLM 크기 (학습)
NVIDIA A100	40 GB	Large	✔	✔	~6B–8B
...

13M LLM 학습은 1,300만(13+ million) 파라미터 모델의 학습을 의미하며, 2B LLM 학습은 20억(2+ billion) 파라미터 모델의 학습을 의미합니다. 데이터 크기는 소형 (small), 중형 (medium), 대형 (large)으로 분류됩니다. 소형 데이터 크기는 약 1GB, 중형 데이터 크기는 약 5GB, 대형 데이터 크기는 약 10GB입니다.

코드베이스는 다음과 같이 구성되어 있습니다:

train-llm-from-scratch/
├── src/
│ ├── models/
...

scripts/ 디렉토리는 데이터셋 다운로드, 데이터 전처리 (preprocessing), 모델 학습 (training), 그리고 학습된 모델을 사용한 텍스트 생성 (generating text)을 위한 스크립트들을 포함하고 있습니다. src/models/ 디렉토리는 트랜스포머 (transformer) 모델, 다층 퍼셉트론 (MLP), 어텐션 메커니즘 (attention mechanisms), 그리고 트랜스포머 블록 (transformer blocks)의 구현을 포함하고 있습니다. config/ 디렉토리는 기본 파라미터 (default parameters)가 포함된 설정 파일 (configuration file)을 포함하고 있습니다. data_loader/ 디렉토리는 데이터 로더/이터레이터 (data loaders/iterators)를 생성하기 위한 함수들을 포함하고 있습니다.

저장소를 클론 (Clone)하고 디렉토리로 이동하세요:

git clone https://github.com/FareedKhan-dev/train-llm-from-scratch.git
cd train-llm-from-scratch

임포트 (imports)와 관련된 문제가 발생하면, PYTHONPATH를 프로젝트의 루트 디렉토리로 변경해야 합니다:

export PYTHONPATH="${PYTHONPATH}:/path/to/train-llm-from-scratch"
# 또는 이미 train-llm-from-scratch 디렉토리에 있는 경우
export PYTHONPATH="$PYTHONPATH:."

필요한 의존성 (dependencies)을 설치하세요:

pip install -r requirements.txt

src/models/transformer.py에서 트랜스포머 아키텍처 (transformer architecture)를 수정할 수 있으며,

config/config.py에서 학습 설정 (training configurations)을 수정할 수 있습니다.

학습 데이터를 다운로드하려면 다음을 실행하세요:

python scripts/data_download.py

이 스크립트는 다음 인자 (arguments)를 지원합니다:

--train_max: 다운로드할 학습 파일의 최대 개수입니다. 기본값은 1이며 (최대 30), 각 파일은 약 11GB입니다.
--train_dir: 학습 데이터를 저장할 디렉토리입니다. 기본값은 data/train입니다.
--val_dir: 검증 (validation) 데이터를 저장할 디렉토리입니다. 기본값은 data/val입니다.

다운로드된 데이터를 전처리하려면 다음을 실행하세요:

python scripts/data_preprocess.py

이 스크립트는 다음 인자를 지원합니다:

--train_dir: 학습 데이터 파일이 저장된 디렉토리입니다 (기본값은 data/train입니다).
--val_dir: 검증 데이터 파일이 저장된 디렉토리입니다 (기본값은 data/val입니다).
--out_train_file

: 처리된 학습 데이터를 HDF5 형식으로 저장할 경로입니다 (기본값은 data/train/pile_train.h5입니다).
--out_val_file

: 처리된 검증 데이터를 HDF5 형식으로 저장할 경로입니다 (기본값은 data/val/pile_dev.h5입니다).
--tokenizer_name

: 데이터 처리에 사용할 토크나이저 (tokenizer)의 이름입니다 (기본값은 r50k_base입니다).
--max_data

: 각 데이터셋(학습 및 검증 모두)에서 처리할 JSON 객체(라인)의 최대 개수입니다. 기본값은 1000입니다.

이제 데이터 전처리가 완료되었으므로, config/config.py의 설정을 다음과 같이 변경하여 1,300만 파라미터의 LLM을 학습시킬 수 있습니다:

# 어휘 사전 크기 및 트랜스포머 (transformer) 설정 정의 (3 Billion)
VOCAB_SIZE = 50304 # 어휘 사전 내 고유 토큰 수
CONTEXT_LENGTH = 128 # 모델의 최대 시퀀스 길이
...

모델을 학습시키려면 다음을 실행하세요:

python scripts/train_transformer.py

모델 학습이 시작되며, 학습된 모델은 models/ 기본 디렉토리 또는 설정 파일에 지정된 디렉토리에 저장됩니다.

학습된 모델을 사용하여 텍스트를 생성하려면 다음을 실행하세요:

python scripts/generate_text.py --model_path models/your_model.pth --input_text hi

이 스크립트는 다음 인자 (arguments)를 지원합니다:

--model_path

: 학습된 모델의 경로입니다.
--input_text

: 새로운 텍스트 생성을 위한 초기 텍스트 프롬프트 (prompt)입니다.
--max_new_tokens

: 생성할 최대 토큰 수입니다 (기본값은 100입니다).

학습된 모델을 사용하여 입력 프롬프트를 기반으로 텍스트를 생성합니다.

이 섹션은 코드를 자세히 이해하고자 하는 분들을 위한 것입니다. 라이브러리 임포트 (importing)부터 모델 학습 및 텍스트 생성에 이르기까지 코드를 단계별로 설명하겠습니다.

이전에 Medium에 Tiny Shakespeare 데이터셋을 사용하여 230만 개 이상의 파라미터를 가진 LLM을 만드는 것에 대한 글을 썼지만, 출력 결과가 의미가 없었습니다. 다음은 샘플 출력입니다:

# 230만 파라미터 LLM 출력
ZELBETH:
Sey solmenter! tis tonguerered if
...

한 가지 생각이 들었습니다. 만약 Transformer (트랜스포머) 아키텍처를 더 작고 덜 복잡하게 만들고, 학습 데이터 (training data)를 더 다양하게 만든다면 어떨까요? 그렇다면 거의 수명이 다해가는 GPU를 사용하는 개인이, 올바른 문법을 구사하고 어느 정도 말이 되는 텍스트를 생성할 수 있는 파라미터 (parameter) 규모의 모델을 얼마나 크게 만들 수 있을까요?

저는 1,300만 개 이상의 파라미터를 가진 모델만으로도 올바른 문법과 구두점 측면에서 말이 되기 시작한다는 것을 발견했으며, 이는 긍정적인 부분입니다. 이는 우리가 이전에 학습시킨 모델을 특정 작업에 맞춰 미세 조정 (fine-tune)하기 위해 매우 특정한 데이터셋을 사용할 수 있음을 의미합니다. 결과적으로 10억 개 (1 billion) 미만, 혹은 심지어 5억 개 (500 million) 정도의 파라미터를 가진 모델을 얻을 수 있으며, 이는 특히 개인 데이터를 안전하게 처리하기 위한 특정 사용 사례에 완벽할 수 있습니다.

먼저 제 GitHub 저장소에 있는 스크립트를 사용하여 1,300만 개 이상의 파라미터를 가진 모델을 먼저 학습시키는 것을 권장합니다. 이렇게 하면 더 오래 기다릴 필요 없이 하루 안에 결과를 얻을 수 있으며, 로컬 GPU 성능이 10억 개 규모의 모델을 학습시키기에 충분하지 않을 경우에도 유용합니다.

이 블로그 전반에 걸쳐 사용될 필수 라이브러리들을 임포트 (import)해 보겠습니다:

# 딥러닝 함수와 텐서 (tensors)를 위한 PyTorch
import torch
import torch.nn as nn
...

우리의 학습 데이터셋 (training dataset)은 다양한 도메인의 정보를 포함하여 다양해야 하며, 이를 위해 The Pile이 적절한 선택입니다. 크기는 825GB에 달하지만, 우리는 그중 아주 작은 부분인 5%~10%만을 사용할 것입니다. 먼저 데이터셋을 다운로드하고 어떻게 작동하는지 살펴보겠습니다. 저는 HuggingFace에서 제공되는 버전을 다운로드하겠습니다.

# 검증 (validation) 데이터셋 다운로드
!wget https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main/val.jsonl.zst
# 학습 데이터셋의 첫 번째 부분 다운로드
...

다운로드하는 데 시간이 좀 걸리겠지만, 학습 데이터셋을 세 개 대신 단 하나의 파일인 00.jsonl.zst로 제한할 수도 있습니다. 이 데이터셋은 이미 학습/검증/테스트 (train/val/test)로 분할되어 있습니다. 다운로드가 완료되면 파일들을 각각의 해당 디렉토리에 올바르게 배치했는지 확인하세요.

import os
import shutil
import glob
...

위 코드의 출력 결과는 다음과 같습니다:

#### OUTPUT ####
Line: 0
{
...

이제 데이터셋을 인코딩 (토큰화 (tokenize))해야 합니다. 우리의 목표는 최소한 적절한 단어를 출력할 수 있는 LLM을 만드는 것입니다. 이를 위해 이미 사용 가능한 토크나이저 (tokenizer)를 사용해야 합니다. 우리는 OpenAI의 오픈 소스 토크나이저인 tiktoken을 사용할 것입니다. 데이터셋을 토큰화하기 위해 ChatGPT (GPT-3) 모델에 사용되는 r50k_base 토크나이저를 사용할 것입니다.

훈련 (train) 데이터셋과 검증 (validation) 데이터셋을 모두 토큰화할 것이므로, 중복을 피하기 위해 이를 위한 함수를 만들어야 합니다.

def process_files(input_dir, output_file):
"""
지정된 입력 디렉토리의 모든 .zst 파일을 처리하고 인코딩된 토큰을 HDF5 파일로 저장합니다.
...

이 함수와 관련하여 두 가지 중요한 사항이 있습니다:

우리는 토큰화된 데이터를 HDF5 파일에 저장하며, 이는 모델을 학습시키는 동안 더 빠른 데이터 접근을 위한 유연성을 제공합니다.

<|endoftext|>

토큰을 추가하는 것은 각 텍스트 시퀀스의 끝을 표시하며, 모델에게 의미 있는 문맥 (context)의 끝에 도달했음을 알려주어 일관된 출력을 생성하는 데 도움을 줍니다.

이제 다음과 같이 훈련 및 검증 데이터셋을 간단히 인코딩할 수 있습니다:

# 토큰화된 데이터 출력 디렉토리 정의
out_train_file = "data/train/pile_train.h5"
out_val_file = "data/val/pile_dev.h5"
...

토큰화된 데이터의 샘플을 살펴보겠습니다:

with h5py.File(out_val_file, 'r') as file:
# 'tokens' 데이터셋에 접근
tokens_dataset = file['tokens']
...

위 코드의 출력 결과는 다음과 같습니다:

#### OUTPUT ####
Dtype of 'tokens' dataset: int32
First few elements of the 'tokens' dataset:
...

이제 학습을 위한 데이터셋 준비가 완료되었습니다. 이제 트랜스포머 (transformer) 아키텍처를 코드로 구현하고 그에 따른 이론을 살펴보겠습니다.

트랜스포머 (transformer) 아키텍처가 텍스트를 처리하고 이해하는 데 어떻게 사용되는지 빠르게 살펴보겠습니다. 트랜스포머는 텍스트를 토큰 (tokens)이라고 불리는 더 작은 조각들로 나누고, 시퀀스 (sequence) 내의 다음 토큰을 예측하는 방식으로 작동합니다. 트랜스포머는 트랜스포머 블록 (transformer blocks)이라고 불리는 많은 레이어 (layers)가 서로 쌓여 있으며, 마지막에는 예측을 수행하기 위한 최종 레이어가 있습니다.

각 트랜스포머 블록은 두 가지 주요 구성 요소를 가집니다:

셀프 어텐션 헤드 (Self-Attention Heads): 모델이 집중해야 할 입력의 어떤 부분이 가장 중요한지를 파악합니다. 예를 들어, 문장을 처리할 때 어텐션 헤드는 대명사가 지칭하는 명사와의 관계와 같이 단어 사이의 관계를 강조할 수 있습니다.

MLP (Multi-Layer Perceptron, 다층 퍼셉트론): 이는 단순한 피드포워드 신경망 (feed-forward neural network)입니다. 어텐션 헤드에 의해 강조된 정보를 받아 추가로 처리합니다. MLP는 어텐션 헤드로부터 데이터를 받는 입력 레이어 (input layer), 처리에 복잡성을 더하는 은닉 레이어 (hidden layer), 그리고 결과를 다음 트랜스포머 블록으로 전달하는 출력 레이어 (output layer)로 구성됩니다.

함께 작용할 때, 어텐션 헤드는 "무엇에 대해 생각할 것인가" 하는 역할을 수행하고, MLP는 "그것에 대해 어떻게 생각할 것인가" 하는 역할을 수행합니다. 많은 트랜스포머 블록을 쌓으면 모델이 텍스트의 복잡한 패턴과 관계를 이해할 수 있게 되지만, 이것이 항상 보장되는 것은 아닙니다.

원문 논문의 다이어그램을 보는 대신, 우리가 코드로 구현할 더 단순하고 이해하기 쉬운 아키텍처 다이어그램을 시각화해 보겠습니다.

우리가 코드로 구현할 아키텍처의 흐름을 읽어봅시다:

입력 토큰 (Input tokens)은 임베딩 (embeddings)으로 변환되며 위치 정보 (position information)와 결합됩니다.

모델은 데이터를 순차적으로 처리하는 64개의 동일한 트랜스포머 블록을 가집니다.

각 블록은 먼저 토큰 간의 관계를 살펴보기 위해 멀티 헤드 어텐션 (multi-head attention)을 실행합니다.

그 다음 각 블록은 데이터를 확장한 후 다시 압축하는 MLP를 통해 데이터를 처리합니다.

각 단계는 정보의 흐름을 돕기 위해 잔차 연결 (residual connections, shortcuts)을 사용합니다.

학습을 안정화하기 위해 전 과정에서 레이어 정규화 (Layer normalization)가 사용됩니다.

어텐션 메커니즘 (Attention mechanism)은 어떤 토큰들이 서로에게 주의를 기울여야 하는지를 계산합니다.

MLP는 데이터를 4배 크기로 확장하고, ReLU를 적용한 다음, 다시 원래 크기로 압축합니다.

모델은 다양한 유형의 관계를 포착하기 위해 16개의 어텐션 헤드 (Attention heads)를 사용합니다.

최종 레이어는 처리된 데이터를 어휘 사전 크기 (Vocabulary-sized)의 예측값으로 변환합니다.

모델은 다음에 올 가능성이 가장 높은 토큰을 반복적으로 예측함으로써 텍스트를 생성합니다.

MLP는 트랜스포머 (Transformer)의 피드포워드 네트워크 (Feed-forward network) 내에서 핵심적인 구성 요소입니다. MLP의 역할은 비선형성 (Non-linearity)을 도입하고 임베딩된 표현 (Embedded representations) 내의 복잡한 관계를 학습하는 것입니다. MLP 모듈을 정의할 때 중요한 파라미터는 입력 임베딩의 차원 (Dimensionality)을 정의하는 n_embed입니다.