moep_gemma / moep_gemma.py

Update moep_gemma.py

b78146b verified 12 months ago

20.2 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
    one_hot: Optional[torch.Tensor] = None

@dataclass
class CausalLMOutputWithAux(CausalLMOutput):
    aux_loss: Optional[torch.Tensor] = None
    one_hot: 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 = False):
        super().__init__()
        self.linear = nn.Linear(in_dim, out_dim, bias=b)

    def forward(self, x: torch.Tensor):
        return self.linear(x)
        

class RMS_norm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5, b=True):
        super().__init__()
        self.rms_norm = nn.RMSNorm(dim, eps=eps, elementwise_affine=b)
        
    def forward(self, x: torch.Tensor):
        return self.rms_norm(x)


class GeGLU(nn.Module):
    def __init__(self, in_dim, hidden_dim, b=False):
        super().__init__()        
        self.linear1 = nn.Linear(in_dim, hidden_dim, bias=b)
        self.linear2 = nn.Linear(in_dim, hidden_dim, bias=b)
        self.linear3 = nn.Linear(hidden_dim, in_dim, bias=b)

    def forward(self, x: torch.Tensor):
        gate = F.gelu(self.linear1(x), approximate='tanh')
        return self.linear3(gate * self.linear2(x))


# --- RoPE --- #

def precompute_theta_pos_frequencies(emb_dim: int, seq_len: int, device: str, theta: float = 10000.0):
    assert emb_dim % 2 == 0, "Dimension must be divisible by 2"
    theta_numerator = torch.arange(0, emb_dim, 2, device=device).float()
    theta_values = 1.0 / (theta ** (theta_numerator / emb_dim))
    m = torch.arange(seq_len, device=device)
    freqs_matrix = torch.outer(m, theta_values).float()
    freqs_complex = torch.polar(torch.ones_like(freqs_matrix, device=device), freqs_matrix)
    return freqs_complex

def RoPE(x: torch.Tensor, freqs: torch.Tensor, index: int = -1):
    ndim = x.ndim
    assert 1 < ndim
    assert freqs.shape[-1] == x.shape[-1]
    
    if index != -1:
        # Ensure that index cannot be less than 1 --> (1 == first token position) 
        # and more than freqs.shape[-2] --> (index > max sequence length)
        index = max(1, min(index, freqs.shape[-2]))
        freqs = freqs[(index -1), :].unsqueeze(0)
        
    if freqs.shape[-2] != x.shape[-2] and index == -1:
        # if in x seq_len != max lenght, drop unneeded m frequence positions 
        freqs = freqs[:x.shape[-2], :]
    
    # from (seq_len, head_dim) to (1, 1, seq_len, head_dim)
    freqs = freqs.unsqueeze(0).unsqueeze(0)
    rotated_x = x * freqs
     
    return rotated_x

def apply_rotary_embeddings(x: torch.Tensor, freqs_complex: torch.Tensor):
    x_reshaped = x.reshape(*x.shape[:-1], -1, 2)  
    x_complex = torch.view_as_complex(x_reshaped.float())
    x_rotated = RoPE(x_complex, freqs_complex)
    x_out = torch.view_as_real(x_rotated)
    x_out = x_out.reshape(*x.shape).type_as(x)
    return x_out


