ONNX (Open Neural Network Exchange) là định dạng mở để lưu một mô hình học máy dưới dạng đồ thị tính toán, kèm trọng số, tách khỏi framework đã huấn luyện ra nó. Mô hình ONNX được chạy bằng ONNX Runtime, một engine suy luận viết bằng C++ có provider riêng cho CPU, GPU và nhiều phần cứng khác. Bài này lấy một mô hình embedding đa ngôn ngữ đang dùng cho tìm kiếm tiếng Việt, xuất sang ONNX, tối ưu đồ thị, lượng tử hoá int8 và đo xem trên CPU được gì, mất gì.
Bài nối tiếp Phục vụ mô hình AI bằng FastAPI: batching, timeout và giám sát: khi dịch vụ đã có batching và giám sát, câu hỏi tiếp theo là làm mỗi lượt suy luận rẻ hơn. Mô hình chính là paraphrase-multilingual-MiniLM-L12-v2 đã dùng ở Embedding và semantic search. Mọi số trong bài là output thật trên Apple M4, 16 GB, Python 3.12, torch==2.8.0, onnxruntime==1.23.2, mã nguồn ở code/s7-48. Máy không có GPU NVIDIA nên bài chỉ nói về CPU; phần GPU (TensorRT, CUDA provider) không chạy trong bài.
Vì sao cần tối ưu suy luận trên CPU
Rất nhiều khối lượng thực tế ở doanh nghiệp là mô hình nhỏ chạy đều đặn: embedding cho tìm kiếm, phân loại ticket, chấm điểm rủi ro, reranker. Những mô hình này vừa trong RAM của một máy chủ thường, và chi phí chính là độ trễ cùng số CPU phải thuê. Với mô hình cỡ này, tuỳ phần cứng và tải, có thể đủ để không cần GPU. Trong bài, bản lượng tử hoá nhỏ hơn bốn lần và nhanh hơn khoảng hai lần ở một câu trên Apple M4; chưa đo trên VPS x86 nên đừng suy ra con số đó cho máy của bạn.
Có ba đòn bẩy khác nhau: đổi runtime (PyTorch sang ONNX Runtime, có tối ưu đồ thị), giảm độ chính xác số (float32 sang int8) và chỉnh cách chạy (số luồng, độ dài đệm). Bài đo từng đòn bẩy riêng để thấy cái nào thực sự có tác dụng.
ONNX là gì và nằm ở đâu trong chuỗi triển khai
Một file ONNX là một protobuf gồm: danh sách input/output có kiểu và hình dạng, danh sách node (mỗi node là một toán tử như MatMul, Add, Softmax), danh sách initializer (trọng số) và các opset cho biết phiên bản đặc tả toán tử. Vì đồ thị đã được đóng băng, runtime có thể phân tích toàn bộ để gộp node và chọn kernel tốt nhất cho CPU, điều mà vòng lặp Python của PyTorch eager không làm được.
Đánh đổi của ONNX: thêm một bước chuyển đổi dễ vỡ (toán tử lạ, điều khiển luồng động, tokenizer vẫn nằm ngoài đồ thị), và bạn phải kiểm chứng rằng kết quả không đổi. Phần còn lại của bài là quy trình kiểm chứng đó.
Bước 1: xuất mô hình sang ONNX
Ta bọc mô hình thành một Embedder gồm cả mean pooling và chuẩn hoá L2, để file ONNX nhận input_ids và attention_mask rồi trả thẳng vector 384 chiều. Đưa phần hậu xử lý vào đồ thị nghĩa là không còn đoạn NumPy rời nào phải giữ cho khớp với lúc huấn luyện.
class Embedder(torch.nn.Module):
def __init__(self, base):
super().__init__(); self.base = base
def forward(self, input_ids, attention_mask):
h = self.base(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
m = attention_mask.unsqueeze(-1).to(h.dtype)
v = (h * m).sum(1) / m.sum(1).clamp(min=1e-9)
return torch.nn.functional.normalize(v, dim=-1)
torch.onnx.export(
model, (enc["input_ids"], enc["attention_mask"]), "embedder_fp32.onnx",
input_names=["input_ids", "attention_mask"], output_names=["embedding"],
dynamic_axes={"input_ids": {0: "batch", 1: "seq"},
"attention_mask": {0: "batch", 1: "seq"},
"embedding": {0: "batch"}},
opset_version=17, dynamo=False)
onnx.checker.check_model(onnx.load("embedder_fp32.onnx"))Hai chi tiết quan trọng. Thứ nhất, trục động: nếu không khai báo dynamic_axes, file ONNX khoá cứng batch và độ dài của ví dụ mẫu, và mọi câu có độ dài khác sẽ lỗi hoặc cho kết quả sai. Thứ hai, chọn bộ xuất: bài dùng PyTorch 2.8, nơi bộ xuất dựa trên TorchScript là mặc định. Theo tài liệu PyTorch, từ bản 2.9 tham số dynamo mặc định là True (bộ xuất dựa trên torch.export) và dynamic_shapes dành cho đường này, còn dynamic_axes dùng với dynamo=False. Vì vậy trong bài tôi ghi tường minh dynamo=False; nếu bạn nâng PyTorch lên 2.9 trở đi mà không đặt tham số này, đường xuất sẽ khác, hãy chạy lại bước kiểm chứng ở dưới. Output thật:
opset: [('ai.onnx', 17)] ir_version: 8
so node: 1038 | initializer: 197
top op: [('Constant', 308), ('Add', 122), ('MatMul', 96), ('Unsqueeze', 70), ('Mul', 51), ('Concat', 50), ('Reshape', 50), ('Transpose', 48)]
kich thuoc: 470.2 MBCon số đáng chú ý nằm ở trọng số. Script step6_dynamic_check.py đếm được 117.653.760 tham số, trong đó riêng bảng embedding từ vựng (word_embeddings, 250.037 token × 384 chiều) chiếm 96.014.208, tức 81,6%. Phần còn lại, nơi chứa các phép MatMul của 12 lớp transformer, chỉ khoảng 21,6 triệu. Bảng tra từ vựng gần như không tốn phép tính nhưng chiếm phần lớn dung lượng file, nên việc lượng tử hoá có đụng tới nó hay không (xem bước 3) quyết định kích thước file.
Bước 2: tối ưu đồ thị trong ONNX Runtime
Theo tài liệu ONNX Runtime, có ba nhóm tối ưu đồ thị, mỗi mức bao gồm các mức trước: basic (gấp hằng số, loại node thừa), extended (gộp node phức tạp hơn như MatMul+Add, LayerNorm, Attention, chỉ cho CPU/CUDA/ROCm) và layout (đổi bố cục dữ liệu, chỉ cho CPU). ORT_ENABLE_ALL bật cả ba và là mặc định. Ta ghi đồ thị sau tối ưu ra file ở từng mức để đếm node:
so = ort.SessionOptions()
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
so.optimized_model_filepath = "embedder_fp32_extended.onnx"
ort.InferenceSession("embedder_fp32.onnx", so, providers=["CPUExecutionProvider"])onnxruntime 1.23.2
disable node= 729 470.2 MB
basic node= 634 470.2 MB
extended node= 574 470.2 MB
all node= 574 470.2 MBSố node giảm từ 1.038 (file xuất ra) xuống 574, gần một nửa, nhưng kích thước file không đổi: tối ưu đồ thị chỉ sắp xếp lại phép tính, không đụng tới trọng số. Ngay ở mức disable, số node đã là 729 chứ không phải 1.038. Mức extended và all ra cùng 574 node vì mô hình chỉ có toán tử không phụ thuộc bố cục NCHWc (loại bố cục dành cho tích chập).
Khi ghi mô hình đã tối ưu ra đĩa, ONNX Runtime in cảnh báo thật sau đây, ở mức lớn hơn extended:
Serializing optimized model with Graph Optimization level greater than ORT_ENABLE_EXTENDED and the NchwcTransformer enabled. The generated model may contain hardware specific optimizations, and should only be used in the same environment the model was optimized in.Hệ quả thực tế: đừng ghi file đã tối ưu mức all rồi mang sang máy khác. Cách an toàn là phân phối file ONNX gốc và để mỗi máy tự tối ưu khi nạp phiên (mất vài trăm mili giây, xem bảng thời gian nạp bên dưới).
Bước 3: lượng tử hoá động int8
Lượng tử hoá thay trọng số float32 (4 byte) bằng số nguyên 8 bit kèm hệ số tỉ lệ. Với ONNX Runtime có hai kiểu. Lượng tử hoá động chỉ lượng tử trọng số lúc chuyển đổi và tính hệ số tỉ lệ của activation ngay khi suy luận. Lượng tử hoá tĩnh dùng dữ liệu hiệu chuẩn để cố định hệ số đó trước. Theo tài liệu lượng tử hoá của ONNX Runtime, nên dùng kiểu động cho transformer và RNN, kiểu tĩnh cho CNN. Kiểu động không cần dữ liệu hiệu chuẩn nên là điểm xuất phát gọn nhất:
from onnxruntime.quantization import quantize_dynamic, QuantType
quantize_dynamic("embedder_fp32.onnx", "embedder_int8.onnx", weight_type=QuantType.QInt8)int8: node=1305 118.1 MB
op lượng tử: [('DynamicQuantizeLinear', 48), ('MatMulInteger', 72)]File nhỏ đi đúng 4 lần: 470,2 MB còn 118,1 MB. Kiểm tra bằng onnx cho thấy ngay cả bảng embedding từ vựng cũng đã thành số nguyên 8 bit (output thật của step6_dynamic_check.py: initializer lớn nhất là base.embeddings.word_embeddings.weight_quantized, kiểu UINT8, hình dạng [250037, 384]), nên toàn bộ trọng số lớn đều co lại. Đồ thị lại có nhiều node hơn (1.305 so với 1.038) vì mỗi MatMul được chèn thêm node lượng tử hoá activation (DynamicQuantizeLinear) và đổi thành MatMulInteger. Công cụ cảnh báo nên chạy bước tiền xử lý trước; bài bỏ qua bước đó nên độ chính xác có thể chưa tối đa.
Bước 4: kiểm chứng ONNX cho kết quả giống PyTorch
Đo tốc độ trước khi đo độ đúng là cách nhanh nhất để tối ưu ra một mô hình hỏng. Script step3_parity.py so ba biến thể trên cùng 12 tài liệu và 8 truy vấn tiếng Việt (có đáp án đúng đã gán tay), đo độ lệch tuyệt đối lớn nhất của vector, cosine nhỏ nhất giữa vector gốc và vector biến thể, rồi top-1 và top-3 truy hồi:
mô hình max|Δ| doc cos min doc top-1 top-3
torch 0.00e+00 1.000000 0.75 1.00
onnx fp32 2.51e-07 1.000000 0.75 1.00
onnx int8 3.53e-02 0.977287 0.75 1.00
fp32: câu riêng vs trong batch có padding, max|Δ| = 0.00e+00
int8: câu riêng vs trong batch có padding, max|Δ| = 1.35e-02
torch điểm cosine cặp đúng: 0.564 0.244 0.853 0.745 0.605 0.630 0.359 0.806
onnx fp32 điểm cosine cặp đúng: 0.564 0.244 0.853 0.745 0.605 0.630 0.359 0.806
onnx int8 điểm cosine cặp đúng: 0.549 0.247 0.847 0.729 0.577 0.624 0.343 0.798Đọc kết quả:
- ONNX fp32 gần như giống hệt PyTorch: độ lệch lớn nhất 2,5×10-7, ở mức sai số làm tròn float32.
- int8 lệch rõ hơn: độ lệch lớn nhất mỗi chiều 0,035, vector lệch nhiều nhất vẫn có cosine 0,977 so với bản gốc, và điểm cosine của cặp đúng giảm khoảng 0,006 đến 0,028 trên 7 trong 8 truy vấn (một truy vấn tăng 0,003).
- Thứ hạng truy hồi không đổi: cả ba cho top-1 = 0,75 và top-3 = 1,00. Nhưng đừng đọc quá nhiều: 8 truy vấn quá ít để kết luận về chất lượng. Một bộ kiểm thử thật cần vài trăm truy vấn lấy từ log của chính hệ thống, như đã làm ở bài đánh giá RAG.
- Kết quả int8 phụ thuộc vào batch: cùng một câu cho vector khác nhau (lệch tới 0,0135) khi đứng riêng và khi nằm trong batch 12 câu có đệm, trong khi fp32 lệch đúng 0. Điều này khớp với cơ chế lượng tử hoá động: hệ số tỉ lệ của activation được tính trên cả tensor lúc suy luận, nên câu bên cạnh ảnh hưởng tới câu của bạn. Tôi chưa cô lập nguyên nhân này bằng thí nghiệm riêng, nhưng hệ quả thì có thật: một hệ thống lưu vector đã tính sẵn nên dùng cố định một kiểu suy luận, và không nên trộn vector từ hai biến thể.
Bước 5: đo độ trễ trên CPU
Đo độ trễ có nhiều cái bẫy hơn vẻ ngoài. Máy tôi dùng chung với nhiều tiến trình khác, và lần đầu tôi chạy lúc load average vượt 20 trên 10 nhân: số tuyệt đối khi đó cao gấp 3 đến 7 lần và thứ tự giữa các cấu hình cũng đổi. Vì vậy step4_bench.py chạy xen kẽ từng cấu hình trong cùng một vòng lặp, có khởi động trước, torch.set_num_threads(2) và intra_op_num_threads = 2 cho cả hai bên, rồi lặp cả khối đo 5 lượt và báo trung vị cùng khoảng dao động của các lượt. Tôi đợi load average xuống dưới 9 mới chạy lần cuối. Output thật:
batch=1 (5 lượt x 40 vòng)
PyTorch fp32 min 5.1 p50 8.5 (lượt: 5.4- 9.3) p95 37.5 ms | so với PyTorch: 1.00x (1.00-1.00)
ORT fp32 (tắt tối ưu) min 6.1 p50 7.1 (lượt: 6.5- 10.5) p95 12.4 ms | so với PyTorch: 0.84x (0.82-1.23)
ORT fp32 (ALL) min 6.1 p50 7.5 (lượt: 7.4- 9.7) p95 14.6 ms | so với PyTorch: 0.90x (0.72-1.14)
ORT int8 động min 2.4 p50 3.7 (lượt: 3.6- 4.3) p95 8.1 ms | so với PyTorch: 2.02x (1.49-2.37)
batch=12 (5 lượt x 15 vòng)
PyTorch fp32 min 19.8 p50 21.1 (lượt: 18.1- 22.8) p95 33.4 ms | so với PyTorch: 1.00x (1.00-1.00)
ORT fp32 (tắt tối ưu) min 59.6 p50 63.9 (lượt: 56.5- 70.9) p95 78.7 ms | so với PyTorch: 0.33x (0.32-0.33)
ORT fp32 (ALL) min 58.2 p50 61.5 (lượt: 54.5- 64.1) p95 97.2 ms | so với PyTorch: 0.34x (0.33-0.36)
ORT int8 động min 21.0 p50 22.5 (lượt: 19.6- 24.1) p95 33.5 ms | so với PyTorch: 0.94x (0.91-1.00)
load average đầu/cuối: 8.8 19.1 21.1 -> 8.3 18.1 20.7Cột "so với PyTorch" là tỉ lệ p50 (lớn hơn 1 là nhanh hơn PyTorch); trong ngoặc là khoảng dao động giữa 5 lượt. Các nhận xét từ số đo:
- Với một câu, ONNX Runtime fp32 không nhanh hơn PyTorch một cách ổn định: p50 7,5 ms so với 8,5 ms, tỉ lệ 0,90x với khoảng dao động 0,72 đến 1,14 nằm cả hai phía của 1. Một người khác chạy lại trên máy ít tải hơn còn thấy fp32 chậm hơn PyTorch. Kết luận an toàn là hai bên ngang nhau trong biên nhiễu; int8 là bản nhanh nhất, p50 3,7 ms, khoảng 2 lần (1,5 đến 2,4 lần giữa các lượt).
- Với batch 12 câu, ONNX Runtime fp32 chậm hơn PyTorch khoảng ba lần (61,5 ms so với 21,1 ms, tỉ lệ 0,34x). Kết quả này lặp lại ở mọi lần đo và mọi lượt, kể cả khi máy bận. Không có đòn bẩy nào tự động đúng: trên máy này, với batch lớn, đường fp32 của ONNX Runtime thua PyTorch. Tôi chưa tìm nguyên nhân nên chỉ báo kết quả đo; trên máy chủ x86 kết quả có thể khác hẳn.
- int8 ở batch 12 ngang PyTorch (22,5 so với 21,1 ms, tỉ lệ 0,94x, khoảng 0,91 đến 1,00): không còn lợi thế tốc độ, chỉ còn lợi thế kích thước. Lợi thế rõ nhất của int8 nằm ở trường hợp một câu mỗi lượt.
- Tối ưu đồ thị ALL so với tắt: không khác trong biên nhiễu ở batch 1 (7,5 so với 7,1 ms) và batch 12 (61,5 so với 63,9 ms).
Về bộ nhớ và thời gian khởi động, step5_size_load.py đo ba lần nạp liên tiếp mỗi file:
embedder_fp32.onnx 470.2 MB nạp phiên: 152 / 107 / 104 ms
embedder_int8.onnx 118.1 MB nạp phiên: 80 / 69 / 61 msThời gian nạp dao động theo bộ nhớ đệm của hệ điều hành và tải máy, nên tôi chỉ rút ra kết luận về kích thước: với pod Kubernetes có nhiều bản sao trên node nhỏ (xem requests và limits), file 118 MB tiết kiệm RAM và băng thông kéo image hơn hẳn file 470 MB.
Dùng ONNX Runtime trong dịch vụ thật
Một phiên (InferenceSession) an toàn để dùng chung giữa nhiều luồng, nên tạo một lần khi khởi động dịch vụ và dùng lại, không tạo mỗi request. Đoạn mã tối thiểu cho hàm embed thay vào endpoint FastAPI ở bài trước:
import onnxruntime as ort
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained(MODEL_ID)
so = ort.SessionOptions()
so.intra_op_num_threads = 2 # khớp số vCPU được cấp cho pod
sess = ort.InferenceSession("embedder_int8.onnx", so, providers=["CPUExecutionProvider"])
def embed(texts: list[str]):
enc = tok(texts, padding=True, truncation=True, max_length=128, return_tensors="np")
return sess.run(None, {"input_ids": enc["input_ids"],
"attention_mask": enc["attention_mask"]})[0]Theo tài liệu về luồng của ONNX Runtime, nếu không đặt intra_op_num_threads thì mặc định bằng số lõi vật lý, và đặt tường minh sẽ tắt việc gắn luồng vào lõi. Trong container, nên đặt tường minh theo CPU limit của pod. Ngoài ra hãy cắt độ dài (truncation=True, max_length=128): độ dài đệm là tham số tốc độ lớn mà nhiều người quên, vì chi phí attention tăng theo bình phương độ dài.
Khi nào không nên dùng ONNX hay lượng tử hoá
- Khi chưa có bộ kiểm thử chất lượng. Không có tập truy vấn có đáp án thì không biết int8 làm hỏng gì. Hãy có bộ kiểm thử trước, như đã làm ở bài đánh giá LLM.
- Khi tải nhỏ. 4 ms hay 15 ms cho một câu không quyết định trải nghiệm nếu mỗi giây chỉ có vài request.
- Khi mô hình có điều khiển luồng động hoặc toán tử riêng, xuất ONNX có thể lỗi hoặc cho đồ thị khác. LLM sinh văn bản dùng KV cache, chạy lặp từng token, cần công cụ khác (llama.cpp ở bài chạy LLM local, hoặc vLLM ở bài sau).
- Khi bạn cần GPU. Bài chỉ đo CPU. Chạy ONNX Runtime trên GPU (CUDA hoặc TensorRT provider) cần phần cứng NVIDIA mà máy thử nghiệm không có; đoạn này cần GPU, chưa chạy trong bài.
Lỗi thường gặp
- Quên trục động: mô hình chỉ chạy đúng với batch và độ dài của ví dụ mẫu. Hãy chạy thử với độ dài khác (kiểm tra ở bài này bằng
step6_dynamic_check.py: câu 4 token và câu 62 token đều ra vector 384 chiều). - Đổi phiên bản PyTorch mà không kiểm lại: từ 2.9, bộ xuất mặc định đổi và cách khai báo trục động cũng đổi (
dynamic_shapesthay chodynamic_axes); ghim phiên bản trongrequirements.txtvà chạy lại bước kiểm chứng.
Checklist tối ưu suy luận trên CPU
- Xuất ONNX kèm tiền/hậu xử lý có thể đưa vào đồ thị (pooling, chuẩn hoá), khai báo trục động, chạy
onnx.checker. - So sánh với PyTorch bằng độ lệch số và, quan trọng hơn, bằng chỉ số nghiệp vụ trên bộ kiểm thử riêng.
- Thử lượng tử hoá động int8 cho transformer; xem kích thước, độ trễ và chất lượng cùng lúc.
- Đặt
intra_op_num_threadstheo CPU limit, tạo phiên một lần, cắt độ dài đầu vào. - Đo p50, p95 trên batch thật và phần cứng thật: x86 và ARM cho kết quả khác nhau.
Cài đặt: python3.12 -m venv .venv && .venv/bin/pip install -r requirements.txt, rồi chạy export_onnx.py, optimize_quantize.py, step3_parity.py, step4_bench.py và step5_size_load.py trong code/s7-48; các file *_output.txt là output đã dán. Tài liệu gốc: lượng tử hoá ONNX Runtime, tối ưu đồ thị và bộ xuất ONNX của PyTorch.
Bài tiếp theo
Phục vụ LLM tự host với vLLM: batching liên tục và KV cache. Mô hình embedding nhỏ chạy tốt trên CPU, nhưng mô hình sinh văn bản có bài toán khác hẳn: bộ nhớ KV cache và lập lịch yêu cầu quyết định thông lượng.
