modeling_gpt_moe_mcts.py

import torch
import torch.nn as nn
from torch.nn import functional as F
from transformers import PreTrainedModel
from .configuration_gpt_moe_mcts import GPTMoEMCTSConfig

class FlashAttention3(nn.Module):
    def __init__(self, d_model, n_heads, block_size_q, block_size_kv, num_blocks_kv, device='cuda'):
        super(FlashAttention3, self).__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.block_size_q = block_size_q
        self.block_size_kv = block_size_kv
        self.num_blocks_kv = num_blocks_kv
        self.device = device

        self.q_proj = nn.Linear(d_model, d_model).to(device)
        self.k_proj = nn.Linear(d_model, d_model).to(device)
        self.v_proj = nn.Linear(d_model, d_model).to(device)
        self.out_proj = nn.Linear(d_model, d_model).to(device)

    def forward(self, x):
        B, T, C = x.size()
        Q = self.q_proj(x).view(B, T, self.n_heads, C // self.n_heads).transpose(1, 2)
        K = self.k_proj(x).view(B, T, self.n_heads, C // self.n_heads).transpose(1, 2)
        V = self.v_proj(x).view(B, T, self.n_heads, C // self.n_heads).transpose(1, 2)

        O = torch.zeros(B, self.n_heads, T, C // self.n_heads).to(self.device)
        L = torch.zeros(B, self.n_heads, T).to(self.device)
        M = torch.full((B, self.n_heads, T), -float('inf')).to(self.device)

        for i in range(0, T, self.block_size_q):
            Q_block = Q[:, :, i:i+self.block_size_q]
            O_block = torch.zeros_like(Q_block).to(self.device)
            L_block = torch.zeros(B, self.n_heads, Q_block.size(2)).to(self.device)
            M_block = torch.full((B, self.n_heads, Q_block.size(2)), -float('inf')).to(self.device)

            for j in range(0, T, self.block_size_kv):
                K_block = K[:, :, j:j+self.block_size_kv]
                V_block = V[:, :, j:j+self.block_size_kv]

                S_block = torch.matmul(Q_block, K_block.transpose(-2, -1))
                M_block_old = M_block
                M_block = torch.max(M_block, S_block.max(dim=-1).values)

                exp_S_block = torch.exp(S_block - M_block.unsqueeze(-1))
                L_block = torch.exp(M_block_old - M_block) * L_block + exp_S_block.sum(dim=-1)

                O_block += torch.matmul(exp_S_block, V_block)

            O_block /= L_block.unsqueeze(-1)
            O[:, :, i:i+self.block_size_q] = O_block

        O = O.transpose(1, 2).contiguous().view(B, T, self.n_heads * (C // self.n_heads))
        O = self.out_proj(O)

        return O

# Define the MLP module
class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
        self.gelu = nn.GELU(approximate='tanh')
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
        self.dropout = nn.Dropout(config.dropout)
        self.c_proj.scale_init = 1

    def forward(self, x):
        x = self.c_fc(x)
        x = self.gelu(x)
        x = self.c_proj(x)
        x = self.dropout(x)
        return x

# Define the MixtureOfExperts module
class MixtureOfExperts(nn.Module):
    def __init__(self, config, num_experts, expert_layers):
        super().__init__()
        self.num_experts = num_experts
        self.expert_layers = expert_layers

        self.experts = nn.ModuleList([self._create_expert(config) for _ in range(num_experts)])
        self.gate = nn.Linear(config.n_embd, num_experts)

    def _create_expert(self, config):
        layers = []
        for _ in range(self.expert_layers):
            layers.append(FlashAttention3(d_model=config.n_embd, n_heads=config.n_head, block_size_q=32, block_size_kv=32, num_blocks_kv=4))
            layers.append(nn.LayerNorm(config.n_embd))
            layers.append(MLP(config))
        return nn.Sequential(*layers)

    def forward(self, x):
        B, T, C = x.size()
        
        gate_scores = self.gate(x)
        gate_probs = F.softmax(gate_scores, dim=-1)
        
        expert_outputs = torch.stack([expert(x) for expert in self.experts], dim=1)

        gate_probs = gate_probs.unsqueeze(-1)
        gate_probs = gate_probs.permute(0, 2, 1, 3)
        
        output = torch.sum(gate_probs * expert_outputs, dim=1)

        return output

# Define the BlockWithMoE module
class BlockWithMoE(nn.Module):
    def __init__(self, config, num_experts=4, expert_layers=2, block_size_q=32, block_size_kv=32, num_blocks_kv=4, device='cuda'):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = FlashAttention3(d_model=config.n_embd, n_heads=config.n_head, block_size_q=block_size_q, block_size_kv=block_size_kv, num_blocks_kv=num_blocks_kv, device=device)
        self.dropout1 = nn.Dropout(config.dropout)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.moe = MixtureOfExperts(config, num_experts, expert_layers)
        self.dropout2 = nn.Dropout(config.dropout)
        self.ln_3 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)
        self.dropout3 = nn.Dropout(config.dropout)

    def forward(self, x):
        B, T, C = x.size()

        attn_output = self.attn(x)
        x = x + attn_output
        x = self.dropout1(x)
        x = x + self.moe(self.ln_2(x))
        x = self.dropout2(x)
        x = x + self.mlp(self.ln_3(x))
        x = self.dropout3(x)
        return x
    
class GPTMoEMCTSPreTrainedModel(PreTrainedModel):
    config_class = GPTMoEMCTSConfig
    base_model_prefix = "gpt_moe_mcts"

    def __init__(self, *inputs, **kwargs):
        super().__init__(*inputs, **kwargs)

class GPTMoEMCTSModel(GPTMoEMCTSPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.config = config

        self.transformer = nn.ModuleDict(dict(
            wte=nn.Embedding(config.vocab_size, config.n_embd),
            wpe=nn.Embedding(config.block_size, config.n_embd),
            h=nn.ModuleList([BlockWithMoE(config) for _ in range(config.n_layer)]),
            ln_f=nn.LayerNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        self.transformer.wte.weight = self.lm_head.weight

        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # Forward pass
        B, T = input_ids.size()
        assert T <= self.config.block_size, f"Cannot forward sequence of length {T}, block size is only {self.config.block_size}"

        pos = torch.arange(0, T, dtype=torch.long, device=input_ids.device)
        pos_emb = self.transformer.wpe(pos)
        tok_emb = self.transformer.wte(input_ids)
        x = tok_emb + pos_emb

        for block in self.transformer.h:
            x = block(x)

        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1))

        if not return_dict:
            output = (logits,) + (loss,) if loss is not None else (logits,)
            return output

        return {
            "logits": logits,
            "loss": loss,
        }