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

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

이 아티클은 Apple Silicon 환경에서 외부 라이브러리 없이 Swift만으로 LLM 학습의 핵심인 행렬 곱셈(Matrix Multiplication)을 최적화하는 과정을 상세히 설명합니다. 초기 Swift 구현의 성능 병목 지점을 파악하고, SIMD 명령어 활용, 루프 언롤링, 멀티스레딩 처리를 통해 성능을 점진적으로 개선합니다. 최종적으로 Apple의 독자적인 AMX 유닛과 Metal GPU 커널을 직접 작성하여 성능을 2.8 Gflop/s에서 1.1 Tflop/s까지 끌어올리는 성과를 보여줍니다. 이는 고성능 수치 연산에서 Swift가 C 언어에 필적하거나 오히려 능가할 수 있음을 입증하는 실전 사례입니다.

배경

Swift 프로그래밍 언어 기초, 행렬 곱셈 및 신경망 학습의 기본 개념, Apple Silicon 하드웨어 구조에 대한 이해

대상 독자

Apple Silicon 환경에서 LLM 성능 최적화에 관심 있는 Swift 개발자 및 ML 엔지니어

의미 / 영향

Swift가 단순한 앱 개발 언어를 넘어 고성능 ML 학습 엔진으로 충분히 활용될 수 있음을 보여줍니다. 특히 프레임워크 없이 하드웨어 레벨의 최적화가 가능하다는 점은 특수 목적의 온디바이스 AI 모델 개발에 중요한 시사점을 제공합니다.

섹션별 상세

기본 Swift 구현에서 배열의 Copy-on-Write 오버헤드와 참조 확인 비용이 성능의 주요 병목으로 작용했습니다. UnsafeMutableBufferPointer를 사용하여 메모리에 직접 접근하고 Span 구조체를 활용해 슬라이싱 비용을 제거함으로써 초기 대비 약 3.5배의 성능 향상을 달성했습니다.

Swift Numerics 라이브러리의 'Relaxed' 제약 조건을 적용하여 컴파일러가 FMA(Fused Multiply-Add)와 같은 SIMD 명령어를 적극적으로 사용하도록 유도했습니다. 이를 통해 부동 소수점 연산의 정밀도 제약을 일부 완화하는 대신 하드웨어 가속을 극대화하여 성능을 대폭 개선했습니다.

루프 구조를 재설계하고 수동으로 루프를 언롤링하여 CPU의 파이프라이닝 효율을 높였습니다. 데이터가 메모리에 저장된 순서(Row-major)에 맞춰 내부 루프를 순회하도록 변경함으로써 캐시 적중률을 최적화하고 SIMD 처리량을 극대화했습니다.

DispatchQueue.concurrentPerform을 도입하여 단일 코어의 한계를 넘는 멀티스레드 병렬 처리를 구현했습니다. 행렬의 행 단위를 여러 스레드에 분산 배치하여 계산함으로써 Apple Silicon의 다중 코어 성능을 온전히 활용할 수 있게 되었습니다.

Apple Silicon 전용 가속 유닛인 AMX(Apple Matrix Extension)를 직접 제어하여 CPU 기반 최적화보다 70% 더 높은 성능을 기록했습니다. 타일링(Tiling) 기법을 적용해 데이터를 AMX 레지스터 크기에 맞춰 패킹하고 연산함으로써 메모리 대역폭 한계를 극복했습니다.

Metal GPU 커널을 직접 작성하여 1.1 Tflop/s라는 최종 성능 목표를 달성했습니다. 스레드 그룹 공유 메모리를 활용한 타일링 커널을 구현하여 GPU의 수천 개 코어가 효율적으로 행렬 연산을 수행하도록 설계했습니다.

이미지 분석

Chart
이 이미지는 llm.c, Basic Swift, Fast Swift, Multithreaded Swift, Accelerate BLAS 등 각 구현 방식에 따른 학습 손실(Training Loss) 변화와 단계별 실행 시간(ms), 초당 반복 횟수(Iter/s)를 상세히 비교합니다. 특히 Accelerate BLAS가 가장 빠른 성능(13.381s)을 보이며, 기본 Swift 구현과의 극명한 성능 차이를 시각적으로 증명합니다.
TinyShakespeare 데이터셋을 사용한 다양한 엔진별 학습 손실 및 성능 비교 그래프와 표

실무 Takeaway

고성능 수치 연산 시 Swift의 안전 장치(참조 확인 등)가 오버헤드가 될 수 있으므로 Unsafe 포인터 계열 API를 적절히 활용해야 합니다.
Apple Silicon의 성능을 극대화하려면 단순 CPU 코어 활용을 넘어 AMX나 Metal과 같은 전용 하드웨어 가속 유닛을 타겟팅한 최적화가 필수적입니다.
행렬 연산 최적화의 핵심은 메모리 접근 패턴 최적화(타일링)와 하드웨어 전용 명령어(SIMD, FMA)의 결합에 있습니다.

언급된 리소스

GitHubCwlLlmSwift GitHub Repository

GitHubllm.c by Andrej Karpathy

튜토리얼Let's build GPT: from scratch, in code, spelled out