MLX
MLX (Apple 머신러닝 프레임워크)
MLX는 Apple의 머신러닝 연구팀이 만든, Apple silicon에서 머신러닝을 돌리기 위한 배열(array) 프레임워크예요. NumPy를 따라가는 익숙한 Python API를 제공하면서, 동시에 모델을 효율적으로 학습·배포할 수 있게 설계됐어요. 연구자들이 새 아이디어를 빠르게 탐구할 수 있도록 프레임워크 자체의 설계도 개념적으로 단순하게 유지하려는 것을 목표로 해요.
출처: MLX 공식 문서 (공식)
핵심 특징
MLX는 여러 독특한 특징을 갖고 있어요.
- 익숙한 API: Python API가 NumPy를 밀접하게 따라가요. C++, C, Swift API도 Python API를 거울처럼 반영해요.
mlx.nn과mlx.optimizers같은 상위 레벨 패키지는 PyTorch를 밀접하게 따라가서 더 복잡한 모델을 쉽게 만들 수 있어요. - 조합 가능한 함수 변환: 자동 미분, 자동 벡터화, 계산 그래프 최적화를 위한 조합 가능한 함수 변환을 지원해요.
- 지연 계산: MLX에서 계산은 지연돼요. 배열은 필요할 때만 메모리에 실체화돼요.
- 동적 그래프 구성: 계산 그래프가 동적으로 구성돼요. 함수 인자의 형태를 바꿔도 느린 컴파일이 트리거되지 않고, 디버깅이 간단해요.
- 멀티 디바이스: 연산이 지원되는 모든 디바이스(현재 CPU와 GPU)에서 실행될 수 있어요.
- 통합 메모리: 다른 프레임워크와 가장 두드러지게 다른 점이 _통합 메모리 모델_이에요. MLX 배열은 공유 메모리에 있어요. 따라서 데이터를 전송하지 않고도 배열 연산을 어떤 지원 디바이스 유형에서든 수행할 수 있어요.
설치하기
macOS에서 MLX를 설치하는 방법은 간단해요.
pip install mlx
Linux에서 CUDA 백엔드로 설치하려면:
pip install mlx[cuda]
CPU 전용 Linux 패키지라면:
pip install mlx[cpu]
예제
MLX 예제 저장소에는 다양한 예제가 있어요.
- Transformer 언어 모델 학습
- LLaMA로 대규모 텍스트 생성, LoRA로 파인튜닝
- Stable Diffusion으로 이미지 생성
- OpenAI의 Whisper로 음성 인식
설계 영감
MLX의 설계는 NumPy, PyTorch, JAX, ArrayFire 같은 프레임워크에서 영감을 받았어요. 익숙하면서도 효율적인, 그리고 개념적으로 단순한 프레임워크를 만드는 게 목표예요.