apple-aiml-research/ml-ane-transformers

Apple Neural Engine (ANE) Transformers

무엇인가요 – ane_transformers는 Apple의 Neural Engine (ANE)에서 효율적으로 실행되도록 최적화된 Transformer 모델의 참조 구현을 제공하는 소규모 오픈소스 라이브러리입니다. 다음 두 가지 패키지를 제공합니다:

• ane_transformers.reference – ANE에서 효율적으로 실행하기 위해 필요한 저수준 변경 사항을 보여주는 독립적이고 이해하기 쉬운 구현입니다.
• ane_transformers.huggingface – DistilBert와 같은 특정 Hugging Face 모델 클래스의 드롭인 대체품으로, 동일한 Python API를 유지하면서 이러한 최적화를 포함합니다.

왜 중요한가요 – Apple의 A14, A15, M1, M2 등에는 전용 신경처리장치(NPU)가 포함되어 있습니다. 계산 그래프를 재구성(연산 통합, 임베딩 조회를 CPU로 이동, ANE의 명령어 집합에 맞게 모델 조정)함으로써, 이 라이브러리는 Core ML을 통해 디바이스에서 실행할 때 일반적인 PyTorch/torch-script 모델과 비교해 최대 10배의 지연 감소와 최대 14배의 피크 메모리 사용량 감소를 주장합니다.

사용 방법 – README는 Hugging Face의 distilbert-base-uncased-finetuned-sst-2-english 모델에 대한 완전한 튜토리얼을 안내합니다:

1. transformers.AutoModelForSequenceClassification로 표준 모델을 로드합니다.
2. ane_transformers.huggingface.distilbert.DistilBertForSequenceClassification로 교체하고 원본 state dict를 복사합니다.
3. 입력을 토큰화하고, torch.jit.trace로 모델을 추적한 후, coremltools.convert를 사용해 추적된 모델을 Core ML mlprogram 패키지로 변환합니다.
4. 생성된 .mlpackage를 Xcode 프로젝트에 드래그 앤 드롭하고, Xcode의 성능 탭을 사용해 iPhone 또는 Mac에서 성능 향상을 확인합니다.

설치 – 패키지는 PyPI에 게시되어 있습니다:

pip install ane_transformers          # 사전 빌드된 웨일 (가장 빠름)
# 또는 개발용:
pip install -e .

tokenizers 종속성에서 빌드 실패 시, README는 Rust 컴파일러를 사용한 알려진 해결 방법을 안내합니다.

지원되는 하드웨어 – 최적화는 ANE를 포함한 Apple Silicon(예: iPhone 12/A14 이상, A14+ 탑재 iPad, M1 이상 Mac)을 대상으로 합니다. 오래된 하드웨어에서 실행 시 단위 테스트는 경고를 표시하지만, 생성된 Core ML 모델은 여전히 실행 가능합니다(단, 홍보된 성능 향상은 없음).

제한 사항 / 주의사항

• 현재는 Hugging Face 모델 중 일부만 래핑되어 있습니다(예: DistilBERT). 다른 아키텍처로 확장하려면 추가 작업이 필요합니다.
• 최적화된 모델은 한 번만 발생하는 컴파일 비용이 추가됩니다. 첫 번째 로드는 느립니다.
• 606개의 연산 중 4개는 CPU에 남아 있습니다(임베딩 조회), 이는 이 모델 크기에서는 의도된 설계입니다.
• 성능 향상은 시퀀스 길이와 배치 크기에 따라 달라집니다. README는 긴 시퀀스(예: 512 토큰, 배치 8)에서 최대 10배의 지연 감소를 보고하고 있습니다.

누가 이 프로젝트를 봐야 할까요 – iOS/macOS 개발자 중 로컬에 NLP 모델을 배포하고 가능한 한 최고의 디바이스 내 지연과 메모리 사용량을 원하는 사람, 그리고 Apple의 ANE를 위한 Transformer 그래프 최적화 방법에 관심 있는 연구자.