How to build & run on-device model

• vision/speech understanding, language/speech/action generation, etc

• 어떤 조건을 온디바이스라고 부를까
서버/클라우드가 아닌 로컬 기기에서 실행되면 전부 온디바이스 모델이라고 부른다.
꼭 cpu로 실행되는 것만을 지칭하는건 아니다.
맥북에도 이미 어느정도 좋은 GPU/NPU 등이 내장되어 있다.

• 일반적으로 on-device가 되기 위해서 필요한 것
1. 로컬 메모리에 올라갈 수 있는 정도의 모델 크기
2. 낮은 병렬 처리 효율성 환경에서도 충분히 기다릴만한 inference 속도

메모리에 올릴 수 있는 모델 크기에 대한 이해.

1. model weight
fp16 → 2 bytes * 30억 = 6GB

정밀도	파라미터 1개당 크기	3B 모델 가중치 크기
FP32	4 bytes	약 12GB
FP16 / BF16	2 bytes	약 6GB
INT8	1 byte	약 3GB
INT4	0.5 byte	약 1.5GB

FP16과 INT4는 당연히 4배 차이가 난다.
하지만 실제로는 quantization을 위해서 추가로 저장해야하는 정보가 있기 때문에 실제로는 1.7 ~ 2.0GB 정도 차지. (아래에서)
 
Transformer block 한개 기준
아주 단순하게
Norm + Attention + Norm + MLP/FFN
라고 생각하면, 일반적으로 MLP의 hidden_dim은 4배 크기 때문에
• Norm : dim * 2 : 1% 이하
• Attention : 4 * dim^2 : 약 33% -> Grouped query attention을 사용하면 더 작아짐
• MLP : 8 * dim^2 : 약 66%
 

2. activation
모델이 계산 중간에 만들어내는 임시 tensor : 대표적인 예시 = attention score
attention score는 단순 구현시 [seq_len, seq_len] 이기 때문에 크기가 꽤 커질 수 있다. 하지만 실제로는 flash attention 같은 프레임워크들이 잘 처리해주기 때문에 그렇게 커지지 않는다.

3. KV cache
저번 주에 얘기 하려던거
LLM에서 주로 사용하는 KV Cache : sequence 길이에 따라 달라진다.
연산 속도를 크게 높여주는 대신 메모리를 차지한다.

num_layers * sequence_length * kv_heads * head_dim * 2(K,V) * bytes

 
3B LLM 가정

num_layers = 28
hidden_size = 3072
num_attention_heads = 24
num_kv_heads = 8
head_dim = 128
KV cache dtype = FP16, 2 bytes

여기서 num_kv_heads는 attention에서 MHA, GQA, MQA에 따라 달라진다. KV Cache를 줄이기 위해서 MQA가 나왔다가 성능을 위해 GQA로 타협
 
→ 이제 토큰 길이가 4096이라고 하면,

28 × 4096 × 8 × 128 × 2 × 2
= 469,762,048 bytes
≈ 448 MiB
≈ 0.45GB

전체 토큰 수	KV cache
1K tokens	약 0.11GB
2K tokens	약 0.22GB
4K tokens	약 0.45GB
8K tokens	약 0.90GB
16K tokens	약 1.8GB
32K tokens	약 3.6GB

단순 곱이기 때문에 선형적으로 증가해서 토큰이 2배 길어지면 메모리도 2배 더 차지함.

4. 기타
• tokenizer, attention mask, input ids 등등
• token embedding

vocab_size * dim * fp16
128,000 * 3072 * 2 bytes = 786MB

마지막에 embedding → token을 처리하는 LM head는 공유하는 경우에는 추가되지 않음
• 일반적으로 앱들이 컴퓨터에서 차지하는 메모리
macOS + 기본 앱 + Chrome + VSCode 등등..
Mac Apple Chip = 16GB = CPU + NPU + GPU
Activity Monitor 사용하면 볼 수 있다. →



이미 Swap이 일어난 상태이지만 남은게 거의 2GB 밖에 없음;
모델을 올리기 위해서는 다른 프로그램들을 Swap해야하니 다른 동작들이 느려진다.
보통 5 ~ 8GB 정도 된다고 한다.

