Quantization-Aware Training (QAT)

Huấn luyện có nhận biết lượng tử hóa (QAT) là kỹ thuật huấn luyện giúp chuẩn bị model cho suy luận độ chính xác thấp bằng cách mô phỏng lượng tử hóa trong quá trình huấn luyện. Model thường được huấn luyện và lưu trọng số dưới dạng số thực dấu phẩy động 32 bit (FP32); lượng tử hóa model chuyển chúng sang các định dạng như số nguyên 8 bit (INT8) để triển khai, giúp trọng số nhỏ hơn 4 lần và cho phép CPU, GPU cùng bộ xử lý thần kinh (NPU) sử dụng phép toán số nguyên tốc độ cao. Quá trình chuyển đổi làm tròn mọi giá trị và cắt các giá trị ngoại lai, gây tổn thất accuracy. QAT đưa phép làm tròn đó vào lượt truyền xuôi khi model vẫn đang học, nhờ vậy gradient descent tìm được các trọng số vẫn hoạt động tốt sau khi lượng tử hóa. QAT thường được chạy như một giai đoạn tinh chỉnh ngắn trên model đã huấn luyện. Phương pháp suy luận chỉ dùng phép toán số nguyên của Google dành cho CPU di động kết hợp suy luận INT8 với quy trình huấn luyện kiểu này; từ đó, QAT trở thành công cụ tiêu chuẩn để triển khai AI biên.

Cách hoạt động của huấn luyện có nhận biết lượng tử hóa#

QAT bổ sung các phép toán lượng tử hóa giả, còn gọi là cặp lượng tử hóa-khử lượng tử hóa (Q/DQ), vào trọng số và activation của model. Mỗi phép toán ánh xạ một giá trị float về mức INT8 gần nhất bằng scale và zero point, sau đó lập tức ánh xạ giá trị đó trở lại float; vì vậy, mạng vẫn tính toán bằng số dấu phẩy động nhưng nhìn thấy các giá trị đã làm tròn giống như engine INT8. Các scale được lấy từ bước hiệu chỉnh, ghi lại phạm vi thường gặp của mỗi tensor trên một số batch dữ liệu; một số phương pháp còn học các scale này trong quá trình huấn luyện.

Đạo hàm của phép làm tròn bằng 0 ở hầu hết mọi điểm, khiến quá trình học bị dừng. QAT dùng bộ ước lượng truyền thẳng (STE), xem bước làm tròn như hàm đồng nhất trong quá trình lan truyền ngược, nhờ vậy gradient truyền tới các trọng số FP32 bên dưới:

import torch

# FP32 weights that training keeps and updates
w = torch.tensor([0.3137, -0.7214, 0.0521, 1.1800], requires_grad=True)

# Symmetric INT8 fake quantization: snap each weight to one of 255 levels, then map back to float
scale = w.detach().abs().max().item() / 127
w_q = torch.fake_quantize_per_tensor_affine(w, scale, 0, -127, 127)
print(f"Fake-quantized: {[round(v, 4) for v in w_q.tolist()]}")

# The forward pass sees the rounded weights; the straight-through estimator still delivers gradients to w
x = torch.tensor([1.0, 2.0, 3.0, 4.0])
loss = (w_q * x).sum()
loss.backward()
print(f"Gradient on FP32 weights: {w.grad.tolist()}")

Lượt truyền xuôi sử dụng các trọng số đã làm tròn 0.3159, −0.7247, 0.0557 và 1.18; mỗi giá trị nằm trong phạm vi nửa bước lượng tử hóa (0.0046) so với giá trị FP32 tương ứng, nhưng gradient truyền tới trọng số FP32 vẫn chính xác là các giá trị đầu vào 1.0, 2.0, 3.0 và 4.0, như thể không có phép làm tròn nào diễn ra. Qua nhiều bước, optimizer dịch chuyển trọng số tới các vị trí mà sai số làm tròn gây ít tác động nhất; model đã huấn luyện được xuất sang INT8 cùng các scale đã học.

QAT và lượng tử hóa sau huấn luyện#

