Train and run transformers directly on Apple's Neural Engine

개요

Espresso는 Apple Silicon의 Neural Engine(ANE)을 직접 활용하여 트랜스포머 모델의 추론 속도를 Core ML 대비 4.76배 향상시키는 도구입니다.

주요 내용

• 직접 ANE 추론: Espresso는 역공학된 비공개 API(_ANEClient, _ANEInMemoryModel)를 통해 MIL(Machine Intermediate Language) 프로그램을 ANE 실리콘으로 직접 컴파일합니다.
• 성능 향상: 6개 레이어 모델에서 토큰당 1.08ms의 추론 속도를 기록하여 Core ML의 5.09ms보다 4.76배 빠릅니다. 이는 6개의 트랜스포머 레이어를 2개의 ANE 디스패치로 융합된 커널을 사용하여 달성됩니다.
• Zero-copy I/O: NEON 벡터화된 읽기, vDSP argmax를 활용하고 데이터 마샬링 없이 IOSurface 버퍼를 직접 사용하여 데이터를 처리합니다.
• ANE에서의 전체 학습 지원: 순방향 및 역방향 패스, 그래디언트 누적 및 Adam 최적화 기능을 포함하여 ANE에서 모델 학습이 가능합니다.
• Swift 6.2 기반: ~Copyable move-only 텐서, strict concurrency, typed throws 등의 기능을 활용하는 순수 Swift로 작성되었으며 외부 의존성이 없습니다.
• 간편한 시작: Git clone 후 ./espresso 명령어로 빌드, 데모 가중치 다운로드, TUI 실행이 가능하며, Swift Package Manager를 통해 프로젝트에 쉽게 통합할 수 있습니다.
• .esp 모델 플랫폼: 포터블 .esp 번들과 번들 인식 런타임 선택을 중심으로 하는 비공개 모델 플랫폼을 제공합니다.
• 다양한 벤치마크 결과: Core ML, llama.cpp(Metal, CPU) 등과 비교했을 때 Espresso ANE의 상당한 성능 우위를 보여줍니다.
• 호환성: M1, M2, M3, M4 시리즈 Apple Silicon SoC의 Neural Engine을 지원하며 macOS 15.0 이상이 필요합니다. Apple A-series 칩(iOS)도 지원하지만, App Store 배포에는 제한이 있습니다.

시사점

Espresso는 Apple Silicon Neural Engine의 잠재력을 최대한 활용하여 트랜스포머 모델의 추론 및 학습 성능을 크게 향상시킬 수 있는 새로운 접근 방식을 제시하며, 이를 통해 온디바이스 AI 애플리케이션의 개발 및 효율성을 높일 수 있습니다.