이 문서는 Apple Neural Engine (ANE)에서 직접 신경망 훈련을 가능하게 하는 연구 프로젝트를 상세히 설명합니다. 이 프로젝트는 Apple의 사설(private), 비공개(undocumented) API를 역설계하여 CoreML의 추론 전용 제약을 우회하고, ANE의 하드웨어 잠재력을 훈련 목적으로 탐색하는 것을 목표로 합니다.

프로젝트 범위 및 목표:
이 프로젝트는 ANE 훈련의 가능성을 보여주는 개념 증명(proof of concept)이며, ANE의 성능 특성(처리량, 전력, SRAM 동작)을 문서화하는 벤치마크를 제공합니다. CoreML 외부에서 ANE에 직접 접근하는 방법을 탐구하는 이들에게 참고 자료가 됩니다. 이는 유지보수되는 프레임워크나 생산용 스택을 대체하는 것이 아니라, 소규모 연구 모델을 위한 연구 코드입니다. 현재 ANE 활용도는 최고 성능의 약 2-3% 수준이며, 많은 요소별(element-wise) 연산이 여전히 CPU로 폴백(fall back)됩니다.

핵심 방법론:

1. 사설 API 역설계 및 사용 (Private API Reverse Engineering and Usage):

_ANEClient 및 _ANECompiler와 같은 Apple의 비공개 API를 런타임 인트로스펙션(introspection)을 통해 식별하고 사용합니다. 이들 API는 일반적으로 CoreML 내부에서 사용됩니다.

2. MIL (Model Intermediate Language) 생성 (MIL Generation):

Objective-C 코드가 런타임에 MIL 프로그램 텍스트를 구성합니다. 이는 선형 계층(linear layers)을 위한 convolution, 어텐션(attention)을 위한 matmul, softmax, 그리고 요소별 연산(element-wise ops)과 같은 ANE의 기본 연산들을 지정합니다.

3. 인메모리 컴파일 (In-Memory Compilation):

디스크에 .mlmodelc 파일을 생성하는 대신, _ANEInMemoryModelDescriptor를 사용하여 MIL 텍스트와 가중치 블롭(weight blobs)을 직접 ANE 프로그램으로 인메모리 컴파일합니다.

4. IOSurface I/O:

입력 및 출력 텐서는 IOSurface 공유 메모리를 통해 CPU와 ANE 간에 전달됩니다. 데이터 형식은 [1, channels, 1, spatial]이며 fp16 (float16) 정밀도를 사용합니다.

5. 가중치 임베딩 (Weight Embedding):

가중치(weights)는 BLOBFILE 상수로 ANE 프로그램 내에 직접 내장(bake)됩니다. 가중치가 변경되면 (예: 훈련 중 업데이트 시) ANE 프로그램은 매 배치(batch)마다 재컴파일되어야 합니다.

6. 경사 흐름 (Gradient Flow):
   • 순방향(Forward Pass): 모든 순방향 연산은 ANE에서 실행됩니다. 역전파에 필요한 중간 값들은 concat 출력을 통해 노출(forward taps)되어 CPU에서의 재계산을 방지합니다.
   • 입력 경사(Input Gradients, dx): dx는 ANE 상에서 계산됩니다. 예를 들어, kFFNBwd는 FFN 역방향 계산($W_2^T + SiLU_{bwd} + W_1^T + W_3^T$)을, kSdpaBwd1, kSdpaBwd2, kQKVb는 SDPA 및 QKV 역방향 계산을 ANE에서 수행합니다.
   • 가중치 경사(Weight Gradients, dW): dW는 CPU의 Accelerate 프레임워크 cblas_sgemm을 통해 계산됩니다. 이는 특히 행렬 곱셈이 필요한 가중치 경사 계산의 부하를 CPU로 분산시키는 전략입니다.

7. 훈련 아키텍처:

단일 트랜스포머 레이어(dim=768, seq=512) 훈련 루프는 단계당 6개의 ANE 커널을 사용합니다:

• kFwdAttn: RMSNorm + QKV 투영 + SDPA + 출력 투영 (순방향 어텐션)
• kFwdFFN: RMSNorm + SwiGLU FFN (순방향 FFN)
• kFFNBwd: FFN 역방향 (가중치 전치 $W_2^T, W_1^T, W_3^T$ 포함)
• kSdpaBwd1: SDPA 역방향 파트 1 (출력 투영 전치 $W_o^T$ 포함)
• kSdpaBwd2: SDPA 역방향 파트 2 (softmax 경사, $dQ, dK$)
• kQKVb: QKV 역방향 ($W_q^T, W_k^T, W_v^T \rightarrow dx$)

CPU는 RMSNorm 역방향, 잔차 연결(residual connections), 손실 계산, dW 경사 누적(cblas\_sgemm), Adam 최적화기 업데이트를 처리합니다.

8. 핵심 최적화:
   • 채널-우선 CPU 레이아웃 (Channel-first CPU layout): ANE IOSurface의 [1,C,1,S] 형식과 일치시켜 전치(transpose) 오버헤드를 제거합니다.
   • vDSP 벡터화 RMSNorm: RMSNorm 계산을 10배 빠르게 만듭니다.
   • GCD 비동기 cblas 오버랩 (GCD async cblas overlap): dW 경사 sgemm을 ANE 평가와 병렬로 실행합니다.
   • 지연된 cblas 대기 (Deferred cblas wait): 최대 오버랩을 위해 다음 단계의 순방향 패스로 대기 시간을 미룹니다.
   • ANE RMSNorm 퓨전 (ANE RMSNorm fusion): RMSNorm을 MIL 연산으로 순방향 커널에 통합합니다.
   • Wo^T 퓨전: 출력 투영 역방향을 SDPA 역방향 커널에 병합하여 커널 수를 줄입니다.
   • exec() 재시작: 프로세스당 약 119개의 ANE 컴파일 제한을 우회합니다.

성능 결과 (M4 칩):
단일 트랜스포머 레이어(dim=768, seq=512) 훈련 기준, 단계당 9.3ms, 11.2%의 ANE 활용도 (1.78 TFLOPS 지속)를 달성했습니다. 이는 초기 3.1% ANE 활용도에서 크게 개선된 결과입니다.

제한 사항:

• SDPA 인과적 마스킹 (SDPA causal masking): ANE 하드웨어는 SDPA 연산에서 attn_mask를 무시하므로, 인과적 어텐션은 별도의 연산들로 분해됩니다.
• 약 119개의 컴파일 제한: ANE 컴파일러가 리소스를 누출하여, 체크포인트를 사용한 exec() 재시작으로 해결했습니다.
• 단일 레이어: 현재는 단일 트랜스포머 레이어만 훈련하며, 다중 레이어는 파이프라인 스케줄링이 필요합니다.
• 합성 데이터 (Synthetic data): 현재는 벤치마킹을 위해 무작위 데이터를 사용하며, 실제 토큰화된 데이터 지원은 진행 중입니다.

면책 조항 (Disclaimer):
이 프로젝트는 Apple의 비공개, 비공식 API를 사용하므로 macOS 업데이트에 따라 작동이 중단될 수 있습니다. 이는 ANE 아키텍처에 대한 독립적인 연구이며, 공정 사용(fair use) 및 상호 운용성 조항에 따라 연구 및 교육 목적으로 런타임 인트로스펙션을 통해 발견된 API를 사용합니다. Apple의 독점 코드나 바이너리는 포함되지 않았으며, Apple Inc.와 제휴하거나 보증하지 않습니다.