MLX - Apple Silicon용 Deep Learning Framework

MLX 관련 소개문서들을 한글로 번역하고, 잠깐 살펴보면서 느낀점들을 이어서 적어보려고 한다. 

 

1. 소개글

1-1. 공식 문서 소개글

https://github.com/ml-explore/mlx

https://ml-explore.github.io/mlx/build/html/index.html

MLX는 애플 머신 러닝 연구에서 제공하는 Apple Silicon용 Machine Learning을 위한 Array Framework입니다.

MLX의 주요 기능은 다음과 같습니다:
- 익숙한 API(Familiar APIs): MLX는 NumPy를 밀접하게 따르는 Python API를 가지고 있습니다. 또한 MLX는 Python API를 밀접하게 반영하는 완전한 기능을 갖춘 C++ API를 가지고 있습니다. MLX는 더 복잡한 모델 구축을 단순화하기 위해 PyTorch를 밀접하게 따르는 API를 가진 상위 레벨 패키지인 mlx.nn과 mlx.optimizers를 가지고 있습니다.
- 조합 가능한 함수 변환(Composable function transformations): MLX는 자동 미분, 자동 벡터화 및 계산 그래프 최적화를 위한 조합 가능한 함수 변환을 지원합니다.
- 느긋한 계산(Lazy computation): MLX에서의 계산은 느긋합니다. 배열은 필요할 때만 구체화됩니다. (실제 느긋하다는게 아니라, Lazy 동작을 의미함)
- 동적 그래프 구성(Dynamic graph construction): MLX의 계산 그래프는 동적으로 구성됩니다. 함수 인수의 형태를 변경해도 느린 컴파일이 발생하지 않으며, 디버깅은 간단하고 직관적입니다.
- 다중 장치(Multi-device): 연산은 지원되는 모든 장치(현재 CPU 및 GPU)에서 실행할 수 있습니다.
- 통합 메모리(Unified memory): MLX와 다른 프레임워크와의 눈에 띄는 차이점은 통합 메모리 모델입니다. MLX의 배열은 공유 메모리에 존재합니다. MLX 배열에 대한 작업은 데이터를 전송하지 않고도 지원되는 모든 장치 유형에서 수행할 수 있습니다.

MLX는 머신 러닝 연구자를 위해 머신 러닝 연구자가 설계했습니다. 이 프레임워크는 사용자 친화적이면서도 모델을 훈련하고 배포하기에 효율적을 목적으로 하고 있습니다. 프레임워크 자체의 디자인도 개념적으로 간단합니다. 새로운 아이디어를 빠르게 탐색하는 목표를 가지고 연구자들이 MLX를 확장하고 개선하기 쉽도록 하려는 것입니다.MLX의 디자인은 NumPy, PyTorch, Jax 및 ArrayFire와 같은 프레임워크에서 영감을 받았습니다.

 

1-2. 배경지식

- Lazy computation은 순차적으로 연산하고 메모리에 저장해두는 방식이 아니라, 필요할때만 연산하고 저장하기 위해 뒤로 미루는 방식이다. 지연된 시간동안 다른 일을 처리할 수 있고, 불필요한 연산을 줄이기 때문에 효율적이다.

 

- Unified memory는 Apple Silicon에 적용된 하드웨어 최적화로, CPU, GPU, 기타 프로세서가 동일한 Memory Pool을 가지기 때문에 서로간에 데이터 복사, 전송이 필요없는 구조다.

예를들어 우리가 쓰는 PC는 엔비디아 그래픽 카드에 8G RAM(VRAM), 메인보드에 16G RAM이 꽃혀있어서

 

CPU ↔ 메모리 ↔  GPU메모리 ↔  GPU

 

이런 식으로 통신해야 했다면, Unified memory는 SOC 칩 안에서

 

CPU ↔ Unified 메모리 ↔ GPU

 

M1 Ultra consists of 114 billion transistors, and supports up to a massive 128GB of fast unified memory.

 

이렇게 데이터 복사가 덜 일어나고, 대용량 데이터일 경우 이러한 장점이 더 빛을 발하게 되는 구조다. Machine Learning이나 Deep Learning이 아닌 동영상 편집, 게임 등에서도 이점을 갖는다. 그래서 그런지 요즘 Apple Silicon에서 게임 돌릴 수 있다는걸 엄청 강조한다.

스냅드래곤 혹은 엑시노스 또한 AP, GPU, NPU가 같은 메모리를 사용하긴 하겠지만, Appli Silicon 처럼 같은 SoC 안에서 메모리를 공유하는 것보다 성능이 떨어진다.

 

- Multi-device는 스마트폰, 테블릿, 노트북 같은 디바이스가 아니다. 여기서는 CPU, GPU, NPU 등의 연산 유닛을 디바이스라고 한다.

 

2. 성능

https://github.com/ml-explore/mlx-examples/tree/main/stable_diffusion

Performance
The following table compares the performance of the UNet in stable diffusion. We report throughput in images per second processed by the UNet for the provided txt2image.py script and the diffusers library using the MPS PyTorch backend.
At the time of writing this comparison convolutions are still some of the least optimized operations in MLX. Despite that, MLX still achieves ~40% higher throughput than PyTorch with a batch size of 16 and ~15% higher when comparing the optimal batch sizes.
Notably, PyTorch achieves almost ~50% higher throughput for the batch size of 1 which is unfortunate as that means that a single image can be computed faster. However, when starting with the models not loaded in memory and PyTorch's MPS graph kernels not cached, the compilation time more than accounts for this speed difference.

 

PyTorch로 구현한 것 보다 Batch Size 16일때 약 40%, 최적의 Batch Size인 경우 약 15% 더 높은 성능을 보인다고 한다. 심지어 Convolution 연산은 아직 최적화가 덜 되었다고 한다.

(M2 Ultra with PyTorch version 2.1, diffusers version 0.21.4 and transformers version 4.33.3)

 

3. CoreML 과의 차이

CoreML은 기존 머신러닝 모델을 Apple 장치에 맞게 변환하고 최적화하는 데 사용되는 프레임워크입니다. CoreML은 주로 기존의 머신 러닝 모델을 Apple 장치에서 실행하기 위해 최적화하는데 초점을 맞추고 있으며, MLX 처럼 처음부터 머신 러닝 모델을 생성, 실행하는 데는 적합하지 않습니다.

 

 

- HyScaler https://hyscaler.com/insights/mlx-apple-new-framework-ai-development/
- Decrypt https://decrypt.co/209003/apple-ai-ml-development-mlx-silicon

사진 출처 : Apple unveils M1 Ultra, the world’s most powerful chip for a personal computer - Apple