Swift로 LLM 학습시키기, 파트 1: 행렬 곱셈을 Gflop/s에서 Tflop/s로 끌어올리기

이 글에서 저는 Swift로 대규모 언어 모델 (LLM)을 학습시키기 위해, 제가 직접 작성한 행렬 곱셈 (matrix multiplication) 코드를 가능한 한 빠르게 실행해 보려고 합니다. 목표는 Swift에서 수학 코드를 최적화하는 핵심 단계에 대한 통찰을 제공하는 것입니다. 또한 이러한 예제들이 Apple Silicon의 다양한 유닛인 CPU, SIMD, AMX 및 GPU의 성능 규모에 대한 감각을 제공하기를 바랍니다.

이 글은 Apple Silicon에서 Swift로 신경망 (neural networks)을 학습시키는 과정을 살펴보는 시리즈의 첫 번째 글입니다. 향후 게시물에서는 Apple이 Mac용 머신러닝 (machine learning)을 위해 제공하는, 어쩌면 너무 많을지도 모르는 프레임워크들을 살펴볼 예정입니다. 행렬 곱셈과 머신러닝을 위해서는 실제로 그 기성 프레임워크들을 사용해야 합니다 (그들은 저보다 몇 년 더 오래 행렬 커널 (matrix kernels)을 최적화하는 데 시간을 보냈으니까요).

하지만 그때까지는, "프레임워크도 라이브러리도 없는" 순수 코드 접근 방식으로 모든 것을 직접 작성하며 즐거움을 느껴보려 합니다.

단순히 행렬 곱셈 커널만 작성하는 것은 아닙니다. 샘플 앱은 전체 LLM 구현의 일부로 이 커널들을 사용할 것이며, 제가 인용할 수치들은 전체 순전파 (forward) 및 역전파 (backward) 학습 반복에 대한 것입니다. 이 시리즈의 참조 구현체는 Andrej Karpathy의 llm.c (GPT2 호환 모델의 순수 C 구현체)가 될 것입니다. 이는 상당히 기본적인 모델이지만, 필요한 모든 구성 요소를 포함하고 있으며 실제 워크로드 (workloads)를 대표합니다.

즉, 제가 가장 좋아하는 게임을 할 시간이라는 뜻입니다: Swift가 C보다 빨라질 때까지 최적화하기.

배경 이야기

약 2년 전, 저는 2000년대 초반에 작성했던 공학 학위 논문을 찾아냈습니다. 그것은 이미지를 분류하기 위해 신경망을 사용하는 C++로 작성된 이미지 인식기였습니다. 예전 코드를 다시 실행해 보고 싶었지만, 오랫동안 ML 코드를 다루지 않았습니다. 짜증이 났고 결국 포기하고 말았습니다.

2024년 초 LLM(Large Language Models)에 대한 수많은 논의가 있었음에도 불구하고, Mac에서 신경망 (Neural Networks)을 *학습(training)*시키는 사람은 아무도 없는 것처럼 느껴졌습니다. 적어도 Swift와 같은 언어를 사용하는 경우는 말이죠. 저는 PyTorch나 TensorFlow 같은 Python 라이브러리들을 다뤄보았지만, Python은 결코 스스로 계산을 수행하지 않습니다. Python은 내부적으로 다른 연산 엔진을 조율하는 오케스트레이터 (Orchestrator)처럼 작동하며, 이러한 분리 구조는 제가 직접 제어하고 있지 않다는 느낌을 주었습니다.

한 달 후, Andrej Karpathy가 llm.c를 공개했습니다. 이 프로젝트는 다른 머신러닝 (Machine Learning) 콘텐츠와는 다른 방식으로 저에게 다가왔는데, 그 이유는 아무것도 숨겨져 있지 않았기 때문입니다. 약 1,000줄의 순수한 C 언어로 작성되어 있으며, (비록 꽤 난해한 변수명들로 가득 차 있긴 하지만) 비교적 읽기 쉽습니다.

그래서 자연스럽게, 저는 즉시 이를 Swift로 다시 작성했습니다. 그리고 그것을 가지고 노는 것은 매우 즐거운 일이었습니다.

물론, 코드를 가지고 놀기 위해서는 빠르게 실행되도록 만드는 작업이 필요했습니다. 미리 암시를 드리자면, 초기 Swift 구현체는 정말로 매우 느렸습니다. 하지만 최적화 (Optimization)는 끊임없는 과정입니다. 항상 시도해 볼 수 있는 무언가가 더 있기 마련이니까요.

이제 드디어 이 글로 이어집니다. 저는 라이브러리를 사용하지 않고 LLM을 상당히 빠르게 학습시키기 위해, 당시 제가 작성했던 다양한 탐구 과정(그리고 지난 한 주 동안 추가한 몇 가지)을 살펴볼 것입니다. 대부분의 코드는 Swift로 작성될 것입니다 (마지막에 Metal 구현을 보여드리긴 하겠지만요).

