On-Device AI with CoreML(feat. DeepSeek-R1)

개요

https://www.youtube.com/watch?v=p_hyo2FRil4&t=536s

 

최근 유행하는 Deepseek-R1 모델(Distill)을 iPad에서 온디바이스로 실행시켜 보기 위해 관련 내용들을 공부하게 되었다

 

iOS 기기에서 PyTorch 기반의 모델을 실행하려면 CoreML이 지원하는 형식으로 변환해야 한다

이때 coremltools를 활용하여 모델을 변환시킬 수 있고, 변환된 모델의 최적화를 위해 경량화 기법을 적용할 수도 있다

 

이렇게 변환시킨 모델은 CoreML을 사용하여 iOS에서 온디바이스로 추론을 수행시킬 수 있다

 

이 글은 위의 사진처럼 학습시킨 모델을 coremltools로 변환시키고, CoreML에서 이용하는 과정에서 알아야 할 것들을 중점적으로 정리했다

 

 

• 배경지식
• CoreML
• coremltools

 

배경지식

모델 경량화

• Quantization: 모델 가중치를 더 적은 비트로 표현하여 메모리 사용량을 줄이고 연산 속도를 높이는 기법
• Pruning: 중요하지 않은 연결을 제거하여 모델을 최적화하는 기법
• Palettization: 모델 가중치를 제한된 값의 집합으로 변환하여 경량화하는 기법

ScriptModule(TorchScript)

ScriptModule은 PyTorch 모델(nn.Module의 하위 클래스)을 변환하여 JIT(Just-In-Time) 컴파일이 가능하도록 만드는 형식이다

 

• 모델을 파일로 저장한 후 Python 없이 실행 가능하며 최적화 적용 가능
• 정확한 변환을 위해. eval()을 호출하여 모델을 평가 모드로 고정해야 함
• torch.jit.script() 또는 torch.jit.trace()를 사용하여 ScriptModule 생성
  • .script(model): 전체 코드 구조를 분석하여 변환 (조건문, 루프 포함)
  • .trace(model, example_input): 입력 데이터를 흘려보내 그래프를 기록 (조건문 반영 X)

Hugging Face

허깅페이스는 머신러닝 모델의 개발, 훈련, 공유를 위한 플랫폼으로, 다양한 사전 학습된 모델을 제공하며, 사용자는 간단하게 이러한 모델을 다운로드하고 사용할 수 있다

 

특히, 허깅페이스의 transformers 프레임워크의 AutoModel과 같은 클래스는 모델의 종류에 상관없이 모델을 자동으로 선택하고 간편하게 로드하는 기능을 제공한다

Stateful Model

Stateful Model은 처리 중인 데이터의 상태를 기억하고, 이를 통해 연산 효율성을 높인다(대표적으로 트랜스포머 모델에서는 KVCache를 활용하여 속도를 향상)

 

PyTorch에서는 이 캐시 정보를 register_buffer를 통해 관리하고 유지하게 하여, 같은 데이터를 반복해서 처리하지 않고도 빠르게 결과를 도출할 수 있게 한다

Apple의 CoreML에서는 MLModelState를 통해 이러한 캐시를 이용할 수 있다(WWDC24 발표)

 

CoreML

• https://developer.apple.com/kr/machine-learning/core-ml/

CoreML은 Apple 플랫폼(iOS, macOS, watchOS, tvOS)에서 최적화된 머신러닝 모델을 실행할 수 있도록 지원하는 프레임워크로, Metal 및 뉴럴 엔진을 활용하여 고속 추론을 제공한다

구성 요소

MLModelConfiguration

CoreML 모델의 설정을 관리하는 구조체로, 실행 시 사용할 장치(CPU/GPU/ANE 등)를 지정할 수 있다

MLFeatureDescription

모델의 입력과 출력 특징(feature)을 설명하는 클래스다

이 클래스는 각 특징의 이름, 데이터 유형, 크기 등의 메타데이터를 포함하여 모델의 인터페이스를 정의하여 모델이 기대하는 데이터 형식과 제약 조건을 명확하게 파악할 수 있다

MLFeatureProvider

CoreML 모델의 입출력 데이터를 제공하는 프로토콜이다

protocol MLFeatureProvider {
	var featureNames: Set<String>
	func featureValue(for: String) -> MLFeatureValue?
}

 

MLFeatureValue

MLFeatureProvider가 제공하는 값(value)을 나타내는 클래스다(MLMultiArray 또는 MLTensor 형태)

MLModel

CoreML 모델을 나타내는 기본 클래스. 실제 모델을 로드하고 예측을 수행할 때 사용된다

예측 수행의 출력은 항상 MLFeatureProvider다

MLState

모델의 state buffers(상태 버퍼)를 관리하는 객체로, 상태를 MLMultiArray와 함께 사용할 수 있다

MLTensor

CoreML에서 딥러닝 연산을 수행하는 데 최적화된 텐서 구조다(MLMultiArray와 유사하지만, 딥러닝 연산에 최적화되어 있다)

 

MLMultiArray

다차원 배열 데이터(이미지, 텍스트 임베딩 등)를 표현하는 클래스로, 모델의 입력 및 출력을 다루는 데 사용된다

모델 타입

.mlmodel

모델 아키텍처(구조)와 학습된 가중치를 포함하고 있는 CoreML 모델의 기본 파일이다

이 파일은 Xcode에 직접 추가하여 사용할 수 있으며, Xcode가 자동으로 이 파일을 컴파일하여 파일명으로 클래스(예시 코드에서 사용된 클래스)를 생성한다

 

.mlmodelc

.mlmodel 파일을 컴파일한 결과물로, 실행 환경에 최적화된 형식이다

