Tokens of Curiosity

Hoang's exploring & learning place 📚

Những kỹ thuật LLM Inference: Parallelism, Flash Attention, GPTQ, AWQ,…

CategorIes:

By

Hoang Tran

·

9–14 phút
  1. Giới thiệu
  2. Song song hóa (Parallelism)
    1. Song song hóa dữ liệu (Data Parallelism)
    2. Song song hóa trong từng phép toán (Intra-operation/Tensor Parallelism)
    3. Song song hóa mô hình (Model Parallelism)
    4. Song song hóa theo pipeline (Pipeline Parallelism)
  3. Tối ưu hóa Tính toán và Biên dịch
    1. Tối ưu hóa đồ thị tính toán (Graph Optimization)
    2. Kết hợp toán tử (Operator Fusion)
    3. Truy cập bộ nhớ GPU trực tiếp (Direct Storage Access)
    4. Tối ưu hóa thuật toán
  4. Kỹ thuật tối ưu hóa dành cho LLM
    1. Flash Attention
    2. GPTQ – GPT Quantization
    3. AWQ – Activation-aware Weight Quantization
  5. Tối ưu hóa bằng Caching và Memoization
    1. Bộ nhớ đệm truy vấn (Query Caching)
    2. Caching dựa trên embeddings và cơ sở dữ liệu vector
  6. Kết luận
  7. 📚 Tham khảo

Giới thiệu

Gần đây ở công ty, mình có làm việc với các tasks liên quan đến fine-tune và serving LLMs cho các tác vụ đặc thù. Qua đó, mình cũng có cơ hội để vọc vạch và hiểu thêm về các thư viện cũng như các kỹ thuật liên quan được dùng phổ biến trong thực tế.

Triển khai mô hình lớn (LLM) một cách hiệu quả không chỉ yêu cầu về phần cứng mạnh mẽ mà còn đòi hỏi các kỹ thuật tối ưu hóa tinh vi. Ở bài viết này, mình tổng hợp lại các kỹ thuật nâng cao được sử dụng trong LLM inference để cải thiện hiệu suất, giảm latency và tối ưu chi phí khi triển khai trong sử dụng thực tế như một cheatsheet cho bản thân, hi vọng sẽ hữu ích cho mọi người.

Song song hóa (Parallelism)

Song song hóa dữ liệu (Data Parallelism)

Đây là phương pháp dễ tiếp cận nhất. Dữ liệu được chia thành các nhóm nhỏ và xử lý đồng thời trên nhiều GPUs khác nhau. Mỗi bản sao của mô hình sẽ được load vào GPU và xử lý một phần dữ liệu riêng biệt, giúp cải thiện tốc độ xử lý và scale một cách linearly theo số lượng đơn vị tính toán.

  • Cơ chế: Tạo bản sao mô hình trên nhiều GPU, xử lý các batch dữ liệu khác nhau song song. Khi training thay vì inference, ta sẽ cần thêm những bước như đồng bộ hóa gradients giữa các GPU qua all-reduce
  • Ưu điểm:
    • Tăng throughput bằng cách xử lý nhiều dữ liệu cùng lúc
    • Phù hợp khi model có thể fit trên bộ nhớ của từng GPU đơn lẻ
  • Hạn chế:
    • Không giúp chạy được model lớn hơn bộ nhớ của GPU
  • Ứng dụng: Thích hợp cho các tác vụ yêu cầu thông lượng cao, chẳng hạn như xử lý hàng loạt dữ liệu lớn.

Một ví dụ thực tế về việc song song hoá dữ liệu là trong LLM evaluation với thư viện lm-evaluation-harness từ EleutherAI. Thư viện này hỗ trợ việc chạy data parallelism trên nhiều GPU bằng việc sử dụng launcher từ thư viện accelerate của Huggingface:

accelerate launch -m lm_eval --model hf \
    --tasks lambada_openai,arc_easy \
    --batch_size 16

