TristanBilot/mlx-benchmark

이 저장소는 모든 Apple Silicon 칩과 일부 GPU에서 Apple의 MLX 연산 및 레이어를 벤치마킹하는 것을 목표로 합니다.

기여: 누구나 이 벤치마크에 기여할 수 있습니다! 누락된 장치가 있거나 빠진 레이어/연산을 추가하고 싶다면, 기여 가이드라인을 읽어주세요.

현재 M 칩: M1, M1 Pro, M1 Max, M1 Ultra, M2, M2 Pro, M2 Max, M2 Ultra, M3, M3 Pro, M3 Max, M3 Ultra, M4, M4 Pro, M4 Max, M5, M5 Max.

현재 CUDA GPU: RTX4090, Tesla V100, A100.

참고

MLX의 새 버전을 사용한다면, 이미 목록에 있는 장치에 대해서도 벤치마크를 제출할 수 있습니다. 이전 벤치마크 테이블을 오버라이드(overriding)하여 PR을 제출하기만 하면 됩니다. 또한, 기존 벤치마크 대부분에는 최근 mlx-benchmark에 추가된 mx.compile 기능이 포함되어 있지 않습니다.

벤치마크는 GPU와 CPU에서 모든 mlx 연산의 실행 시간을 측정하고, 이를 mps, cpu, 그리고 cuda 백엔드를 사용하는 pytorch의 해당 연산과 비교하여 생성됩니다. MLX와 GPU를 사용할 경우, mx.compile으로 컴파일된 연산은 기본적으로 벤치마크에 포함됩니다. 컴파일된 함수를 벤치마킹하고 싶지 않다면 --compile=False를 설정하세요.

각 연산에 대해 여러 실험의 실행 시간을 측정합니다. 저희는 이 실험들을 기반으로 두 가지 벤치마크를 제안합니다:

• 상세 벤치마크 (Detailed benchmark): 각 실험의 실행 시간을 제공합니다.
• 평균 실행 시간 벤치마크 (Average runtime benchmark): 실험들의 평균을 계산하며, 탐색하기 쉽고 세부 정보가 적습니다.

로컬에서 벤치마크를 실행하는 것은 간단합니다. osx-arm64 아키텍처로 새 환경을 만들고 의존성을 설치하세요.

CONDA_SUBDIR=osx-arm64 conda create -n mlx_benchmark python=3.10 numpy pytorch torchvision scipy requests -c conda-forge
pip install -r requirements.txt

macOS 외의 다른 운영 체제에서는 CPU 또는 CUDA 장치를 사용하여 torch 실험만 실행할 수 있습니다. CONDA_SUBDIR=osx-arm64 없이 새 환경을 설치하세요.

CUDA 버전을 지원하는 torch 패키지를 프리픽스(prefix)하고 설치하세요. 그런 다음 requirements.txt 내의 모든 요구사항을 설치하되, mlx는 제외합니다.

마지막으로, config.py 파일을 열고 다음을 설정하여 mlx 패키지 가져오기(importing)를 방지하세요. 이 패키지는 Mac 장치가 아닌 곳에서는 설치할 수 없습니다:

USE_MLX = False

mps, mlx 및 CPU에서 벤치마크를 실행하려면 다음 명령을 사용하세요:
python run_benchmark.py --include_mps=True --include_mlx_gpu=True --include_mlx_cpu=True --include_cpu=True

CUDA 및 CPU에서 torch 벤치마크를 실행하려면 다음 명령을 사용하세요:
python run_benchmark.py --include_mps=False --include_mlx_gpu=False --include_mlx_cpu=False --include_cuda=True --include_cpu=True

만약 mx.compile로 컴파일된 연산에 대한 벤치마킹만 관심이 있다면, 다음 명령을 실행할 수 있습니다:
python run_benchmark.py --include_mps=False --include_cpu=False --include_mlx_cpu=False

프로젝트를 계속 업데이트해 주신 모든 분들께 감사드립니다!