참고로, 저는 신경망 (Neural Network)이나 LLM이 어떻게 작동하는지는 설명하지 않을 것입니다. 만약 관심이 있다면, Karpathy의 영상인 "Let’s build GPT: from scratch, in code, spelled out."는 GPT와 같은 LLM이 어떻게 작동하는지 배우기 위한 사실상의 결정적인 가이드이며, 더 기초적인 강의를 원하신다면 "The spelled-out intro to language modeling: building makemore"로 시작하는 그의 이전 시리즈가 5개의 영상 시리즈를 통해 많은 입문 개념을 다루고 있습니다. 물론 두 영상 모두 Python으로 작성되어 있으므로, Swift로 어떻게 구현할 수 있는지 볼 준비가 되었을 때 다시 이곳으로 돌아와 주시기 바랍니다.

llm.c

머신러닝 (Machine learning)은 본질적으로 입력 데이터에 모델 가중치 (model weights)를 적용하고 (이를 순전파 (forward pass) 또는 추론 (inference)이라고 부릅니다), 그 다음 오차 기울기 (error gradients)를 계산하여 해당 가중치를 업데이트하는 것 (역전파 (backward pass))입니다.

우리는 보통 이러한 계산들을 하나로 묶어 가능한 한 빠르게 실행되도록 만듭니다. 이러한 연산 패키지는 "선형 텐서 투영 (linear tensor projection)", "행렬 곱셈 (matrix multiplication)", 또는 작업 단위를 얼마나 크거나 작게 나누느냐에 따라 일련의 "벡터 내적 (vector dot products)"이라고 불릴 수도 있습니다. 궁극적으로 이는 z += x * y를 아주 많이 반복하는 루프입니다.

이러한 행렬 곱셈 (matrix multiplications)이 머신러닝 작업의 상당 부분을 차지하기 때문에, 저는 이 작업을 수행하는 코드에 집중할 것입니다. 진행하면서 나머지 구현 부분도 업데이트하겠지만, 행렬 곱셈에 보여줄 것과 동일한 개선 사항들만 적용할 것입니다.

먼저 순전파 (forward pass)에서 사용되는 핵심 행렬 곱셈인 llm.c의 matmul_forward를 살펴보겠습니다. 이 함수는 입력 (inp)에 대해 반복하며, 모델 가중치 (weight)를 곱하고, 그 결과를 누적 합계 (val)에 더합니다.

