- MLX는 Apple Silicon 기반 기계학습 연구자를 위한 어레이 프레임워크로 NumPy 및 PyTorch와 유사한 사용성을 제공
- MLX에 CUDA 백엔드를 추가하는 작업이 진행되고 있음
- 주요 목적은 통합 메모리(unified memory) 지원과 NVIDIA 하드웨어의 폭넓은 활용
- 현재 튜토리얼 예제만 동작 가능하며, 빌드와 테스트는 Ubuntu 22.04 + CUDA 11.6 환경에서 확인됨
- 초기 성능 문제와 병목 개선 과정을 거치면서 최적화 및 리팩토링이 반복적으로 진행 중
- MLX의 CUDA 백엔드는 Apple의 후원을 받아 진행되고 있으며, 통합 메모리와 NVIDIA 하드웨어 지원을 통해 Mac에서 개발하고 대형 컴퓨팅 환경에 배포할 때 일관된 경험을 제공하는 것이 목표
MLX 프레임워크 소개
- MLX는 Apple의 머신러닝 연구팀이 개발한 기계학습용 배열 프레임워크임
- 엔지니어와 연구자가 효과적으로 실험하고 빠르게 아이디어를 검증할 수 있도록 개발됨
- Apple Silicon을 지원하는 것이 특징이지만, 이제 CUDA 백엔드 추가로 NVIDIA GPU 환경에서도 동작 가능성을 확장하려고 하고 있음
주요 특징
- 친숙한 API: Python API는 NumPy를, 상위 패키지인
mlx.nn,mlx.optimizers는 PyTorch와 유사한 인터페이스 제공- C++ , C, Swift API도 함께 제공되며, Python API와 일관된 사용성 보장
- 합성 가능한 함수 변환: 자동 미분, 자동 벡터화, 컴퓨팅 그래프 최적화 등의 기능이 포함됨
- 지연 계산: 여러 연산을 묶어 실제로 필요할 때만 결과를 계산하는 구조로 효율적인 리소스 활용 가능
- 동적 그래프 생성: 입력 데이터의 shape이 바뀌어도 느린 컴파일 없이 즉각적인 실행, 개발 및 디버깅이 용이함
- 멀티 디바이스 및 통합 메모리 모델:
- 기존에는 CPU와 Apple GPU만 지원했으나, CUDA 백엔드로 NVIDIA GPU 환경 지원 추진 중
- 통합 메모리 구조 채택으로, 메모리 복사 없이 모든 지원 장치에서 동일 객체 연산 가능
- 간결하고 확장 가능한 설계: 연구자들이 프레임워크를 손쉽게 확장 및 개선 가능
타 프레임워크와의 차이점 및 영감
- NumPy, PyTorch, Jax, ArrayFire 등에서 설계 영감 받음
- 특징적으로 통합 메모리 모델 및 간결한 인터페이스, 즉시 실행/디버깅 환경을 강조함
대표적 예시 및 활용 분야
- MLX Examples 레포에는 다양한 실습 코드가 존재함
- Transformer 언어 모델 학습
- LLaMA 기반 대규모 텍스트 생성 및 LoRA 활용한 파인튠
- Stable Diffusion으로 이미지 생성
- OpenAI Whisper 기반 음성 인식 등 최신 모델 구현 사례 포함
MLX의 의미와 CUDA 지원
- Apple Silicon 중심의 프레임워크였던 MLX가 CUDA 백엔드 도입으로 NVIDIA GPU 환경에서도 범용적으로 활용될 가능성 확장 중
- 최신 CUDA 지원을 통해, Apple과 NVIDIA 양쪽 하드웨어에서 최신 연구 및 산업 활용 모두에 적합한 프레임워크로 자리매김할 전망