Tăng tốc LLM gấp 3 lần: Giải mã cơ chế Speculative Decoding

Mọi agent đều đã chạy theo một vòng lặp. "Loop engineering" (kỹ thuật vòng lặp) bổ sung thêm một vòng lặp bao quanh chính agent đó, cho phép nó tự đánh giá kết quả đầu ra, thử lại khi công việc chưa đạt yêu cầu và tinh chỉnh các hướng dẫn khi những lỗi tương tự tái diễn. Hiện nay, bạn đang đóng vai trò đó: xem xét công việc, chẩn đoán sai sót và đưa ra gợi ý (prompt) lại cho agent. Bài viết này sẽ chỉ ra cách tự động hóa quy trình đó với một ví dụ thực tế, đồng thời khám phá những khía cạnh mà sự phán đoán của con người vẫn là cần thiết.

Đọc bài viết

Một mô hình 70 tỷ tham số yêu cầu đọc khoảng 140 GB trọng số từ bộ nhớ GPU. Trên một GPU trung tâm dữ liệu hiện đại, quá trình truyền tải này có thể mất hàng chục mili giây. Các phép tính thực tế áp dụng lên những trọng số này chỉ chiếm một phần nhỏ thời gian đó. Điều này có nghĩa là các đơn vị tính toán của bộ xử lý không được sử dụng trong phần lớn thời gian dành cho bước tạo token.

Speculative decoding (giải mã suy đoán) là một kỹ thuật chuyển đổi công suất dư thừa này thành kết quả đầu ra. Một mô hình thứ hai, nhỏ hơn nhiều, sẽ tạo ra trước một vài token dự kiến. Mô hình lớn sẽ đánh giá tất cả các token đó trong một lần truyền tiến (forward pass) duy nhất thay vì mỗi lần truyền cho một token, giúp tốc độ tạo văn bản nhanh gấp 2-3 lần. Hơn nữa, văn bản được tạo ra vẫn giữ nguyên tính thống kê giống hệt với kết quả từ mô hình lớn khi chạy độc lập.

Trong bài viết này, chúng ta sẽ tìm hiểu cách thức hoạt động của speculative decoding. Dưới đây là những nội dung chúng ta sẽ đề cập:

Tại sao việc tạo token lại chạy từng bước một

GPU dành thời gian vào việc gì trong quá trình tạo văn bản

Cách đánh giá nhiều token dự kiến trong một lần truyền duy nhất

Vòng lặp chấp nhận và từ chối, và điều gì xảy ra khi một token dự kiến bị sai

Tại sao chất lượng đầu ra được bảo toàn chính xác

Tỷ lệ chấp nhận và tại sao nó thay đổi tùy theo khối lượng công việc

Bốn nguồn gốc có thể tạo ra bản nháp (draft)

Khi nào speculative decoding không còn hiệu quả

Tuyên bố miễn trừ trách nhiệm: Bài viết này dựa trên các chi tiết được chia sẻ công khai từ nhiều nguồn khác nhau. Tài liệu tham khảo ở cuối bài. Vui lòng để lại bình luận nếu bạn phát hiện bất kỳ điểm nào không chính xác.

Việc tạo văn bản hoạt động theo từng token một.

Mô hình đọc tất cả những gì đã được tạo ra cho đến thời điểm hiện tại, tính toán phân phối xác suất trên từ vựng của nó, chọn token tiếp theo, thêm token đó vào đầu vào và lặp lại chu trình. Mỗi chu trình được gọi là một lần truyền tiến (forward pass), và mỗi lần truyền tiến sẽ đưa dữ liệu đầu vào đi qua tất cả các lớp của mô hình.

Ví dụ, token 50 phụ thuộc vào việc token 49 có mặt trong đầu vào, và token 49 lại phụ thuộc vào token 48, cứ như vậy. Việc tính toán chúng đồng thời sẽ phá vỡ chuỗi phụ thuộc tạo nên sự mạch lạc của kết quả đầu ra.

Điều này có nghĩa là một phản hồi dài 500 token đòi hỏi 500 lần truyền tiến tuần tự, mỗi lần phải hoàn thành trước khi lần tiếp theo bắt đầu. Vì thời gian của một lần truyền phụ thuộc vào kích thước mô hình, tổng thời gian tạo văn bản sẽ bằng số lượng token đầu ra nhân với thời gian cho mỗi lần truyền tiến.

