Tối ưu hóa LLM cho Tốc độ và Bộ nhớ

Các Mô hình Ngôn ngữ Lớn (LLM) như GPT3/4, Falcon và Llama đang tiến bộ nhanh chóng trong khả năng giải quyết các tác vụ mang tính con người, khẳng định vị thế là công cụ thiết yếu trong các ngành công nghiệp dựa trên tri thức hiện đại. Tuy nhiên, việc triển khai các mô hình này trong các tác vụ thực tế vẫn còn nhiều thách thức:

Để thể hiện khả năng hiểu và tạo văn bản gần như con người, các LLM hiện tại cần được cấu thành từ hàng tỷ tham số (xem Kaplan et al, Wei et. al). Điều này dẫn đến việc tăng đáng kể yêu cầu về bộ nhớ cho suy luận.
Trong nhiều tác vụ thực tế, các LLM cần được cung cấp thông tin ngữ cảnh chi tiết. Điều này đòi hỏi mô hình phải có khả năng xử lý các chuỗi đầu vào rất dài trong quá trình suy luận.

Điểm mấu chốt của những thách thức này nằm ở việc tăng cường khả năng tính toán và bộ nhớ của LLM, đặc biệt khi xử lý các chuỗi đầu vào có quy mô lớn.

Trong hướng dẫn này, chúng ta sẽ xem xét các kỹ thuật hiệu quả để triển khai LLM hiệu quả:

Độ chính xác thấp hơn: Nghiên cứu đã chỉ ra rằng việc vận hành ở độ chính xác số thấp hơn, cụ thể là 8-bit và 4-bit, có thể mang lại lợi thế về mặt tính toán mà không gây suy giảm đáng kể hiệu suất của mô hình.
Flash Attention: Flash Attention là một biến thể của thuật toán attention, không chỉ cung cấp cách tiếp cận tiết kiệm bộ nhớ hơn mà còn đạt được hiệu suất cao hơn nhờ tối ưu hóa việc sử dụng bộ nhớ GPU.
Sáng kiến kiến trúc: Vì LLM luôn được triển khai theo cùng một cách trong suy luận, cụ thể là tạo văn bản tự hồi quy với ngữ cảnh đầu vào dài, các kiến trúc mô hình chuyên biệt đã được đề xuất để cho phép suy luận hiệu quả hơn. Những tiến bộ quan trọng nhất trong kiến trúc mô hình ở đây là Alibi, Rotary embeddings, Multi-Query Attention (MQA) và Grouped-Query-Attention (GQA).

Trong suốt hướng dẫn này, chúng ta sẽ phân tích quá trình sinh tự hồi quy (auto-regressive generation) từ góc nhìn của tensor. Chúng ta sẽ đi sâu vào những ưu và nhược điểm của việc áp dụng độ chính xác thấp hơn, khám phá toàn diện các thuật toán attention mới nhất và thảo luận về các kiến trúc LLM được cải tiến. Trong quá trình đó, chúng ta sẽ chạy các ví dụ thực tế để minh họa cho từng cải tiến tính năng.

1. Độ chính xác thấp hơn

Yêu cầu bộ nhớ của LLM có thể được hiểu rõ nhất bằng cách xem LLM như một tập hợp các ma trận và vector trọng số, và các đầu vào văn bản như một chuỗi các vector. Trong phần tiếp theo, định nghĩa trọng số sẽ được sử dụng để chỉ tất cả các ma trận và vector trọng số của mô hình.

Tại thời điểm viết hướng dẫn này, các LLM bao gồm ít nhất vài tỷ tham số. Mỗi tham số do đó được tạo thành từ một số thập phân, ví dụ 4.5689, thường được lưu trữ ở định dạng float32, bfloat16 hoặc float16. Điều này cho phép chúng ta dễ dàng tính toán yêu cầu bộ nhớ để nạp LLM vào bộ nhớ:

Việc nạp trọng số của một mô hình có X tỷ tham số yêu cầu khoảng 4 * X GB bộ nhớ VRAM ở độ chính xác float32

Tuy nhiên, ngày nay các mô hình hiếm khi được huấn luyện ở độ chính xác float32 đầy đủ, mà thường ở độ chính xác bfloat16 hoặc ít phổ biến hơn là float16. Do đó, quy tắc kinh nghiệm trở thành:

Việc nạp trọng số của một mô hình có X tỷ tham số yêu cầu khoảng 2 * X GB bộ nhớ VRAM ở độ chính xác bfloat16/float16

Đối với các đầu vào văn bản ngắn hơn (ít hơn 1024 token), yêu cầu bộ nhớ cho suy luận phần lớn bị chi phối bởi yêu cầu bộ nhớ để tải trọng số. Do đó, hiện tại, hãy giả sử rằng yêu cầu bộ nhớ cho suy luận bằng với yêu cầu bộ nhớ để tải mô hình vào VRAM của GPU.

Để đưa ra một số ví dụ về lượng VRAM cần thiết để nạp một mô hình ở bfloat16:

GPT3 cần 2 * 175 GB = 350 GB VRAM
Bloom cần 2 * 176 GB = 352 GB VRAM
Llama-2-70b cần 2 * 70 GB = 140 GB VRAM
Falcon-40b cần 2 * 40 GB = 80 GB VRAM
MPT-30b cần 2 * 30 GB = 60 GB VRAM
bigcode/starcoder yêu cầu 2 * 15.5 = 31 GB VRAM

Tính đến thời điểm viết tài liệu này, chip GPU lớn nhất trên thị trường là A100 & H100 với 80GB VRAM. Hầu hết các mô hình được liệt kê ở trên cần hơn 80GB chỉ để tải, do đó bắt buộc phải sử dụng tensor parallelism và/hoặc pipeline parallelism.

🤗 Transformers hiện hỗ trợ tensor parallelism cho các model được hỗ trợ có base_tp_plan trong các config class tương ứng của chúng. Tìm hiểu thêm về Tensor Parallelism tại đây. Hơn nữa, nếu bạn quan tâm đến việc viết model theo cách thân thiện với tensor parallelism, hãy xem thư viện text-generation-inference.

Pipeline parallelism đơn giản được hỗ trợ ngay từ đầu. Để thực hiện, chỉ cần tải model với device="auto", hệ thống sẽ tự động phân bổ các lớp khác nhau lên các GPU khả dụng như đã mô tả tại đây. Tuy nhiên, mặc dù rất hiệu quả, pipeline parallelism đơn giản này không giải quyết được vấn đề GPU nằm chờ. Để xử lý vấn đề này, cần đến pipeline parallelism nâng cao hơn như đã mô tả tại đây.

