Core ML로 Mistral 7B 실행

Hugging Face는 WWDC 24에서 소개된 새로운 Core ML 기능을 활용하여 Mac 하드웨어에서 Mistral 7B 모델을 실행하는 워크플로우를 자세히 설명했습니다. 상태 버퍼, 새로운 Swift Tensor 타입, 고급 양자화를 결합함으로써 70억 파라미터 모델을 4GB 미만의 메모리로 실행할 수 있습니다.

WWDC 24에서 소개된 새로운 Core ML 기능

Apple의 최신 Core ML 업데이트는 온디바이스 대형 언어 모델(LLM)을 위한 여러 중요한 최적화를 제공합니다.

Swift Tensor (MLTensor)

새로운 MLTensor 타입은 Swift에서 다차원 데이터 구조에 대한 고수준 추상화를 제공하며, numpy 배열이나 torch 텐서의 기능을 모방합니다. MLMultiArray 또는 MLShapedArray와 같은 저수준 조작이 필요했던 것을 대체합니다. 이전에는 기본 저장소에 대한 불투명 포인터에 접근해야 했습니다. MLTensor는 softmax와 같은 내장 연산을 포함하여 언어 모델의 전처리 및 후처리 파이프라인을 단순화합니다.

Stateful Buffers

전통적으로 Core ML 모델은 무상태 함수로 동작했습니다. Stateful buffers는 모델이 GPU에 메모리 블록을 예약하여 반복 사이에 상태를 유지하도록 합니다. LLM의 경우, 이는 키-값(KV) 캐시를 구현하는 데 사용됩니다. KV 캐시를 GPU에 유지함으로써, 매 토큰 생성 단계마다 CPU와 GPU 간에 대용량 상태 데이터를 주고받는 오버헤드를 피하고, 메모리 대역폭 병목을 줄여 성능을 크게 향상시킵니다.

블록 단위 양자화

품질 손실을 크게 줄이지 않으면서 모델 크기를 줄이기 위해, Core ML은 이제 블록 단위 양자화를 지원합니다. 이 기법은 전체 텐서에 대한 단일 테이블 대신 동일 텐서의 서로 다른 영역에 대해 여러 개의 조회 테이블(LUT)을 생성합니다. 이를 통해 모델을 4비트 정밀도로 압축할 수 있으며, float32에 비해 크기가 8배, float16에 비해 4배 감소합니다.

다기능 지원

다기능 지원을 통해 개발자는 LoRA(저랭크 적응) 어댑터를 생성 모델에 패키징할 수 있습니다. 이를 통해 전체 모델을 다시 로드할 필요 없이 작은 추가 파라미터(어댑터) 세트를 교체함으로써 단일 기본 모델을 다양한 작업이나 스타일에 사용할 수 있습니다.

Mistral 7B를 Core ML로 변환하기

Mistral 7B를 효율적으로 실행하려면 전체 KV 캐시 버퍼를 사전 할당하고 제자리에서 업데이트하는 맞춤형 어텐션 구현이 필요하며, 이는 Core ML의 stateful buffer 요구사항을 충족합니다.

추적 및 변환 과정

  1. Tracing: 모델은 StatefulMistralForCausalLM의 패치된 구현을 사용해 로드되고, torch.jit.trace를 이용해 예시 입력으로 추적됩니다.
  2. Input Definition: coremltools를 통해 범위 차원을 사용하여 입력 시퀀스가 단일 토큰에서 최대 컨텍스트 길이인 2048까지 확장될 수 있도록 합니다.
  3. State Preparation: keyCache와 valueCache 버퍼는 ct.StateType을 사용해 정의되어 stateful Core ML 버퍼로 변환됩니다.
  4. Conversion: 모델은 iOS 18 또는 macOS 15 이상의 최소 배포 대상으로 Core ML에 변환되어 새로운 stateful 기능에 접근할 수 있습니다.

모델 압축

OpLinearQuantizerConfig를 사용하여 모델을 블록 크기 32의 4비트 선형 대칭 양자화로 압축합니다. 이를 통해 최종 mlpackage 크기가 약 14GB(float16 기준)에서 약 3.8GB로 감소합니다.

실행 및 구현

사용자는 swift-transformers 저장소의 preview 브랜치를 사용해 변환된 Mistral 7B 모델을 실행할 수 있습니다.

Swift로 실행하기

preview 브랜치는 전체 MLTensor 지원과 Stateful API의 Swift 대응 구현을 포함합니다. 추론은 명령줄을 통해 실행할 수 있습니다:

git clone -b preview https://github.com/huggingface/swift-transformers
swift run transformers "Prompt text" --max-length 128 Examples/Mistral7B/StatefulMistral7BInstructInt4.mlpackage

Python으로 실행하기

추론은 coremltools를 사용한 Python에서도 지원됩니다:

python3 generate.py Examples/Mistral7B/StatefulMistral7BInstructInt4.mlpackage --prompt "Prompt text"

향후 로드맵

Hugging Face는 이러한 실험적 방법들을 exporters에 통합할 계획이며, exporters는 transformers 모델을 Core ML로 변환하는 Python 도구입니다. 또한 현재 Mistral 7B 구현이 주로 Mac GPU에 최적화되어 있기 때문에, OpenELM이나 DCLM과 같은 더 작은 모델을 탐색하여 iPhone의 Apple Neural Engine(ANE) 성능을 최적화하는 것을 목표로 하고 있습니다.

Sources