Mở Đầu

Mixtral 8x7B đã trở thành một cái tên quen thuộc trong cộng đồng người sử dụng Large Language Model (LLM). Tuy nhiên, nhiều người vẫn chưa rõ tại sao tên gọi của nó lại là "8x7B" thay vì những con số dễ hiểu như 56B hay 7B. Bài viết này sẽ giải thích ý nghĩa của tên gọi Mixtral 8x7B và công nghệ chính mà nó sử dụng, đó là Mixture of Experts (MoE).

Giới Thiệu về Transformer

Hầu hết mọi người đều đã quen thuộc với kiến trúc Transformer, và đặc biệt, khái niệm Self-Attention chính là thành phần nổi bật nhất. Tuy nhiên, mô hình Transformer không chỉ đơn giản là Self-Attention mà còn bao gồm nhiều lớp Transformer Block (Tầng Transformer). Mỗi Transformer Block có cấu trúc cơ bản bao gồm Self-Attention và Feed Forward Network (FFN).

FFN có chức năng nhận một chuỗi đầu vào có chiều dài dim, chuyển đổi nó thành chiều dài hidden_dim, và sau đó đưa nó trở lại chiều dài dim. Khi nâng cấp kích thước một LLM, chẳng hạn từ 7B lên 70B, không chỉ số lượng Transformer Block tăng lên mà cả dim và hidden_dim cũng được mở rộng. Điều này làm gia tăng số lượng parameters trong FFN một cách đáng kể; do đó, cần thiết phải giảm thiểu số lượng parameters này.

Hơn nữa, tất cả các token trong một sequence đều đi qua cùng một FFN, điều này có nghĩa là tất cả token sẽ được xử lý giống nhau. Do vậy, việc thay thế một FFN bằng nhiều FFN để phục vụ cho nhiều đầu vào trở nên hợp lý hơn.

Giới Thiệu về Mixture of Experts (MoE)

Tổng Quan

Mixture of Experts (MoE) là giải pháp cho hai vấn đề đã đề cập ở trên. Để giảm số lượng parameters trong FFN, khái niệm sparsity được áp dụng. Điều này có nghĩa là với một đầu vào, không cần thiết phải sử dụng toàn bộ các parameters của mô hình trong quá trình suy diễn, mà chỉ cần sử dụng một phần parameters. Điều này được thể hiện thực tế qua ví dụ có 4 FFN, tuy nhiên chỉ một FFN được chọn để thực hiện suy diễn.

Tăng số lượng FFN để tạo ra nhiều heads, mỗi FFN sẽ được gọi là một expert. Để định hướng token nào sẽ được chuyển đến expert nào, một router sẽ đảm nhận vai trò này. Kiến trúc của một MoELayer sẽ cấu thành từ nhiều expert như vậy.

Chi Tiết Kỹ Thuật

Ý tưởng của MoE được trình bày ngắn gọn ở phần trước và giờ đây sẽ được đi sâu hơn vào khía cạnh toán học của nó. Trong Mixtral 8x7B, FFN Layer đã được thay thế bằng MoELayer. Thông số kiến trúc của Mixtral có thể được tham khảo ở bảng chi tiết.

Đầu ra từ MoELayer khi nhận một đầu vào xxx sẽ là tổng trọng số (weighted sum) các đầu ra của các experts, với trọng số này được xác định bởi một Gating Network. Với n experts {E0, Ei,..., En−1}, đầu ra của expert layer được tính như sau:

[\sum_{i=0}^{n-1} G(x)_i \cdot E_i(x)]

Trong đó, G(x)_i là đầu ra của Gating Network cho expert thứ i và E_i(x) là đầu ra của expert thứ i.

Nếu vector đầu ra của Gating Network là sparse (chỉ một số thành phần khác 0), có thể bỏ qua các expert với Gating score bằng 0. Cách tạo G(x) có thể được thực hiện theo nhiều cách, nhưng Mixtral chọn cách đơn giản là sử dụng softmax trên top-k mảng logits từ một lớp FC:

[G(x) = \text{Softmax}(\text{TopK}(x . W_g))]

Trong đó TopK(li) = li nếu li nằm trong K giá trị lớn nhất và TopK(li) = −∞ nếu không. Ở đây K là một hyper-parameter kiểm soát số lượng experts sử dụng cho một token trong sequence. Nếu tăng n mà vẫn giữ K không đổi, tổng số parameters của mô hình sẽ tăng nhưng số lượng parameters sử dụng cho tính toán sẽ không thay đổi. Như đã nêu ở bảng thông số, Mixtral chỉ sử dụng K=2 experts cho từng token và số lượng n=8 experts. Vậy nên tổng số parameters của Mixtral 8x7B không phải 56B mà chỉ có 47B, và số parameters chỉ sử dụng trong quá trình suy diễn chỉ còn 13B.

Kỹ Thuật Expert Parallelism

Đối với các máy có nhiều hơn 1 GPU, kỹ thuật model parallelism sẽ được áp dụng cho các expert. Chúng ta sẽ phân chia các expert trên các GPU khác nhau. Khi tính toán cho mỗi expert, đầu ra từ Gating Network sẽ phân chia các token sang GPU chứa expert đó, và sau khi tính toán xong, sẽ trả lại token về GPU ban đầu.

Kết Luận

Trên đây là cách thức hoạt động của Mixture of Expert (MoE) và sự áp dụng của nó trong Mixtral 8x7B. Hy vọng bài viết này giúp bạn hiểu rõ hơn về công nghệ tiên tiến này.
source: viblo