5. Runtime / framework overhead
같은 모델이라도 어떤걸 사용해서 돌리느냐에 따라 사용되는 메모리가 다르다.
PyTorch, Tensorflow, Jax는 학습을 위한 프레임워크로 inference를 위해서는 주로 다른 것들이 사용된다.

llama.cpp
Ollama
MLX
PyTorch
ONNX Runtime
Core ML
Transformers
TensorFlow.js
WebGPU runtime

일반적으로 on-device 모델을 위해서 하는 것들

경량화 = model size 줄이기

1. Pruning



 
중요하지 않은 가중치를 drop하는걸 의미한다. 단순히 neuron 단위로 drop할 수도 있는데, 이렇게되면 모델의 가중치 수가 줄어들어 가벼워지기는 하지만 latency는 줄어들지 않을 수도 있다.
왜냐하면 kernel들이 어짜피 병렬 처리에 최적화되어있어 부분적으로 연산이 비더라도 한 단위에 영향이 없을 수 있기 때문에.
이론적으로 sparsity는 높아지지만, 실제 속도 이득은 sparse 상태를 지원하는 kernel이 있어야함.
Structured pruning은 channel, attention head, layer 같은 덩어리를 제거한다.
이건 실제 행렬 크기 자체가 줄어들기 때문에 온디바이스에서는 더 실용적일 때가 많다.
 
SlimQwen, Nvidia minitron에서는 depth(d_model), width(layer num), expert 단위로 pruning을 적용하는 실험을 했다.



width를 줄이는게 제일 좋았다고 하는데, 큰 의미는 모르겠다.
 

2. Knowledge Distillation
잘 학습되었지만 큰 크기의 모델(Teacher)이 있을 때 이걸 활용해서 모델(Student)을 더 효율적으로 학습하는 방법.
우선 모델의 크기를 줄인다음, 줄어든 크기의 모델도 좋은 성능을 낼 수 있게 학습하는 방법이다.
 
LLM의 경우에는 기존 LM Loss(정답 token과의 cross entropy) 대신 teacher model의 output logits 분포 자체를 따라가도록 학습한다.
이렇게되면 한번의 학습 안에 정보량이 더 많다.
 
Diffusion에서는 모델의 크기보다도 더 적은 step으로 inference하기위해 teacher model을 사용해서 distillation한다.
 
한번의 학습마다 teacher/student 모델 둘다 inference해야하기 때문에 학습 자체는 더 오래 걸릴 수 있다.
 

• 실제로는 LM Loss + KD Loss까지 하는게 더 좋았고, Multi-token prediction으로도 KD를 해서 speculative decoding 성능도 최적화했다.

• 또 어짜피 4배를 줄여서 400B token을 학습한다고 할 때, 2배로 줄이고 40B로 학습하고 다시 2배 줄여서 360B를 학습하는 식으로 progressive pruning을 하는게 성능이 더 좋았다고 한다.

Model quantization

1. weight만 줄일 것인가, activation(= 처리되는 데이터)도 줄일 것인가

2. scale/zero point를 미리 정할 것인가, 실행 중에 계산할 것인가

3. 학습 없이 바꿀 것인가, 다시 학습하면서 적응시킬 것인가

4. CPU, GPU, NPU, Apple Neural Engine 중 어디에서 돌릴 것인가

• int 양자화는 대충 이런 변환입니다.

float 값 ≈ scale × (int 값 - zero_point)

예를 들어 원래 activation 값이 -3.2 ~ 5.7 범위에 있다면, 이 범위를 INT8의 -128 ~ 127 또는 UINT8의 0 ~ 255 공간에 대응시킵니다.
즉 zero_point로 위치를 만들고 scale로 size를 만든다. 이 때 일반적으로 weight에서는 channel 단위로 서로 다른 scale, zero_point를 가짐.
 
 
하지만 여기서 문제는 어떻게 이 float 값의 범위를 미리 알고 결정하지?(특히 activation)
→ dynamic quantization과 static quantization이 갈립니다.
• dynamic quantization : 사전에는 weight만 양자화 한 다음, inference 단계에서 input tensor를 보고 그 값 범위로 양자화 한다.
• static quantization : 사전에 dataset 일부를 이용해 tensor 범위를 계산한 뒤 scale, zero_point를 고정해둔다.