Nếu bạn có quyền truy cập vào một node 8 x 80GB A100, bạn có thể nạp BLOOM như sau

!pip install transformers accelerate bitsandbytes optimum

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("bigscience/bloom", device_map="auto", pad_token_id=0)

Khi sử dụng device_map="auto", các lớp attention sẽ được phân bổ đều trên tất cả các GPU khả dụng.

Trong hướng dẫn này, chúng ta sẽ sử dụng bigcode/octocoder vì nó có thể chạy trên một chip GPU A100 40 GB đơn lẻ. Lưu ý rằng tất cả các tối ưu hóa bộ nhớ và tốc độ mà chúng ta sẽ áp dụng từ đây đều có thể áp dụng tương tự cho các mô hình yêu cầu song song hóa mô hình hoặc tensor.

Vì mô hình được nạp ở độ chính xác bfloat16, theo quy tắc ước lượng ở trên, chúng ta có thể kỳ vọng yêu cầu bộ nhớ để chạy suy luận với bigcode/octocoder là khoảng 31 GB VRAM. Hãy thử xem.

Đầu tiên, chúng ta nạp mô hình và tokenizer, sau đó truyền cả hai vào đối tượng pipeline của Transformers.

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", dtype=torch.bfloat16, device_map="auto", pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer:"

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Đầu ra:

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

Tuyệt, giờ chúng ta có thể trực tiếp sử dụng kết quả để chuyển đổi byte sang Gigabyte.

def bytes_to_giga_bytes(bytes):
  return bytes / 1024 / 1024 / 1024

Hãy gọi torch.accelerator.memory.max_memory_allocated để đo mức phân bổ bộ nhớ accelerator cao nhất.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

Đầu ra:

29.0260648727417

Đủ gần với phép tính ước lượng của chúng ta! Chúng ta có thể thấy con số không hoàn toàn chính xác vì việc chuyển đổi từ byte sang kilobyte yêu cầu nhân với 1024 thay vì 1000. Do đó, công thức ước lượng cũng có thể được hiểu là phép tính “tối đa X GB”. Lưu ý rằng nếu chúng ta cố gắng chạy mô hình ở độ chính xác float32 đầy đủ, sẽ cần tới 64 GB VRAM.

Hầu hết các mô hình hiện nay đều được huấn luyện ở bfloat16, không có lý do gì để chạy mô hình ở độ chính xác float32 đầy đủ nếu GPU của bạn hỗ trợ bfloat16. Float32 sẽ không cho kết quả suy luận tốt hơn so với độ chính xác đã được sử dụng để huấn luyện mô hình.

Nếu bạn không chắc trọng số mô hình được lưu ở định dạng nào trên Hub, bạn luôn có thể xem cấu hình của checkpoint tại "dtype", ví dụ ở đây. Khuyến nghị đặt mô hình cùng loại độ chính xác như đã ghi trong cấu hình khi tải bằng from_pretrained(..., dtype=...), ngoại trừ khi loại gốc là float32, trong trường hợp đó có thể dùng cả float16 hoặc bfloat16 cho suy luận.

Hãy định nghĩa một hàm flush(...) để giải phóng toàn bộ bộ nhớ đã cấp phát, giúp chúng ta đo lường chính xác bộ nhớ GPU đã cấp phát ở mức đỉnh.

del pipe
del model

import gc
import torch

def flush():
  gc.collect()
  torch.accelerator.empty_cache()
  torch.accelerator.reset_peak_memory_stats()

Hãy gọi nó ngay bây giờ cho thí nghiệm tiếp theo.

flush()

Từ thư viện Accelerate, bạn cũng có thể sử dụng một phương tiện tiện ích độc lập với thiết bị có tên là release_memory, phương tiện này có tính đến các backend phần cứng khác nhau như XPU, MLU, NPU, MPS và nhiều hơn nữa.

from accelerate.utils import release_memory
# ...

release_memory(model)

Vậy nếu GPU của bạn không có 32 GB VRAM thì sao? Người ta đã phát hiện ra rằng trọng số mô hình có thể được lượng tử hóa xuống 8-bit hoặc 4-bit mà không gây suy giảm hiệu suất đáng kể (xem Dettmers et al.). Mô hình thậm chí có thể được lượng tử hóa xuống 3 hoặc 2 bit với mức suy giảm hiệu suất chấp nhận được, như được thể hiện trong bài báo GPTQ gần đây 🤯.

Không đi sâu vào quá nhiều chi tiết, các phương pháp lượng tử hóa nhằm mục đích giảm độ chính xác của trọng số trong khi cố gắng giữ cho kết quả suy luận của mô hình chính xác nhất có thể (còn gọi là gần nhất có thể với bfloat16). Lưu ý rằng lượng tử hóa hoạt động đặc biệt tốt cho việc tạo văn bản vì tất cả những gì chúng ta quan tâm là chọn bộ các token tiếp theo có xác suất cao nhất và không thực sự quan tâm đến các giá trị chính xác của phân phối logit của token tiếp theo. Tất cả những gì quan trọng là phân phối logit của token tiếp theo phải giữ nguyên xấp xỉ để một thao tác argmax hoặc topk cho ra kết quả giống hệt nhau.

Có nhiều kỹ thuật lượng tử hóa, mà chúng ta sẽ không thảo luận chi tiết ở đây, nhưng nhìn chung, tất cả các kỹ thuật lượng tử hóa đều hoạt động như sau:

1. Lượng tử hóa tất cả trọng số xuống độ chính xác đích
1. Tải trọng số đã lượng tử hóa, và truyền vào chuỗi đầu vào các vector với độ chính xác bfloat16
1. Giải nén lượng tử động trọng số sang bfloat16 để thực hiện phép tính với các vector đầu vào của chúng ở độ chính xác bfloat16

Nói ngắn gọn, điều này có nghĩa là ma trận nhân với đầu vào các phép nhân, trong đó $X$ là đầu vào, $W$ là ma trận trọng số và $Y$ là đầu ra: $Y = X * W$

được thay đổi thành $Y = X * \text{dequantize}(W)$

cho mỗi phép nhân ma trận. Việc khử lượng tử hóa và lượng tử hóa lại được thực hiện tuần tự cho tất cả các ma trận trọng số khi các đầu vào chạy qua đồ thị mạng.

