단일 GPU로 트랜스포머 LLM을 처음부터 학습하는 오픈소스 레포지토리 공개

개발자 FareedKhan-dev가 트랜스포머 기반 대규모언어모델(LLM)을 처음부터 직접 학습할 수 있는 오픈소스 코드 'train-llm-from-scratch'를 깃허브에 공개했다. 이 프로젝트는 머신러닝 프레임워크 파이토치(PyTorch)로 작성됐으며, 저자는 자신의 스크립트를 이용하면 단일 GPU만으로 수백만에서 수십억 파라미터 규모의 LLM을 직접 학습할 수 있다고 설명했다.

모델 구조는 논문 'Attention is All You Need'에 기반한 트랜스포머를 처음부터 구현한 것이다. 저자는 이 코드로 학습한 1,300만 파라미터 모델이 생성한 텍스트 예시도 함께 공개했다.

학습 데이터로는 The Pile 데이터셋이 사용된다. The Pile은 책, 기사, 웹사이트 등에서 모은 22개의 다양한 데이터셋을 묶은 오픈소스 대규모 데이터셋으로, 전체 용량은 825GB에 달한다.

모델 학습에는 GPU가 필요하다. 저자는 1,300만 파라미터 이상 모델은 Colab이나 Kaggle의 T4 GPU로도 학습할 수 있지만, 수십억 파라미터 모델 학습에는 이들 GPU로는 부족하다고 안내했다. 함께 제공된 비교표에 따르면 40GB 메모리의 NVIDIA A100은 20억 파라미터 모델 학습이 가능하며 실질적으로 약 60억~80억 파라미터까지 다룰 수 있다. 24GB의 RTX 4090도 20억 파라미터 학습이 가능하고 최대 약 40억 파라미터까지 가능하다. 반면 16GB의 V100은 20억 파라미터 학습은 어렵고 1,300만 파라미터급 학습과 최대 약 20억 파라미터 수준에 머문다.

데이터 규모는 소·중·대로 구분된다. 작은 데이터는 약 1GB, 중간은 약 5GB, 큰 데이터는 약 10GB 수준이다.

코드베이스는 기능별 디렉터리로 정리돼 있다. src/models 디렉터리에는 다층 퍼셉트론(MLP), 어텐션 메커니즘, 트랜스포머 블록, 트랜스포머 모델 본체 구현이 들어 있다. 이 밖에 기본 설정값을 담은 config, 데이터 로더를 담은 data_loader, 데이터 다운로드·전처리·학습·텍스트 생성 스크립트를 모은 scripts 디렉터리로 구성된다.

사용 절차는 데이터 다운로드부터 시작한다. 다운로드 스크립트는 기본적으로 학습 파일 1개를 받으며 최대 30개까지 지정할 수 있고, 파일 하나의 크기는 약 11GB다. 이어 전처리 스크립트가 데이터를 HDF5 형식으로 변환하는데, 기본 토크나이저는 r50k_base이고 각 데이터셋에서 처리하는 JSON 객체 수는 기본 1,000개로 설정돼 있다.

1,300만 파라미터 모델을 학습하려면 설정 파일에서 어휘 크기 50,304, 컨텍스트 길이 128, 임베딩 차원 128, 어텐션 헤드 8개, 트랜스포머 블록 1개로 값을 지정하면 된다. 학습이 끝난 모델은 별도 디렉터리에 저장되며, 텍스트 생성 스크립트에 입력 프롬프트를 주면 새 텍스트를 만들어 낸다. 생성 토큰 수는 기본 100개다.

저자는 코드를 이해하려면 객체지향 프로그래밍(OOP), 신경망(NN), 파이토치에 대한 기본 지식이 필요하다며 관련 학습 자료도 함께 안내했다.