• 양자화 방법도 학습 후 양자화를 적용(PTQ)시키는 방법이 있고 학습 단계에서 미리 양자화를 고려해서 학습(QAT)하는 방법이 있다. - DeepSeek V4에서는 QAT

• Weight를 양자화할지 Activation을 양자화할지에 따라 보통 W8A8, W4A16과 같이 표기한다.
어짜피 하는거 둘다 하지않을 이유는?
1. 학습 후 weight는 고정되기 때문에 양자화가 더 쉽다. 하지만 activation은 계속 바뀌니 양자화가 더 어렵다.
→ outlier 때문에 LLM에서 양자화가 더 어렵다.
작은 값들: -1.2, 0.3, 1.5
큰 outlier: 31.0
큰 outlier까지 살리려고 scale을 크게 잡으면, 대부분의 작은 값들은 INT8에서 너무 거칠게 표현된다
2. activation이 해상도가 높은게 정확도 유지가 더 잘된다. 어쨌든 다음 레이어로, output으로 전달되는 값 자체에 손실이 더 적다는거니까
3. 목적이 약간 다르다.
weight를 양자화하는건 주로 모델 크기를 줄이고 메모리 점유를 줄이기 위함이고, activation을 양자화하는건 메모리 이득은 크지 않지만 연산 속도를 높이기 위함이다.
그러니 모델을 줄이는게 중요한 경우에는 weight를 양자화하고 정확도를 유지하기 위해 activation은 그대로 둔다.
 
그래서 모델을 올리는게 문제인 LLM에서는 로컬 LLM을 위해 GGUF, GPTQ, AWQ 같은 quantized model들이 보통 W4A16, W4A16 비슷한 형태를 많이 쓴다.
 
4. 대신 W8A8과 같이 둘다 int8로 만들게 되면, 진짜 INT8 연산만 하면되니 최적화된 matmul method를 쓸 수 있기 때문에 연산속도 이득이 크다.

Inference Framework

모델 포맷:
- ONNX - 범용적인 포맷
- GGUF - LLM 전용
- .mlmodel / .mlpackage - Core ML 용 포맷
- .tflite - tensorRT용 포맷

Runtime / Framework:
- ONNX Runtime
- llama.cpp
- Core ML - Apple 생태계에서 모델을 실행하고 싶은 경우
- LiteRT
- TensorRT
- OpenVINO

ONNX - 범용적인 모델 파일 포맷

PyTorch / TensorFlow에서 만든 모델
→ ONNX라는 공통 그래프 표현으로 export
→ ONNX Runtime, TensorRT, OpenVINO, DirectML 같은 실행 엔진에서 돌림

Python 코드
nn.Module class 구조
autograd graph
optimizer
loss function
training-only branch
dropout / batchnorm training behavior
custom Python control flow

입력 tensor가 들어오면
어떤 연산들을 어떤 순서로 하고
어떤 weight를 써서
출력 tensor를 만들 것인가

• GGUF : LLM용 quantized 모델 파일 포맷
quantized LLM을 로컬에서 빠르게 로드/실행하기 위한 파일 포맷

모델 graph를 어디까지 TensorRT로 넘길지
어떤 precision을 쓸지
dynamic shape를 어떻게 처리할지
어떤 op는 fallback할지
batch size / sequence length / memory layout을 어떻게 고정할지

결론

500M도 메모리에 올라간다
= 실행 가능

99M이면 훨씬 낫다
= 다운로드, 로딩, latency, 배터리, 브라우저/모바일 배포가 쉬워진다

99M params  × 2 bytes ≈ 198MB
500M params × 2 bytes ≈ 1GB

• Nvidia Compact Language Models via Pruning and Knowledge Distillation
arxiv.org
https://arxiv.org/pdf/2407.14679

• Qwen Exploring the Pruning and Distillation in Large MoE Model Pre-training
SlimQwen: Exploring the Pruning and Distillation in Large MoE...
Structured pruning and knowledge distillation (KD) are typical techniques for compressing large language models, but it remains unclear how they should be applied at pretraining scale, especially...

https://arxiv.org/abs/2605.08738