Do đó, thời gian suy luận thường không giảm khi sử dụng trọng lượng đã lượng tử hóa, mà thậm chí còn tăng lên. Đủ lý thuyết rồi, hãy thử ngay thôi! Để lượng tử hóa trọng lượng với Transformers, bạn cần đảm bảo rằng thư viện bitsandbytes đã được cài đặt.

!pip install bitsandbytes

Sau đó, chúng ta có thể nạp các mô hình ở định lượng 8-bit bằng cách đơn giản thêm cờ load_in_8bit=True vào from_pretrained.

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_8bit=True), pad_token_id=0)

Bây giờ, hãy chạy lại ví dụ của chúng ta và đo lường mức sử dụng bộ nhớ.

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Đầu ra:

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

Tốt, chúng ta đang đạt được kết quả giống như trước đó, vậy nên không có sự suy giảm về độ chính xác! Hãy xem xét lượng bộ nhớ đã được sử dụng lần này.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

Đầu ra:

15.219234466552734

Ít hơn đáng kể! Chúng ta chỉ còn lại hơn 15 GB một chút và do đó có thể chạy mô hình này trên các GPU tiêu dùng như 4090. Chúng ta thấy được sự cải thiện rất tốt về hiệu quả sử dụng bộ nhớ và hầu như không có sự suy giảm nào đối với đầu ra của mô hình. Tuy nhiên, chúng ta cũng có thể nhận thấy một chút chậm lại trong quá trình suy luận.

Chúng ta xóa các mô hình và làm sạch bộ nhớ một lần nữa.

del model
del pipe

flush()

Hãy xem mức tiêu thụ bộ nhớ GPU đỉnh mà lượng tử hóa 4-bit mang lại. Việc lượng tử hóa mô hình xuống 4-bit có thể thực hiện bằng cùng một API như trước - lần này bằng cách truyền load_in_4bit=True thay vì load_in_8bit=True.

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_4bit=True), pad_token_id=0)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

Đầu ra:

