Create mini_moe.py

mini_moe.py

@@ -0,0 +1,374 @@
+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)  # batch, seq_len, dim
+            last_logits = logits[:, -1]  # batch, dim
+            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