이 파일은 Xcode 프로젝트에 포함되지 않지만, Xcode가 .mlmodel을 컴파일할 때 자동으로 생성된다

 

.mlpackage

.mlpackage는 최신 CoreML 모델 파일 포맷으로, 모델의 가중치와 구조를 분리하여 저장할 수 있다(모델을 업데이트하거나 성능을 최적화하는 데 유리한 포맷)

예시

모델을 불러오는 것은 너무나도 간단하다

이 모델에 입력과 출력을 위해 MLFeatureProvider를 통해 데이터를 전달하고 전달받는 것이 중요하다

 

import CoreML
import Foundation

// 1. 모델 로드 (YourModel.mlmodel 또는 YourModel.mlpackage)
guard let model = try? YourModel(configuration: MLModelConfiguration()) else {
    fatalError("모델 로드 실패")
}

// 2. 입력 데이터 준비
let inputData: MLMultiArray = ...

// 3. 입력 데이터를 MLDictionaryFeatureProvider로 래핑
let inputFeature = try? MLDictionaryFeatureProvider(dictionary: ["input" : MLFeatureValue(multiArray: inputData)])

// 4. 예측 실행
guard let prediction = try? model.prediction(from: inputFeature!) else {
    fatalError("예측 실패")
}

// 5. 예측 결과
guard let output = prediction.featureValue(for: "output")?.multiArrayValue else {
    fatalError("예측 결과 추출 실패")
} 
 
print(output)

 

coremltools

coremltools는 PyTorch, TensorFlow, JAX 등에서 학습된 모델을 CoreML 형식으로 변환하고 최적화하는 Python 라이브러리다

데이터 정의 타입

TensorType

다차원 배열 데이터를 정의하는 기본적인 타입이며, CoreML에서 MLMultiArray로 변환되어 사용된다

 

StateType

Transformer와 같은 상태 유지(Stateful) 모델에서 내부 상태(KV Cache 등)를 저장하기 위해 사용된다

• PyTorch에서는 register_buffer 메서드를 사용하며 등록 및 관리
• CoreML에서는 MLModelState에 대한 read_state 및 update_state API를 사용하여 관리

모델 변환

• coremltools.convert

import coremltools as ct

# 입력 형태
inputs: list[ct.TensorType] = [
    ct.TensorType(
        shape=(args.batch_size, query_size), dtype=np.int32, name="input_ids"
    ),
    ct.TensorType(
        shape=(args.batch_size, 1, query_size, final_step),
        dtype=np.float16,
        name="causal_mask",
    ),
]
# 상태모델의 경우, 상태 형태
states: list[ct.StateType] = [
    ct.StateType(
        wrapped_type=ct.TensorType(shape=model.kv_cache_shape, dtype=np.float16),
        name="key_cache",
    ),
    ct.StateType(
        wrapped_type=ct.TensorType(shape=model.kv_cache_shape, dtype=np.float16),
        name="value_cache",
    ),
]
# 출력 형태
outputs: list[ct.TensorType] = [ct.TensorType(dtype=np.float16, name="logits")]

# 각 형태들 정보를 전달하여 모델 변환
mlmodel: ct.models.MLModel = ct.convert(
    model, 
    inputs=inputs, 
    outputs=outputs, 
    states=states,
)

 

이때 모델 입출력 데이터 타입을 나타내기위해 TensorType가 사용되고, Stateful Model의 상태를 나타내기 위해 StateType을 사용한다

모델 최적화

• coremltools.optimize.coreml.optimize_model

Core ML 모델 변환 후, 추가적인 최적화를 수행할 수 있다

 

여러 기법 중, WWDC24에서 발표한 선형 양자화를 통한 최적화 예시는 다음과 같다(선형 양자화는 원래 부동소수점 값과 정수 값 사이에 일정한 비율을 적용하여 정수로 변환한다)

 

선형 양자화는 먼저 구성 설정을 하고, 이를 바탕으로 최적화를 진행시킨다

 

# 선형 양자화 구성 설정
op_config = ct.optimize.coreml.OpLinearQuantizerConfig(
    mode="linear_symmetric",  # 대칭 선형 양자화
    dtype="int8",             # 8비트 정수형
    granularity="per_block",  # 블록 단위 양자화
    block_size=32             # 블록 크기 32
)

# 최적화 구성 생성
config = ct.optimize.coreml.OptimizationConfig(global_config=op_config)

# 가중치 양자화 적용
quantized_mlmodel = ct.optimize.coreml.linear_quantize_weights(mlmodel, config=config)

모델 저장

...
mlmodel: ct.models.MLModel = ct.convert(
    model, 
    inputs=inputs, 
    outputs=outputs, 
    states=states,
)

mlmodel.save("경로")

모델 컴파일

기본적으로 변환된 모델을 저장하면, .mlpackage 형태다

 

만약 이를 .mlmodelc로 변환하려면, 다음과 같이 coremlcopiler를 통해 컴파일을 수행하면 된다

.mlmodelc는 Core ML 엔진이 직접 실행할 수 있는 최적화된 바이너리 형태의 모델이다

 

xcrun coremlcompiler compile MyModel.mlmodel ./MyCompiledModel

 

마무리

• https://github.com/finnvoor/CoreMLPipelines/tree/main

이러한 내용들을 통해 모델을 불러오는 과정을 다음 레포지터리의 코드들에서 볼 수 있다

 

4비트 양자화 모델은 아직 iOS에서 지원하지 못해서 iPad에서 돌리지 못했으나, 일단 노트북에서 돌려보고 많이 신기했었다

1.5GB 정도 되는 크기의 모델이 다양한 질문에 빠르게 적절하게 괜찮은 응답하는 모습이 상당히 인상적이다..