# makeMoE: Triển Khai Mô Hình Ngôn Ngữ Sparse Mixture of Experts từ Đầu ## Giới Thiệu Trong bài viết này, chúng ta sẽ triển khai một mô hình ngôn ngữ **Sparse Mixture of Experts (MoE)** từ đầu, lấy cảm hứng từ dự án **makemore** của Andrej Karpathy. Mục tiêu là xây dựng một mô hình có thể hiểu và tái tạo được, đồng thời minh họa các khái niệm cốt lõi của kiến trúc MoE. --- ## Mixture of Experts là gì? **Mixture of Experts** là một kỹ thuật trong học máy trong đó nhiều mạng con chuyên biệt ("experts") được kết hợp lại. Thay vì kích hoạt toàn bộ mô hình cho mỗi token đầu vào, **Sparse MoE** chỉ kích hoạt một tập con nhỏ các experts, giúp: - ✅ Tăng dung lượng mô hình mà không tăng chi phí tính toán tương ứng - ✅ Mỗi expert chuyên biệt hóa cho các loại đầu vào khác nhau - ✅ Hiệu quả tính toán cao hơn so với mô hình dày đặc (dense model) --- ## Kiến Trúc Tổng Quan ``` Input Tokens │ ▼ Token Embedding + Positional Embedding │ ▼ ┌─────────────────────────────┐ │ Transformer Block │ │ ┌───────────────────────┐ │ │ │ Multi-Head Attention │ │ │ └───────────────────────┘ │ │ │ │ │ ┌───────────────────────┐ │ │ │ Sparse MoE Layer │ │ │ │ ┌─────────────────┐ │ │ │ │ │ Router/Gate │ │ │ │ │ └────────┬────────┘ │ │ │ │ ┌──────┴──────┐ │ │ │ │ Expert1 Expert2 ... │ │ │ └───────────────────────┘ │ └─────────────────────────────┘ │ ▼ Language Model Head │ ▼ Output Logits ``` --- ## Các Thành Phần Chính ### 1. Noisy Top-K Gating (Cổng Top-K có Nhiễu) Đây là trái tim của kiến trúc MoE. Router quyết định expert nào được kích hoạt: ```python import torch import torch.nn as nn import torch.nn.functional as F class NoisyTopKGating(nn.Module): """ Noisy Top-K Gating theo paper 'Outrageously Large Neural Networks' của Shazeer et al., 2017 """ def __init__(self, n_embed, num_experts, top_k, noise_std=1.0): super().__init__() self.num_experts = num_experts self.top_k = top_k self.noise_std = noise_std # Ma trận trọng số cho gate self.W_gate = nn.Linear(n_embed, num_experts, bias=False) # Ma trận trọng số cho nhiễu self.W_noise = nn.Linear(n_embed, num_experts, bias=False) def forward(self, x): """ Args: x: tensor shape (batch_size * seq_len, n_embed) Returns: gates: trọng số cho mỗi expert load: tải trọng của mỗi expert (dùng cho auxiliary loss) """ # Tính logits sạch clean_logits = self.W_gate(x) # (N, num_experts) # Thêm nhiễu trong quá trình training if self.training: noise = torch.randn_like(clean_logits) noise_logits = self.W_noise(x) noise_std = F.softplus(noise_logits) + 1e-2 noisy_logits = clean_logits + noise * noise_std * self.noise_std else: noisy_logits = clean_logits # Lấy top-k experts top_k_logits, top_k_indices = noisy_logits.topk( self.top_k, dim=-1 ) # Tạo mask: chỉ giữ top-k, còn lại là -inf zeros = torch.full_like(noisy_logits, float('-inf')) sparse_logits = zeros.scatter(-1, top_k_indices, top_k_logits) # Softmax để chuẩn hóa gates = F.softmax(sparse_logits, dim=-1) return gates, top_k_indices ``` ### 2. Expert Network Mỗi expert là một mạng Feed-Forward đơn giản: ```python class Expert(nn.Module): """ Một expert đơn lẻ - về cơ bản là một FFN (Feed-Forward Network) """ def __init__(self, n_embed, dropout=0.1): super().__init__() self.net = nn.Sequential( nn.Linear(n_embed, 4 * n_embed), nn.ReLU(), nn.Linear(4 * n_embed, n_embed), nn.Dropout(dropout), ) def forward(self, x): return self.net(x) ``` ### 3. Sparse MoE Layer Kết hợp router và các experts: ```python class SparseMoE(nn.Module): """ Sparse Mixture of Experts Layer """ def __init__(self, n_embed, num_experts, top_k, dropout=0.1): super().__init__() self.num_experts = num_experts self.top_k = top_k # Khởi tạo các experts self.experts = nn.ModuleList([ Expert(n_embed, dropout) for _ in range(num_experts) ]) # Khởi tạo router self.gate = NoisyTopKGating(n_embed, num_experts, top_k) def forward(self, x): """ Args: x: (batch_size, seq_len, n_embed) Returns: output: (batch_size, seq_len, n_embed) aux_loss: auxiliary loss để cân bằng tải """ batch_size, seq_len, n_embed = x.shape # Reshape để xử lý từng token độc lập x_flat = x.view(-1, n_embed) # (B*T, n_embed) num_tokens = x_flat.shape[0] # Lấy trọng số từ router gates, top_k_indices = self.gate(x_flat) # gates: (num_tokens, num_experts) - sparse # Khởi tạo output output = torch.zeros_like(x_flat) # Xử lý từng expert for expert_idx, expert in enumerate(self.experts): # Tìm các token được gán cho expert này token_mask = (top_k_indices == expert_idx).any(dim=-1) if not token_mask.any(): continue # Lấy các token tương ứng expert_input = x_flat[token_mask] # Chạy qua expert expert_output = expert(expert_input) # Lấy trọng số gate cho expert này expert_gates = gates[token_mask, expert_idx].unsqueeze(-1) # Cộng có trọng số vào output output[token_mask] += expert_gates * expert_output # Reshape về kích thước ban đầu output = output.view(batch_size, seq_len, n_embed) # Tính auxiliary loss để cân bằng tải giữa các experts aux_loss = self._compute_auxiliary_loss(gates) return output, aux_loss def _compute_auxiliary_loss(self, gates): """ Auxiliary loss để đảm bảo các experts được sử dụng đều nhau. Tránh hiện tượng 'expert collapse' - chỉ một vài experts được dùng. L_aux = num_experts * sum(f_i * P_i) trong đó: f_i = tỷ lệ token được gán cho expert i P_i = xác suất trung bình expert i được chọn """ # Tỷ lệ token được gán cho mỗi expert # gates > 0 nghĩa là expert đó được kích hoạt expert_usage = (gates > 0).float().mean(dim=0) # (num_experts,) # Xác suất trung bình expert_probs = gates.mean(dim=0) # (num_experts,) # Auxiliary loss aux_loss = self.num_experts * (expert_usage * expert_probs).sum() return aux_loss ``` ### 4. Multi-Head Self-Attention ```python class MultiHeadAttention(nn.Module): """ Multi-Head Self-Attention với causal masking """ def __init__(self, n_embed, n_heads, block_size, dropout=0.1): super().__init__() assert n_embed % n_heads == 0 self.n_heads = n_heads self.head_size = n_embed // n_heads # Projection cho Q, K, V self.qkv_proj = nn.Linear(n_embed, 3 * n_embed, bias=False) self.out_proj = nn.Linear(n_embed, n_embed, bias=False) self.attn_dropout = nn.Dropout(dropout) self.resid_dropout = nn.Dropout(dropout) # Causal mask self.register_buffer( 'mask', torch.tril(torch.ones(block_size, block_size)) .view(1, 1, block_size, block_size) ) def forward(self, x): B, T, C = x.shape # Tính Q, K, V qkv = self.qkv_proj(x) q, k, v = qkv.split(C, dim=-1) # Reshape cho multi-head q = q.view(B, T, self.n_heads, self.head_size).transpose(1, 2) k = k.view(B, T, self.n_heads, self.head_size).transpose(1, 2) v = v.view(B, T, self.n_heads, self.head_size).transpose(1, 2) # Attention scores scale = self.head_size ** -0.5 attn = (q @ k.transpose(-2, -1)) * scale # Áp dụng causal mask attn = attn.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf')) attn = F.softmax(attn, dim=-1) attn = self.attn_dropout(attn) # Weighted sum out = attn @ v # (B, n_heads, T, head_size) out = out.transpose(1, 2).contiguous().view(B, T, C) return self.resid_dropout(self.out_proj(out)) ``` ### 5. Transformer Block với MoE ```python class MoETransformerBlock(nn.Module): """ Transformer Block thay thế FFN bằng Sparse MoE Layer """ def __init__(self, n_embed, n_heads, num_experts, top_k, block_size, dropout=0.1): super().__init__() # Layer Normalization self.ln1 = nn.LayerNorm(n_embed) self.ln2 = nn.LayerNorm(n_embed) # Multi-Head Attention self.attention = MultiHeadAttention( n_embed, n_heads, block_size, dropout ) # Sparse MoE thay cho FFN thông thường self.moe = SparseMoE(n_embed, num_experts, top_k, dropout) def forward(self, x): # Self-attention với residual connection x = x + self.attention(self.ln1(x)) # MoE với residual connection moe_out, aux_loss = self.moe(self.ln2(x)) x = x + moe_out return x, aux_loss ``` ### 6. Mô Hình MoE Hoàn Chỉnh ```python class MakeMoE(nn.Module): """ Mô hình ngôn ngữ Sparse Mixture of Experts hoàn chỉnh """ def __init__( self, vocab_size, n_embed=256, n_heads=8, n_layers=6, num_experts=8, top_k=2, block_size=256, dropout=0.1 ): super().__init__() self.block_size = block_size # Embeddings self.token_embedding = nn.Embedding(vocab_size, n_embed) self.position_embedding = nn.Embedding(block_size, n_embed) self.dropout = nn.Dropout(dropout) # Transformer blocks với MoE self.blocks = nn.ModuleList([ MoETransformerBlock( n_embed, n_heads, num_experts, top_k, block_size, dropout ) for _ in range(n_layers) ]) # Final layer norm self.ln_final = nn.LayerNorm(n_embed) # Language model head self.lm_head = nn.Linear(n_embed, vocab_size, bias=False) # Weight tying self.token_embedding.weight = self.lm_head.weight # Khởi tạo trọng số self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: nn.init.zeros_(module.
Repo Github tại đây cung cấp triển khai đầu cuối đến đầu cuối: https://github.com/AviSoori1x/makeMoE/tree/main
Với sự ra mắt của Mixtral và những thông tin cho rằng GPT-4 có thể là một mô hình ngôn ngữ mixture of experts, kiến trúc mô hình này đang nhận được sự quan tâm đáng kể. Tuy nhiên, trong các mô hình ngôn ngữ sparse mixture of experts, phần lớn các thành phần được chia sẻ với các transformer truyền thống. Bất chấp vẻ đơn giản bề ngoài, bằng chứng thực nghiệm cho thấy tính ổn định trong quá trình huấn luyện là một trong những vấn đề chính của các mô hình này. Các triển khai quy mô nhỏ dễ tùy chỉnh như thế này có thể giúp thử nghiệm nhanh chóng các phương pháp mới.
Trong triển khai này, tôi thực hiện một số thay đổi đáng kể so với kiến trúc makemore:
Tuy nhiên, những điểm sau đây vẫn được giữ nguyên từ makemore:
Các mô hình ngôn ngữ sparse mixture of experts, như dự đoán, phụ thuộc vào self-attention để hiểu ngữ cảnh. Tiếp theo, chúng ta sẽ khám phá những chi tiết phức tạp của khối mixture of experts. Trước tiên, hãy cùng tìm hiểu lại về self-attention để làm mới hiểu biết của chúng ta.
Đoạn code được cung cấp minh họa cơ chế và các khái niệm cơ bản của self-attention, tập trung cụ thể vào scaled dot product self-attention cổ điển. Trong biến thể này, các ma trận query, key và value đều xuất phát từ cùng một chuỗi đầu vào. Để đảm bảo tính toàn vẹn của quá trình sinh ngôn ngữ tự hồi quy, đặc biệt trong mô hình chỉ có decoder, code triển khai kỹ thuật masking. Kỹ thuật masking này rất quan trọng vì nó che khuất mọi thông tin đứng sau vị trí của token hiện tại, từ đó hướng sự chú ý của mô hình chỉ vào các phần trước đó của chuỗi. Cơ chế attention như vậy được gọi là causal self-attention. Cần lưu ý rằng mô hình Sparse Mixture of Experts không bị giới hạn ở các kiến trúc Transformer chỉ có decoder. Trên thực tế, phần lớn các công trình quan trọng trong lĩnh vực này, đặc biệt là của Shazeer và cộng sự, xoay quanh kiến trúc T5, bao gồm cả thành phần encoder lẫn decoder trong mô hình Transformer.
Code cho causal self-attention và multi-head causal self-attention có thể được tổ chức như sau. Multi-head self-attention áp dụng nhiều attention head song song, mỗi head tập trung vào một phần riêng biệt của channel (chiều embedding). Multi-head self-attention về cơ bản cải thiện quá trình học và nâng cao hiệu quả huấn luyện mô hình nhờ vào triển khai song song vốn có. Lưu ý rằng tôi đã sử dụng dropout xuyên suốt triển khai này để regularization, tức là ngăn chặn overfitting.
Multi-head self-attention được triển khai như sau:
Trong kiến trúc Sparse Mixture of Experts (MoE), cơ chế self-attention bên trong mỗi khối transformer vẫn không thay đổi. Tuy nhiên, một sự thay đổi đáng chú ý xảy ra trong cấu trúc của mỗi khối: mạng nơ-ron feed-forward tiêu chuẩn được thay thế bằng nhiều mạng feed-forward được kích hoạt thưa thớt, được gọi là các expert. "Kích hoạt thưa thớt" (sparse activation) đề cập đến quá trình trong đó mỗi token trong chuỗi chỉ được định tuyến đến
# Bản Dịch
một số lượng hạn chế các chuyên gia này – thường là một hoặc hai – trong tổng số chuyên gia hiện có. Điều này giúp tăng tốc độ huấn luyện và suy luận, vì chỉ một số ít chuyên gia được kích hoạt trong mỗi lần forward pass. Tuy nhiên, tất cả các chuyên gia đều phải nằm trong bộ nhớ GPU, từ đó tạo ra những vấn đề triển khai thú vị khi tổng số lượng tham số đạt đến hàng trăm tỷ hoặc thậm chí hàng nghìn tỷ.
Mạng gating, còn được gọi là router, xác định mạng chuyên gia nào sẽ nhận đầu ra cho mỗi token từ multi-head attention. Hãy xem xét một ví dụ đơn giản: giả sử có 4 chuyên gia, và token sẽ được định tuyến đến top 2 chuyên gia. Ban đầu, chúng ta đưa token vào mạng gating thông qua một linear layer. Layer này chiếu tensor đầu vào từ shape (2, 4, 32) — đại diện cho (Batch size, Tokens, n_embed, trong đó n_embed là chiều kênh của đầu vào) — sang shape mới (2, 4, 4), tương ứng với (Batch size, Tokens, num_experts), trong đó num_experts là số lượng mạng chuyên gia. Tiếp theo, chúng ta xác định top k=2 giá trị cao nhất và các chỉ số tương ứng của chúng dọc theo chiều cuối cùng.
Thu được đầu ra gating thưa bằng cách chỉ giữ lại top k giá trị tại chỉ số tương ứng của chúng dọc theo chiều cuối cùng. Điền phần còn lại bằng '-inf' và đưa qua hàm kích hoạt softmax. Điều này đẩy các giá trị '-inf' về không, làm cho hai giá trị hàng đầu trở nên nổi bật hơn và tổng bằng 1. Việc tổng bằng 1 này hỗ trợ cho việc tính trọng số của các đầu ra chuyên gia.
Hãy kiểm tra chức năng với một số đầu vào mẫu:
Mặc dù bài báo về Mixtral được công bố gần đây không đề cập đến điều này, tôi tin rằng Noisy top-k Gating là một công cụ quan trọng trong việc huấn luyện các mô hình MoE. Về cơ bản, bạn không muốn tất cả các token đều được gửi đến cùng một tập hợp các chuyên gia 'được ưu tiên'. Bạn muốn có sự cân bằng tinh tế giữa khai thác và khám phá. Vì mục đích này, để cân bằng tải, việc thêm nhiễu chuẩn tắc vào các logit từ linear layer của gating là rất hữu ích. Điều này giúp quá trình huấn luyện hiệu quả hơn.
Hãy thử nghiệm lại cách triển khai này.
Khía cạnh chính của quá trình này liên quan đến đầu ra của mạng gating. Sau khi thu được các kết quả này, top k giá trị được nhân có chọn lọc với các đầu ra từ các chuyên gia top-k tương ứng cho một token nhất định. Phép nhân có chọn lọc này tạo thành một tổng có trọng số, cấu thành đầu ra của khối SparseMoe. Phần quan trọng và đầy thách thức của quá trình này là tránh các phép nhân không cần thiết. Điều thiết yếu là chỉ thực hiện forward pass cho top_k chuyên gia, sau đó tính tổng có trọng số này. Việc thực hiện forward pass cho từng chuyên gia sẽ đánh mất mục đích của việc sử dụng sparse MoE, vì khi đó nó sẽ không còn thưa nữa.
Việc kiểm tra với các đầu vào mẫu để xác nhận cách triển khai trên có hoạt động hay không là rất hữu ích. Khi chạy đoạn code sau, chúng ta có thể thấy nó hoạt động tốt!
Để nhấn mạnh, điều quan trọng cần nhận ra là độ lớn của đầu ra top_k chuyên gia từ Router/mạng gating, như được minh họa trong đoạn code trên, cũng rất có ý nghĩa. Các chỉ số top_k này
Xác định các chuyên gia được kích hoạt, và độ lớn của các giá trị trong những chiều top_k đó xác định trọng số tương ứng của chúng. Khái niệm tổng có trọng số này được làm nổi bật hơn trong sơ đồ bên dưới.
Multi-head self attention và sparse mixture of experts được kết hợp để tạo thành một khối transformer sparse mixture of experts. Giống như trong một khối transformer thông thường, các skip connection được thêm vào để đảm bảo quá trình huấn luyện ổn định và tránh các vấn đề như vanishing gradient. Ngoài ra, layer normalization được sử dụng để ổn định quá trình học hơn nữa.
Cuối cùng, kết hợp tất cả lại để tạo ra một mô hình ngôn ngữ sparse mixture of experts.
Khởi tạo tham số đóng vai trò quan trọng cho việc huấn luyện hiệu quả các mạng nơ-ron sâu. Khởi tạo Kaiming He được sử dụng ở đây vì sự hiện diện của các hàm kích hoạt ReLU trong các chuyên gia. Bạn có thể thử nghiệm với khởi tạo Glorot, vốn phổ biến hơn trong các transformer. Phần 2 khóa học Fastai của Jeremy Howard có một bài giảng xuất sắc triển khai những phương pháp này từ đầu: https://course.fast.ai/Lessons/lesson17.html. Trong tài liệu có ghi nhận rằng khởi tạo Glorot thường được sử dụng trong các mô hình transformer, vì vậy đây là cơ hội để có thể cải thiện hiệu suất mô hình.
Tôi đã sử dụng mlflow để theo dõi và ghi lại các chỉ số quan trọng cũng như các siêu tham số huấn luyện. Vòng lặp huấn luyện tôi trình bày ở đây bao gồm đoạn code này. Nếu bạn muốn huấn luyện mà không dùng mlflow, các notebook trong repo makeMoE trên GitHub có các khối code không sử dụng MLFlow. Cá nhân tôi thấy việc theo dõi các tham số và chỉ số rất tiện lợi, đặc biệt khi thử nghiệm.
Ghi lại loss trên tập huấn luyện và tập kiểm định cho bạn thấy rõ quá trình huấn luyện đang diễn ra như thế nào. Biểu đồ cho thấy tôi có lẽ nên dừng lại ở khoảng bước 4500 (khi validation loss tăng nhẹ).
Bây giờ chúng ta có thể tạo văn bản bằng mô hình này theo từng ký tự một, theo cách tự hồi quy. Với một mô hình ~9 triệu tham số được kích hoạt thưa thớt, tôi không thể phàn nàn gì.
Tôi hy vọng phần giải thích này đã giúp bạn xây dựng được sự hiểu biết của mình.