Here is a Python function that transforms bytes to Giga bytes:\n\n```\ndef bytes_to_gigabytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single argument

Chúng ta gần như đang thấy cùng một văn bản đầu ra như trước - chỉ thiếu python ngay trước đoạn mã. Hãy xem cần bao nhiêu bộ nhớ.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

Đầu ra:

9.543574333190918

Chỉ 9,5GB! Đó thực sự không phải là một con số lớn đối với một mô hình có hơn 15 tỷ tham số.

Mặc dù chúng ta thấy độ chính xác của mô hình ở đây hầu như không bị suy giảm, nhưng trên thực tế, lượng tử hóa 4-bit thường có thể dẫn đến kết quả khác so với lượng tử hóa 8-bit hoặc suy luận bfloat16 đầy đủ. Người dùng có thể tự thử nghiệm.

Cũng lưu ý rằng suy luận ở đây lại chậm hơn một chút so với lượng tử hóa 8-bit, điều này là do phương pháp lượng tử hóa mạnh hơn được sử dụng cho lượng tử hóa 4-bit dẫn đến $\text{quantize}$ và $\text{dequantize}$ mất nhiều thời gian hơn trong quá trình suy luận.

del model
del pipe

flush()

Nhìn chung, chúng ta thấy rằng việc chạy OctoCoder ở độ chính xác 8-bit đã giảm bộ nhớ VRAM GPU cần thiết từ 32G xuống chỉ còn 15GB, và việc chạy mô hình ở độ chính xác 4-bit tiếp tục giảm bộ nhớ VRAM GPU cần thiết xuống chỉ hơn 9GB một chút.

Lượng tử hóa 4-bit cho phép mô hình chạy trên các GPU như RTX3090, V100 và T4, những phần cứng khá dễ tiếp cận với hầu hết mọi người.

Để biết thêm thông tin về lượng tử hóa và xem cách lượng tử hóa mô hình để yêu cầu ít bộ nhớ GPU VRAM hơn cả 4-bit, chúng tôi khuyến nghị bạn tìm hiểu về triển khai GPT-QModel.

Tóm lại, cần ghi nhớ rằng lượng tử hóa mô hình đánh đổi hiệu quả bộ nhớ được cải thiện với độ chính xác và trong một số trường hợp là thời gian suy luận.

Nếu bộ nhớ GPU không phải là ràng buộc cho trường hợp sử dụng của bạn, thường thì không cần phải tìm hiểu về lượng tử hóa. Tuy nhiên, nhiều GPU đơn giản là không thể chạy LLM mà không có các phương pháp lượng tử hóa, và trong trường hợp này, các lược đồ lượng tử hóa 4-bit và 8-bit là những công cụ cực kỳ hữu ích.

Để có thông tin chi tiết hơn về cách sử dụng, chúng tôi rất khuyến khích bạn tham khảo Tài liệu Lượng tử hóa Transformers. Tiếp theo, hãy cùng tìm hiểu cách chúng ta có thể cải thiện hiệu quả tính toán và bộ nhớ bằng cách sử dụng các thuật toán tốt hơn và kiến trúc mô hình được cải tiến.

2. Flash Attention

Các LLM hiệu suất cao hàng ngày đều chia sẻ phần lớn cùng một kiến trúc cơ bản, bao gồm các lớp feed-forward, các lớp kích hoạt, các lớp chuẩn hóa lớp và quan trọng nhất là các lớp self-attention.

Các lớp self-attention đóng vai trò trung tâm trong các Large Language Models (LLM) vì chúng cho phép mô hình hiểu được mối quan hệ ngữ cảnh giữa các token đầu vào. Tuy nhiên, mức tiêu thụ bộ nhớ GPU đỉnh của các lớp self-attention tăng bậc hai về cả độ phức tạp tính toán lẫn bộ nhớ theo số lượng token đầu vào (còn được gọi là chiều dài chuỗi) mà chúng tôi ký hiệu là $N$ . Mặc dù điều này không thực sự đáng kể đối với các chuỗi đầu vào ngắn (tối đa 1000 token đầu vào), nhưng nó trở thành một vấn đề nghiêm trọng đối với các chuỗi đầu vào dài (khoảng 16000 token đầu vào).

Hãy xem xét kỹ hơn. Công thức để tính toán đầu ra $\mathbf{O}$ của một lớp self-attention cho một đầu vào $\mathbf{X}$ có độ dài $N$ là: $\textbf{O} = \text{Attn}(\mathbf{X}) = \mathbf{V} \times \text{Softmax}(\mathbf{QK}^T) \text{ with } \mathbf{Q} = \mathbf{W}_q \mathbf{X}, \mathbf{V} = \mathbf{W}_v \mathbf{X}, \mathbf{K} = \mathbf{W}_k \mathbf{X}$

$\mathbf{X} = (\mathbf{x}_1, ... \mathbf{x}_{N})$ do đó là chuỗi đầu vào của lớp attention. Các phép chiếu $\mathbf{Q}$ và $\mathbf{K}$ sẽ mỗi cái gồm $N$ vector, dẫn đến $\mathbf{QK}^T$ có kích thước $N^2$ .

LLM thường có nhiều đầu chú ý, do đó thực hiện nhiều phép tính self-attention song song. Giả sử LLM có 40 đầu chú ý và chạy ở độ chính xác bfloat16, chúng ta có thể tính toán bộ nhớ cần thiết để lưu trữ $\mathbf{QK^T}$ các ma trận $40 * 2 * N^2$ là $N=1000$ chỉ cần khoảng 50 MB VRAM, tuy nhiên, với $N=16000$ chúng ta sẽ cần 19 GB VRAM, và với $N=100,000$ chúng ta sẽ cần gần 1TB chỉ để lưu trữ $\mathbf{QK}^T$ các ma trận.

Nói ngắn gọn, thuật toán self-attention mặc định nhanh chóng trở nên tốn bộ nhớ đến mức không thể chấp nhận được với các ngữ cảnh đầu vào lớn.

Khi các LLM ngày càng nâng cao khả năng hiểu và tạo văn bản, chúng được ứng dụng vào những tác vụ ngày càng phức tạp. Trong khi trước đây các mô hình chỉ xử lý việc dịch hoặc tóm tắt một vài câu, thì nay chúng phải quản lý cả trang văn bản, đòi hỏi khả năng xử lý độ dài đầu vào lớn.

Làm thế nào để loại bỏ yêu cầu bộ nhớ quá lớn cho các độ dài đầu vào lớn? Chúng ta cần một cách mới để tính toán cơ chế self-attention, loại bỏ được $\mathbf{QK}^T$ ma trận. Tri Dao và cộng sự đã phát triển chính xác một thuật toán mới như vậy và đặt tên là Flash Attention.

Nói ngắn gọn, Flash Attention tách phép tính $\mathbf{V} \times \text{Softmax}(\mathbf{QK}^T)$ ra và thay vào đó tính các phần nhỏ hơn của kết quả bằng cách lặp qua nhiều bước tính softmax: $\textbf{O}_i \leftarrow s^a_{ij} * \textbf{O}_i + s^b_{ij} * \mathbf{V}_{j} \times \text{Softmax}(\mathbf{QK}^T_{i,j}) \text{ for multiple } i, j \text{ iterations}$

với $s^a_{ij}$ và $s^b_{ij}$ là một số thống kê chuẩn hóa softmax cần được tính lại cho mỗi $i$ và $j$ .

Xin lưu ý rằng toàn bộ cơ chế Flash Attention phức tạp hơn một chút và được đơn giản hóa đáng kể ở đây vì việc đi sâu vào chi tiết nằm ngoài phạm vi của hướng dẫn này. Người đọc được khuyến khích tham khảo bài báo Flash Attention được viết rất tốt để biết thêm chi tiết.

Điểm chính cần rút ra ở đây là:

Bằng cách theo dõi các thống kê chuẩn hóa softmax và sử dụng một số phép toán thông minh, Flash Attention cho ra kết quả giống hệt về mặt số học so với lớp self-attention mặc định, với chi phí bộ nhớ chỉ tăng tuyến tính theo $N$ .

Nhìn vào công thức, người ta có thể trực giác cho rằng Flash Attention chắc chắn phải chậm hơn nhiều so với công thức self-attention mặc định vì cần thực hiện nhiều phép tính hơn. Quả thực, Flash Attention yêu cầu nhiều FLOPs hơn so với attention thông thường vì các thống kê chuẩn hóa softmax phải liên tục được tính lại (xem bài báo để biết thêm chi tiết nếu bạn quan tâm)

Tuy nhiên, Flash Attention nhanh hơn đáng kể trong suy luận so với cơ chế chú ý mặc định, nhờ khả năng giảm đáng kể yêu cầu đối với bộ nhớ băng thông cao nhưng chậm hơn của GPU (VRAM), thay vào đó tập trung vào bộ nhớ trên chip nhanh hơn (SRAM).

Về cơ bản, Flash Attention đảm bảo rằng tất cả các thao tác ghi và đọc trung gian đều có thể được thực hiện bằng bộ nhớ SRAM trên chip nhanh thay vì phải truy cập bộ nhớ VRAM chậm hơn để tính toán vector đầu ra $\mathbf{O}$ .

Trong thực tế, hiện tại hoàn toàn không có lý do nào để không sử dụng Flash Attention nếu có sẵn. Thuật toán này cho ra kết quả giống hệt nhau về mặt toán học, đồng thời nhanh hơn và tiết kiệm bộ nhớ hơn.

3. Đổi mới kiến trúc

Cho đến nay, chúng ta đã xem xét việc cải thiện hiệu quả tính toán và bộ nhớ bằng cách:

Đổ trọng số sang định dạng độ chính xác thấp hơn
Thay thế thuật toán tự chú ý bằng một phiên bản hiệu quả hơn về bộ nhớ và tính toán

Bây giờ, hãy cùng xem xét cách chúng ta có thể thay đổi kiến trúc của một LLM để nó hiệu quả và tối ưu nhất cho các tác vụ yêu cầu đầu vào văn bản dài, ví dụ:

Trả lời câu hỏi có tăng cường truy xuất,
Tóm tắt,
Chat

Lưu ý rằng chat không chỉ yêu cầu LLM xử lý các đầu vào văn bản dài, mà còn đòi hỏi LLM phải có khả năng xử lý hiệu quả cuộc đối thoại qua lại giữa người dùng và trợ lý (chẳng hạn như ChatGPT).

Sau khi được huấn luyện, kiến trúc LLM cơ bản rất khó thay đổi, do đó việc xem xét các tác vụ của LLM trước khi huấn luyện và tối ưu hóa kiến trúc mô hình cho phù hợp là rất quan trọng. Có hai thành phần quan trọng trong kiến trúc mô hình nhanh chóng trở thành nút thắt cổ chai về bộ nhớ và/hoặc hiệu suất đối với các chuỗi đầu vào lớn.

Các vector nhúng vị trí
Bộ nhớ đệm khóa-giá trị

Hãy xem xét chi tiết hơn từng thành phần

3.1 Cải thiện nhúng vị trí của các LLM

Self-attention đặt mỗi token vào mối quan hệ với các token khác. Ví dụ, ma trận $\text{Softmax}(\mathbf{QK}^T)$ của chuỗi đầu vào văn bản ““Hello”, “I”, “love”, “you” có thể có dạng như sau:

Mỗi token từ được gán một khối lượng xác suất cho việc chú ý đến tất cả các token từ khác, do đó thiết lập mối quan hệ với tất cả các token từ khác. Ví dụ, từ “love” chú ý đến từ “Hello” với 5%, đến “I” với 30%, và đến chính nó với 65%.

Một LLM dựa trên self-attention nhưng không có position embeddings sẽ gặp rất nhiều khó khăn trong việc hiểu vị trí tương đối của các đầu vào văn bản. Điều này là do điểm xác suất được tính bởi $\mathbf{QK}^T$ liên kết mỗi token từ với mọi token từ khác trong $O(1)$ phép tính bất kể khoảng cách vị trí tương đối giữa chúng. Do đó, đối với LLM không có position embeddings, mỗi token dường như có cùng khoảng cách đến tất cả các token khác, e.g. việc phân biệt giữa ““Hello I love you” và ““You love I hello” sẽ rất khó khăn.

Để LLM hiểu được thứ tự các câu, cần có thêm một đầu mối và thường được áp dụng dưới dạng mã hóa vị trí (còn được gọi là nhúng vị trí). Mã hóa vị trí mã hóa vị trí của từng token thành một biểu diễn số mà LLM có thể khai thác để hiểu tốt hơn thứ tự các câu.

Các tác giả của bài báo Attention Is All You Need đã giới thiệu các embedding vị trí dạng sin $\mathbf{P} = \mathbf{p}_1, \ldots, \mathbf{p}_N$ . trong đó mỗi vector $\mathbf{p}_i$ được tính như một hàm sin của vị trí $i$ . Các mã hóa vị trí sau đó đơn giản được cộng vào các vector của chuỗi đầu vào $\mathbf{\hat{X}} = \mathbf{\hat{x}}_1, \ldots, \mathbf{\hat{x}}_N$ = $\mathbf{x}_1 + \mathbf{p}_1, \ldots, \mathbf{x}_N + \mathbf{p}_N$ qua đó giúp mô hình học tốt hơn thứ tự câu.

Thay vì sử dụng position embeddings cố định, các nghiên cứu khác (chẳng hạn như Devlin et al.) đã sử dụng positional encodings được học, trong đó các positional embeddings $\mathbf{P}$ được học trong quá trình huấn luyện.

Các vector nhúng vị trí hình sin và vector nhúng vị trí học được từng là các phương pháp chủ đạo để mã hóa thứ tự câu vào LLM, nhưng một số vấn đề liên quan đến các phép mã hóa vị trí này đã được phát hiện:

Cả position embedding dạng sinusoidal và position embedding học được đều là position embedding tuyệt đối, i.e. mã hóa một embedding duy nhất cho mỗi position id: $0, \ldots, N$ . Như đã được chỉ ra bởi Huang et al. và Su et al., position embedding tuyệt đối dẫn đến hiệu suất kém của LLM với các đầu vào văn bản dài. Với các đầu vào văn bản dài, sẽ có lợi nếu mô hình học được khoảng cách vị trí tương đối giữa các token đầu vào thay vì vị trí tuyệt đối của chúng.
Khi sử dụng position embeddings đã học, LLM phải được huấn luyện trên một độ dài đầu vào cố định $N$, điều này khiến việc ngoại suy sang độ dài đầu vào dài hơn so với độ dài đã huấn luyện trở nên khó khăn.

Gần đây, các nhúng vị trí tương đối có thể giải quyết các vấn đề được nêu ở trên đã trở nên phổ biến hơn, nổi bật nhất là:

Cả RoPE và ALiBi đều cho rằng cách tốt nhất là đưa thông tin về thứ tự câu vào trực tiếp trong thuật toán self-attention, vì đó là nơi các token từ được đặt vào mối quan hệ với nhau. Cụ thể hơn, thứ tự câu nên được đưa vào bằng cách sửa đổi phép tính $\mathbf{QK}^T$ trong phép nhân ma trận.

Không đi vào quá nhiều chi tiết, RoPE chỉ ra rằng thông tin vị trí có thể được mã hóa vào các cặp query-key, e.g. $\mathbf{q}_i$ và $\mathbf{x}_j$ bằng cách xoay mỗi vector một góc $\theta * i$ và $\theta * j$ tương ứng với $i, j$ mô tả vị trí câu của mỗi vector: $\mathbf{\hat{q}}_i^T \mathbf{\hat{x}}_j = \mathbf{{q}}_i^T \mathbf{R}_{\theta, i -j} \mathbf{{x}}_j.$

$\mathbf{R}_{\theta, i - j}$ do đó biểu diễn một ma trận quay. $\theta$ là không được học trong quá trình huấn luyện, mà được đặt thành một giá trị định trước, phụ thuộc vào độ dài tối đa của chuỗi đầu vào trong quá trình huấn luyện.

Nhờ đó, điểm xác suất giữa $\mathbf{q}_i$ và $\mathbf{q}_j$ chỉ bị ảnh hưởng nếu $i \ne j$ và chỉ phụ thuộc vào khoảng cách tương đối $i - j$ bất kể vị trí cụ thể của từng vector $i$ và $j$ .

RoPE được sử dụng trong nhiều LLM quan trọng nhất hiện nay, chẳng hạn như:

Một phương án thay thế, ALiBi đề xuất một lược đồ mã hóa vị trí tương đối đơn giản hơn nhiều. Khoảng cách tương đối giữa các token đầu vào được thêm vào dưới dạng số nguyên âm được nhân với một giá trị định trước m vào mỗi mục query-key của ma trận $\mathbf{QK}^T$ ngay trước khi tính softmax.

Như đã trình bày trong bài báo ALiBi, phương pháp mã hóa vị trí tương đối đơn giản này cho phép mô hình duy trì hiệu suất cao ngay cả với các chuỗi đầu vào văn bản rất dài.

ALiBi được sử dụng trong nhiều LLM quan trọng nhất hiện nay, chẳng hạn như:

Cả RoPE và ALiBi position encodings đều có thể ngoại suy đến các độ dài đầu vào chưa từng thấy trong quá trình huấn luyện, trong khi đã được chứng minh rằng ngoại suy hoạt động tốt hơn nhiều ngay từ đầu với ALiBi so với RoPE. Đối với ALiBi, chỉ cần tăng các giá trị của ma trận vị trí tam giác dưới để khớp với độ dài của chuỗi đầu vào. Đối với RoPE, việc giữ nguyên $\theta$ được sử dụng trong quá trình huấn luyện sẽ dẫn đến kết quả kém khi xử lý các đầu vào văn bản dài hơn nhiều so với những gì đã thấy trong quá trình huấn luyện, c.f Press et al.. Tuy nhiên, cộng đồng đã tìm ra một số mẹo hiệu quả để điều chỉnh $\theta$, từ đó cho phép RoPE position embeddings hoạt động tốt cho các chuỗi đầu vào văn bản ngoại suy (xem tại đây).

Cả RoPE và ALiBi đều là các nhúng vị trí tương đối không được học trong quá trình huấn luyện, mà thay vào đó dựa trên các trực giác sau:

Các tín hiệu vị trí về đầu vào văn bản nên được cung cấp trực tiếp cho $\mathbf{QK}^T$ ma trận của lớp self-attention.
LLM nên được khuyến khích học mã hóa vị trí khoảng cách tương đối không đổi.
Các token đầu vào văn bản càng cách xa nhau thì xác suất query-value của chúng càng thấp. Cả RoPE và ALiBi đều làm giảm xác suất query-key của các token nằm xa nhau. RoPE làm giảm bằng cách giảm tích vector của chúng thông qua việc tăng góc giữa các vector query-key. ALiBi làm giảm bằng cách cộng các số âm lớn vào tích vector.

Tóm lại, các LLM được triển khai cho các tác vụ yêu cầu xử lý đầu vào văn bản lớn sẽ được huấn luyện tốt hơn với các embedding vị trí tương đối, chẳng hạn như RoPE và ALiBi. Ngoài ra, hãy lưu ý rằng ngay cả khi một LLM sử dụng RoPE và ALiBi chỉ được huấn luyện trên một độ dài cố định, ví dụ như $N_1 = 2048$ nó vẫn có thể được sử dụng trong thực tế với các đầu vào văn bản lớn hơn nhiều so với $N_1$, chẳng hạn như $N_2 = 8192 > N_1$, bằng cách ngoại suy các embedding vị trí.

3.2 Bộ đệm khóa-giá trị

Việc tạo văn bản tự hồi quy bằng LLM hoạt động bằng cách lặp lại các bước: đưa vào một chuỗi đầu vào, lấy mẫu token tiếp theo, thêm token đó vào chuỗi đầu vào, và tiếp tục cho đến khi LLM tạo ra một token cho biết quá trình tạo đã hoàn tất.

Hãy xem Transformer’s Generate Text Tutorial để có lời giải thích trực quan hơn về cách hoạt động của quá trình sinh tự hồi quy.

Hãy chạy một đoạn mã ngắn để xem cách tự hồi quy hoạt động trên thực tế. Chúng ta sẽ đơn giản lấy token tiếp theo có khả năng cao nhất thông qua torch.argmax.

input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)

for _ in range(5):
  next_logits = model(input_ids)["logits"][:, -1:]
  next_token_id = torch.argmax(next_logits,dim=-1)

  input_ids = torch.cat([input_ids, next_token_id], dim=-1)
  print("shape of input_ids", input_ids.shape)

generated_text = tokenizer.batch_decode(input_ids[:, -5:])
generated_text

Đầu ra:

shape of input_ids torch.Size([1, 21])
shape of input_ids torch.Size([1, 22])
shape of input_ids torch.Size([1, 23])
shape of input_ids torch.Size([1, 24])
shape of input_ids torch.Size([1, 25])
[' Here is a Python function']

Như chúng ta có thể thấy, mỗi lần chúng ta tăng số token đầu vào văn bản lên bằng chính token vừa được lấy mẫu.

Ngoại trừ một số rất ít trường hợp, các LLM được huấn luyện bằng mục tiêu mô hình hóa ngôn ngữ nhân quả và do đó che đi ma trận tam giác trên của điểm chú ý - đây là lý do tại sao trong hai sơ đồ ở trên, các điểm chú ý được để trống (còn gọi là có xác suất bằng 0). Để ôn lại nhanh về mô hình hóa ngôn ngữ nhân quả, bạn có thể tham khảo blog Illustrated Self Attention.

Do đó, các token không bao giờ phụ thuộc vào các token sau, cụ thể hơn là vector $\mathbf{q}_i$ không bao giờ được đặt quan hệ với bất kỳ vector key, values nào $\mathbf{k}_j, \mathbf{v}_j$ nếu $j > i$. Thay vào đó, $\mathbf{q}_i$ chỉ chú ý đến các vector key-value trước đó $\mathbf{k}{m < i}, \mathbf{v}{m < i} \text{ , for } m \in {0, \ldots i - 1}$. Để giảm thiểu tính toán không cần thiết, ta có thể cache các vector key-value của từng lớp cho tất cả các bước thời gian trước đó.

Trong phần tiếp theo, chúng ta sẽ yêu cầu LLM sử dụng bộ nhớ đệm key-value bằng cách truy xuất và chuyển tiếp nó cho mỗi lần truyền tiến. Trong Transformers, chúng ta có thể truy xuất bộ nhớ đệm key-value bằng cách truyền cờ use_cache vào lời gọi forward và sau đó truyền nó cùng với token hiện tại.

past_key_values = None # past_key_values is the key-value cache
generated_tokens = []
next_token_id = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)

for _ in range(5):
  next_logits, past_key_values = model(next_token_id, past_key_values=past_key_values, use_cache=True).to_tuple()
  next_logits = next_logits[:, -1:]
  next_token_id = torch.argmax(next_logits, dim=-1)

  print("shape of input_ids", next_token_id.shape)
  print("length of key-value cache", past_key_values.get_seq_length())  # past_key_values are of shape [num_layers, 0 for k, 1 for v, batch_size, length, hidden_dim]
  generated_tokens.append(next_token_id.item())

generated_text = tokenizer.batch_decode(generated_tokens)
generated_text

Đầu ra:

shape of input_ids torch.Size([1, 1])
length of key-value cache 20
shape of input_ids torch.Size([1, 1])
length of key-value cache 21
shape of input_ids torch.Size([1, 1])
length of key-value cache 22
shape of input_ids torch.Size([1, 1])
length of key-value cache 23
shape of input_ids torch.Size([1, 1])
length of key-value cache 24
[' Here', ' is', ' a', ' Python', ' function']

Như có thể thấy, khi sử dụng bộ nhớ đệm key-value, các token đầu vào văn bản không tăng về độ dài, mà vẫn là một vector đầu vào duy nhất. Ngược lại, độ dài của bộ nhớ đệm key-value tăng thêm một ở mỗi bước giải mã.

Việc tận dụng bộ nhớ đệm key-value có nghĩa là $\mathbf{QK}^T$ về cơ bản được rút gọn thành $\mathbf{q}_c\mathbf{K}^T$ với $\mathbf{q}_c$ là phép chiếu query của token đầu vào hiện tại, vốn luôn chỉ là một vector duy nhất.

Việc sử dụng bộ nhớ đệm khóa-giá trị có hai lợi thế:

Hiệu quả tính toán tăng đáng kể vì ít phép tính hơn so với việc tính toán toàn bộ $\mathbf{QK}^T$ ma trận. Điều này dẫn đến tốc độ suy luận tăng lên
Bộ nhớ tối đa cần thiết không tăng theo cấp số nhân với số lượng token được tạo ra, mà chỉ tăng tuyến tính.

Người dùng luôn nên sử dụng bộ nhớ đệm key-value vì nó cho ra kết quả giống hệt và tăng tốc đáng kể với các chuỗi đầu vào dài hơn. Transformers bật bộ nhớ đệm key-value theo mặc định khi sử dụng pipeline văn bản hoặc phương thức generate. Chúng tôi có một hướng dẫn riêng dành cho bộ nhớ đệm tại đây.

Lưu ý rằng, mặc dù chúng tôi khuyên bạn nên sử dụng bộ nhớ đệm key-value, đầu ra của LLM của bạn có thể hơi khác khi sử dụng chúng. Đây là đặc tính của chính các kernel nhân nhân ma trận — bạn có thể đọc thêm về điều này tại đây.

3.2.1 Trò chuyện nhiều vòng

Bộ nhớ đệm key-value đặc biệt hữu ích cho các ứng dụng như chat, nơi cần nhiều lần giải mã tự hồi quy. Hãy xem một ví dụ.

User: How many people live in France?
Assistant: Roughly 75 million people live in France
User: And how many are in Germany?
Assistant: Germany has ca. 81 million inhabitants

Trong cuộc trò chuyện này, LLM chạy giải mã tự hồi quy hai lần:

Lần đầu tiên, bộ nhớ đệm key-value trống và prompt đầu vào là "User: How many people live in France?", mô hình tự hồi quy sinh ra văn bản "Roughly 75 million people live in France" trong khi tăng bộ nhớ đệm key-value ở mỗi bước giải mã.
Lần thứ hai, prompt đầu vào là "User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many in Germany?". Nhờ vào cache, tất cả các vector key-value cho hai câu đầu tiên đã được tính toán sẵn. Do đó, prompt đầu vào chỉ bao gồm "User: And how many in Germany?". Trong khi xử lý prompt đầu vào rút gọn này, các vector key-value được tính toán của nó sẽ được nối vào key-value cache của lần giải mã đầu tiên. Sau đó, câu trả lời thứ hai của Assistant "Germany has ca. 81 million inhabitants" được tạo theo kiểu tự hồi quy với key-value cache bao gồm các vector key-value đã mã hóa của "User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many are in Germany?".

Có hai điều cần lưu ý ở đây:

Việc giữ lại toàn bộ ngữ cảnh là rất quan trọng đối với các LLM được triển khai trong chat để LLM có thể hiểu được toàn bộ ngữ cảnh trước đó của cuộc trò chuyện. Ví dụ, với ví dụ ở trên, LLM cần hiểu rằng người dùng đang đề cập đến dân số khi hỏi "And how many are in Germany".
Bộ nhớ đệm key-value cực kỳ hữu ích cho chat vì nó cho phép chúng ta liên tục mở rộng lịch sử chat đã mã hóa thay vì phải mã hóa lại lịch sử chat từ đầu (như trường hợp khi sử dụng kiến trúc encoder-decoder, chẳng hạn).

Trong transformers, một lệnh gọi generate sẽ trả về past_key_values khi truyền return_dict_in_generate=True, bên cạnh use_cache=True mặc định. Lưu ý rằng tính năng này hiện chưa khả dụng qua giao diện pipeline.

# Generation as usual
prompt = system_prompt + "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(**model_inputs, max_new_tokens=60, return_dict_in_generate=True)
decoded_output = tokenizer.batch_decode(generation_output.sequences)[0]

# Piping the returned `past_key_values` to speed up the next conversation round
prompt = decoded_output + "\nQuestion: How can I modify the function above to return Mega bytes instead?\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(
  **model_inputs,
  past_key_values=generation_output.past_key_values,
  max_new_tokens=60,
  return_dict_in_generate=True
)
tokenizer.batch_decode(generation_output.sequences)[0][len(prompt):]

Đầu ra:

 is a modified version of the function that returns Mega bytes instead.

def bytes_to_megabytes(bytes):
   return bytes / 1024 / 1024

Answer: The function takes a number of bytes as input and returns the number of

Tuyệt vời, không mất thêm thời gian để tính lại cùng các key và values cho lớp attention! Tuy nhiên, có một điểm cần lưu ý. Trong khi bộ nhớ đỉnh cần thiết cho $\mathbf{QK}^T$ ma trận được giảm đáng kể, việc giữ cache key-value trong bộ nhớ có thể trở nên rất tốn kém về mặt bộ nhớ đối với các chuỗi đầu vào dài hoặc chat nhiều lượt. Hãy nhớ rằng cache key-value cần lưu trữ các vector key-value cho tất cả các vector đầu vào trước đó $\mathbf{x}_i \text{, for } i \in \{1, \ldots, c - 1\}$ cho tất cả các lớp self-attention và tất cả các attention heads.

Hãy tính số lượng giá trị float cần được lưu trữ trong bộ nhớ đệm key-value cho LLM bigcode/octocoder mà chúng ta đã sử dụng trước đó. Số lượng giá trị float bằng hai lần nhân với độ dài chuỗi, nhân với số đầu chú ý, nhân với chiều kích đầu chú ý và nhân với số lớp. Tính toán cho LLM của chúng ta với độ dài chuỗi đầu vào giả định là 16000 cho kết quả:

config = model.config
2 * 16_000 * config.n_layer * config.n_head * config.n_embd // config.n_head

Đầu ra:

7864320000

Khoảng 8 tỷ giá trị float! Lưu trữ 8 tỷ giá trị float ở độ chính xác float16 đòi hỏi khoảng 15 GB RAM, tức là xấp xỉ một nửa so với chính trọng số của mô hình! Các nhà nghiên cứu đã đề xuất hai phương pháp cho phép giảm đáng kể chi phí bộ nhớ khi lưu trữ bộ đệm key-value, được khám phá trong các tiểu mục tiếp theo.

3.2.2 Multi-Query-Attention (MQA)

Multi-Query-Attention được đề xuất trong bài báo Fast Transformer Decoding: One Write-Head is All You Need của Noam Shazeer. Như tiêu đề đã nói, Noam phát hiện ra rằng thay vì sử dụng n_head trọng số chiếu key-value, người ta có thể sử dụng một cặp trọng số chiếu head-value duy nhất được chia sẻ giữa tất cả các đầu chú ý mà hiệu suất của mô hình không bị suy giảm đáng kể.

Bằng cách sử dụng một cặp trọng số chiếu head-value duy nhất, các vector key-value $\mathbf{k}_i, \mathbf{v}_i$ phải giống nhau trên tất cả các attention head, điều này có nghĩa là chúng ta chỉ cần lưu trữ 1 cặp chiếu key-value trong cache thay vì n_head cặp.

Vì hầu hết các LLM sử dụng từ 20 đến 100 đầu chú ý, MQA giảm đáng kể mức tiêu thụ bộ nhớ của bộ đệm key-value. Đối với LLM được sử dụng trong notebook này, chúng ta có thể giảm mức tiêu thụ bộ nhớ cần thiết từ 15 GB xuống dưới 400 MB ở độ dài chuỗi đầu vào là 16000.

Ngoài việc tiết kiệm bộ nhớ, MQA còn giúp cải thiện hiệu quả tính toán như được giải thích bên dưới. Trong quá trình giải mã tự hồi quy (auto-regressive decoding), các vector key-value lớn cần được nạp lại, nối với cặp vector key-value hiện tại để sau đó đưa vào $\mathbf{q}_c\mathbf{K}^T$ tính toán ở mỗi bước. Đối với giải mã tự hồi quy, băng thông bộ nhớ cần thiết cho việc nạp lại liên tục có thể trở thành nút thắt cổ chai nghiêm trọng về thời gian. Bằng cách giảm kích thước của các vector key-value, lượng bộ nhớ cần truy cập sẽ giảm, từ đó giảm nút thắt cổ chai về băng thông bộ nhớ. Để biết thêm chi tiết, hãy xem bài báo của Noam.

Phần quan trọng cần hiểu ở đây là việc giảm số lượng đầu chú ý key-value xuống còn 1 chỉ có ý nghĩa nếu có sử dụng bộ nhớ đệm key-value. Mức tiêu thụ bộ nhớ đỉnh của mô hình cho một lần truyền tiến đơn lẻ mà không có bộ nhớ đệm key-value vẫn không thay đổi, vì mỗi đầu chú ý vẫn có một vector truy vấn riêng biệt, do đó mỗi đầu chú ý vẫn có một $\mathbf{QK}^T$ ma trận.

MQA đã được cộng đồng áp dụng rộng rãi và hiện đang được sử dụng bởi nhiều LLM phổ biến nhất:

Ngoài ra, checkpoint được sử dụng trong notebook này - bigcode/octocoder - có sử dụng MQA.

3.2.3 Grouped-Query-Attention (GQA)

Grouped-Query-Attention, do Ainslie et al. từ Google đề xuất, nhận thấy rằng việc sử dụng MQA thường có thể dẫn đến suy giảm chất lượng so với việc sử dụng các phép chiếu đầu key-value đa đầu thông thường. Bài báo lập luận rằng hiệu suất mô hình có thể được duy trì tốt hơn bằng cách giảm số lượng trọng số phép chiếu đầu query ít triệt để hơn. Thay vì chỉ sử dụng một trọng số phép chiếu key-value duy nhất, nên sử dụng n < n_head trọng số phép chiếu key-value. Bằng cách chọn n là một giá trị nhỏ hơn đáng kể so với n_head, chẳng hạn như 2, 4 hoặc 8, gần như toàn bộ các lợi ích về bộ nhớ và tốc độ từ MQA có thể được giữ lại trong khi hy sinh ít hơn dung lượng mô hình và do đó có thể nói là ít suy giảm hiệu suất hơn.

Hơn nữa, các tác giả của GQA đã phát hiện ra rằng các checkpoint mô hình hiện có có thể được uptrain để có kiến trúc GQA với chỉ 5% lượng tính toán pre-training ban đầu. Mặc dù 5% lượng tính toán pre-training ban đầu vẫn có thể là một con số rất lớn, nhưng uptraining GQA cho phép các checkpoint hiện có tiếp tục hữu ích cho các chuỗi đầu vào dài hơn.

GQA mới chỉ được đề xuất gần đây, vì vậy tại thời điểm viết notebook này, việc áp dụng công nghệ này còn hạn chế. Ứng dụng nổi bật nhất của GQA là Llama-v2.

Tóm lại, nếu LLM được triển khai với giải mã tự hồi quy và cần xử lý các chuỗi đầu vào lớn, chẳng hạn như trong trường hợp chat, thì rất nên sử dụng GQA hoặc MQA.

Kết luận

Cộng đồng nghiên cứu liên tục đưa ra những cách thức mới, tinh tế để rút ngắn thời gian suy luận cho các LLM ngày càng lớn. Một ví dụ về hướng nghiên cứu đầy triển vọng là speculative decoding, trong đó các “token dễ” được tạo ra bởi các mô hình ngôn ngữ nhỏ hơn, nhanh hơn, và chỉ các “token khó” mới được tạo ra bởi chính LLM. Việc đi sâu vào chi tiết nằm ngoài phạm vi của notebook này, nhưng bạn có thể tìm hiểu thêm trong bài blog hay này.

Lý do các LLM khổng lồ như GPT3/4, Llama-2-70b, Claude, PaLM có thể chạy nhanh đến vậy trong các giao diện chat như Hugging Face Chat hay ChatGPT là một phần lớn nhờ vào những cải tiến đã đề cập ở trên về độ chính xác, thuật toán và kiến trúc. Trong tương lai, các bộ tăng tốc như GPU, TPU, v.v. sẽ ngày càng nhanh hơn và cho phép sử dụng nhiều bộ nhớ hơn, nhưng người dùng vẫn luôn nên đảm bảo sử dụng các thuật toán và kiến trúc tốt nhất có sẵn để đạt được hiệu quả cao nhất 🤗

Update on GitHub