|
|
from torch import nn as nn |
|
|
import torch |
|
|
from torch.nn import functional as F |
|
|
from transformers import PretrainedConfig, PreTrainedModel |
|
|
|
|
|
|
|
|
class MiniMoEConfig(PretrainedConfig): |
|
|
model_type = "mini-moe" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
vocab_size=32000, |
|
|
num_layers=12, |
|
|
dim=1024, |
|
|
rope_base=10000, |
|
|
num_attention_q_heads=16, |
|
|
num_attention_kv_heads=8, |
|
|
num_expert=8, |
|
|
top_k=4, |
|
|
qkv_bias=False, |
|
|
drop_rate=0.0, |
|
|
use_aux_loss=True, |
|
|
**kwargs, |
|
|
): |
|
|
super().__init__(**kwargs) |
|
|
self.vocab_size = vocab_size |
|
|
self.num_layers = num_layers |
|
|
self.dim = dim |
|
|
self.rope_base = rope_base |
|
|
self.num_attention_q_heads = num_attention_q_heads |
|
|
self.num_attention_kv_heads = num_attention_kv_heads |
|
|
self.qkv_bias = qkv_bias |
|
|
self.drop_rate = drop_rate |
|
|
self.num_expert = num_expert |
|
|
self.top_k = top_k |
|
|
self.use_aux_loss = use_aux_loss |
|
|
self.auto_map = { |
|
|
"AutoConfig": "mini_moe.MiniMoEConfig", |
|
|
"AutoModelForCausalLM": "mini_moe.MiniMoE", |
|
|
} |
|
|
|
|
|
|
|
|
class RMSNorm(nn.Module): |
|
|
def __init__(self, dim): |
|
|
super().__init__() |
|
|
self.weight = nn.Parameter(torch.ones(dim)) |
|
|
|
|
|
def forward(self, x: torch.Tensor): |
|
|
norm_x = x / torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + 1e-8) |
|
|
output = self.weight * norm_x |
|
|
return output |
|
|
|
|
|
|
|
|
class RopePositionEmbedding(nn.Module): |
|
|
def __init__(self, dim: int, base=10000): |
|
|
super().__init__() |
|
|
inv_freq = 1 / base ** (torch.arange(0, dim, 2).float() / dim) |
|
|
inv_freq = inv_freq.unsqueeze(0) |
|
|
self.register_buffer("inv_freq", inv_freq) |
|
|
|
|
|
def rotate_half(self, x: torch.Tensor): |
|
|
odd = x[..., 1::2] |
|
|
even = x[..., 0::2] |
|
|
return torch.stack((-odd, even), dim=-1).flatten(-2) |
|
|
|
|
|
def apply_rope(self, x: torch.Tensor): |
|
|
x_len = x.shape[2] |
|
|
t = torch.arange(0, x_len, device=x.device, dtype=torch.float32).unsqueeze(1) |
|
|
freq = t * self.inv_freq |
|
|
freq = torch.repeat_interleave(freq, repeats=2, dim=-1)[None, None, :, :] |
|
|
xf = x.float() |
|
|
y = xf * freq.cos() + self.rotate_half(xf) * freq.sin() |
|
|
return y.to(x.dtype) |
|
|
|
|
|
def forward(self, q: torch.Tensor, k: torch.Tensor): |
|
|
return self.apply_rope(q), self.apply_rope(k) |
|
|
|
|
|
|
|
|
class GroupQueryAttention(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
num_attention_q_heads, |
|
|
num_attention_kv_heads, |
|
|
dim, |
|
|
qkv_bias, |
|
|
drop_rate, |
|
|
rope_base, |
|
|
): |
|
|
super().__init__() |
|
|
|
|
|
self.head_dim = dim // num_attention_q_heads |
|
|
|
|
|
assert dim % num_attention_q_heads == 0, "dim 必须被 Q 头数整除" |
|
|
assert ( |
|
|
num_attention_q_heads % num_attention_kv_heads == 0 |
|
|
), "Q头数必须是KV头数的整数倍" |
|
|
assert self.head_dim % 2 == 0, "head_dim 必须为偶数以应用 RoPE" |
|
|
|
|
|
self.q_proj = nn.Linear(dim, dim, bias=qkv_bias) |
|
|
self.k_proj = nn.Linear( |
|
|
dim, self.head_dim * num_attention_kv_heads, bias=qkv_bias |
|
|
) |
|
|
self.v_proj = nn.Linear( |
|
|
dim, self.head_dim * num_attention_kv_heads, bias=qkv_bias |
|
|
) |
|
|
self.out_proj = nn.Linear(dim, dim, bias=qkv_bias) |
|
|
|
|
|
self.num_repeat_kv = num_attention_q_heads // num_attention_kv_heads |
|
|
self.drop = nn.Dropout(drop_rate) |
|
|
|
|
|
self.position_embedding = RopePositionEmbedding(self.head_dim, rope_base) |
|
|
|
|
|
self.num_attention_q_heads = num_attention_q_heads |
|
|
self.num_attention_kv_heads = num_attention_kv_heads |
|
|
self.drop_rate = drop_rate |
|
|
|
|
|
def repeat_kv(self, k: torch.Tensor, v: torch.Tensor): |
|
|
k = k.repeat_interleave(self.num_repeat_kv, dim=1) |
|
|
v = v.repeat_interleave(self.num_repeat_kv, dim=1) |
|
|
return k, v |
|
|
|
|
|
def forward(self, x: torch.Tensor): |
|
|
batch_size, seq_len, dim = x.shape |
|
|
Q = ( |
|
|
self.q_proj(x) |
|
|
.reshape(batch_size, seq_len, self.num_attention_q_heads, self.head_dim) |
|
|
.transpose(1, 2) |
|
|
) |
|
|
K = ( |
|
|
self.k_proj(x) |
|
|
.reshape(batch_size, seq_len, self.num_attention_kv_heads, self.head_dim) |
|
|
.transpose(1, 2) |
|
|
) |
|
|
V = ( |
|
|
self.v_proj(x) |
|
|
.reshape(batch_size, seq_len, self.num_attention_kv_heads, self.head_dim) |
|
|
.transpose(1, 2) |
|
|
) |
|
|
|
|
|
Q, K = self.position_embedding(Q, K) |
|
|
|
|
|
K, V = self.repeat_kv(K, V) |
|
|
|
|
|
out = F.scaled_dot_product_attention( |
|
|
Q, K, V, dropout_p=self.drop_rate if self.training else 0.0, is_causal=True |
|
|
) |
|
|
out = out.transpose(1, 2).reshape(batch_size, seq_len, dim) |
|
|
out = self.out_proj(out) |
|
|
out = self.drop(out) |
|
|
return out |
|
|
|
|
|
|
|
|
class Expert(nn.Module): |
|
|
def __init__(self, dim, drop_rate): |
|
|
super().__init__() |
|
|
self.ffn = nn.Sequential( |
|
|
nn.Linear(dim, dim * 4), |
|
|
nn.SiLU(), |
|
|
nn.Linear(dim * 4, dim), |
|
|
nn.Dropout(drop_rate), |
|
|
) |
|
|
|
|
|
def forward(self, x): |
|
|
return self.ffn(x) |
|
|
|
|
|
|
|
|
class NoiseRouter(nn.Module): |
|
|
def __init__(self, num_expert, top_k, dim): |
|
|
super().__init__() |
|
|
self.gate = nn.Linear(dim, num_expert) |
|
|
self.noise_gate = nn.Linear(dim, num_expert) |
|
|
self.top_k = top_k |
|
|
|
|
|
def forward(self, x): |
|
|
gate = self.gate(x) |
|
|
logits = gate + torch.randn_like(gate) + self.noise_gate(x) |
|
|
|
|
|
top_k_val, top_k_ids = torch.topk(logits, k=self.top_k, dim=-1) |
|
|
scores = torch.full_like(logits, -torch.inf) |
|
|
scores.scatter_(dim=-1, index=top_k_ids, src=top_k_val) |
|
|
scores = scores.softmax(dim=-1) |
|
|
return scores, top_k_ids |
|
|
|
|
|
|
|
|
class SparseMoe(nn.Module): |
|
|
def __init__(self, num_expert, top_k, dim, drop_rate, use_aux_loss=True): |
|
|
super().__init__() |
|
|
self.route = NoiseRouter(num_expert=num_expert, top_k=top_k, dim=dim) |
|
|
self.experts = nn.ModuleList( |
|
|
[Expert(dim=dim, drop_rate=drop_rate) for _ in range(num_expert)] |
|
|
) |
|
|
self.use_aux_loss = use_aux_loss |
|
|
self.num_expert = num_expert |
|
|
|
|
|
def forward(self, x: torch.Tensor): |
|
|
batch_size, seq_len, dim = x.shape |
|
|
|
|
|
scores, indices = self.route(x) |
|
|
flatten_x = x.reshape(-1, dim) |
|
|
flatten_scores = scores.reshape(-1, scores.shape[-1]) |
|
|
|
|
|
final_out = torch.zeros_like(flatten_x) |
|
|
|
|
|
for i, expert in enumerate(self.experts): |
|
|
expert_mask = (indices == i).any(dim=-1) |
|
|
expert_mask = expert_mask.reshape(-1) |
|
|
if expert_mask.any(): |
|
|
expert_in = flatten_x[expert_mask] |
|
|
expert_out = expert(expert_in) |
|
|
expert_weight = flatten_scores[expert_mask, i].unsqueeze(1) |
|
|
expert_out = expert_weight * expert_out |
|
|
|
|
|
final_out[expert_mask] += expert_out |
|
|
|
|
|
final_out = final_out.reshape(batch_size, seq_len, dim) |
|
|
|
|
|
if self.use_aux_loss: |
|
|
importance = flatten_scores.mean(dim=0).float() |
|
|
uniform = torch.full_like(importance, fill_value=1.0 / self.num_expert).float() |
|
|
|
|
|
importance_log = (importance + 1e-8).log() |
|
|
uniform_log = uniform.log() |
|
|
|
|
|
aux_loss = F.kl_div( |
|
|
input=importance_log, target=uniform_log, log_target=True, reduction="sum" |
|
|
) |
|
|
return final_out, aux_loss |
|
|
return final_out |
|
|
|
|
|
|
|
|
class DecoderLayer(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
num_attention_q_heads, |
|
|
num_attention_kv_heads, |
|
|
dim, |
|
|
qkv_bias, |
|
|
drop_rate, |
|
|
rope_base, |
|
|
num_expert, |
|
|
top_k, |
|
|
use_aux_loss, |
|
|
): |
|
|
super().__init__() |
|
|
self.norm1 = RMSNorm(dim=dim) |
|
|
self.attn = GroupQueryAttention( |
|
|
num_attention_q_heads=num_attention_q_heads, |
|
|
num_attention_kv_heads=num_attention_kv_heads, |
|
|
dim=dim, |
|
|
qkv_bias=qkv_bias, |
|
|
drop_rate=drop_rate, |
|
|
rope_base=rope_base, |
|
|
) |
|
|
self.norm2 = RMSNorm(dim=dim) |
|
|
self.moe = SparseMoe( |
|
|
num_expert=num_expert, |
|
|
top_k=top_k, |
|
|
dim=dim, |
|
|
drop_rate=drop_rate, |
|
|
use_aux_loss=use_aux_loss, |
|
|
) |
|
|
self.use_aux_loss = use_aux_loss |
|
|
|
|
|
def forward(self, x): |
|
|
x = x + self.attn(self.norm1(x)) |
|
|
hidden_state = self.moe(self.norm2(x)) |
|
|
if self.use_aux_loss: |
|
|
x = x + hidden_state[0] |
|
|
aux_loss = hidden_state[1] |
|
|
|
|
|
return x, aux_loss |
|
|
else: |
|
|
x = x + hidden_state |
|
|
return x |
|
|
|
|
|
|
|
|
class MiniMoE(PreTrainedModel): |
|
|
model_type = "mini-moe" |
|
|
config_class = MiniMoEConfig |
|
|
|
|
|
def __init__(self, config: MiniMoEConfig, pretrain_ckpt=None): |
|
|
super().__init__(config) |
|
|
self.embedding = nn.Embedding(config.vocab_size, config.dim) |
|
|
self.layers = nn.ModuleList([]) |
|
|
for _ in range(config.num_layers): |
|
|
self.layers.append( |
|
|
DecoderLayer( |
|
|
num_attention_q_heads=config.num_attention_q_heads, |
|
|
num_attention_kv_heads=config.num_attention_kv_heads, |
|
|
dim=config.dim, |
|
|
qkv_bias=config.qkv_bias, |
|
|
drop_rate=config.drop_rate, |
|
|
rope_base=config.rope_base, |
|
|
num_expert=config.num_expert, |
|
|
top_k=config.top_k, |
|
|
use_aux_loss=config.use_aux_loss, |
|
|
) |
|
|
) |
|
|
self.norm = RMSNorm(dim=config.dim) |
|
|
self.head = nn.Linear(config.dim, config.vocab_size, bias=False) |
|
|
self.apply(self.init_weight) |
|
|
self.head.weight = self.embedding.weight |
|
|
self.use_aux_loss = config.use_aux_loss |
|
|
if pretrain_ckpt is not None: |
|
|
self.load_ckpt(pretrain_ckpt) |
|
|
|
|
|
def load_ckpt(self, ckpt_path): |
|
|
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) |
|
|
state_dict = ckpt["state_dict"] |
|
|
new_state_dict = {} |
|
|
for k, v in state_dict.items(): |
|
|
new_k = k[len("net._orig_mod.") :] |
|
|
new_state_dict[new_k] = v |
|
|
self.load_state_dict(new_state_dict, strict=True) |
|
|
print(f"load state dict from {ckpt_path}") |
|
|
|
|
|
def init_weight(self, m): |
|
|
if isinstance(m, nn.Linear): |
|
|
nn.init.normal_(m.weight, mean=0, std=0.02) |
|
|
if m.bias is not None: |
|
|
nn.init.constant_(m.bias, 0) |
|
|
elif isinstance(m, RMSNorm): |
|
|
nn.init.constant_(m.weight, 1) |
|
|
elif isinstance(m, nn.Embedding): |
|
|
nn.init.normal_(m.weight, mean=0, std=0.02) |
|
|
|
|
|
def forward(self, input_ids: torch.Tensor): |
|
|
hidden_state = self.embedding(input_ids) |
|
|
aux_loss = None |
|
|
for layer in self.layers: |
|
|
hidden_state = layer(hidden_state) |
|
|
if self.use_aux_loss: |
|
|
if aux_loss is None: |
|
|
aux_loss = hidden_state[1] |
|
|
else: |
|
|
aux_loss += hidden_state[1] |
|
|
hidden_state = hidden_state[0] |
|
|
|
|
|
hidden_state = self.norm(hidden_state) |
|
|
logits = self.head(hidden_state) |
|
|
if self.use_aux_loss: |
|
|
return logits, aux_loss |
|
|
return logits |
|
|
|
|
|
def top_k_sample(self, logits, top_k=5): |
|
|
|
|
|
weights, indices = torch.topk(logits, k=top_k, dim=-1) |
|
|
|
|
|
probs = torch.softmax(weights, dim=-1) |
|
|
chosssed_id = torch.multinomial(probs, num_samples=1) |
|
|
new_token = torch.gather(indices, dim=-1, index=chosssed_id) |
|
|
return new_token |
|
|
|
|
|
@torch.no_grad() |
|
|
def chat(self, conversations, tokenizer, max_new_token=256, top_k=5): |
|
|
ids = tokenizer.apply_chat_template( |
|
|
conversations, add_generation_prompt=True, tokenize=True |
|
|
) |
|
|
eos_ids = tokenizer.eos_token_id |
|
|
input_ids = torch.tensor(ids, dtype=torch.long).unsqueeze(0) |
|
|
for _ in range(max_new_token): |
|
|
|
|
|
logits = self(input_ids) |
|
|
last_logits = logits[:, -1] |
|
|
new_token = self.top_k_sample(last_logits, top_k=top_k) |
|
|
input_ids = torch.cat((input_ids, new_token), dim=-1) |
|
|
|
|
|
if new_token.detach()[0].cpu().item() == eos_ids: |
|
|
break |
|
|
|
|
|
output_id = input_ids.detach().cpu()[0].tolist() |
|
|
output_id = output_id[len(ids) :] |
|
|
answer = tokenizer.decode(output_id, skip_special_tokens=True) |
|
|
return answer |
