MLX 대 PyTorch 비교 벤치마크
요약
본 문서는 Apple Silicon 환경에서 인기 AI 프레임워크인 MLX와 PyTorch의 성능을 비교하는 벤치마크 결과를 제시합니다. 트랜스포머 학습, BERT 파인튜닝, Whisper 추론 등 다양한 시나리오를 테스트했으며, M3 Max 같은 최신 장비에서 MLX가 전반적으로 더 빠른 성능을 보여주었습니다.
핵심 포인트
- Apple Silicon 환경에 특화된 AI 프레임워크 비교 분석 자료입니다.
- MLX는 PyTorch 대비 트랜스포머 학습 및 추론 등 다수 벤치마크에서 우수한 속도를 보였습니다.
- M3 Max와 같은 고성능 장비일수록 MLX의 성능 격차가 더욱 두드러집니다.
이 저장소는 Apple Silicon 장치에서 작동하는 두 가지 인기 있는 인공지능 프레임워크인 MLX와 PyTorch를 비교하기 위한 벤치마크를 담고 있습니다.
이 간단한 프로젝트의 아이디어는 Apple 컴퓨터에서 AI 프로젝트를 시작할 때 현명한 선택을 할 수 있도록 하는 것입니다.
우리는 일상적인 사용을 에뮬레이션하기 위해 다섯 가지 벤치마크를 여러 번 실행했습니다. 각 벤치마크에 대한 자세한 정보는 '각 벤치마크에 대한 세부 정보' 섹션을 참조해 주십시오.
- 트랜스포머 언어 모델 학습 (
lm_train.py) - BERT 학습/파인튜닝 (
bert_fine_tune.py) - OpenAI의 whisper 모델을 사용한 추론 (
whisper_inference.py) - TinyLLama를 사용한 언어 모델 추론 (
llm_inference.py) - 행렬 곱셈을 위해 CPU와 GPU 간에 데이터를 이동시키는 합성 벤치마크 (
switch_test.py)
언어 모델 학습과 BERT 학습의 경우 시간이 더 오래 걸려 세 번만 실행했으며, 나머지 테스트는 각각 열 번씩 실행했습니다.
아래 표의 결과는 우리가 실행한 반복 횟수에 대한 평균 시간을 보여줍니다. 각 벤치마크 실행 시간의 중앙값에 대한 정보는 raw_results.txt를 참조하십시오.
M1 Pro (CPU 코어 10개, GPU 코어 16개, RAM 32 GB) |||
| 벤치마크 | PyTorch 시간 (초) | MLX 시간 (초) |
|---|---|---|
| 트랜스포머 언어 모델 학습 | 1806.63 | 1157.00 |
| BERT 학습 | 751.02 | 718.35 |
| Whisper 추론 | 31.99 | 8.50 |
| ... | ||
| M1 Max (CPU 코어 10개, GPU 코어 32개, RAM 64 GB) | ||
| 벤치마크 | PyTorch 시간 (초) | MLX 시간 (초) |
| --- | --- | --- |
| 트랜스포머 언어 모델 학습 | 1106.75 | 752.25 |
| BERT 학습 | 793.67 | 499.34 |
| Whisper 추론 | 21.28 | 6.95 |
| ... | ||
| M3 Max (CPU 코어 16개, GPU 코어 40개, RAM 48 GB) | ||
| 벤치마크 | PyTorch 시간 (초) | MLX 시간 (초) |
| --- | --- | --- |
| 트랜스포머 언어 모델 학습 | 912.52 | 426.00 |
| BERT 학습 | 550.29 | 408.45 |
| Whisper 추론 | 17.90 | 4.85 |
| ... | ||
| 먼저, 저장소를 설정할 수 있도록 git LFS가 설치되어 있는지 확인하세요: |
pip3 install -r requirements.txt
cd pytorch_models
./configure.sh
...
루트 폴더의 모든 Python 파일은 다른 벤치마크를 나타냅니다. 이들 모두 두 가지 인자, 즉 벤치마크 실행 횟수와 프레임워크가 필요합니다. 예를 들어 PyTorch를 사용하여 TinyLLama 추론 벤치마크를 열 번 실행하고 싶다면 다음을 실행하세요:
python3 llm_inference.py --framework pytorch --iter 10
명령이 완료되면 터미널에 열 번 반복한 평균 시간과 중앙값을 출력합니다.
lm_train.py 벤치마크는 PyTorch와 함께 사용될 때 PYTORCH_MPS_HIGH_WATERMARK_RATIO 환경 변수가 0으로 설정되어야 합니다.
whisper_inference 벤치마크는 PyTorch 저장소의 최신 커밋에서만 작동하므로, 이 벤치마크를 실행하려면 소스에서 빌드해야 합니다.
이 벤치마크를 위해 MLX의 TransformerLM 예제에서 모델을 복사했습니다. PyTorch 버전의 경우, 다른 프레임워크에서 모델을 제대로 재현하기 위해 사용 가능한 가장 유사한 함수들을 활용했습니다. 사용된 데이터셋은 PTB 코퍼스입니다. 모델 크기, 에포크 및 기타 하이퍼파라미터에 대한 자세한 정보는 lm_train.py를 참조하십시오.
저희는 Conneau 등이 제시한 모델을 활용했으며, 해당 BERT 블록에는 BERT-tiny 모델을 사용했습니다. 이 모델은 문장 쌍이 모순(contradiction), 함의(entailment) 또는 중립 관계(neutral relation)를 갖는지 분류합니다. 각각 순수 PyTorch와 순수 MLX로 구현되었습니다. 사전 학습된 가중치로 초기화하지 않았기 때문에, 이 벤치마크는 순수한 학습으로 볼 수 있습니다. 학습 데이터셋은 NLI 데이터셋이었습니다.
여기서의 유일한 조정 사항은 토크나이저 라이브러리와 호환되도록 MLX 모델에도 PyTorch dataloader를 사용했다는 점입니다. 데이터 로더가 각 입력에 대해 PyTorch 텐서를 생성하더라도, 추가적인 복사 없이 NumPy 배열로 변환할 수 있으므로 이 설정이 MLX 결과에 해를 끼치지는 않았습니다.
PyTorch 설정의 경우, HuggingFace transformers 라이브러리를 사용하여 tiny whisper 모델을 다운로드하고 실행했습니다. MLX 벤치마크의 경우, float32를 내부 데이터 타입으로 사용하고 MLX examples tools를 사용하여 tiny whisper를 다운로드하고 MLX 형식으로 변환했습니다(PyTorch와 일치하도록). (mlx_models/configure.sh 참조). MLX용 추론 코드는 mlx_whisper 라이브러리를 활용합니다.
PyTorch의 경우, HuggingFace 저장소에서 TinyLlama-1.1B-Chat-v1.0 모델을 다운로드하고 (pytorch_models/configure.sh 참조), transformers 라이브러리를 사용하여 모델을 로드하고 실행했습니다.
MLX의 경우, MLX examples tools를 사용하여 모델을 MLX 형식으로 변환하고, PyTorch와 일치하도록 float32를 데이터 타입으로 사용했습니다. 적절한 프롬프트 포맷팅 및 실행 제약 조건을 처리하기 위해 여러 가지 조정이 가해진 MLX examples 저장소의 실행 스크립트를 활용했습니다.
이 벤치마크에서는 루프 내에서 행렬 곱셈을 수행합니다. 먼저, CPU에서 행렬을 곱한 다음, 그 결과로 나온 행렬들을 GPU에서 곱합니다. 마지막으로, 후자의 결과를 다음 반복의 CPU 곱셈에 대한 입력으로 재사용합니다.
이 벤치마크의 아이디어는 각 프레임워크가 실행 장치(execution units) 간 데이터 이동을 얼마나 효과적으로 처리하는지 평가하는 것입니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 GitHub ML Hardware의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기