Upload main.py

main.py

@@ -0,0 +1,231 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from gguf import GGUFWriter
+
+class ModelConfig:
+    def __init__(self):
+        # Core parameters
+        self.vocab_size = 32000
+        self.hidden_size = 768
+        self.num_hidden_layers = 4
+        self.num_attention_heads = 8
+        self.intermediate_size = 3072
+        
+        # Expert parameters
+        self.num_experts = 4
+        
+        # Efficiency parameters
+        self.chunk_size = 256
+        self.compression_ratio = 4
+        
+        # Reasoning parameters
+        self.max_graph_nodes = 512
+        self.node_dim = self.hidden_size // 4  # 192
+        
+        # Regularization
+        self.hidden_dropout_prob = 0.1
+        self.initializer_range = 0.02
+
+class CoATGraphManager(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        
+        # Node initialization
+        self.base_nodes = nn.Parameter(torch.randn(config.max_graph_nodes, config.node_dim))
+        self.projection = nn.Linear(config.hidden_size, config.node_dim)
+        
+        # Update mechanism
+        self.update_gate = nn.Sequential(
+            nn.Linear(config.node_dim * 2, config.hidden_size),
+            nn.GELU(),
+            nn.Linear(config.hidden_size, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, hidden_states, current_nodes):
+        batch_size = hidden_states.size(0)
+        
+        # Clone to prevent in-place errors
+        current_nodes = current_nodes.clone()
+        
+        # Aggregate sequence information
+        seq_aggregated = hidden_states.mean(dim=1)  # [batch, hidden_size]
+        
+        # Project to node space
+        projected = self.projection(seq_aggregated)  # [batch, node_dim]
+        
+        # Calculate similarity scores
+        similarity = torch.matmul(projected.unsqueeze(1), current_nodes.transpose(1, 2))  # [batch, 1, max_nodes]
+        
+        # Get top-2 nodes
+        _, topk_indices = torch.topk(similarity.squeeze(1), k=2, dim=-1)  # [batch, 2]
+        
+        # Gather relevant nodes
+        selected_nodes = torch.gather(
+            current_nodes,
+            1,
+            topk_indices.unsqueeze(-1).expand(-1, -1, self.config.node_dim)
+        )  # [batch, 2, node_dim]
+        
+        # Calculate updates
+        combined = torch.cat([
+            selected_nodes,
+            self.base_nodes[topk_indices]
+        ], dim=-1)
+        update_weights = self.update_gate(combined)
+        updated_nodes = selected_nodes * update_weights + self.base_nodes[topk_indices] * (1 - update_weights)
+        
+        # Safe scatter update
+        current_nodes.scatter_(
+            1,
+            topk_indices.unsqueeze(-1).expand(-1, -1, self.config.node_dim),
+            updated_nodes
+        )
+        
+        return current_nodes
+
+class ChunkKVAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        
+        # Projections
+        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size)
+        self.k_proj = nn.Linear(config.hidden_size, config.hidden_size)
+        self.v_proj = nn.Linear(config.hidden_size, config.hidden_size)
+        
+        # Compression
+        self.k_compress = nn.Linear(config.chunk_size, config.chunk_size//config.compression_ratio)
+        self.v_compress = nn.Linear(config.chunk_size, config.chunk_size//config.compression_ratio)
+
+    def forward(self, hidden_states):
+        batch_size, seq_len, _ = hidden_states.size()
+        
+        # Process queries
+        q = self.q_proj(hidden_states)
+        
+        # Process keys/values in chunks
+        k = self._process_chunk(self.k_proj, self.k_compress, hidden_states)
+        v = self._process_chunk(self.v_proj, self.v_compress, hidden_states)
+
+        # Reshape for attention
+        q = q.view(batch_size, -1, self.config.num_attention_heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch_size, -1, self.config.num_attention_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch_size, -1, self.config.num_attention_heads, self.head_dim).transpose(1, 2)
+
+        # Attention calculation
+        attn = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim))
+        attn = F.softmax(attn, dim=-1)
+        output = torch.matmul(attn, v)
+        
+        return output.transpose(1, 2).flatten(2)
+
+    def _process_chunk(self, proj, compress, x):
+        chunks = []
+        for i in range(0, x.size(1), self.config.chunk_size):
+            chunk = proj(x[:, i:i+self.config.chunk_size])
+            compressed = compress(chunk.transpose(1, 2)).transpose(1, 2)
+            chunks.append(compressed)
+        return torch.cat(chunks, dim=1)
+
+class SelfMoA(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.experts = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(config.hidden_size, config.intermediate_size),
+                nn.GELU(),
+                nn.Linear(config.intermediate_size, config.hidden_size)
+            ) for _ in range(config.num_experts)
+        ])
+        self.gate = nn.Linear(config.hidden_size, config.num_experts)
+
+    def forward(self, x):
+        gate = F.gumbel_softmax(self.gate(x), hard=True, dim=-1)
+        return sum(expert(x) * gate[..., i].unsqueeze(-1) for i, expert in enumerate(self.experts))
+
+class DeepSeekLiteBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.attention = ChunkKVAttention(config)
+        self.moa = SelfMoA(config)
+        self.coat = CoATGraphManager(config)
+        self.norm = nn.LayerNorm(config.hidden_size)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, x, nodes):
+        # Attention path
+        attn_out = self.attention(self.norm(x))
+        x = x + self.dropout(attn_out)
+        
+        # Update graph nodes
+        updated_nodes = self.coat(x, nodes)
+        
+        # MOA path
+        moa_out = self.moa(self.norm(x))
+        x = x + self.dropout(moa_out)
+        
+        return x, updated_nodes
+
+class DeepSeekLite(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([DeepSeekLiteBlock(config) for _ in range(config.num_hidden_layers)])
+        self.final_norm = nn.LayerNorm(config.hidden_size)
+        
+        # Initialize graph nodes with cloning
+        self.graph_nodes = nn.ParameterList([
+            nn.Parameter(torch.randn(config.max_graph_nodes, config.node_dim).clone().detach().requires_grad_(True))
+            for _ in range(config.num_hidden_layers)
+        ])
+
+    def forward(self, input_ids):
+        x = self.embedding(input_ids)
+        batch_size = input_ids.size(0)
+        
+        for layer_idx, layer in enumerate(self.layers):
+            # Clone and expand nodes for each layer
+            nodes = self.graph_nodes[layer_idx].unsqueeze(0).expand(batch_size, -1, -1).clone()
+            x, _ = layer(x, nodes)
+        
+        return self.final_norm(x)
+
+def save_gguf(model, filename):
+    writer = GGUFWriter(filename, "deepseek-lite")
+    
+    # Add model configuration
+    writer.add_uint32("vocab_size", model.config.vocab_size)
+    writer.add_uint32("hidden_size", model.config.hidden_size)
+    writer.add_uint32("num_hidden_layers", model.config.num_hidden_layers)
+    writer.add_uint32("num_attention_heads", model.config.num_attention_heads)
+    writer.add_uint32("num_experts", model.config.num_experts)
+    writer.add_uint32("max_graph_nodes", model.config.max_graph_nodes)
+    
+    # Add all parameters
+    for name, param in model.named_parameters():
+        writer.add_tensor(name, param.detach().cpu().numpy())
+    
+    writer.write_header_to_file()
+    writer.write_kv_data_to_file()
+    writer.write_tensors_to_file()
+    writer.close()
+
+if __name__ == "__main__":
+    config = ModelConfig()
+    model = DeepSeekLite(config)
+    
+    # Test forward pass
+    inputs = torch.randint(0, config.vocab_size, (2, 1024))
+    with torch.no_grad():
+        outputs = model(inputs)
+        print(f"Successful execution! Output shape: {outputs.shape}")
+        print(f"Parameter count: {sum(p.numel() for p in model.parameters())/1e6:.1f}M")
+    
+    # Save model
+    save_gguf(model, "deepseek-lite.gguf")
+    print("Model saved in GGUF format")