Tiny hackable CUDA language model implementation

Who Should Read

LLM의 내부 동작 원리(어텐션, 역전파, 옵티마이저 등)를 코드 레벨에서 직접 이해하고 싶은 ML 엔지니어나 CUDA 프로그래밍을 배우고 싶은 시스템 개발자.

Core Mechanics

• 이 프로젝트는 C + CUDA로 작성된 GPT 아키텍처의 미니멀 구현체로, 외부 딥러닝 프레임워크(PyTorch, TensorFlow 등) 없이 트랜스포머 모델의 학습과 추론을 처음부터 구현했다.
• 입력 단위로 텍스트의 단어나 서브워드가 아닌 바이트(8비트 토큰, 0~255)를 사용하기 때문에 별도의 토크나이저가 필요 없고, DNA/RNA 시퀀스, 압축 데이터, 이미지, 오디오, 실행 바이너리 등 어떤 바이트 스트림도 학습할 수 있다.
• 트랜스포머 레이어는 Causal Self-Attention(현재 위치보다 앞선 위치만 참조하는 어텐션)과 Feed-Forward Network로 구성되며, 각각 Residual Connection(입력을 그대로 더해주는 우회 경로)으로 감싸져 있다.
• 위치 인코딩으로 RoPE(Rotary Positional Encoding, 쿼리와 키에 회전 행렬을 곱해 상대적 위치를 인코딩하는 방식)를 사용한다. 절대 위치 임베딩 대신 상대 위치 관계를 더 잘 포착한다는 장점이 있다.
• 활성화 함수로 Swish(입력에 sigmoid를 곱하는 부드럽고 비단조적인 함수)를 사용하며, 이는 ReLU보다 그래디언트 흐름이 더 매끄럽다.
• 옵티마이저로 AdamW를 사용하는데, 기존 Adam에서 Weight Decay(가중치 감쇠, 과적합 방지를 위해 가중치 크기에 패널티를 주는 기법)를 그래디언트 업데이트에서 분리(decoupling)해 정규화 효과를 더 정확하게 적용한다.
• 출력 레이어에서는 256개 바이트 값에 대한 logit을 계산하고, softmax로 확률로 변환한 뒤 Cross-Entropy Loss로 학습한다. 즉, 매 스텝마다 '다음 바이트가 무엇일지'를 예측하는 방식이다.
• 파일 구조가 gpt.c, gpt.h, train.c, infer.c로 단순하게 분리되어 있어 코드베이스 전체를 빠르게 파악하고 직접 수정하기 좋다.

Evidence

• Backward pass(역전파) 구현의 정확성을 검증하는 Numerical Gradient Check(수치 미분으로 실제 그래디언트와 비교하는 테스트)가 없다는 지적이 있었다. 이 테스트가 없으면 역전파 코드에 버그가 있어도 학습 손실이 줄어드는 것처럼 보일 수 있어 디버깅이 어렵다는 우려다.
• LoRA(적은 파라미터만 학습하는 파인튜닝 기법) 지원을 추가할 계획이 있는지 묻는 댓글이 있었다. 구현이 너무 low-level이라 LoRA 같은 확장을 얹기가 쉽지 않을 수 있다는 뉘앙스도 함께 담겨 있었다.
• VRAM이 8GB인 RTX 3050(876MiB 사용 중)에서 `make run -j 10` 실행 시 `CUDA error in attention.c:91: out of memory`가 발생했다는 실패 사례가 공유됐다. CUDA 12.4 버전도 'partially supported' 경고가 떴다고 하며, 기본 설정으로는 소비자용 GPU 메모리가 부족할 수 있음을 시사한다.
• 데이터 준비 방법, 커스텀 데이터셋 생성, 토크나이저 분리, 컨텍스트 크기 설정, 메모리 요구량, 학습 중단 방법, LoRA 적용, 양자화(quantization) 방법 등 실용적인 사용 가이드가 전혀 없다는 불만 댓글이 있었다. Ruby 바인딩 추가도 제안됐다.
• ARM 아키텍처(예: Apple Silicon, Raspberry Pi 등)에서 동작하는지 묻는 댓글이 있었지만, CUDA 기반 구현이므로 NVIDIA GPU 없이는 실행이 불가능하다는 점이 암시된다.