class SelfAttention(nn.Module):
    def __init__(self, dim: int, n_heads: int):
        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)
        self.K = nn.Linear(dim, dim)
        self.V = nn.Linear(dim, dim)
        self.O = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor, *, causal_mask: Optional[torch.Tensor] = None, key_padding_mask: Optional[torch.Tensor] = None, freqs: 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)

        q = apply_rotary_embeddings(q, freqs)
        k = apply_rotary_embeddings(k, freqs)

        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)
        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):
        super().__init__()
        self.norm_1 = RMS_norm(dim)
        self.norm_2 = RMS_norm(dim)
        self.attention = SelfAttention(dim, n_heads)
        self.hidden_dim = int(dim * 4)
        self.feedforward = GeGLU(dim, self.hidden_dim)
    
    def forward(self, x: torch.Tensor, *, causal_mask=None, key_padding_mask=None, freqs=None):
        residual = x
        x = self.norm_1(x)
        x = self.attention(x, causal_mask=causal_mask, key_padding_mask=key_padding_mask, freqs=freqs)
        x = residual + x

        residual = x
        x = self.norm_2(x)
        x = self.feedforward(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, one_hot


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, one_hot = 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, one_hot.view(B, S, self.n_experts)


class ParallelLayer(nn.Module):
    def __init__(self, dim: int, n_heads: int, n_blocks: int, k: int):
        super().__init__()
        self.dim = dim
        self.n_blocks = n_blocks
        self.blocks = nn.ModuleList([Layer_Block(dim, n_heads) 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, freqs=None):
        B, S, D = x.shape
        T = B * S
        outs = [blk(x, causal_mask=causal_mask, key_padding_mask=key_padding_mask, freqs=freqs) 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, one_hot = self.router(x_flat)

        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, one_hot.view(B, S, self.n_blocks)


class MoEPgemmaLMConfig(PretrainedConfig):
    model_type = "moep_gemma_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_context_length: 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,
        # Norms
        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_context_length = max_context_length
        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

        # norms
        self.layer_norm_epsilon = layer_norm_epsilon

        # Tying
        self.tie_word_embeddings = tie_word_embeddings

class MoEPgemmaModel(PreTrainedModel):
    config_class = MoEPgemmaLMConfig
    base_model_prefix = "model"

    def __init__(self, config: MoEPgemmaLMConfig):
        super().__init__(config)
        self.config = config

        # --- embeddings ---
        self.embed_tokens = nn.Embedding(config.vocab_size, config.dim, padding_idx=config.pad_token_id)
     
        self.layers = nn.ModuleList()

        # Large Layer
        self.layers.append(Layer_Block(config.dim, 
                                       config.n_heads))

        # 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))

        # 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))

        self.norm = RMS_norm(config.dim)
        self.post_init()

        self.sliding_window = config.sliding_window

        self.freqs_p = None

        self.freqs = None


    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

        if self.freqs is None or self.freqs.device != device:
            self.freqs = precompute_theta_pos_frequencies(
                int(self.config.dim / self.config.n_heads),
                self.config.max_context_length,
                device=device
            )
        
        if self.freqs_p is None or self.freqs_p.device != device:
            self.freqs_p = precompute_theta_pos_frequencies(
                int(self.config.parallel_dim / self.config.parallel_n_heads),
                self.config.max_context_length,
                device=device
            )

        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.sliding_window, device=self.device, dtype=mask_dtype)
        sw_mask = None if sw is None else torch.minimum(causal, sw)
        
      
        aux_total = None
        all_one_hots = []
        
        
        freqs = self.freqs

        # First large
        x = self.layers[0](x, causal_mask=causal, key_padding_mask=key_padding, freqs=freqs)

        # MoE shrink
        x, aux, one_hot = self.layers[1](x)
 
        if aux is not None:
            aux_total = aux

        if one_hot is not None:
            all_one_hots.append(one_hot)

        freqs_p = self.freqs_p
        
        idx = 2
        for i in range(self.config.n_layers - 2):
            
            use_sw = (self.sliding_window is not None) and (i % 2 == 1)
            mask = sw_mask if (use_sw and sw_mask is not None) else causal
            
            x, aux, one_hot = self.layers[idx](x, causal_mask=mask, key_padding_mask=key_padding, freqs=freqs_p)
            if aux is not None:
                aux_total += aux
            if one_hot is not None:
                all_one_hots.append(one_hot)
            idx += 1

        # MoE grow
        x, aux, one_hot = self.layers[idx](x)
        
        if aux is not None:
                aux_total += aux

        if one_hot is not None:
            all_one_hots.append(one_hot)
            
        idx += 1

        # Second large
        x = self.layers[idx](x, causal_mask=causal, key_padding_mask=key_padding, freqs=freqs)

        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,
                one_hot=all_one_hots,
            )
        else:
            return BaseModelOutput(
                last_hidden_state=last_hidden,
                hidden_states=None,
                attentions=None,
            )


class MoEPgemmaForCausalLM(PreTrainedModel):
    config_class = MoEPgemmaLMConfig
    base_model_prefix = "model"

    def __init__(self, config: MoEPgemmaLMConfig):
        super().__init__(config)
        self.config = config
        self.model = MoEPgemmaModel(config)
        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
        self.post_init()
        if config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight

    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)
 
    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 return_aux:
            return CausalLMOutputWithAux(
                loss=loss,
                logits=logits,
                hidden_states=base_out.hidden_states,
                attentions=base_out.attentions,
                aux_loss=getattr(base_out, "aux_loss", None),
                one_hot=getattr(base_out, "one_hot", None),
            )
        else:
            return CausalLMOutput(
                loss=loss,
                logits=logits,
                hidden_states=base_out.hidden_states,
                attentions=base_out.attentions,
            )

