RustGPT: Rust로 처음부터 완전히 구현한 순수 트랜스포머 LLM

1 month ago 20

  • RustGPT는 외부 머신러닝 프레임워크 없이, 순수 Rust와 ndarray만으로 구현된 트랜스포머 기반 언어 모델
  • 사전 학습(Pre-training)지시 튜닝(Instruction tuning) 을 통해 사실 기반 지식과 대화형 패턴을 학습하도록 설계됨
  • 구조는 토크나이저 → 임베딩 → 트랜스포머 블록 → 출력 프로젝션으로 이어지는 전형적인 LLM 아키텍처를 따름
  • 모듈화된 소스 구조테스트 코드를 제공하여 학습, 추론, 최적화 과정을 세부적으로 이해할 수 있음
  • 러스트 생태계에서 프레임워크 의존 없이 LLM을 처음부터 구현해보고 싶은 개발자나 학습자에게 중요한 참고 자료

프로젝트 개요

  • RustGPT는 외부 머신러닝 프레임워크나 복잡한 의존성 없이 순수 Rust 언어와 선형대수 연산 라이브러리(ndarray) 만으로 LLM을 구현한 오픈소스 프로젝트
  • 주요 목표는 현대 LLM의 핵심 구성 요소(트랜스포머, 어텐션, 임베딩, 최적화 등)를 직접 구현하고 학습 과정을 이해하는 것
  • 다른 주류 LLM들과 달리, 직접 트랜스포머 구조와 백프로파게이션, 토크나이저, 옵티마이저 등을 모두 Rust 코드로 설계해, Rust 개발자와 연구자들이 딥러닝의 원리를 처음부터 직접 이해하고 확장할 수 있는 것이 큰 장점
  • ndarray로 행렬 연산을 처리하고, PyTorch나 TensorFlow 같은 외부 머신러닝 패키지에 의존하지 않는다는 것이 차별점
  • 모듈화와 테스트 커버리지가 탄탄해, 다양한 실험과 개선에 적합하고 "처음부터 직접 만드는 LLM(From Scratch)" 교육용 목적에도 적합함

주요 특징 및 구현 방식

  • 트랜스포머 아키텍처: 입력 텍스트 → 토크나이즈 → 임베딩 → 트랜스포머 블록 → 최종 예측
    • 입력 텍스트는 토큰화 과정을 거쳐 임베딩 벡터로 변환됨
    • 임베딩은 Transformer Block(멀티헤드 어텐션 + 피드포워드 네트워크)을 통과함
    • 마지막으로 Output Projection Layer에서 어휘 확률 분포를 생성하여 예측을 수행함

구현 구조

  • main.rs: 학습 파이프라인, 데이터 준비, 인터랙티브 모드 실행
  • llm.rs: LLM 전반의 순전파·역전파 및 학습 로직
  • transformer.rs, self_attention.rs, feed_forward.rs: 핵심 트랜스포머 블록
  • embeddings.rs, output_projection.rs: 임베딩 및 최종 출력 계층
  • adam.rs: Adam 옵티마이저 구현
  • 각 모듈은 대응되는 테스트 코드(tests/)를 포함하여 기능 검증이 가능함

학습·테스트 방법 및 데이터 흐름

  • 학습 과정
    • 어휘집 생성 → 사전학습(100epoch, 사실 문장 데이터) → Instruction 튜닝(100epoch, 대화 데이터)
    • 사전학습 예시: "The sun rises in the east and sets in the west"
    • Instruction 튜닝 예시: "User: How do mountains form? Assistant: ..."
  • 인터랙티브 모드 지원
    • 학습 완료 후, 프롬프트-응답 기반 대화 테스트 가능
    • 예시: "How do mountains form?" → "Mountains are formed through tectonic forces or volcanism..."

기술적 세부 구성

  • 어휘크기: 트레이닝 데이터 기반 동적 설정
  • 임베딩 차원: 128, 히든 레이어: 256
  • 최대 시퀀스 길이: 80 토큰
  • 아키텍처: 3개 트랜스포머 블록 + 임베딩 + 출력 계층
  • 학습 알고리듬: Adam 옵티마이저, gradient clipping (L2 norm 5.0 제한)
  • 학습률: pre-training 0.0005, instruction tuning 0.0001
  • 손실 함수: cross-entropy loss

모델 및 코드 특징

  • 맞춤형 토크나이저 (구두점 처리)
  • 그리디 디코딩 기반 텍스트 생성
  • 모듈형 계층 구조와 명료한 인터페이스
  • 테스트 커버리지: 각 계층·기능 단위 테스트 코드 탑재
  • 의존성: ndarray(행렬연산), rand/ rand_distr(난수 초기화)만 사용 (PyTorch/TensorFlow 등 외부 ML 미사용)
  • 교육적 가치: 대표적인 현대 LLM 내부 구조/훈련 원리 학습에 최적

발전 가능성

  • 고급 아키텍처 도입: 다중 헤드 어텐션, RoPE, 위치 인코딩 등
  • 성능 최적화: SIMD, 병렬 학습, 메모리 효율 개선
  • 모델 저장/로드 지원
  • 향상된 샘플링(빔서치, Top-k/Top-p) 및 평가 지표 추가

의의

  • Python 기반 PyTorch, TensorFlow 프레임워크에 의존하지 않고도 Rust만으로 LLM을 직접 구현할 수 있음을 보여주는 학습용·실험용 프로젝트
  • LLM의 내부 작동 원리를 이해하고 Rust 환경에서 ML 시스템을 만들고자 하는 개발자에게 유용한 레퍼런스

Read Entire Article