| |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import math |
|
|
| class SparseMoE(nn.Module): |
| """Sparse Mixture-of-Experts layer.""" |
| def __init__(self, d_model, num_experts, top_k, routing_algorithm, d_ff): |
| super().__init__() |
| self.d_model = d_model |
| self.num_experts = num_experts |
| self.top_k = top_k |
| self.routing_algorithm = routing_algorithm |
| |
| self.experts = nn.ModuleList([ |
| nn.Sequential( |
| nn.Linear(d_model, d_ff), |
| nn.ReLU(), |
| nn.Linear(d_ff, d_model) |
| ) for _ in range(num_experts) |
| ]) |
|
|
| if self.routing_algorithm == 'top_k': |
| self.gate = nn.Linear(d_model, num_experts) |
| |
| self.load_balancing_loss = 0.0 |
|
|
| def hash_routing(self, x): |
| token_hashes = x.sum(dim=-1).long().abs() |
| expert_indices = token_hashes % self.num_experts |
| return F.one_hot(expert_indices, num_classes=self.num_experts).float() |
|
|
| def top_k_routing(self, x): |
| gate_logits = self.gate(x) |
| top_k_logits, top_k_indices = torch.topk(gate_logits, self.top_k, dim=-1) |
| gate_scores = F.softmax(top_k_logits, dim=-1) |
| |
| router_mask = torch.zeros_like(gate_logits).scatter_(-1, top_k_indices, gate_scores) |
|
|
| if self.training: |
| probs_per_expert = gate_logits.softmax(dim=-1) |
| tokens_per_batch_seq = router_mask.shape[0] |
| fraction_tokens_per_expert = router_mask.sum(dim=0) / tokens_per_batch_seq |
| mean_prob_per_expert = probs_per_expert.mean(dim=0) |
| self.load_balancing_loss = self.num_experts * torch.sum(fraction_tokens_per_expert * mean_prob_per_expert) |
|
|
| return router_mask |
|
|
| def forward(self, x): |
| batch_size, seq_len, _ = x.shape |
| x_flat = x.view(-1, self.d_model) |
|
|
| if self.routing_algorithm == 'top_k': |
| router_output = self.top_k_routing(x_flat) |
| elif self.routing_algorithm == 'hash': |
| router_output = self.hash_routing(x_flat) |
| else: |
| raise ValueError(f"Unknown routing algorithm: {self.routing_algorithm}") |
| |
| final_output = torch.zeros_like(x_flat) |
| for i, expert in enumerate(self.experts): |
| expert_mask = router_output[:, i].unsqueeze(1) |
| active_tokens_indices = torch.where(expert_mask.squeeze() > 0)[0] |
| if active_tokens_indices.numel() > 0: |
| active_tokens = x_flat[active_tokens_indices] |
| expert_out = expert(active_tokens) |
| weighted_out = expert_out * expert_mask[active_tokens_indices] |
| final_output.index_add_(0, active_tokens_indices, weighted_out) |
|
|
| return final_output.view(batch_size, seq_len, self.d_model) |
|
|
| class GroupedQueryAttention(nn.Module): |
| """ |
| Implements Grouped-Query Attention (GQA). |
| - MHA is a special case of GQA where num_kv_heads == num_heads. |
| - MQA is a special case of GQA where num_kv_heads == 1. |
| """ |
| def __init__(self, d_model, num_heads, num_kv_heads): |
| super().__init__() |
| assert d_model % num_heads == 0, "d_model must be divisible by num_heads" |
| assert num_heads % num_kv_heads == 0, "num_heads must be divisible by num_kv_heads" |
|
|
| self.d_model = d_model |
| self.num_heads = num_heads |
| self.num_kv_heads = num_kv_heads |
| self.num_key_value_groups = num_heads // num_kv_heads |
| self.d_k = d_model // num_heads |
| |
| self.W_q = nn.Linear(d_model, d_model) |
| self.W_k = nn.Linear(d_model, self.num_kv_heads * self.d_k) |
| self.W_v = nn.Linear(d_model, self.num_kv_heads * self.d_k) |
| self.W_o = nn.Linear(d_model, d_model) |
| |
| def scaled_dot_product_attention(self, Q, K, V, mask=None): |
| attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) |
| if mask is not None: |
| attn_scores = attn_scores.masked_fill(mask == 0, -1e9) |
| attn_probs = F.softmax(attn_scores, dim=-1) |
| output = torch.matmul(attn_probs, V) |
| return output |
|
|
| def forward(self, q, k, v, mask=None): |
| batch_size = q.size(0) |
| |
| Q = self.W_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) |
| K = self.W_k(k).view(batch_size, -1, self.num_kv_heads, self.d_k).transpose(1, 2) |
| V = self.W_v(v).view(batch_size, -1, self.num_kv_heads, self.d_k).transpose(1, 2) |
| |
| if self.num_key_value_groups > 1: |
| K = K.repeat_interleave(self.num_key_value_groups, dim=1) |
| V = V.repeat_interleave(self.num_key_value_groups, dim=1) |
| |
| context = self.scaled_dot_product_attention(Q, K, V, mask) |
| context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) |
| |
| output = self.W_o(context) |
| return output |
|
|
| class PositionalEncoding(nn.Module): |
| def __init__(self, d_model, dropout=0.1, max_len=5000): |
| super().__init__() |
| self.dropout = nn.Dropout(p=dropout) |
| pe = torch.zeros(max_len, d_model) |
| position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) |
| div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) |
| pe[:, 0::2] = torch.sin(position * div_term) |
| pe[:, 1::2] = torch.cos(position * div_term) |
| self.register_buffer('pe', pe.unsqueeze(0)) |
|
|
| def forward(self, x): |
| x = x + self.pe[:, :x.size(1)] |
| return self.dropout(x) |
|
|
| class EncoderLayer(nn.Module): |
| def __init__(self, d_model, num_heads, num_kv_heads, d_ff, num_experts, top_k, routing_algorithm, dropout): |
| super().__init__() |
| self.self_attn = GroupedQueryAttention(d_model, num_heads, num_kv_heads) |
| self.moe_ffn = SparseMoE(d_model, num_experts, top_k, routing_algorithm, d_ff) |
| self.norm1 = nn.LayerNorm(d_model) |
| self.norm2 = nn.LayerNorm(d_model) |
| self.dropout = nn.Dropout(dropout) |
|
|
| def forward(self, x, mask): |
| attn_output = self.self_attn(x, x, x, mask) |
| x = self.norm1(x + self.dropout(attn_output)) |
| |
| moe_output = self.moe_ffn(x) |
| x = self.norm2(x + self.dropout(moe_output)) |
| return x |
|
|
| class DecoderLayer(nn.Module): |
| def __init__(self, d_model, num_heads, num_kv_heads, d_ff, num_experts, top_k, routing_algorithm, dropout): |
| super().__init__() |
| self.self_attn = GroupedQueryAttention(d_model, num_heads, num_kv_heads) |
| self.cross_attn = GroupedQueryAttention(d_model, num_heads, num_kv_heads) |
| self.moe_ffn = SparseMoE(d_model, num_experts, top_k, routing_algorithm, d_ff) |
| self.norm1 = nn.LayerNorm(d_model) |
| self.norm2 = nn.LayerNorm(d_model) |
| self.norm3 = nn.LayerNorm(d_model) |
| self.dropout = nn.Dropout(dropout) |
|
|
| def forward(self, x, enc_output, src_mask, tgt_mask): |
| attn_output = self.self_attn(x, x, x, tgt_mask) |
| x = self.norm1(x + self.dropout(attn_output)) |
| |
| cross_attn_output = self.cross_attn(x, enc_output, enc_output, src_mask) |
| x = self.norm2(x + self.dropout(cross_attn_output)) |
| |
| moe_output = self.moe_ffn(x) |
| x = self.norm3(x + self.dropout(moe_output)) |
| return x |
|
|
| class MoETransformer(nn.Module): |
| def __init__(self, config, vocab_size): |
| super().__init__() |
| self.config = config |
| self.encoder_embedding = nn.Embedding(vocab_size, config['d_model']) |
| self.decoder_embedding = nn.Embedding(vocab_size, config['d_model']) |
| self.positional_encoding = PositionalEncoding(config['d_model'], config['dropout']) |
|
|
| self.encoder_layers = nn.ModuleList([ |
| EncoderLayer(config['d_model'], config['num_heads'], config['num_kv_heads'], config['d_ff'], config['num_experts'], config['top_k'], config['routing_algorithm'], config['dropout']) |
| for _ in range(config['num_encoder_layers']) |
| ]) |
| self.decoder_layers = nn.ModuleList([ |
| DecoderLayer(config['d_model'], config['num_heads'], config['num_kv_heads'], config['d_ff'], config['num_experts'], config['top_k'], config['routing_algorithm'], config['dropout']) |
| for _ in range(config['num_decoder_layers']) |
| ]) |
|
|
| self.fc_out = nn.Linear(config['d_model'], vocab_size) |
|
|
| def generate_mask(self, src, tgt, pad_idx): |
| src_mask = (src != pad_idx).unsqueeze(1).unsqueeze(2) |
| tgt_pad_mask = (tgt != pad_idx).unsqueeze(1).unsqueeze(2) |
| seq_len = tgt.size(1) |
| tgt_sub_mask = torch.tril(torch.ones((seq_len, seq_len), device=tgt.device)).bool() |
| tgt_mask = tgt_pad_mask & tgt_sub_mask |
| return src_mask, tgt_mask |
|
|
| def forward(self, src, tgt, pad_idx=0): |
| src_mask, tgt_mask = self.generate_mask(src, tgt, pad_idx) |
| |
| src_emb = self.positional_encoding(self.encoder_embedding(src) * math.sqrt(self.config['d_model'])) |
| tgt_emb = self.positional_encoding(self.decoder_embedding(tgt) * math.sqrt(self.config['d_model'])) |
|
|
| enc_output = src_emb |
| for layer in self.encoder_layers: |
| enc_output = layer(enc_output, src_mask) |
| |
| dec_output = tgt_emb |
| for layer in self.decoder_layers: |
| dec_output = layer(dec_output, enc_output, src_mask, tgt_mask) |
| |
| return self.fc_out(dec_output) |
|
|
| def get_total_load_balancing_loss(self): |
| total_loss = 0 |
| for layer in self.encoder_layers + self.decoder_layers: |
| total_loss += layer.moe_ffn.load_balancing_loss |
| return total_loss |
|
|
| @torch.no_grad() |
| def generate(self, src, max_length, start_symbol, pad_idx=0): |
| self.eval() |
| device = next(self.parameters()).device |
| src = src.to(device) |
| batch_size = src.shape[0] |
|
|
| src_mask = (src != pad_idx).unsqueeze(1).unsqueeze(2) |
| src_emb = self.positional_encoding(self.encoder_embedding(src) * math.sqrt(self.config['d_model'])) |
| enc_output = src_emb |
| for layer in self.encoder_layers: |
| enc_output = layer(enc_output, src_mask) |
|
|
| tgt = torch.full((batch_size, 1), start_symbol, dtype=torch.long, device=device) |
|
|
| for _ in range(max_length - 1): |
| _, tgt_mask = self.generate_mask(src, tgt, pad_idx) |
| tgt_emb = self.positional_encoding(self.decoder_embedding(tgt) * math.sqrt(self.config['d_model'])) |
| dec_output = tgt_emb |
| for layer in self.decoder_layers: |
| dec_output = layer(dec_output, enc_output, src_mask, tgt_mask) |
| |
| logits = self.fc_out(dec_output[:, -1]) |
| next_token = torch.argmax(logits, dim=-1).unsqueeze(1) |
| tgt = torch.cat([tgt, next_token], dim=1) |
|
|
| if (tgt == 1).any(dim=-1).all(): |
| break |
|
|
| return tgt |
