Initial commit of GPT-MoE-MCTS model

cnfiguration_gpt_moe_mcts.py

@@ -0,0 +1,32 @@
+from transformers import PretrainedConfig
+
+class GPTMoEMCTSConfig(PretrainedConfig):
+    model_type = "gpt_moe_mcts"
+
+    def __init__(
+        self,
+        vocab_size=50257,
+        block_size=512,
+        n_layer=6,
+        n_head=4,
+        n_embd=256,
+        dropout=0.2,
+        num_experts=3,
+        expert_layers=3,
+        block_size_q=32,
+        block_size_kv=32,
+        num_blocks_kv=4,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.block_size = block_size
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_embd = n_embd
+        self.dropout = dropout
+        self.num_experts = num_experts
+        self.expert_layers = expert_layers
+        self.block_size_q = block_size_q
+        self.block_size_kv = block_size_kv
+        self.num_blocks_kv = num_blocks_kv

example_usage.py

@@ -0,0 +1,29 @@
+from transformers import GPT2Tokenizer
+from modeling_gpt_moe_mcts import GPTMoEMCTSModel
+from configuration_gpt_moe_mcts import GPTMoEMCTSConfig
+
+# Initialize configuration
+config = GPTMoEMCTSConfig()
+
+# Initialize model
+model = GPTMoEMCTSModel(config)
+
+# Initialize tokenizer (using GPT2Tokenizer as a base)
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+
+# Prepare input
+text = "Hello, how are you?"
+inputs = tokenizer(text, return_tensors="pt")
+
+# Forward pass
+outputs = model(**inputs)
+
+# Get the predicted next token
+next_token_logits = outputs.logits[0, -1, :]
+next_token = next_token_logits.argmax()
+
+# Decode the predicted token
+predicted_text = tokenizer.decode(next_token)
+
+print(f"Input: {text}")
+print(f"Predicted next token: {predicted_text}")

modeling_gpt_moe_mcts.py

@@ -0,0 +1,208 @@
+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,
+        }

readme.md

@@ -0,0 +1,141 @@
+# GPT-MoE-MCTS: GPT with Mixture of Experts and Monte Carlo Tree Search
+
+## Table of Contents
+1. [Introduction](#introduction)
+2. [Key Features](#key-features)
+3. [Model Architecture](#model-architecture)
+4. [Installation](#installation)
+5. [Usage](#usage)
+6. [Training](#training)
+7. [Evaluation](#evaluation)
+8. [MCTS Decoding](#mcts-decoding)
+9. [Contributing](#contributing)
+10. [License](#license)
+
+## Introduction
+
+GPT-MoE-MCTS is an advanced language model that combines the power of GPT (Generative Pre-trained Transformer) with Mixture of Experts (MoE) and Monte Carlo Tree Search (MCTS) decoding. This model is designed to provide high-quality text generation with improved efficiency and performance.
+
+## Key Features
+
+- **GPT-based Architecture**: Utilizes the powerful GPT architecture for language modeling.
+- **Mixture of Experts**: Incorporates a dynamic routing system to specialize different parts of the network for different inputs.
+- **FlashAttention3**: Implements an optimized attention mechanism for improved efficiency.
+- **Monte Carlo Tree Search Decoding**: Uses MCTS during inference for higher quality text generation.
+- **Hugging Face Compatible**: Easily integrates with the Hugging Face Transformers library.
+
+## Model Architecture
+
+The GPT-MoE-MCTS model consists of the following key components:
+
+1. **Token and Positional Embeddings**: Converts input tokens into embeddings and adds positional information.
+2. **Transformer Blocks with MoE**: Multiple layers of transformer blocks, each incorporating:
+   - FlashAttention3: An optimized attention mechanism.
+   - Mixture of Experts Layer: A dynamic routing system for specialized processing.
+   - Feed-Forward Network: Standard MLP for additional processing.
+3. **Output Layer**: Final layer normalization and projection to vocabulary logits.
+
+## Installation
+
+To install the GPT-MoE-MCTS model, follow these steps:
+
+```bash
+git clone https://github.com/yourusername/gpt-moe-mcts.git
+cd gpt-moe-mcts
+pip install -r requirements.txt
+```
+
+## Usage
+
+Here's a basic example of how to use the GPT-MoE-MCTS model:
+
+```python
+from transformers import GPT2Tokenizer
+from modeling_gpt_moe_mcts import GPTMoEMCTSModel
+from configuration_gpt_moe_mcts import GPTMoEMCTSConfig
+
+# Initialize configuration and model
+config = GPTMoEMCTSConfig()
+model = GPTMoEMCTSModel(config)
+
+# Initialize tokenizer (using GPT2Tokenizer as a base)
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+
+# Prepare input
+text = "Hello, how are you?"
+inputs = tokenizer(text, return_tensors="pt")
+
+# Forward pass
+outputs = model(**inputs)
+
+# Get the predicted next token
+next_token_logits = outputs.logits[0, -1, :]
+next_token = next_token_logits.argmax()
+
+# Decode the predicted token
+predicted_text = tokenizer.decode(next_token)
+
+print(f"Input: {text}")
+print(f"Predicted next token: {predicted_text}")
+```
+
+## Training
+
+To train the GPT-MoE-MCTS model on your own data:
+
+1. Prepare your dataset in the format of tokenized .npy files.
+2. Adjust the hyperparameters in the `train_model()` function in `train.py`.
+3. Run the training script:
+
+```bash
+python train.py
+```
+
+The script will automatically save checkpoints and display training progress.
+
+## Evaluation
+
+To evaluate the model's performance:
+
+```python
+from eval_utils import evaluate_model
+
+perplexity, accuracy = evaluate_model(model, eval_dataloader)
+print(f"Perplexity: {perplexity}, Accuracy: {accuracy}")
+```
+
+## MCTS Decoding
+
+The GPT-MoE-MCTS model uses Monte Carlo Tree Search for decoding during inference. To use MCTS decoding:
+
+```python
+from mcts_decode import mcts_decode
+
+generated_text = mcts_decode(model, input_text, max_length=50, num_simulations=100)
+print(f"Generated text: {generated_text}")
+```
+
+## Contributing
+
+We welcome contributions to the GPT-MoE-MCTS project! Please see our [CONTRIBUTING.md](CONTRIBUTING.md) file for guidelines on how to contribute.
+
+## License
+
+This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.
+
+---
+
+For more detailed information about the model architecture, training process, and advanced usage, please refer to our [documentation](docs/index.md).
+
+If you use GPT-MoE-MCTS in your research, please cite:
+
+```
+@misc{GPT-MoE-MCTS,
+  author = {Your Name},
+  title = {GPT-MoE-MCTS: GPT with Mixture of Experts and Monte Carlo Tree Search},
+  year = {2023},
+  publisher = {GitHub},
+  journal = {GitHub repository},
+  howpublished = {\url{https://github.com/yourusername/gpt-moe-mcts}}
+}
+```

tokenizer_gpt_moe_mcts.py

@@ -0,0 +1,14 @@
+from transformers import GPT2Tokenizer
+
+class GPTMoEMCTSTokenizer(GPT2Tokenizer):
+    def __init__(
+        self,
+        vocab_file,
+        merges_file,
+        **kwargs
+    ):
+        super().__init__(
+            vocab_file,
+            merges_file,
+            **kwargs
+        )