How to Apply

• LLM 아키텍처를 프레임워크 추상화 없이 처음부터 이해하고 싶다면, gpt.c와 gpt.h를 읽으면서 어텐션 연산, RoPE 인코딩, AdamW 업데이트 루프가 CUDA 커널 레벨에서 어떻게 구현되는지 직접 추적해볼 수 있다.
• DNA 시퀀스나 바이너리 파일처럼 기존 텍스트 토크나이저를 쓰기 어려운 바이트 스트림 데이터에 언어 모델을 실험적으로 적용하고 싶다면, 이 구현체를 기반으로 별도의 전처리 없이 바로 학습 파이프라인을 구성할 수 있다.
• VRAM이 8GB 이하인 환경(RTX 3050 등)에서 실행하면 OOM이 발생할 수 있으므로, train.c나 Makefile에서 모델 크기(레이어 수, 헤드 수, 히든 차원)와 배치 크기, 컨텍스트 길이를 줄여서 실험해야 한다.
• 역전파 구현이 올바른지 확인하고 싶다면, 커뮤니티에서 지적된 대로 Numerical Gradient Check를 직접 추가해볼 수 있다. 특정 파라미터에 작은 perturbation(±ε)을 줬을 때 loss 변화량과 backward에서 계산된 그래디언트를 비교하면 버그를 잡을 수 있다.

Code Example

# 빌드 및 실행 (NVIDIA GPU + CUDA 필요)
git clone https://github.com/markusheimerl/gpt
cd gpt
make run -j 4   # -j 뒤 숫자는 병렬 작업 수 (8GB VRAM이면 낮게 설정)

# 추론만 실행
# make infer (별도 infer.c 참고)

Terminology

Related Papers

• Is One Layer Enough? A Single Transformer Layer Matches Full-Parameter RL Train

  LLM의 RL 후처리 학습(post-training)에서 성능 향상의 대부분이 중간 레이어 소수에 집중되며, 단 하나의 레이어만 학습해도 전체 파라미터 학습과 비슷하거나 더 나은 결과를 낼 수 있다는 연구 결과. 이는 RL 학습 비용을 대폭 줄일 수 있는 가능성을 시사한다.

• Knowledge Distillation of Black-Box Large Language Models (2024)

  GPT-4 같은 내부 구조에 접근할 수 없는 독점 LLM에서 작은 모델로 지식을 효과적으로 전달하는 Proxy-KD 기법을 소개하는 논문으로, 전통적인 White-Box 방식보다 성능이 높다는 점에서 주목할 만하다.

• Show HN: NanoEuler – GPT-2 scale model in pure C/CUDA from scratch

  PyTorch나 autograd 없이 C와 CUDA만으로 GPT-2 수준의 LLM을 처음부터 구현한 교육용 프로젝트로, 역전파·BPE 토크나이저·FlashAttention까지 직접 손으로 작성했다.

• Show HN: Neural Particle Automata

  고정된 격자 대신 움직이는 파티클 위에서 동작하는 Neural Cellular Automata의 확장 버전으로, 형태 생성·포인트 클라우드 분류·텍스처 합성 등 다양한 작업에서 자기조직화 동작을 학습할 수 있다.

• The annotated PyTorch training loop

  PyTorch 학습 루프의 각 코드 줄이 왜 그 위치에 있어야 하는지, 순서를 바꾸거나 빠뜨렸을 때 어떤 문제가 생기는지를 단계별로 설명한 심층 가이드다.

• When Good Verifiers Go Bad: Self-Improving VLMs Can Regress on New Tasks

  VLM 자가학습 루프에서 verifier가 특정 태스크에 맞지 않으면 학습할수록 오히려 성능이 떨어지는데, DPO 손실값은 멀쩡히 내려가서 눈치채기도 어렵다.

Related Resources

• markusheimerl/gpt - GitHub
• transformer 디렉토리 (attention.c 등 핵심 CUDA 커널)