Song song hóa trong từng phép toán (Intra-operation/Tensor Parallelism)

  • Cơ chế: Chia tensor trong các phép toán như nhân ma trận, convolution, attention thành những phần nhỏ, được xử lý cùng lúc trên nhiều GPUs, sau đó kết hợp kết quả qua phép all-gather. Ví dụ: một phép nhân ma trận WX có thể được chia thành W1X trên GPU1 và W2X trên GPU2
  • Ứng dụng:
    • Xử lý các layer lớn (FFN, attention heads) trong transformers
    • Về hiệu quả, có thể giúp tăng 57-80% throughput khi dùng 2 GPU so với single-GPU
Hình ảnh minh hoạ cho việc song song hoá phép nhân ma trận. Nguồn: Huggingface

Bạn có thể tìm đọc thêm về Tensor Parallelism tại documentation của transformers về distributed training.

Song song hóa mô hình (Model Parallelism)

  • Cơ chế: Thay vì duplicate mô hình trên nhiều GPUs khác nhau, kỹ thuật này chia mô hình lớn, không vừa để fit trên một GPU đơn lẻ, thành các phần nhỏ và triển khai trên nhiều thiết bị. Tuỳ theo architecture của mô hình mà ta có nhiều cách để chia nhỏ mô hình ra.
  • Phân loại:
    • Expert Parallelism: Phân phối các expert networks trong kiến trúc MoE
    • Layer Parallelism: Merge các transformer layers và chỉnh sửa để cho phép chúng chạy song song (ví dụ attention và FFN layers). Về hiệu quả, kỹ thuật này có thể giúp giảm độ sâu mô hình đi 18-21% và tăng tốc inference 1.2-1.38x
    • Pipeline Parallelism: đề cập rõ hơn ở bên dưới
  • Ứng dụng: Bên cạnh để tăng tốc quá trình inference, kỹ thuật này cũng được sử dụng khi mô hình quá lớn để load trên một GPU đơn lẻ

Một ví dụ về kỹ thuật này là trong thư viện lm-evaluation-harness mình đã đề cập phía trên. Ta có thể chia weights của một model lớn lên các GPU đơn lẻ và thực hiện đánh giá bằng cách thêm model argument parallelize=True:

lm_eval --model hf \
    --tasks lambada_openai,arc_easy \
    --model_args parallelize=True \
    --batch_size 16

Song song hóa theo pipeline (Pipeline Parallelism)

  • Cơ chế: model được chia thành nhiều stages và các phần khác nhau sẽ được xử lý trên các GPUs khác nhau theo kiểu pipeline. Cụ thể:
    1. Chia model thành N stage (vd: encoder/decoder)
    2. Mỗi stage xử lý trên GPU riêng
    3. Kết hợp micro-batching để lấp đầy pipeline
  • Ứng dụng: Thích hợp với các mô hình lớn cần tối ưu cả bộ nhớ lẫn tốc độ xử lý.

Tương như như Tensor Parallelism, đây cũng là một phương pháp được hỗ trợ trong thư viện transformers.

Tối ưu hóa Tính toán và Biên dịch

Tối ưu hóa đồ thị tính toán (Graph Optimization)

  • Cơ chế: Việc tái cấu trúc computation graph giúp loại bỏ các phép toán dư thừa và tối ưu quá trình thực thi. Ví dụ, trong kiến trúc Transformer, các phép toán như Softmax và Layernorm thường đi kèm với các lớp Linear. Bằng cách phân tách và sắp xếp lại các phép toán này, ta có thể song song hóa quá trình tính toán và giảm độ trễ tổng thể.
  • Thư viện liên quan: TensorRT, ONNX Runtime, TorchScript,…

