Add moe.py

moe.py

@@ -0,0 +1,558 @@
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple, List, Union 
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers import AutoConfig, AutoModelForCausalLM
+
+
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    CausalLMOutput,
+)
+
+
+@dataclass
+class BaseModelOutputWithAux(BaseModelOutput):
+    aux_loss: Optional[torch.Tensor] = None
+
+
+def create_causal_mask(seq_len: int, device=None, dtype=torch.float32):
+    mask = torch.full((seq_len, seq_len), torch.finfo(dtype).min, device=device, dtype=dtype)
+    mask = torch.triu(mask, diagonal=1)
+    return mask
+
+
+def create_sw_mask(seq_len: int, window_size: Optional[int], device=None, dtype=torch.float32):
+    if not window_size:
+        return None
+    i = torch.arange(seq_len, device=device)
+    j = torch.arange(seq_len, device=device)
+    allowed = (j[None, :] <= i[:, None]) & (j[None, :] >= i[:, None] - (window_size - 1))
+    mask = torch.zeros((seq_len, seq_len), device=device, dtype=dtype)
+    mask = mask.masked_fill(~allowed, torch.finfo(dtype).min)
+    return mask
+
+
+class Linear(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, b: bool = True):
+        super().__init__()
+        self.linear = nn.Linear(in_dim, out_dim, bias=b)
+
+    def forward(self, x: torch.Tensor):
+        return self.linear(x)
+
+
+class Layer_norm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5, b: bool = True):
+        super().__init__()
+        self.layer_norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=b)
+
+    def forward(self, x: torch.Tensor):
+        return self.layer_norm(x)
+
+
+class MLP(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int, dim_out: Optional[int] = None, b: bool = True):
+        super().__init__()
+        self.linear1 = nn.Linear(in_dim, hidden_dim, bias=b)
+        out_dim = in_dim if dim_out is None else dim_out
+        self.linear2 = nn.Linear(hidden_dim, out_dim, bias=b)
+
+    @staticmethod
+    def gelu_new(x: torch.Tensor):
+        return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * x ** 3)))
+
+    def forward(self, x: torch.Tensor):
+        x = self.linear1(x)
+        x = self.gelu_new(x)
+        x = self.linear2(x)
+        return x
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, n_heads: int, attn_dropout: float = 0.1):
+        super().__init__()
+        assert dim % n_heads == 0, "dim must be divisible by n_heads"
+        self.dim = dim
+        self.n_heads = n_heads
+        self.head_dim = dim // n_heads
+        self.Q = nn.Linear(dim, dim, bias=True)
+        self.K = nn.Linear(dim, dim, bias=True)
+        self.V = nn.Linear(dim, dim, bias=True)
+        self.O = nn.Linear(dim, dim, bias=True)
+        self.attn_dropout = nn.Dropout(attn_dropout)
+
+    def forward(self, x: torch.Tensor, *, causal_mask: Optional[torch.Tensor] = None, key_padding_mask: Optional[torch.Tensor] = None):
+        # x: [B,S,D]
+        B, S, D = x.shape
+        H = self.n_heads
+        Dh = self.head_dim
+
+        q = self.Q(x).view(B, S, H, Dh).transpose(1, 2)  # [B,H,S,Dh]
+        k = self.K(x).view(B, S, H, Dh).transpose(1, 2)
+        v = self.V(x).view(B, S, H, Dh).transpose(1, 2)
+
+        attn = (q @ k.transpose(-2, -1)) / math.sqrt(Dh)  # [B,H,S,S]
+        neg_inf = torch.finfo(attn.dtype).min
+
+        if causal_mask is not None:
+            attn = attn + causal_mask[None, None, :, :]
+
+        if key_padding_mask is not None:  # True at PAD keys
+            attn = attn.masked_fill(key_padding_mask[:, None, None, :], neg_inf)
+
+        attn = F.softmax(attn, dim=-1)
+        attn = self.attn_dropout(attn)
+        out = (attn @ v).transpose(1, 2).contiguous().view(B, S, D)
+
+        if key_padding_mask is not None:
+            query_mask = (~key_padding_mask).to(out.dtype).unsqueeze(-1)  # [B,S,1]
+            out = out * query_mask
+
+        return self.O(out)
+
+
+class Layer_Block(nn.Module):
+    def __init__(self, dim: int, n_heads: int, ffn_multp: int, dropout: float = 0.1, attn_dropout: float = 0.1):
+        super().__init__()
+        self.norm_1 = Layer_norm(dim)
+        self.norm_2 = Layer_norm(dim)
+        self.attention = SelfAttention(dim, n_heads, attn_dropout)
+
+        ### 
+        self.feedforward = MLP(dim, dim * ffn_multp)
+
+        ####
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, *, causal_mask=None, key_padding_mask=None):
+        residual = x
+        x = self.norm_1(x)
+        x = self.attention(x, causal_mask=causal_mask, key_padding_mask=key_padding_mask)
+        x = self.dropout1(x)
+        x = residual + x
+
+        residual = x
+        x = self.norm_2(x)
+        x = self.feedforward(x)
+        x = self.dropout2(x)
+        x = residual + x
+        return x
+
+
+class Top_k_Router(nn.Module):
+    def __init__(self, dim: int, n_experts: int, k: int):
+        super().__init__()
+        self.n_experts = n_experts
+        self.k = k
+        self.router = Linear(dim, n_experts)
+
+    def forward(self, x: torch.Tensor):
+        # x: [T, D]
+        logits = self.router(x)
+        probs = F.softmax(logits, dim=-1)
+        topk_prob, topk_idx = torch.topk(probs, k=self.k, dim=-1)
+        topk_prob = topk_prob / (topk_prob.sum(dim=-1, keepdim=True) + 1e-9)
+
+        p_mean = probs.mean(dim=0)
+        T, E = probs.size()
+        one_hot = F.one_hot(topk_idx, num_classes=E).sum(dim=1)
+        f_mean = one_hot.float().mean(dim=0) / float(self.k)
+        aux_loss = (E * (p_mean * f_mean).sum())
+
+        return topk_idx, topk_prob, aux_loss
+
+class MoE_Block(nn.Module):
+    def __init__(self, dim_in: int, n_experts: int, k: int, dim_hidden: int, dim_out: Optional[int] = None):
+        super().__init__()
+        self.n_experts = n_experts
+        self.k = k
+        self.dim_out = dim_in if dim_out is None else dim_out
+        self.router = Top_k_Router(dim_in, n_experts, k)
+        self.experts = nn.ModuleList([nn.Linear(dim_in, self.dim_out, bias=True) for _ in range(n_experts)])
+
+    def forward(self, x: torch.Tensor):
+        # x: [B,S,D]
+        B, S, D = x.shape
+        T = B * S
+        x_flat = x.view(T, D)
+
+        topk_idx, topk_prob, aux_loss = self.router(x_flat)  # [T,k], [T,k]
+        y_flat = x_flat.new_zeros((T, self.dim_out))
+
+        for slot in range(self.k):
+            e_ids = topk_idx[:, slot]
+            e_prob = topk_prob[:, slot].unsqueeze(-1)  # [T,1]
+            for e in range(self.n_experts):
+                idx = (e_ids == e).nonzero(as_tuple=False).squeeze(1)
+                if idx.numel() == 0:
+                    continue
+                y_e = self.experts[e](x_flat.index_select(0, idx))  # [Ne, dim_out]
+                y_flat.index_add_(0, idx, y_e * e_prob.index_select(0, idx))
+
+        return y_flat.view(B, S, self.dim_out), aux_loss
+
+
+
+class MoELayer(nn.Module):
+    def __init__(self, dim: int, n_heads: int, n_blocks: int, k: int, ffn_multp: int, dropout: float = 0.1, attn_dropout: float = 0.1):
+        super().__init__()
+
+        self.norm_1 = Layer_norm(dim)
+        self.norm_2 = Layer_norm(dim)
+        self.attention = SelfAttention(dim, n_heads, attn_dropout)
+
+        self.feedforward = nn.ModuleList([MLP(dim, dim * ffn_multp) for _ in range(n_blocks)])
+        
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        self.dim = dim
+        self.n_blocks = n_blocks
+        
+        self.router = Top_k_Router(dim, n_blocks, k)
+
+    def forward(self, x: torch.Tensor, *, causal_mask=None, key_padding_mask=None):
+        B, S, D = x.shape
+        T = B * S
+
+        residual = x
+        x = self.norm_1(x)
+        x = self.attention(x, causal_mask=causal_mask, key_padding_mask=key_padding_mask)
+        x = self.dropout1(x)
+        x = residual + x
+
+        residual = x
+        x = self.norm_2(x)
+        
+        outs = [blk(x) for blk in self.feedforward]
+        
+        Y = torch.stack([y.view(T, D) for y in outs], dim=1)  # [T, n_blocks, D]
+
+        x_flat = x.view(T, D)
+        topk_idx, topk_prob, aux_loss = self.router(x_flat)  # [T,k], [T,k]
+
+        gather_idx = topk_idx.unsqueeze(-1).expand(-1, -1, D)  # [T,k,D]
+        Y_topk = torch.gather(Y, dim=1, index=gather_idx)      # [T,k,D]
+        y_flat = (Y_topk * topk_prob.unsqueeze(-1)).sum(dim=1) # [T,D]
+
+        y = y_flat.view(B, S, D)
+   
+        x = self.dropout2(y)
+        x = residual + x
+        
+        return x, aux_loss
+
+
+class MoELMConfig(PretrainedConfig):
+    model_type = "moe_lm"
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        bos_token_id = 1,
+        eos_token_id = 2,
+        pad_token_id = 3,
+        mask_token_id = 4,
+        max_position_embeddings: int = 1024,
+        n_layers: int = 12,
+        ffn_multp: int = 4,
+        dim: int = 348,
+        n_heads: int = 6,
+        # Parallel specifics
+        n_parallel_blocks: int = 4,
+        parallel_dim: int = 192,
+        parallel_n_heads: int = 3,
+        parallel_k: int = 2,
+        sliding_window: Optional[int] = None,
+        # MoE specifics
+        n_experts: int = 4,
+        moe_k: int = 2,
+        # Dropouts / norms
+        attn_pdrop: float = 0.1,
+        embd_pdrop: float = 0.1,
+        resid_pdrop: float = 0.1,
+        layer_norm_epsilon: float = 1e-5,
+        tie_word_embeddings: bool = True,
+        **kwargs,
+    ):
+        super().__init__(
+            vocab_size=vocab_size,
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )
+
+        # Special tokens
+        self.mask_token_id = mask_token_id
+
+        # Core dims
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.n_layers = n_layers
+        self.ffn_multp = ffn_multp
+        self.dim = dim
+        self.n_heads = n_heads
+
+        # Parallel specifics
+        self.n_parallel_blocks = n_parallel_blocks
+        self.parallel_dim = parallel_dim
+        self.parallel_n_heads = parallel_n_heads
+        self.parallel_k = parallel_k
+        self.sliding_window = sliding_window
+
+        # MoE specifics
+        self.n_experts = n_experts
+        self.moe_k = moe_k
+
+        # Dropouts / norms
+        self.attn_pdrop = attn_pdrop
+        self.embd_pdrop = embd_pdrop
+        self.resid_pdrop = resid_pdrop
+        self.layer_norm_epsilon = layer_norm_epsilon
+
+        # Tying
+        self.tie_word_embeddings = tie_word_embeddings
+
+class MoEModel(PreTrainedModel):
+    config_class = MoELMConfig
+
+    def __init__(self, config: MoELMConfig):
+        super().__init__(config)
+        self.config = config
+
+        # --- embeddings ---
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.dim, padding_idx=config.pad_token_id)
+        self.embed_positions = nn.Embedding(config.max_position_embeddings, config.dim)
+        self.embd_dropout = nn.Dropout(config.embd_pdrop)
+
+    
+        self.layers = nn.ModuleList()
+
+        # Large Layer
+        self.layers.append(Layer_Block(config.dim, 
+                                       config.n_heads,
+                                       config.ffn_multp,
+                                       dropout=config.resid_pdrop, 
+                                       attn_dropout=config.attn_pdrop))
+
+        # MoE Linear
+        self.layers.append(MoE_Block(config.dim, 
+                                     config.n_experts, 
+                                     config.moe_k, 
+                                     dim_hidden=config.dim,
+                                     dim_out=config.parallel_dim))
+
+        # Parallel Layers
+        for _ in range(config.n_layers - 2):
+            self.layers.append(MoELayer(config.parallel_dim, 
+                                             config.parallel_n_heads,
+                                             config.n_parallel_blocks, 
+                                             config.parallel_k,
+                                             config.ffn_multp,
+                                             dropout=config.resid_pdrop, 
+                                             attn_dropout=config.attn_pdrop))
+
+        # MoE Linear
+        self.layers.append(MoE_Block(config.parallel_dim, 
+                                     config.n_experts, 
+                                     config.moe_k, 
+                                     dim_hidden=config.parallel_dim,
+                                     dim_out=config.dim))
+        
+        # Large Layer
+        self.layers.append(Layer_Block(config.dim, 
+                                       config.n_heads,
+                                       config.ffn_multp,
+                                       dropout=config.resid_pdrop, 
+                                       attn_dropout=config.attn_pdrop))
+
+        self.norm = Layer_norm(config.dim)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, new_embeddings):
+        self.embed_tokens = new_embeddings
+
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+        return_aux: Optional[bool] = False,
+        return_aux_dict: Optional[bool] = False,
+        **unused,
+    ) -> Union[BaseModelOutputWithAux, BaseModelOutput, Tuple]:
+
+        device = input_ids.device
+        B, S = input_ids.shape
+
+        key_padding = input_ids.eq(self.config.pad_token_id) if attention_mask is None else (attention_mask == 0)
+
+        # embeddings
+        x = self.embed_tokens(input_ids)
+
+        mask_dtype = x.dtype
+        causal = create_causal_mask(S, device=device, dtype=mask_dtype)
+        sw = create_sw_mask(S, self.config.sliding_window, device=device, dtype=mask_dtype)
+        sw_mask = causal if sw is None else torch.minimum(causal, sw)
+
+        
+        pos = torch.arange(S, device=device)    
+        pos = pos.unsqueeze(0).expand(B, S)
+        x = x + self.embed_positions(pos)
+        x = self.embd_dropout(x)
+
+        aux_total = None
+
+
+        # First large
+        x = self.layers[0](x, causal_mask=causal, key_padding_mask=key_padding)
+
+        # MoE shrink
+        x, aux = self.layers[1](x)
+ 
+        if aux is not None:
+            aux_total = aux
+
+        idx = 2
+        for i in range(self.config.n_layers - 2):
+            use_sw = (self.config.sliding_window is not None) and (i % 2 == 1)
+            mask = sw_mask if use_sw else causal
+            x, aux = self.layers[idx](x, causal_mask=mask, key_padding_mask=key_padding)
+            if aux is not None:
+                aux_total += aux
+            idx += 1
+
+        # MoE grow
+        x, aux = self.layers[idx](x)
+        
+        if aux is not None:
+                aux_total += aux
+            
+        idx += 1
+
+        # Second large
+        x = self.layers[idx](x, causal_mask=causal, key_padding_mask=key_padding)
+
+        last_hidden = self.norm(x)
+
+        if not return_dict:
+            return (last_hidden, aux_total) if return_aux else (last_hidden,)
+
+        if return_aux_dict:
+            return BaseModelOutputWithAux(
+                last_hidden_state=last_hidden,
+                hidden_states=None,
+                attentions=None,
+                aux_loss=aux_total,
+            )
+        else:
+            return BaseModelOutput(
+                last_hidden_state=last_hidden,
+                hidden_states=None,
+                attentions=None,
+            )
+
+
+class MoEForCausalLM(PreTrainedModel):
+    config_class = MoELMConfig
+
+    def __init__(self, config: MoELMConfig):
+        super().__init__(config)
+        self.config = config
+        self.model = MoEModel(config)
+        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, new_embeddings):
+        self.model.set_input_embeddings(new_embeddings)
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            nn.init.zeros_(module.bias)
+            
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        return_dict: bool = True,
+        return_aux: Optional[bool] = False,
+        aux_loss_weight: Optional[float] = 0.2,
+        **unused,
+    ) -> Union[CausalLMOutput, Tuple]:
+    
+
+        if return_aux:
+            base_out = self.model(input_ids=input_ids, 
+                                  attention_mask=attention_mask, 
+                                  return_aux_dict=True,
+                                  return_aux=True)
+        else: 
+            base_out = self.model(input_ids=input_ids, 
+                                  attention_mask=attention_mask)
+    
+        logits = self.lm_head(base_out.last_hidden_state).float()
+
+        loss = None
+        
+        if labels is not None:
+            
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+
+            pad_id = self.config.pad_token_id
+            if pad_id is not None:
+                shift_labels = shift_labels.masked_fill(shift_labels == pad_id, -100)
+
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+            
+            aux = getattr(base_out, "aux_loss", None)
+            if aux is not None:
+                loss = loss + aux_loss_weight * aux
+
+        if not return_dict:
+            return (loss, logits) if loss is not None else (logits,)
+
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=base_out.hidden_states,
+            attentions=base_out.attentions,
+        )
+
+AutoConfig.register("moe_lm", MoELMConfig)
+AutoModelForCausalLM.register(MoELMConfig, MoEForCausalLM)
+
+
+