beautifulsoups.dev
Published on

파이썬 ONNX Runtime로 BERT 4배 경량·2배 가속 튜닝

Authors

서빙 비용을 줄이려면 결국 모델 크기와 지연시간(latency) 을 동시에 잡아야 합니다. BERT 계열은 정확도는 좋지만 파라미터 수가 많고, 특히 CPU 추론에서 병목이 자주 발생합니다. 이 글에서는 파이썬 환경에서 ONNX Runtime(ORT) 으로 BERT를 내보낸 뒤, 그래프 최적화 + 동적 양자화 + 실행 옵션 튜닝 을 조합해

  • 모델 파일 크기: 대략 4배 경량화(FP32 -> INT8)
  • CPU 추론 속도: 대략 2배 가속(환경에 따라 편차)

을 노리는 실전 레시피를 다룹니다.

전제: 예시는 bert-base-uncased 기반의 문장 분류(Sequence Classification)로 설명합니다. 토크나이저는 Hugging Face를 사용하고, 추론 엔진만 ORT로 교체합니다.

왜 ONNX Runtime인가

ONNX Runtime은 단순히 모델 포맷 변환 도구가 아니라, 추론 최적화 엔진 입니다.

  • 그래프 최적화(연산 병합, 상수 폴딩, 불필요 노드 제거)
  • Execution Provider(EP)로 하드웨어별 최적화 경로 선택
  • INT8 양자화 도구 내장
  • 세션 옵션/스레딩/메모리 패턴 등 런타임 튜닝 포인트 제공

특히 CPU에서 BERT를 빠르게 돌리고 싶다면 PyTorch eager 실행보다 ORT가 유리한 경우가 많습니다.

목표와 측정 기준부터 고정하기

튜닝을 시작하기 전에 무엇을 측정할지 정해야 합니다.

  • p50, p95 latency
  • throughput(QPS)
  • 모델 파일 크기(MB)
  • 정확도 하락 허용치(예: -0.3% 이내)

서빙 환경에서 병목이 애매하면, 운영 관점에서 리소스/오토스케일링까지 같이 보게 됩니다. 예를 들어 HPA가 기대만큼 줄지 않아 비용이 새는 경우도 흔합니다. 이런 케이스는 Kubernetes HPA가 0으로 안 줄 때 - PDB·윈도우·종료 같은 체크리스트가 같이 도움이 됩니다.

준비: 패키지 설치

아래는 CPU 기준입니다. GPU는 뒤에서 별도 언급합니다.

pip install -U transformers torch onnx onnxruntime onnxruntime-tools

환경에 따라 onnxruntime 대신 onnxruntime-gpu 를 선택합니다.

1) PyTorch BERT를 ONNX로 Export

핵심은 동적 축(dynamic axes) 을 제대로 지정해 배치/시퀀스 길이가 바뀌어도 재사용 가능한 그래프를 만드는 것입니다.

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_id = "bert-base-uncased"
num_labels = 2

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=num_labels)
model.eval()

text = "onnx runtime makes bert faster"
inputs = tokenizer(text, return_tensors="pt")

# ONNX export
onnx_path = "bert_cls_fp32.onnx"

dynamic_axes = {
    "input_ids": {0: "batch", 1: "seq"},
    "attention_mask": {0: "batch", 1: "seq"},
    "logits": {0: "batch"}
}

with torch.no_grad():
    torch.onnx.export(
        model,
        (inputs["input_ids"], inputs["attention_mask"]),
        onnx_path,
        input_names=["input_ids", "attention_mask"],
        output_names=["logits"],
        dynamic_axes=dynamic_axes,
        opset_version=17,
        do_constant_folding=True,
    )

print("saved:", onnx_path)

Export 체크 포인트

  • opset_version 는 가능하면 최신(예: 17)을 권장합니다.
  • 모델에 따라 token_type_ids 가 필요할 수 있습니다. 토크나이저 출력에 해당 키가 있으면 입력으로 함께 export하세요.

2) ONNX Runtime 세션 생성과 기본 최적화

ORT는 세션 옵션에서 그래프 최적화 레벨을 올릴 수 있습니다.

import onnxruntime as ort

sess_options = ort.SessionOptions()
# ORT_ENABLE_ALL: 가장 공격적인 그래프 최적화
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

# CPU 스레딩 튜닝(환경에 맞게 조정)
sess_options.intra_op_num_threads = 4
sess_options.inter_op_num_threads = 1

providers = ["CPUExecutionProvider"]