AutoConfig.register("moep_gemma_lm", MoEPgemmaLMConfig)
AutoModelForCausalLM.register(MoEPgemmaLMConfig, MoEPgemmaForCausalLM)
AutoModel.register(MoEPgemmaLMConfig, MoEPgemmaModel)

	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
	one_hot: Optional[torch.Tensor] = None

	@dataclass
	class CausalLMOutputWithAux(CausalLMOutput):
	aux_loss: Optional[torch.Tensor] = None
	one_hot: 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 = False):
	super().__init__()
	self.linear = nn.Linear(in_dim, out_dim, bias=b)

	def forward(self, x: torch.Tensor):
	return self.linear(x)


	class RMS_norm(nn.Module):
	def __init__(self, dim: int, eps: float = 1e-5, b=True):
	super().__init__()
	self.rms_norm = nn.RMSNorm(dim, eps=eps, elementwise_affine=b)

	def forward(self, x: torch.Tensor):
	return self.rms_norm(x)


	class GeGLU(nn.Module):
	def __init__(self, in_dim, hidden_dim, b=False):
	super().__init__()
	self.linear1 = nn.Linear(in_dim, hidden_dim, bias=b)
	self.linear2 = nn.Linear(in_dim, hidden_dim, bias=b)
	self.linear3 = nn.Linear(hidden_dim, in_dim, bias=b)

	def forward(self, x: torch.Tensor):
	gate = F.gelu(self.linear1(x), approximate='tanh')
	return self.linear3(gate * self.linear2(x))


	# --- RoPE --- #

	def precompute_theta_pos_frequencies(emb_dim: int, seq_len: int, device: str, theta: float = 10000.0):
	assert emb_dim % 2 == 0, "Dimension must be divisible by 2"
	theta_numerator = torch.arange(0, emb_dim, 2, device=device).float()
	theta_values = 1.0 / (theta ** (theta_numerator / emb_dim))
	m = torch.arange(seq_len, device=device)
	freqs_matrix = torch.outer(m, theta_values).float()
	freqs_complex = torch.polar(torch.ones_like(freqs_matrix, device=device), freqs_matrix)
	return freqs_complex

	def RoPE(x: torch.Tensor, freqs: torch.Tensor, index: int = -1):
	ndim = x.ndim
	assert 1 < ndim
	assert freqs.shape[-1] == x.shape[-1]

	if index != -1:
	# Ensure that index cannot be less than 1 --> (1 == first token position)
	# and more than freqs.shape[-2] --> (index > max sequence length)
	index = max(1, min(index, freqs.shape[-2]))
	freqs = freqs[(index -1), :].unsqueeze(0)

	if freqs.shape[-2] != x.shape[-2] and index == -1:
	# if in x seq_len != max lenght, drop unneeded m frequence positions
	freqs = freqs[:x.shape[-2], :]

	# from (seq_len, head_dim) to (1, 1, seq_len, head_dim)
	freqs = freqs.unsqueeze(0).unsqueeze(0)
	rotated_x = x * freqs

	return rotated_x

	def apply_rotary_embeddings(x: torch.Tensor, freqs_complex: torch.Tensor):
	x_reshaped = x.reshape(*x.shape[:-1], -1, 2)
	x_complex = torch.view_as_complex(x_reshaped.float())
	x_rotated = RoPE(x_complex, freqs_complex)
	x_out = torch.view_as_real(x_rotated)
	x_out = x_out.reshape(*x.shape).type_as(x)
	return x_out


	class SelfAttention(nn.Module):
	def __init__(self, dim: int, n_heads: int):
	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)
	self.K = nn.Linear(dim, dim)
	self.V = nn.Linear(dim, dim)
	self.O = nn.Linear(dim, dim)

	def forward(self, x: torch.Tensor, *, causal_mask: Optional[torch.Tensor] = None, key_padding_mask: Optional[torch.Tensor] = None, freqs: 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)

	q = apply_rotary_embeddings(q, freqs)
	k = apply_rotary_embeddings(k, freqs)

	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)
	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):
	super().__init__()
	self.norm_1 = RMS_norm(dim)
	self.norm_2 = RMS_norm(dim)
	self.attention = SelfAttention(dim, n_heads)
	self.hidden_dim = int(dim * 4)
	self.feedforward = GeGLU(dim, self.hidden_dim)

	def forward(self, x: torch.Tensor, *, causal_mask=None, key_padding_mask=None, freqs=None):
	residual = x
	x = self.norm_1(x)
	x = self.attention(x, causal_mask=causal_mask, key_padding_mask=key_padding_mask, freqs=freqs)
	x = residual + x

	residual = x
	x = self.norm_2(x)
	x = self.feedforward(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, one_hot


	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, one_hot = 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, one_hot.view(B, S, self.n_experts)


	class ParallelLayer(nn.Module):
	def __init__(self, dim: int, n_heads: int, n_blocks: int, k: int):
	super().__init__()
	self.dim = dim
	self.n_blocks = n_blocks
	self.blocks = nn.ModuleList([Layer_Block(dim, n_heads) 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, freqs=None):
	B, S, D = x.shape
	T = B * S
	outs = [blk(x, causal_mask=causal_mask, key_padding_mask=key_padding_mask, freqs=freqs) 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, one_hot = self.router(x_flat)

	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, one_hot.view(B, S, self.n_blocks)


	class MoEPgemmaLMConfig(PretrainedConfig):
	model_type = "moep_gemma_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_context_length: 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,
	# Norms
	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_context_length = max_context_length
	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

	# norms
	self.layer_norm_epsilon = layer_norm_epsilon

	# Tying
	self.tie_word_embeddings = tie_word_embeddings

	class MoEPgemmaModel(PreTrainedModel):
	config_class = MoEPgemmaLMConfig
	base_model_prefix = "model"

	def __init__(self, config: MoEPgemmaLMConfig):
	super().__init__(config)
	self.config = config

	# --- embeddings ---
	self.embed_tokens = nn.Embedding(config.vocab_size, config.dim, padding_idx=config.pad_token_id)

	self.layers = nn.ModuleList()

	# Large Layer
	self.layers.append(Layer_Block(config.dim,
	config.n_heads))

	# 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))

	# 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))

	self.norm = RMS_norm(config.dim)
	self.post_init()

	self.sliding_window = config.sliding_window

	self.freqs_p = None

	self.freqs = None


	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

	if self.freqs is None or self.freqs.device != device:
	self.freqs = precompute_theta_pos_frequencies(
	int(self.config.dim / self.config.n_heads),
	self.config.max_context_length,
	device=device
	)

	if self.freqs_p is None or self.freqs_p.device != device:
	self.freqs_p = precompute_theta_pos_frequencies(
	int(self.config.parallel_dim / self.config.parallel_n_heads),
	self.config.max_context_length,
	device=device
	)

	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.sliding_window, device=self.device, dtype=mask_dtype)
	sw_mask = None if sw is None else torch.minimum(causal, sw)


	aux_total = None
	all_one_hots = []


	freqs = self.freqs

	# First large
	x = self.layers[0](x, causal_mask=causal, key_padding_mask=key_padding, freqs=freqs)

	# MoE shrink
	x, aux, one_hot = self.layers[1](x)

	if aux is not None:
	aux_total = aux

	if one_hot is not None:
	all_one_hots.append(one_hot)

	freqs_p = self.freqs_p

	idx = 2
	for i in range(self.config.n_layers - 2):

	use_sw = (self.sliding_window is not None) and (i % 2 == 1)
	mask = sw_mask if (use_sw and sw_mask is not None) else causal

	x, aux, one_hot = self.layers[idx](x, causal_mask=mask, key_padding_mask=key_padding, freqs=freqs_p)
	if aux is not None:
	aux_total += aux
	if one_hot is not None:
	all_one_hots.append(one_hot)
	idx += 1

	# MoE grow
	x, aux, one_hot = self.layers[idx](x)

	if aux is not None:
	aux_total += aux

	if one_hot is not None:
	all_one_hots.append(one_hot)

	idx += 1

	# Second large
	x = self.layers[idx](x, causal_mask=causal, key_padding_mask=key_padding, freqs=freqs)

	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,
	one_hot=all_one_hots,
	)
	else:
	return BaseModelOutput(
	last_hidden_state=last_hidden,
	hidden_states=None,
	attentions=None,
	)


	class MoEPgemmaForCausalLM(PreTrainedModel):
	config_class = MoEPgemmaLMConfig
	base_model_prefix = "model"

	def __init__(self, config: MoEPgemmaLMConfig):
	super().__init__(config)
	self.config = config
	self.model = MoEPgemmaModel(config)
	self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
	self.post_init()
	if config.tie_word_embeddings:
	self.lm_head.weight = self.model.embed_tokens.weight

	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)

	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 return_aux:
	return CausalLMOutputWithAux(
	loss=loss,
	logits=logits,
	hidden_states=base_out.hidden_states,
	attentions=base_out.attentions,
	aux_loss=getattr(base_out, "aux_loss", None),
	one_hot=getattr(base_out, "one_hot", None),
	)
	else:
	return CausalLMOutput(
	loss=loss,
	logits=logits,
	hidden_states=base_out.hidden_states,
	attentions=base_out.attentions,
	)

	AutoConfig.register("moep_gemma_lm", MoEPgemmaLMConfig)
	AutoModelForCausalLM.register(MoEPgemmaLMConfig, MoEPgemmaForCausalLM)
	AutoModel.register(MoEPgemmaLMConfig, MoEPgemmaModel)