Điều này giải thích tại sao tốc độ phản hồi vẫn tương đối ổn định cho dù câu trả lời là một câu trả lời ngắn gọn hay một đoạn mã dài, bởi vì chi phí cho mỗi token vẫn giữ nguyên trong cả hai trường hợp. Nó cũng giải thích tại sao một mô hình lớn hơn lại tạo văn bản chậm hơn trên cùng một phần cứng.

Các hệ thống suy luận hiện đại sử dụng KV cache, giúp lưu trữ trạng thái chú ý (attention state) cho các token đã xử lý để mỗi lần truyền mới chỉ cần tính toán sự chú ý cho vị trí mới nhất. Điều này cắt giảm đáng kể khối lượng công việc trong mỗi lần truyền, mặc dù yêu cầu mỗi lần truyền cho một token vẫn còn đó.

Vì số lượng lần truyền được cố định bởi lượng văn bản chúng ta muốn tạo, điều này dẫn đến nửa còn lại của phương trình. Một lần truyền tiến thực sự dành thời gian để làm gì?

Nói một cách đơn giản, một lần truyền tiến dành phần lớn thời gian để di chuyển dữ liệu thay vì thực hiện các phép tính số học.

Trọng số của mô hình nằm trong bộ nhớ GPU, thường được gọi là VRAM. Để tính toán bất cứ điều gì với các trọng số đó, GPU phải chuyển chúng vào các đơn vị tính toán nơi phép nhân diễn ra. Đối với một mô hình 70 tỷ tham số được lưu trữ ở độ chính xác 16-bit, quá trình truyền tải này lên tới khoảng 140 GB cho mỗi token.

Các phép tính số học thực hiện trên 140 GB đó khá nhỏ nếu so sánh. Một token nghĩa là một vector hẹp chạy qua mỗi ma trận trọng số. GPU tải một ma trận khổng lồ từ bộ nhớ, nhân nó với vector đó, loại bỏ nó và tải ma trận tiếp theo.

Hệ quả là trong quá trình xử lý prompt, hiệu suất tính toán đạt khoảng 90 đến 95 phần trăm. Tuy nhiên, trong quá trình tạo token, con số này giảm xuống còn khoảng 20 đến 40 phần trăm. Các đơn vị tính toán không được sử dụng trong phần lớn thời gian của mỗi bước trong khi bus bộ nhớ hoạt động gần hết công suất.

Sự khác biệt này được thúc đẩy bởi khối lượng công việc mà mỗi lần đọc trọng số hỗ trợ:

Xử lý prompt đọc các trọng số một lần và áp dụng chúng cho hàng nghìn token đầu vào cùng một lúc.

Tạo token đọc cùng các trọng số đó và áp dụng chúng cho đúng một token.

Đây là công suất đã được chi trả nhưng chưa được tận dụng tối đa.

Nhưng tại sao điều này lại quan trọng trên thực tế?

Một GPU có băng thông bộ nhớ cao hơn sẽ cải thiện tốc độ tạo văn bản tốt hơn so với một GPU có khả năng tính toán thô cao hơn.

Tuy nhiên, công suất dư thừa chỉ hữu ích nếu có công việc hữu ích để đưa vào đó. Câu hỏi đặt ra là liệu một lần truyền tiến có thể tạo ra kết quả đầu ra nhiều hơn một token hay không.

Một lần truyền tiến có thể đánh giá nhiều vị trí cùng một lúc.

Các mô hình Transformer xử lý toàn bộ chuỗi song song. Khi chúng ta đưa vào một chuỗi token, mô hình sẽ tính toán dự đoán token tiếp theo tại mọi vị trí trong chuỗi đó trong cùng một lần truyền. Ví dụ, đầu vào gồm năm token sẽ tạo ra năm dự đoán.

Những dự đoán này vẫn hợp lệ nhờ vào cơ chế "causal masking" (che mặt nhân quả). Bên trong cơ chế chú ý (attention mechanism), vị trí thứ 5 có thể truy cập các vị trí từ 1 đến 5 trong khi các vị trí từ 6 trở đi bị che đi, và vị trí thứ 3 chỉ có thể truy cập các vị trí từ 1 đến 3. Do đó, mỗi vị trí được điều kiện hóa chính xác bởi các token đứng trước nó, giống hệt với cách điều kiện hóa mà nó sẽ nhận được nếu chúng ta tạo chuỗi theo từng bước một.

Đây chính là đặc tính giúp việc xử lý prompt trở nên nhanh chóng. Một prompt dài 2.000 token chạy qua mô hình trong một lần truyền thay vì 2.000 lần, bởi vì tất cả 2.000 vị trí được tính toán cùng nhau.