moep / moep.py

Update moep.py

d6ed2d3 verified about 1 year ago

18.4 kB

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, AutoModel


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 ParallelLayer(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.dim = dim
        self.n_blocks = n_blocks
        self.blocks = nn.ModuleList([Layer_Block(dim, n_heads, ffn_multp, dropout, attn_dropout) for _ in range(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
        outs = [blk(x, causal_mask=causal_mask, key_padding_mask=key_padding_mask) for blk in self.blocks]
        
        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]

        return y_flat.view(B, S, D), aux_loss


class MoEPLMConfig(PretrainedConfig):
    model_type = "moep_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 MoEPModel(PreTrainedModel):
    config_class = MoEPLMConfig
    base_model_prefix = "model"

    def __init__(self, config: MoEPLMConfig):
        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 * 2,
                                     dim_out=config.parallel_dim))

        # Parallel Layers
        for _ in range(config.n_layers - 2):
            self.layers.append(ParallelLayer(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 * config.ffn_multp,
                                     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 MoEPForCausalLM(PreTrainedModel):
    config_class = MoEPLMConfig
    base_model_prefix = "model"

    def __init__(self, config: MoEPLMConfig):
        super().__init__(config)
        self.config = config
        self.model = MoEPModel(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("moep_lm", MoEPLMConfig)
AutoModelForCausalLM.register(MoEPLMConfig, MoEPForCausalLM)
AutoModel.register(MoEPLMConfig, MoEPModel)

	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, AutoModel


	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 ParallelLayer(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.dim = dim
	self.n_blocks = n_blocks
	self.blocks = nn.ModuleList([Layer_Block(dim, n_heads, ffn_multp, dropout, attn_dropout) for _ in range(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
	outs = [blk(x, causal_mask=causal_mask, key_padding_mask=key_padding_mask) for blk in self.blocks]

	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]

	return y_flat.view(B, S, D), aux_loss


	class MoEPLMConfig(PretrainedConfig):
	model_type = "moep_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 MoEPModel(PreTrainedModel):
	config_class = MoEPLMConfig
	base_model_prefix = "model"

	def __init__(self, config: MoEPLMConfig):
	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 * 2,
	dim_out=config.parallel_dim))

	# Parallel Layers
	for _ in range(config.n_layers - 2):
	self.layers.append(ParallelLayer(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 * config.ffn_multp,
	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 MoEPForCausalLM(PreTrainedModel):
	config_class = MoEPLMConfig
	base_model_prefix = "model"

	def __init__(self, config: MoEPLMConfig):
	super().__init__(config)
	self.config = config
	self.model = MoEPModel(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("moep_lm", MoEPLMConfig)
	AutoModelForCausalLM.register(MoEPLMConfig, MoEPForCausalLM)
	AutoModel.register(MoEPLMConfig, MoEPModel)