void matmul_forward(float* out, const float* inp, const float* weight, const float* bias, int B, int T, int C, int OC) {
for (int b = 0; b < B; b++) {
for (int t = 0; t < T; t++) {
...

4중 루프 구조가 시각적인 복잡함을 더하지만, 실제로 val += inp[bt * C + i] * weight[o*C + i]; 라인이 신경망 (neural network)의 심장입니다. 앞서 말했듯이, z += x * y를 아주 많이 수행하는 것이죠.

얼마나 많이 할까요?

val 라인은 2개의 부동 소수점 연산 (floating point operations)을 포함하지만, Karpathy는 전체 학습 반복 (training iteration)에서의 부동 소수점 연산 횟수가 대략 6 x N x D가 되어야 한다고 말합니다. 여기서 N은 모델의 가중치 수(우리의 경우 124,439,808개)이고, D는 우리 앱의 경우 B * T = 4 * 64 = 256입니다. 따라서 우리는 6 x 124,439,808 x 256 ≈ 1.911×10¹¹ ≈ 0.2조(trillion) 번의 학습 반복당 부동 소수점 연산에 대해 이야기하고 있는 것입니다.

그러므로 매우 빠르게 실행되어야만 합니다.

모델	Tokens/s	Training iterations/s
llm.c	0.92	0.174

순수한 C 코드는 Swift Package에서 쉽게 실행됩니다. 저는 C 구현체가 (Xcode 설정과 관계없이) 항상 -O3 최적화 레벨에서 실행되도록 수정했습니다.

이 최적화 레벨에서도 C 구현체는 7초마다 단 한 번의 학습 반복 (training iteration)을 수행하며, 추론 (inference)은 초당 1 토큰 미만으로 처리합니다. 훌륭한 개념 증명 (proof of concept) 이지만, 실제로 유용하게 사용될 수 있는 속도보다 10배는 느립니다.

기본 Swift (Basic Swift)

저는 기본 Swift 버전을 최대한 C 버전과 동일하게 유지하려고 최선을 다했습니다:

static func matmul_forward(out: inout [Float], inp: [Float], weight: [Float], bias: [Float]?, B: Int, T: Int, C: Int, OC: Int) {
for b in 0..<B {
for t in 0..<T {
...

C 코드는 본질적으로 "unsafe"하기 때문에, 저도 Swift 코드에 -remove-runtime-asserts를 설정하여 동일한 이점을 부여했습니다 (배열 인덱스에 대한 런타임 검사 제거). 그리고 항상 앱이 "Release" 구성으로 실행되도록 했습니다. 그렇다면 Swift와 C 구현체는 상당히 대등해야 하지 않을까요?

Debug 모드에서 실행하지 마세요. 저는 오직 Release 구성의 수치만을 인용할 것입니다. Debug 모드에서 이 작업의 많은 부분을 실행해 보긴 했지만, Debug 모드에서 20회의 전체 학습 반복이 완료될 때까지 기다려 본 적은 없습니다. 저는 보통 디버깅 중에도 Xcode의 Scheme을 "Release"로 설정해 둡니다.

배경 설명을 읽으셨다면 이미 언급했듯이, 이것은 "극도로 느렸습니다".

모델	Tokens/s	Training iterations/s	Training versus llm.c
llm.c	0.926	0.175	100%
Basic Swift	0.054	0.014	7.3%

Swift 코드는 15배에서 20배 정도 더 느립니다. 이는 LLM이 19초마다 1 토큰을 생성한다는 의미입니다. 이 엔진으로 20회의 학습 반복을 실행하는 데 거의 30분이 걸립니다. 도대체 무슨 일이 일어나고 있는 걸까요?

이 성능은 약 2.8 Gflop/s를 나타냅니다. 1999년에 Apple은 PowerMac G4의 1 Gflop/s 성능이 미국 군대의 관점에서도 강력한 무기가 될 것이라고 주장하며 광고를 내보냈습니다. 이제 2.8 Gflop/s는 완전히 받아들일 수 없는 수준입니다.

Span, Egg and Span

Instruments로 조사해 본 결과, 이전 실행에서 압도적으로 가장 큰 성능 비용을 차지하는 것은 _ArrayBuffer.beginCOWMutation()입니다.

Swift는 다른 누군가가 우리의 Array를 사용하고 있을지도 모른다고 생각하며, 비록 배열들이 고유하여 (따라서 배열 복사가 발생하지 않음에도 불구하고) 고유성 검사(uniqueness checks) 그 자체만으로도 가장 큰 오버헤드(overhead)가 되고 있습니다.

허? 때로는 단순히 버그일 수도 있는 문제에 직면하곤 합니다. 이것이 그중 하나일지도 모릅니다. 제가 2024년에 처음 이 코드를 작업했을 때의 기억으로는 이것이 문제가 되지 않았습니다. 회귀(regression)가 발생했는지, 아니면 안전 구멍(safety hole)이 메워지면서 _ArrayBuffer.beginCOWMutation()이 성능을 저하시키게 된 것인지는 모르겠습니다. 이 문제는 @inline(none)을 사용하여 함수 인라이닝(function inlining)을 비활성화할 때도 나타나므로, 옵티마이저(optimizer)가 제 역할을 제대로 수행하지 못하고 있는 것처럼 느껴집니다.

어쨌든, 우리는 Array를 사용하면서 필요한 성능을 얻을 수 없습니다. 다행히 Swift 6.2는 기본적으로 오버헤드가 거의 없는 신뢰할 수 있는 해결책인 MutableSpan을 제공합니다.

static func matmul_forward(out: inout [Float], inp: [Float], weight: [Float], bias: [Float]?, B: Int, T: Int, C: Int, OC: Int) {
var out = out.mutableSpan
for b in 0..<B {
...

제가 추가한 것은 맨 윗줄에 var out = out.mutableSpan을 추가하여, out을 자체적인 mutableSpan으로 섀도잉(shadowing)한 것뿐입니다. 파일 전체에 걸쳐 동일한 패턴을 적용했습니다.

모델	Tokens/s	Training iterations/s	Training versus llm.c
llm.c	0.926	0.175	100%
...

흥미롭게도, 이 변경 사항은 순전파(forward pass)에는 큰 영향을 미치지 않았지만, 학습 반복(training iterations, 순전파 + 역전파 + 업데이트) 속도를 3배 이상 빠르게 만들었습니다.

Chill vibes

하지만 우리는 왜 순전파 (forward pass)가 느린지에 대해 주의를 기울여야 합니다. Instruments는 우리가 이미 알고 있는 사실을 확인해 줍니다. 순전파에서 가장 뜨거운(가장 많은 연산이 발생하는) 라인은 우리 루프의 중심부인 value += inp[bt * C + i] * weight[o * C + i]입니다.

이제 냉혹한 진실을 마주할 시간입니다. C에는 Swift에는 없는 몇 가지 컴파일러 최적화 플래그 (compiler optimization flags)가 있습니다. 이 구체적인 사례에서 C는 -ffast-math를 가지고 있는데, 이는 C가 단일 명령어로 부동 소수점 곱셈과 덧셈을 수행하는 결합 곱셈-덧셈 (fused-multiply-addition, FMA) 명령을 사용할 수 있게 해주며, 일반적으로 정밀한 정확도에 대해 크게 걱정하지 않습니다.

+0xa34 fmadd s0, s17, s16, s0
+0xa38 ldr s17, [x20, #0x4]
+0xa3c fmadd s7, s17, s16, s7
...

C의 내부 루프는 대부분 fmadd (결합 곱셈-덧셈 명령어)를 8번 언롤링 (unrolled)하여 적용한 것뿐입니다.

재미있는 사실 하나: 우리는 미래에 살고 있습니다. 따라서 어셈블리 언어 (assembly language) 같은 것이 이해되지 않는다면, 즐겨 사용하는 LLM에 넣으면 번역해 줄 것입니다.

Swift에는 -ffast-math가 없기 때문에, 대신 별도의 곱셈과 덧셈을 수행하게 됩니다.

+0x164 fmul.4s v1, v1, v5
+0x168 mov s5, v1[3]
+0x16c mov s17, v1[2]
...

Swift는 4배 루프 언롤링 (loop unroll)을 시도하고 있습니다. 저 fmul.4s는 4개의 곱셈을 수행하기 위한 SIMD 연산입니다. 하지만 저 모든 mov 명령어들과 마지막의 별도 덧셈 연산들이 우리의 성능을 저하시키고 있습니다.

우리는 C가 사용하는 것처럼 결합 곱셈-덧셈 (fused-multiply-add)을 사용해야 합니다.

다행히도 우리에게는 우리만의 fast math 버전을 제공하는 Relaxed가 포함된 Swift-Numerics가 있습니다. 은하수를 여행하는 히치하이커를 위한 안내서에 나오는 "당황하지 마시오 (Don't panic)"처럼, 저도 "Relax"라는 단어가 우리 모두를 더 차분하게 만들고 느긋한 분위기 (chill vibes)를 즐기게 하기 위해 여기 있는 것이라고 확신합니다. 또는 결과의 반올림 (rounding)에 관한 규칙을 완화 (relax)하여 FMA가 가능하도록 해주는 것일 수도 있습니다.

우리의 모든 a += b * c 및 x = y + z 연산은 Relaxed.multiplyAdd와 Relaxed.sum을 통해 개선될 수 있습니다. 하지만 저는 gelu_backward 함수에는 이 함수들을 적용하는 것을 명시적으로 피할 것입니다. 왜냐하면 C 구현체는 해당 함수 주변에서 -ffast-math를 명시적으로 비활성화하기 때문입니다.

static func matmul_forward(out: inout [Float], inp: [Float], weight: [Float], bias: [Float]?, B: Int, T: Int, C: Int, OC: Int) {
var out = out.mutableSpan
for b in 0..<B {
...

중간에 있는 단 한 가지 변경 사항은 다음과 같습니다: val = Relaxed.multiplyAdd(inp[bt * C + i], weight[o * C + i], val)

읽기에 아주 깔끔하지는 않지만, 성능을 확인해 봅시다.

이제 우리의 어셈블리(Assembly)는 다음과 같이 보입니다:

+0x178 fmla.4s v1, v16, v4
+0x17c fmla.4s v0, v17, v5
+0x180 fmla.4s v2, v18, v6
...

흥미롭게도, Swift는 fmadd의 SIMD 벡터화(Vectorized) 버전인 fmla를 선택하고 있지만, 전제 조건은 동일합니다.

모델	Tokens/s	Training iterations/s	Training versus llm.c
llm.c	0.926	0.175	100%
...

이는 초당 토큰(Tokens per second) 수에서 거의 10배의 속도 향상을 의미합니다. 하지만 우리의 학습 성능은 여전히 C보다 15% 느립니다. 어떻게 이 마지막 격차를 줄일 수 있을까요?

They see me rollin'

한 가지 인정할 것이 있습니다. 저는 지금까지 C의 "나이브(Naive)"한 행렬 곱셈 버전을 보여드려 왔습니다. 실제 C 함수는 컴파일러가 가장 안쪽 루프를 언롤링(Unrolling)하기를 기대하며, 외부 루프를 한 번에 8단계씩 건너뛰는 방식으로 조금 더 복잡하게 구현되어 있습니다. 그리고 C의 -O3 옵션은 8배 루프 언롤링(Loop unrolling)을 제공함으로써 이에 부응합니다.