session = ort.InferenceSession("bert_cls_fp32.onnx", sess_options=sess_options, providers=providers)
print("inputs:", [i.name for i in session.get_inputs()])
print("outputs:", [o.name for o in session.get_outputs()])

스레딩 튜닝 팁

  • intra_op_num_threads: 한 연산 내부 병렬성. CPU 코어 수에 맞춰 올립니다.
  • inter_op_num_threads: 연산 간 병렬성. 보통 BERT는 1 이 유리한 경우가 많습니다.
  • 컨테이너/쿠버네티스 환경에서는 CPU limit에 맞춰 조절하지 않으면 오히려 느려질 수 있습니다.

3) 동적 양자화로 4배 경량화 노리기

BERT는 Linear 계층이 많아 동적 양자화(dynamic quantization) 가 잘 먹습니다.

  • 가중치(weight): INT8로 저장
  • 활성값(activation): 런타임에 동적으로 스케일링

정확도 하락이 비교적 적고, 특히 CPU에서 이득이 큽니다.

from onnxruntime.quantization import quantize_dynamic, QuantType

fp32_model = "bert_cls_fp32.onnx"
int8_model = "bert_cls_int8.onnx"

quantize_dynamic(
    model_input=fp32_model,
    model_output=int8_model,
    weight_type=QuantType.QInt8,
)

print("saved:", int8_model)

정말 4배 줄었는지 확인

대부분 FP32 -> INT8이면 가중치 저장 공간이 1/4 로 줄어드는 것이 정상입니다(메타데이터/일부 노드로 인해 정확히 4배는 아닐 수 있음).

import os

def size_mb(p):
    return os.path.getsize(p) / 1024 / 1024

print("fp32 MB:", size_mb(fp32_model))
print("int8 MB:", size_mb(int8_model))

4) 추론 코드: ORT로 교체

토크나이징은 그대로 두고, 모델 forward만 ORT로 바꿉니다.

import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

sess_options = ort.SessionOptions()
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
sess_options.intra_op_num_threads = 4
sess_options.inter_op_num_threads = 1

session = ort.InferenceSession("bert_cls_int8.onnx", sess_options=sess_options, providers=["CPUExecutionProvider"])

def predict(texts):
    enc = tokenizer(texts, padding=True, truncation=True, return_tensors="np")

    ort_inputs = {
        "input_ids": enc["input_ids"].astype(np.int64),
        "attention_mask": enc["attention_mask"].astype(np.int64),
    }

    logits = session.run(["logits"], ort_inputs)[0]
    probs = softmax(logits)
    return probs

def softmax(x):
    x = x - np.max(x, axis=-1, keepdims=True)
    e = np.exp(x)
    return e / np.sum(e, axis=-1, keepdims=True)

print(predict(["this is great", "this is bad"]))

5) 벤치마크: 2배 가속을 재현하는 방법

성능 비교는 반드시 워밍업 + 반복 측정 으로 합니다.

import time
import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer

texts = ["onnx runtime benchmark"] * 32

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
enc = tokenizer(texts, padding=True, truncation=True, return_tensors="np")
ort_inputs = {
    "input_ids": enc["input_ids"].astype(np.int64),
    "attention_mask": enc["attention_mask"].astype(np.int64),
}

def bench(model_path, threads=4, iters=200, warmup=20):
    so = ort.SessionOptions()
    so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    so.intra_op_num_threads = threads
    so.inter_op_num_threads = 1

    sess = ort.InferenceSession(model_path, sess_options=so, providers=["CPUExecutionProvider"])

    for _ in range(warmup):
        _ = sess.run(["logits"], ort_inputs)

    t0 = time.perf_counter()
    for _ in range(iters):
        _ = sess.run(["logits"], ort_inputs)
    t1 = time.perf_counter()

    ms = (t1 - t0) * 1000 / iters
    return ms

print("fp32 ms:", bench("bert_cls_fp32.onnx"))
print("int8 ms:", bench("bert_cls_int8.onnx"))

2배 가속이 안 나오면 체크할 것

  • 배치 크기가 너무 작으면(예: batch=1) 토크나이저/파이썬 오버헤드가 지배적일 수 있습니다.
  • intra_op_num_threads 를 과하게 키우면 컨텍스트 스위칭으로 느려질 수 있습니다.
  • CPU가 AVX2/AVX512를 지원하는지, 컨테이너가 해당 ISA를 활용하는지 확인하세요.
  • 입력이 짧으면 BERT 자체 연산량이 적어 이득이 제한됩니다.

6) 더 빠르게: I/O Binding으로 복사 비용 줄이기

