Training an LLM in Swift, Part 1: Taking matrix mult from Gflop/s to Tflop/s

개요

Swift를 사용하여 LLM 훈련을 위한 행렬 곱셈 연산을 Gflop/s에서 Tflop/s로 최적화하는 방법을 탐구하며, Apple Silicon의 CPU, SIMD, AMX, GPU 성능을 비교합니다.

주요 내용

• 초기 Swift 구현의 낮은 성능: Andrej Karpathy의 llm.c를 Swift로 재작성한 초기 버전은 C 구현보다 15-20배 느린 성능을 보였습니다. 이는 _ArrayBuffer.beginCOWMutation()으로 인한 오버헤드 때문이었습니다.
• MutableSpan 도입: MutableSpan을 사용하여 배열 복사 및 고유성 검사 오버헤드를 줄여 훈련 반복 속도를 3배 이상 향상시켰으나, 순방향 연산 성능에는 큰 영향을 미치지 못했습니다.
• Relaxed.multiplyAdd 활용: Swift-Numerics의 Relaxed.multiplyAdd를 사용하여 FMA(Fused Multiply-Add) 연산을 활용함으로써 C의 -ffast-math 플래그와 유사한 성능을 달성하여 순방향 연산 속도를 약 10배 향상시켰습니다.
• InlineArray를 사용한 루프 언롤링: C의 LOOP_UNROLL 최적화와 유사하게, InlineArray를 사용하여 스택에 할당되는 배열을 구현함으로써 C와 동일한 수준의 추론 및 훈련 성능을 달성했습니다.
• 멀티스레딩 구현: DispatchQueue.concurrentPerform를 사용하여 matmul_forward, matmul_backward, attention_forward, attention_backward의 네 가지 핵심 루프를 병렬 처리하여 훈련 성능을 5.5배 이상 향상시켰으나, 코드 복잡성이 증가했습니다.
• AMX(Apple Matrix Coprocessor) 활용: Apple Silicon의 AMX 유닛을 직접 사용하지는 않지만, 공개된 지침을 기반으로 AMX_MATFP 명령어를 활용한 행렬 곱셈 구현을 통해 훈련 성능을 약 1.67배 더 향상시켰습니다.
• Metal을 이용한 GPU 연산: Metal 커널을 사용하여 GPU에서 행렬 곱셈을 수행하였으며, 기본 Metal 구현은 AMX보다 약간 느린 성능을 보였습니다. threadsPerThreadgroup 최적화를 통해 성능을 개선하였고, 타일링(Tiling) 기법을 적용한 Metal 커널은 훈련 성능을 3000% 이상 향상시켜 TFLOP/s 수준의 컴퓨팅 성능을 달성했습니다.

시사점

Apple Silicon 환경에서 LLM 훈련을 위한 행렬 곱셈 연산의 성능은 MutableSpan, Relaxed.multiplyAdd, InlineArray, 멀티스레딩, AMX, Metal과 같은 다양한 최적화 기법을 통해 순수한 Swift 코드만으로도 Gflop/s 수준에서 Tflop/s 수준으로 크게 향상될 수 있음을 보여줍니다.