Kết hợp toán tử (Operator Fusion)

  • Cơ chế: Kỹ thuật này gộp các phép toán thành một kernel duy nhất, giảm số lần gọi kernel, tăng hiệu suất tính toán và giảm chi phí truy cập bộ nhớ trên GPU. Ví dụ, trong paper LLM Inference Acceleration via Efficient Operation Fusion:
    • Fused Layernorm: Thay vì tính toán Layernorm và lớp tuyến tính riêng biệt, tác giả kết hợp chúng với nhau bằng cách tận dụng các tính chất đại số một cách thông minh
    • Softmax Fusion: Tương tự, phép Softmax được tính toán song song với các lớp tuyến tính, tận dụng tài nguyên phần cứng hiệu quả hơn

Truy cập bộ nhớ GPU trực tiếp (Direct Storage Access)

  • Cơ chế: Cho phép dữ liệu tải trực tiếp từ disk vào GPU mà không qua CPU, giảm đáng kể thời gian tải mô hình. Ví dụ:
    • Sử dụng unified memory để tránh sao chép dữ liệu giữa CPU và GPU.
    • Tận dụng zero-copy access cho dữ liệu đầu vào, cho phép GPU truy xuất trực tiếp từ bộ nhớ hệ thống mà không cần qua bộ đệm trung gian

Tối ưu hóa thuật toán

  • Cơ chế: Các thuật toán sử dụng trong mô hình có thể được triển khai bằng ngôn ngữ cấp thấp như C++ hoặc CUDA để cải thiện tốc độ xử lý và hạn chế các giới hạn về performance của Python.

Kỹ thuật tối ưu hóa dành cho LLM

Flash Attention

Flash Attention là thuật toán tối ưu hóa cơ chế self-attention trong các mô hình transformer, giúp giảm thời gian huấn luyện và độ trễ khi suy luận (inference latency) cho LLM so với cơ chế attention truyền thống, đặc biệt với input có sequence length dài.

  • Cơ chế:
    • Tối ưu bộ nhớ GPU: Flash Attention sử dụng kỹ thuật tiling để chia ma trận attention thành các khối nhỏ, giảm số lần đọc/ghi dữ liệu giữa bộ nhớ GPU (HBM) và SRAM
    • Kết hợp toán tử (Operator Fusion): Thay vì tải dữ liệu nhiều lần, Flash Attention tải toàn bộ queries, keys, values một lần, thực hiện tính toán liên tục rồi ghi kết quả cuối cùng, giảm loading overhead
  • Các thư viện fine-tuning và serving đa số sẽ hỗ trợ flash attention, ví dụ trong Axolotl:
pip install axolotl[flash-attn]

GPTQ – GPT Quantization

GPTQ (GPT Quantization) là kỹ thuật lượng tử hóa 4-bit giúp chạy các LLM lớn trên consumer hardware như GPU RTX 3090 mà không giảm đáng kể độ chính xác của LLM.

  • Cơ chế:
    • Arbitrary Order Insight: Thay vì chọn parameters để quantize theo thứ tự cố định nào đó Optimal Brain Quantization, GPTQ xử lý tất cả weights trong matrix cùng lúc, tối ưu hóa tốc độ tính toán
    • Lazy Batch-Updates: Nhóm các cột ma trận thành batch (ví dụ: 128 cột), update đồng thời để tận dụng hiệu năng GPU
  • Ưu điểm:
    • Giảm lượng VRAM cần thiết, cho phép chạy LLaMA-30B trên GPU RTX 3090 với 24G VRAM
    • Duy trình độ chính xác tương đương với model 16-bit
  • Ứng dụng: Giảm tiêu thụ bộ nhớ, cho phép triển khai trên phần cứng yếu hơn.

AWQ – Activation-aware Weight Quantization

AWQ (Activation-aware Weight Quantization) là phương pháp lượng tử hóa adaptive dựa trên input activation. 

  • Cơ chế:
    • Bảo toàn weight quan trọng: AWQ xác định và giữ nguyên các weight có độ nhạy cao với dữ liệu đầu vào, chỉ lượng tử hóa những trọng số ít quan trọng hơn.
    • Điều này giúp giảm lỗi do quantization.
  • Ứng dụng: Thích hợp cho việc deploy mô hình LLM có yêu cầu cao về độ chính xác.