배치가 커지거나 호출이 빈번하면, 입력/출력 텐서 복사 비용이 무시 못 할 때가 있습니다. ORT의 I/O Binding 은 메모리 바인딩을 통해 불필요한 복사를 줄이는 데 도움됩니다(특히 GPU에서 체감이 큼).

CPU에서도 상황에 따라 이득이 날 수 있어, p95 튜닝 단계에서 시도할 가치가 있습니다.

import numpy as np
import onnxruntime as ort

so = ort.SessionOptions()
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
sess = ort.InferenceSession("bert_cls_int8.onnx", sess_options=so, providers=["CPUExecutionProvider"])

# 더미 입력
input_ids = np.ones((8, 128), dtype=np.int64)
attn = np.ones((8, 128), dtype=np.int64)

io = sess.io_binding()
io.bind_cpu_input("input_ids", input_ids)
io.bind_cpu_input("attention_mask", attn)
# 출력도 바인딩(이름은 모델에 맞게)
io.bind_output("logits")

sess.run_with_iobinding(io)
outputs = io.copy_outputs_to_cpu()
logits = outputs[0]
print(logits.shape)

7) GPU라면: TensorRT/CUDA EP와 정밀도 전략

GPU에서 ORT를 쓰는 이유는 보통 CUDA EP 또는 TensorRT EP 때문입니다.

  • FP16: 정확도 유지하면서 속도 이득
  • INT8: 캘리브레이션 필요, 정확도/속도 트레이드오프 큼

다만 이 글의 목표(4배 경량, 2배 가속)는 CPU 동적 양자화 조합에서 가장 재현성이 좋습니다. GPU는 환경 변수(드라이버, 커널, TensorRT 버전, 엔진 빌드 캐시) 영향이 커서, 운영에서는 재현성과 장애 대응 절차가 중요합니다.

8) 운영 배포 관점 체크리스트

모델 튜닝이 끝나도 운영에서 성능이 흔들리면 의미가 없습니다.

  • 컨테이너 이미지 빌드 시 ORT 버전 고정, 재현 가능한 빌드
  • 스레드 수를 CPU limit에 맞춰 설정
  • 모델 다운로드/로딩 시간을 스타트업 단계에서 처리
  • 캐시 무효화로 매번 의존성을 다시 받는 문제 방지

이미지 빌드가 자주 느려지거나 레이어 캐시가 예상과 다르게 깨진다면 Docker BuildKit 캐시 무효화 원인·해결 8가지도 같이 점검해보세요.

또한 트래픽 변동이 큰 API라면, 성능 튜닝과 함께 오토스케일링이 원하는 방향으로 동작하는지 확인해야 합니다. 특히 scale-to-zero나 안정화 윈도우 설정은 비용에 직접 영향을 줍니다. 이 부분은 앞서 언급한 Kubernetes HPA가 0으로 안 줄 때 - PDB·윈도우·종료 글이 보완재가 됩니다.

9) 자주 겪는 함정과 해결책

함정 1: ONNX 입력 이름이 예상과 다름

Export 시 input_names 를 명시했는데도 모델에 따라 내부적으로 이름이 바뀌는 경우가 있습니다. session.get_inputs() 로 실제 이름을 확인하고 그대로 넣으세요.

함정 2: 토크나이저가 병목

BERT 추론이 빨라질수록 토크나이저 비용이 상대적으로 커집니다.

  • 배치 토크나이징 사용
  • padding=True 와 truncation=True 를 명시
  • 가능하면 토크나이저를 프로세스 시작 시 1회 로드

함정 3: 정확도 하락

동적 양자화는 대체로 안전하지만 태스크/데이터에 따라 하락이 생길 수 있습니다.

  • 양자화 전후로 validation 셋에서 지표 비교
  • 문제가 크면 FP16(가능 환경) 또는 부분 양자화(특정 노드 제외) 고려

결론: 가장 재현성 좋은 조합

CPU에서 BERT를 실전적으로 튜닝할 때, 재현성과 효과를 동시에 챙기려면 아래 조합이 가장 무난합니다.

  • ONNX export 시 dynamic axes 정확히 지정
  • ORT 그래프 최적화 ORT_ENABLE_ALL
  • 동적 양자화로 INT8 모델 생성(대개 4배 내외 경량화)
  • 스레딩 옵션을 CPU limit에 맞춰 튜닝
  • p95가 중요하면 I/O binding까지 시도

이 과정을 그대로 따라 하면, 많은 환경에서 모델 용량은 크게 줄고, CPU 추론은 의미 있게 빨라지는 결과를 얻을 수 있습니다.

Discuss on Twitter • View on GitHub