Lượng tử hóa sau huấn luyện (PTQ) chuyển đổi trực tiếp model FP32 đã hoàn tất, sử dụng vài trăm ảnh hiệu chỉnh để thiết lập các phạm vi. Quá trình này chỉ mất vài phút và không cần pipeline huấn luyện, nhưng trọng số không có cơ hội thích ứng với phép làm tròn, nên model nhạy cảm có thể mất accuracy đáng kể. QAT cần một lượt tinh chỉnh bằng toàn bộ dữ liệu huấn luyện, đổi lại có thể khôi phục phần lớn accuracy đã mất. Các nhóm thường thử PTQ trước và chỉ chuyển sang QAT khi mức giảm accuracy của INT8 quá lớn:

Ultralytics YOLO chạy QAT khi nhận quantize=8 trong quá trình huấn luyện: công cụ hiệu chỉnh lượng tử hóa giả INT8 thông qua NVIDIA Model Optimizer, tinh chỉnh model trong khi áp dụng quá trình này và lưu các phạm vi mà khi xuất sang ONNX và TensorRT sẽ được ghi thành node Q/DQ mà không cần dữ liệu hiệu chỉnh riêng. Trên engine TensorRT, với COCO val2017, PTQ khiến model YOLO26 cỡ medium giảm từ 0.527 xuống 0.470 mAP_50-95, trong khi QAT giữ mức 0.514; model nano chỉ giảm 0.010 với PTQ nên QAT không đem lại lợi ích nào.

from ultralytics import YOLO

# Bắt đầu với checkpoint YOLO26 FP32 đã huấn luyện
model = YOLO("yolo26n.pt")

# Tinh chỉnh với lượng tử hóa giả INT8 trong vòng lặp, dùng learning rate nhỏ
model.train(data="coco8.yaml", epochs=5, quantize=8, optimizer="AdamW", lr0=0.00001, cos_lr=True)

# Xuất sang ONNX; các phạm vi INT8 đã học trở thành node Q/DQ nên không cần dữ liệu hiệu chỉnh
model.export(format="onnx", quantize=8)

QAT được sử dụng ở đâu#

• Thị giác biên: Detector INT8 trên điện thoại di động, camera thông minh, drone và NPU, nơi các bộ tăng tốc chỉ dùng số nguyên hoàn toàn không thể chạy model dấu phẩy động.
• Máy chủ suy luận: Engine GPU như TensorRT chạy kernel Tensor Core INT8 để tăng thông lượng; QAT giúp accuracy gần với FP32.
• Model ngôn ngữ: QAT hiện có thể giảm xuống dưới 8 bit. Google phát hành checkpoint Gemma 3 QAT ở định dạng 4 bit Q4_0, giữ chất lượng gần với bfloat16 trong khi chỉ dùng một phần nhỏ bộ nhớ.

Các framework lớn đều cung cấp công cụ QAT, bao gồm torchao của PyTorch, Bộ công cụ tối ưu hóa model của TensorFlow và NNCF của Intel dành cho OpenVINO.

QAT và các thuật ngữ liên quan#

• Lượng tử hóa model: Thuật ngữ bao quát chỉ việc giảm độ chính xác số học. PTQ và QAT là hai quy trình chính.
• Huấn luyện độ chính xác hỗn hợp: Huấn luyện bằng FP16 hoặc BF16 để tiết kiệm bộ nhớ và thời gian, nhưng không chuẩn bị model cho suy luận INT8.
• Độ chính xác bán phần: Suy luận FP16; hầu hết model có thể sử dụng định dạng này mà không cần huấn luyện lại, nên hiếm khi cần QAT.
• Cắt tỉa model: Loại bỏ trọng số hoặc channel thay vì giảm độ chính xác của chúng; hai phương pháp nén này có thể kết hợp.
• Chưng cất tri thức: Huấn luyện kiến trúc nhỏ hơn để mô phỏng kiến trúc lớn hơn. Phương pháp này thường được kết hợp với QAT, trong đó model độ chính xác đầy đủ đóng vai trò teacher cho bản sao đã lượng tử hóa của chính nó.