Tối ưu hóa bằng Caching và Memoization

Để tối ưu hóa hiệu suất và chi phí trong LLM inference, caching và memoization đóng vai trò quan trọng bằng cách tái sử dụng kết quả đã tính toán. Dưới đây là hai kỹ thuật chính:

Bộ nhớ đệm truy vấn (Query Caching)

Kỹ thuật này lưu trữ các cặp prompt-response để tránh xử lý lặp lại cùng một truy vấn. Các phương pháp triển khai phổ biến bao gồm:

  • In-Memory Caching: Dùng bộ nhớ RAM để lưu kết quả, phù hợp cho ứng dụng nhỏ hoặc môi trường đơn server.
    • Ưu điểm: tốc độ cao, triển khai đơn giản.
    • Nhược điểm: dữ liệu mất khi hệ thống khởi động lại
  • Distributed Caching (Redis/Memcached): Tối ưu cho hệ thống phân tán hoặc traffic cao, cho phép nhiều server truy cập cache lúc. Có thể thiết lập TTL (time-to-live) để tự động làm mới dữ liệu
  • Database Caching: Lưu trữ lâu dài trong SQL/NoSQL database, hỗ trợ các truy vấn phức tạp. Phù hợp khi cần phân tích lịch sử hoặc yêu cầu persist dữ liệu. Tuy nhiên, độ trễ cao hơn so với in-memory cache

Caching dựa trên embeddings và cơ sở dữ liệu vector

Không chỉ cache văn bản đầu ra, một cách tiếp cận thông minh hơn là lưu trữ embedding của các truy vấn trong một cơ sở dữ liệu vector như FAISS. Khi có truy vấn mới, ta có thể tìm kiếm trong không gian embedding và tái sử dụng kết quả gần nhất.

  • Cơ chế: Kỹ thuật này mở rộng khả năng cache bằng cách xử lý các truy vấn tương tự về ngữ nghĩa (semantically similar), thay vì chỉ trùng khớp chính xác (exact match). Quy trình thực hiện:
  1. Tạo embeddings: chuyển đổi prompt thành multi-dimensional vectors thông qua mô hình embeddings (ví dụ: BERT, OpenAI Embeddings)
  2. Lưu trữ trong Vector Database: các vector và response tương ứng được lưu trong cơ sở dữ liệu chuyên dụng (ví dụ: FAISS, Milvus,…). Khi có truy vấn mới, hệ thống tính  cosine similarity (hoặc Euclid distance) để tìm kết quả phù hợp
  3. Chọn decision threshold: Nếu độ tương tự vượt một ngưỡng định trước (ví dụ: > 0.85), response được trả về từ cache. Ngược lại, truy vấn sẽ được gửi đến LLM

Kết luận

Việc triển khai mô hình LLM tối ưu đòi hỏi sự kết hợp của nhiều chiến lược và kỹ thuật khác nhau. Song song hóa, tối ưu hóa tính toán, lượng tử hóa và caching đều đóng vai trò quan trọng trong việc nâng cao hiệu suất sử dụng LLM trong thực tế. Tùy theo bài toán cụ thể, chúng ta có thể lựa chọn những phương pháp phù hợp để đạt được cân bằng giữa tốc độ, chi phí và chất lượng đầu ra của mô hình.

Cảm ơn mọi người đã đọc bài. Hẹn gặp bạn ở bài viết tiếp theo!

📚 Tham khảo

  1. A Brief Overview of Parallelism Strategies in Deep Learning
  2. What is Inference Parallelism and How it Works?
  3. Leveraging the true depth of LLMs (González et al. 2025)
  4. Towards accelerating model parallelism in distributed deep learning systems
  5. 4-bit Quantization with GPTQ
  6. AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration (Lin et al. 2023)

Bình luận về bài viết này