ANE — Apple Neural Engine에서 신경망 훈련
GitHub: https://github.com/maderix/ANE
🔍 개요
ANE는 Apple Silicon 칩(M1/M2/M4 시리즈)에 내장된 전용 AI 가속기인 Apple의 **Neural Engine(ANE)**에서 전체 **신경망 훈련(순방향 + 역방향 전달)**을 직접 실행할 수 있음을 최초로 공개적으로 시연하는 연구 프로젝트입니다.
Apple은 CoreML 프레임워크를 통해 ANE를 추론 전용으로 공식적으로 제한합니다. 이 프로젝트는 Apple의 비공개 API(_ANEClient, _ANECompiler, _ANEInMemoryModelDescriptor)를 역공학하여 해당 제한을 우회하고 역전파를 포함한 사용자 지정 계산 그래프를 ANE 하드웨어에서 네이티브로 실행합니다.
이것은 프로덕션 프레임워크가 아닙니다. 이것은 개념 증명이자 벤치마크 참조이며, 제한이 하드웨어 제약이 아닌 소프트웨어에 의해 부과되었음을 증명합니다.
🎯 중요성
Apple Silicon 칩(특히 M4)에는 15.8 TFLOPS 등급의 ANE가 포함되어 있습니다. 이는 Apple이 추론에만 잠가두는 막대한 양의 전용 AI 컴퓨팅입니다. CoreML을 실행할 때마다 추론만 가능합니다. 훈련은 항상 GPU 또는 CPU로 대체됩니다.
이 프로젝트는 **"ANE에서 훈련할 수 있는가?"**라는 질문에 답합니다.
답은 예이며, 이 저장소가 그 증거입니다.
⚙️ 작동 방식
이 프로젝트는 다음을 통해 완전한 트랜스포머 레이어 훈련 루프를 구현합니다.
-
MIL 프로그램 생성 — Objective-C에서 런타임에 Apple의 MIL(Model Intermediate Language) 프로그램을 구성하여 컨볼루션(선형 레이어용), 행렬 곱셈(어텐션용), 소프트맥스 및 요소별 연산을 정의합니다.
-
인메모리 컴파일 —
_ANEInMemoryModelDescriptor를 사용하여 MIL 텍스트 + 가중치 블롭을 ANE 프로그램으로 직접 컴파일하며,.mlmodelc파일을 디스크에 쓸 필요가 없습니다. -
IOSurface I/O —
[1, channels, 1, spatial]fp16 형식의IOSurface공유 메모리를 통해 입력/출력 텐서를 전달합니다. 이 형식은 ANE 하드웨어가 네이티브로 예상하는 형식입니다. -
가중치 임베딩 — 가중치는
BLOBFILE상수로 ANE 프로그램에 구워지며, 가중치가 업데이트될 때마다 배치마다 다시 컴파일됩니다. -
그래디언트 흐름 — 순방향 "탭"은 역방향 전달에 필요한 중간 활성화를 노출합니다. 역방향 커널은 ANE에서 입력 그래디언트(
dx)를 계산하고, 가중치 그래디언트(dW)는cblas_sgemm을 통해 CPU에서 계산됩니다.
🧱 아키텍처: 훈련 단계당 6개의 ANE 커널
| 커널 | 기능 |
|---|---|
kFwdAttn |
RMSNorm + QKV 프로젝션 + SDPA + 출력 프로젝션 |
kFwdFFN |
RMSNorm + SwiGLU FFN (W1, W3, SiLU, W2) |
kFFNBwd |
FFN 역방향 (W2ᵀ + SiLU_bwd + W1ᵀ + W3ᵀ) |
kSdpaBwd1 |
Woᵀ + SDPA 역방향 파트 1 (dV, probs, dp) |
kSdpaBwd2 |
SDPA 역방향 파트 2 (softmax grad, dQ, dK) |
kQKVb |
QKV 역방향 (Wqᵀ + Wkᵀ + Wvᵀ → dx) |
CPU 처리: RMSNorm 역방향, 잔차 연결, 손실 계산, dW 그래디언트 누적(cblas_sgemm), Adam 옵티마이저 업데이트.
📊 벤치마크 결과 (M4 Mac, 단일 트랜스포머 레이어, dim=768, seq=512)
| 최적화 | ms/step | ANE 활용률 |
|---|---|---|
| 기준선 (vDSP 전치) | 33.5 | 3.1% |
| 채널 우선 레이아웃 | 20.3 | 5.2% |
| vDSP 벡터화 RMSNorm | 14.2 | 7.4% |
| GCD 비동기 cblas 오버랩 | 11.4 | 9.2% |
| ANE RMSNorm 융합 | 11.4 | 9.2% |
| Wo^T 융합 (7→6 커널) | 11.4 | 9.2% |
| 지연된 cblas 대기 | 9.3 | 11.2% |
최고 결과: 9.3 ms/step, 1.78 TFLOPS 유지 (M4 ANE의 15.8 TFLOPS 피크의 11.2%).
🔑 주요 최적화
- 채널 우선 CPU 레이아웃 — ANE의 네이티브
[1,C,1,S]IOSurface 형식과 일치하여 모든 전치 오버헤드를 제거합니다. - vDSP 벡터화 RMSNorm — 일반 구현보다 10배 빠름 (6.7ms → 0.7ms).
- GCD 비동기 cblas 오버랩 — 가중치 그래디언트
sgemm연산은 직렬 디스패치 큐에서 ANE 평가와 병렬로 실행됩니다. - 지연된 cblas 대기 —
dWsgemm에 대한 대기는 최대 계산 오버랩을 위해 다음 단계의 순방향 전달로 푸시됩니다. - ANE RMSNorm 융합 — RMSNorm은 MIL 연산으로 순방향 ANE 커널에 직접 통합됩니다.
- 순방향 탭 — Q, K, V, 어텐션 점수 및 은닉 상태는 연결 출력으로 노출되어 역방향 전달 중 CPU 재계산을 방지합니다.
- exec() 재시작 — 체크포인트/재개 기능을 사용하여 프로세스당 약 119개의 ANE 컴파일 제한을 우회합니다.
📁 저장소 구조
├── api_exploration.m # 초기 ANE API 검색 실험
├── inmem_basic.m # 인메모리 MIL 컴파일 개념 증명
├── inmem_bench.m # ANE 디스패치 지연 시간 벤치마크
├── inmem_peak.m # 피크 TFLOPS 측정 (2048×2048 행렬 곱셈)
├── sram_bench.m # ANE SRAM 대역폭 프로빙
├── sram_probe.m # SRAM 크기/레이아웃 탐색
└── training/
├── ane_runtime.h # ANE 비공개 API 래퍼 (컴파일, 평가, IOSurface)
├── ane_mil_gen.h # MIL 프로그램 생성 도우미
├── model.h # 모델 가중치 초기화 및 블롭 빌더
├── forward.h # 순방향 전달 MIL 생성기
├── backward.h # 역방향 전달 MIL 생성기
├── train.m # 최소 훈련 루프 (초기 프로토타입)
├── tiny_train.m # 2계층 미니 모델 훈련
├── train_large.m # 메인: 단일 계층 dim=768 훈련 (최적화됨)
├── test_*.m # 개별 커널에 대한 단위 테스트
└── Makefile
🛠️ 빌드 및 실행
요구 사항: Apple Silicon이 설치된 macOS 15 이상 (M4에서 테스트됨).
# 빌드
xcrun clang -O2 -framework Foundation -framework IOSurface \
-framework CoreML -framework Accelerate -ldl -lobjc \
-o train_large training/train_large.m
# 실행
./train_large
외부 종속성 없음 — 시스템 프레임워크와 런타임에 objc_msgSend를 통해 해결되는 비공개 ANE API만 사용합니다.
⚠️ 알려진 제한 사항
| 제한 사항 | 세부 정보 |
|---|---|
| 낮은 ANE 활용률 | 피크의 약 11.2%; 많은 요소별 연산은 여전히 CPU로 대체됩니다. |
| 약 119개의 컴파일 제한 | ANE 컴파일러가 리소스를 누수합니다. exec() 재시작으로 해결되었습니다. |
| 단일 트랜스포머 레이어 | 다중 계층 파이프라인 스케줄링은 아직 구현되지 않았습니다. |
| 합성 데이터만 | 실제 토큰화된 데이터 지원은 작업 중입니다. |
| SDPA 인과 마스킹 | ANE는 SDPA에서 attn_mask를 무시합니다. 수동 분해로 해결합니다. |
| 비공개 API | 모든 macOS 업데이트로 인해 중단될 수 있는 문서화되지 않은 API를 사용합니다. |
📚 관련 연구 논문
🏷️ 기술 스택
- 언어: Objective-C, C
- 플랫폼: macOS 15 이상 / Apple Silicon (M1/M2/M4)
- 프레임워크: Foundation, IOSurface, CoreML, Accelerate
- 비공개 API:
_ANEClient,_ANECompiler,_ANEInMemoryModelDescriptor - 형식: MIL (Model Intermediate Language), fp16
⚖️ 법적 고지
이 프로젝트는 Apple의 비공개, 문서화되지 않은 API를 사용합니다. 이러한 API는 안정성을 보장하지 않으며 모든 macOS 업데이트로 인해 중단될 수 있습니다. 이 프로젝트는 공정 사용 및 상호 운용성 조항(Sega v. Accolade, 1992; DMCA §1201(f))에 따른 독립적인 연구입니다. Apple의 독점 코드나 바이너리는 포함되지 않았습니다. Apple Inc.와 제휴하거나 보증받지 않았습니다.