microgpt를 Futhark로 포팅하기, 제1부

저는 데이터 병렬 언어인 Futhark를 시험해 볼 프로젝트를 찾고 싶었습니다. Futhark 팀은 제가 수년 동안 팔로우해 온 매우 좋은 블로그를 운영하고 있지만, 저는 실제로 그것을 사용하여 무언가를 작성해 본 적이 없습니다.

Andrej Karpathy의 microgpt는 200줄의 Python 코드로 구현된 GPT-2와 유사한 신경망의 독립적인 구현체로, 마침내 제가 시도해 볼 구실을 제공했습니다. 저는 microgpt를 좋아하지만, 이것은 전혀 확장성 (scaling)이 없습니다. 분명히 이 구현의 목적은 효율성이 아니지만, 단순히 느리기만 한 것이 아닙니다. Python의 재귀 깊이 (recursion depth) 오류에 빠르게 도달하기 때문에 약간 더 큰 네트워크로 확장하는 것조차 불가능합니다.

그래서 저는 최대한 1대1로 포팅하면서도 간결함을 크게 잃지 않고 훨씬 더 나은 확장성을 얻을 수 있을지 궁금했습니다. 결과적으로 대답은 '어느 정도는 그렇다'였습니다. 포팅된 버전은 확장성이 훨씬 뛰어나지만, 이전만큼 간결하지는 않습니다. 하지만 일부 부분은 꽤 매끄럽게 번역됩니다.

이 포스트인 제1부에서는 순전파 (forward pass)부터 시작하겠습니다. Karpathy의 원래 Python 버전 코드와 저의 Futhark 번역 코드를 번갈아 가며 제시하며, Futhark의 병렬 프리미티브 (parallel primitives)를 활용하면서도 최대한 유사하게 유지하려고 노력할 것입니다.

먼저, LLM 파라미터(가중치, weights)를 담는 데이터 구조입니다. 우리는 이것들이 사전 학습되었다고 가정하겠습니다 (학습 코드는 제2부에서 다룹니다).

Python:

n_layer = 1
n_embd = 16
block_size = 16
...

Futhark:

def n_layer : i64 = 1
def n_embd : i64 = 16
def block_size : i64 = 16
...

다음은 Karpathy가 선택한 특정 모델 아키텍처의 기본 구성 요소입니다.

Python:

def linear(x, w):
return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]
def softmax(logits):
...

Futhark:

def linear [n][m] (x: [n]f32) (w: [m][n]f32) : [m]f32 =
map (\w_row -> reduce (+) 0f32 (map2 (*) w_row x)) w
def softmax [n] (logits: [n]f32) : [n]f32 =
...

이 세 가지 함수가 얼마나 깔끔하게 번역되는지 매우 만족스럽습니다. 명시적 타이핑 (explicit typing) 때문에 구문 노이즈 (syntax noise)가 약간 추가되기는 합니다. 논쟁의 여지는 있겠지만, reduce (+) 0f32 또한 sum을 표현하는 아주 깔끔한 방식은 아닙니다. 하지만 이러한 종류의 함수 조합자 (functional combinators)에 이미 익숙하다면 일반적으로 읽기 쉽습니다. 코드 라인 수는 정확히 동일하게 유지되었습니다.

(업데이트 2026-05-02: Philip Munksgaard는 reduce (+) 0f32 대신 그냥 f32.sum이라고 쓸 수도 있었다고 지적했습니다. 제 머릿속에 APL이 들어있었나 봅니다. 최종 버전에서는 그 부분을 단순화하겠습니다.)

그리고 마지막으로, KV 캐시 (KV cache)를 포함한 GPT 순전파 (forward pass)입니다.

Karpathy의 원본 Python 코드:

def gpt(token_id, pos_id, keys, values):
tok_emb = state_dict['wte'][token_id] # token embedding
pos_emb = state_dict['wpe'][pos_id] # position embedding
...

저의 Futhark 포팅 버전:

def gpt [v]
(p: params [v])
(token_id: i64) (pos_id: i64)
...

for 루프는 대부분 tabulate로 번역되는데, 이는 본질적으로 병렬 for 루프라고 생각하면 됩니다. 예외는 레이어 (layers)에 대한 외부 루프로, 이는 순차적 (sequential)으로 유지되어야 하므로 loop가 됩니다. 마지막의 MLP 블록 또한 꽤 직접적으로 번역됩니다. 어텐션 (attention) 블록은 일부 명령형 (imperative) Python 기능과 파괴적 업데이트 (destructively updated) 데이터 구조를 사용했기 때문에 Futhark의 제약 조건에 맞춰 조정하는 것이 다소 까다로웠습니다. 하지만 아주 나쁘지는 않았습니다. 주요 변경 사항은 KV 캐시를 고정 크기 배열 (크기 [n_layer][block_size][n_embd])로 사전 할당 (preallocate)한 것이며, 그 후 attn_logits 계산 시 모델의 인과성 (causal)을 유지하기 위해 "미래" 토큰들을 마스킹 (mask out)하는 것이었습니다 (Python 버전에서는 리스트가 인과적 순서로 생성되었기 때문에 이 과정이 필요하지 않았습니다).

이 함수에 대한 총 코드 라인 수(주석 및 빈 줄 제외)는 33줄에서 51줄로 늘어났지만, 이는 부분적으로 제가 원본보다 문장을 더 자유롭게 여러 줄로 나누어 작성했기 때문입니다.

저 또한 이러한 스타일의 함수형 프로그래밍 (functional programming)을 좋아함에도 불구하고, 최종 결과물은 확실히 가독성이 떨어질 수 있다는 점을 인정합니다. 특히 깊은 중첩 (deep nesting) 구조는 따라가기가 다소 어렵습니다. flatten 안에 있는 tabulate 안에 있는 또 다른 tabulate 안에 있는 reduce 안에 있는 tabulate 안에 있는 람다 (lambda) 함수 내부로 들어가게 되면, 현재 무엇이 일어나고 있는지 파악하기가 쉬워질 수 있습니다. 아마도 가독성을 높이기 위해 리팩토링 (refactoring)할 수 있겠지만, 지금은 가능한 한 원문에 가깝게 번역하는 것에 집중했습니다. 따라서 Karpathy가 중첩 루프 (nested loops) 안에 다수의 리스트 컴프리헨션 (list comprehensions)을 사용했기 때문에, 저도 동일한 구조를 유지했습니다.

또한 크기 타이핑 (size-typing) 시스템을 만족시키기 위한 사소한 번거로움이 하나 있습니다. Futhark는 flatten의 결과가 크기 n_head * head_dim인 1차원 (1d) 배열이라고 추론하지만 (크기가 *[n_head][head_dim]*인 배열을 flatten 했기 때문), 이것이 크기 n_embd인 1차원 (1d) 배열과 동일하다는 것까지는 추가로 추론하지 못합니다. 따라서 우리는 :> 크기 강제 변환 (size coercion) 연산자를 사용해야 합니다. 반면에 사소한 가독성 개선 사항도 있습니다. 두 리스트에 대해 2개 매개변수 함수를 요소별로 매핑하는 map2는 제 생각에 zip을 사용하는 버전보다 더 직관적입니다.

물론 여기서 빠진 것은 이번 작업의 주인공인 모델 학습 (training)입니다! 그 내용은 벤치마크 (benchmarks)와 함께 제2부에서 다루도록 하겠습니다.