| import torch |
| import torch.nn.functional as F |
| from torch import nn |
| from transformers.activations import ACT2FN |
|
|
|
|
| class FeedForward(nn.Module): |
| def __init__(self, config: "LMConfig", intermediate_size: int = None): |
| super().__init__() |
| intermediate_size = intermediate_size or config.intermediate_size |
| self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False) |
| self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False) |
| self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False) |
| self.act_fn = ACT2FN[config.hidden_act] |
|
|
| def forward(self, x): |
| return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) |
|
|
|
|
| class MOEFeedForward(nn.Module): |
| def __init__(self, config: "LMConfig"): |
| super().__init__() |
| self.config = config |
| self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) |
| self.experts = nn.ModuleList([FeedForward(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.num_experts)]) |
| self.act_fn = ACT2FN[config.hidden_act] |
|
|
| def forward(self, x): |
| batch_size, seq_len, hidden_dim = x.shape |
| x_flat = x.view(-1, hidden_dim) |
| scores = F.softmax(self.gate(x_flat), dim=-1) |
| topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False) |
| if self.config.norm_topk_prob: |
| topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20) |
| y = torch.zeros_like(x_flat) |
| for i, expert in enumerate(self.experts): |
| mask = (topk_idx == i) |
| if mask.any(): |
| token_idx = mask.any(dim=-1).nonzero().flatten() |
| weight = topk_weight[mask].view(-1, 1) |
| y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype)) |
| elif self.training: |
| y[0, 0] += 0 * sum(p.sum() for p in expert.parameters()) |
| if self.training and self.config.router_aux_loss_coef > 0: |
| load = F.one_hot(topk_idx, self.config.num_experts).float().mean(0) |
| self.aux_loss = (load * scores.mean(0)).sum() * self.config.num_experts * self.config.router_aux_loss_coef |
| else: |
| self.aux_loss = scores.new_zeros(1).squeeze() |
| return y.view(batch_size, seq_len, hidden_dim) |
|
|
