File size: 3,891 Bytes
6c728c0 0cbc2ad 6c728c0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.models.qwen2.modeling_qwen2 import Qwen2ForCausalLM, Qwen2MLP
class Qwen2ExpertMLP(nn.Module):
def __init__(self, config):
super().__init__()
self.hidden_size = config.hidden_size
self.intermediate_size = config.intermediate_size
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
self.act_fn = F.silu
nn.init.zeros_(self.gate_proj.weight)
nn.init.zeros_(self.up_proj.weight)
nn.init.zeros_(self.down_proj.weight)
def forward(self, x):
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
class Qwen2MoEMLP(nn.Module):
def __init__(self, config, shared_mlp: Qwen2MLP, num_experts: int, top_k: int, alpha_values):
super().__init__()
object.__setattr__(self, "shared_mlp", shared_mlp)
self.hidden_size = config.hidden_size
self.intermediate_size = config.intermediate_size
self.num_experts = num_experts
self.top_k = top_k
# Keep the dense/shared MLP weights under their original parameter names.
self.gate_proj = shared_mlp.gate_proj
self.up_proj = shared_mlp.up_proj
self.down_proj = shared_mlp.down_proj
self.act_fn = shared_mlp.act_fn
self.router = nn.Linear(self.hidden_size, self.num_experts, bias=False)
nn.init.zeros_(self.router.weight)
self.experts = nn.ModuleList([Qwen2ExpertMLP(config) for _ in range(self.num_experts)])
alpha_tensor = torch.tensor(alpha_values, dtype=torch.float32)
self.register_buffer("alpha_values", alpha_tensor, persistent=True)
def forward(self, x):
original_shape = x.shape
x_flat = x.reshape(-1, original_shape[-1])
shared_out = self.shared_mlp(x_flat)
logits = self.router(x_flat).float()
probs = F.softmax(logits, dim=-1)
top_k = min(self.top_k, probs.shape[-1])
weights, indices = torch.topk(probs, k=top_k, dim=-1)
weights = weights / weights.sum(dim=-1, keepdim=True).clamp_min(1e-9)
weights = weights.to(shared_out.dtype)
expert_out = torch.zeros_like(shared_out)
for expert_idx in torch.unique(indices).tolist():
assignment = indices == expert_idx
token_mask = assignment.any(dim=-1)
if not token_mask.any():
continue
selected_weights = (weights * assignment.to(weights.dtype)).sum(dim=-1, keepdim=True)
expert_result = self.experts[expert_idx](x_flat[token_mask])
alpha = self.alpha_values[expert_idx].to(expert_result.dtype)
expert_out[token_mask] += selected_weights[token_mask] * alpha * expert_result
return (shared_out + expert_out).reshape(original_shape)
class OutlierMoEForCausalLM(Qwen2ForCausalLM):
def __init__(self, config):
super().__init__(config)
moe_layer_indices = [int(layer) for layer in getattr(config, "moe_layer_indices", [])]
experts_per_layer = int(getattr(config, "experts_per_layer", getattr(config, "n_experts", 0)))
top_k = int(getattr(config, "top_k", 2))
alpha_values = getattr(config, "alpha_values", {})
for layer_idx in moe_layer_indices:
layer = self.model.layers[layer_idx]
per_layer_alpha = alpha_values.get(str(layer_idx), [0.0] * experts_per_layer)
layer.mlp = Qwen2MoEMLP(
config=config,
shared_mlp=layer.mlp,
num_experts=experts_per_layer,
top_k=top_k,
alpha_values=per_layer_alpha,
)
|