codon-model 0.0.1__py3-none-any.whl

codon/__init__.py

@@ -0,0 +1,5 @@
+from typing import Optional
+
+__version__ = '0.0.1'
+
+__seed__: Optional[int] = None

codon/base.py

@@ -0,0 +1,167 @@
+import torch
+import torch.nn as nn
+
+from typing import Callable, Any, Iterator, Union
+
+from safetensors.torch import save_model as safe_save_model
+from safetensors.torch import load_model as safe_load_model
+
+
+class BasicModel(nn.Module):
+    '''
+    Base class for all models, providing common functionality like gradient checkpointing and parameter counting.
+    '''
+    def __init__(self):
+        '''
+        Initialize the BasicModel.
+        '''
+        super(BasicModel, self).__init__()
+        self.gradient_checkpointing: bool = False
+    
+    @property
+    def device(self) -> torch.device:
+        '''
+        Get the device of the model.
+
+        Returns:
+            torch.device: The device where the model parameters are located.
+                          Returns 'cpu' if the model has no parameters.
+        '''
+        try: return next(self.parameters()).device
+        except StopIteration: return torch.device('cpu')
+    
+    def set_checkpoint(self, value:bool) -> None:
+        '''
+        Enable or disable gradient checkpointing for the model and its sub-modules.
+
+        Args:
+            value (bool): True to enable gradient checkpointing, False to disable.
+        '''
+        self.gradient_checkpointing = value
+        for model in self.modules():
+            if isinstance(model, BasicModel) and model is not self:
+                model.gradient_checkpointing = value
+
+    def checkpoint(self, function:Callable, *args, **kwargs) -> Any:
+        '''
+        Apply gradient checkpointing to a function if enabled and in training mode.
+
+        Args:
+            function (Callable): The function to be checkpointed.
+            *args: Positional arguments for the function.
+            **kwargs: Keyword arguments for the function.
+
+        Returns:
+            Any: The output of the function.
+        '''
+        if self.gradient_checkpointing and self.training:
+            return torch.utils.checkpoint.checkpoint(
+                function, *args, use_reentrant=False, **kwargs
+            )
+        return function(*args, **kwargs)
+    
+    def get_params(self, trainable_only:bool=False) -> Iterator[torch.nn.Parameter]:
+        '''
+        Get an iterator over the model parameters.
+
+        Args:
+            trainable_only (bool, optional): If True, only yield parameters that require gradients.
+                                             Defaults to False.
+
+        Returns:
+            Iterator[torch.nn.Parameter]: An iterator over the model parameters.
+        '''
+        if trainable_only:
+            return (p for p in self.parameters() if p.requires_grad)
+        return self.parameters()
+    
+    def count_params(self, trainable_only:bool=False, active_only:bool=False, human_readable:bool=False, seen:set=None) -> Union[int, str]:
+        '''
+        Count the number of parameters in the model.
+
+        Args:
+            trainable_only (bool, optional): If True, count only trainable parameters.
+                                             Defaults to False.
+            active_only (bool, optional): If True, count only active parameters (e.g. for MoE).
+                                          Defaults to False.
+            human_readable (bool, optional): If True, return a string representation with units (e.g. M, B).
+                                             Defaults to False.
+            seen (set, optional): A set of already counted parameters to avoid duplicates.
+                                  Defaults to None.
+
+        Returns:
+            Union[int, str]: The total number of parameters.
+        '''
+        if seen is None:
+            seen = set()
+
+        if not active_only:
+            total = 0
+            for p in self.get_params(trainable_only):
+                if p not in seen:
+                    seen.add(p)
+                    total += p.numel()
+        else:
+            total = self._count_params_recursive(self, trainable_only, active_only, seen)
+        
+        if human_readable:
+            if total >= 1e9:
+                return f'{total / 1e9:.2f}B'
+            elif total >= 1e6:
+                return f'{total / 1e6:.2f}M'
+            elif total >= 1e3:
+                return f'{total / 1e3:.2f}K'
+            return str(total)
+
+        return total
+
+    @staticmethod
+    def _count_params_recursive(module: nn.Module, trainable_only: bool, active_only: bool, seen: set) -> int:
+        total = 0
+        for p in module.parameters(recurse=False):
+            if p not in seen:
+                if not trainable_only or p.requires_grad:
+                    seen.add(p)
+                    total += p.numel()
+        
+        for child in module.children():
+            if isinstance(child, BasicModel):
+                total += child.count_params(trainable_only, active_only, seen=seen)
+            else:
+                total += BasicModel._count_params_recursive(child, trainable_only, active_only, seen)
+        
+        return total
+    
+    def load_pretrained(self, path: str) -> None:
+        '''
+        Load a pretrained model from a file.
+
+        Args:
+            path (str): The path to the model file.
+        '''
+        if path.endswith('.safetensors'):
+            safe_load_model(self, path)
+            return
+
+        state_dict = torch.load(path, map_location=self.device)
+
+        if isinstance(state_dict, dict):
+            if 'model_state_dict' in state_dict:
+                state_dict = state_dict['model_state_dict']
+            elif 'state_dict' in state_dict:
+                state_dict = state_dict['state_dict']
+        
+        self.load_state_dict(state_dict)
+    
+    def save_pretrained(self, path: str) -> None:
+        '''
+        Save the model to a file.
+
+        Args:
+            path (str): The path to save the model file.
+        '''
+        if path.endswith('.safetensors'):
+            safe_save_model(self, path)
+        else:
+            state_dict = self.state_dict()
+            torch.save(state_dict, path)

codon/exp/moe.py

@@ -0,0 +1,307 @@
+import torch.nn.functional as F
+
+from codon.base import *
+from codon.block.mlp import MLP
+from codon.block.moe import *
+
+import math
+
+
+class ParallelExpert(nn.Module):
+    '''
+    A module that computes multiple expert outputs in parallel.
+
+    This module manages weights for multiple experts and processes inputs efficiently using batch matrix multiplication.
+
+    Attributes:
+        use_gate (bool): Whether to use Gated Linear Unit (GLU) variants.
+        num_experts (int): The number of experts.
+        weight1 (nn.Parameter): First weight matrix with shape (num_experts, in_features, inter_dim).
+        weight2 (nn.Parameter): Second weight matrix with shape (num_experts, hidden_features, out_features).
+        act (nn.Module): Activation function.
+        dropout (nn.Dropout): Dropout layer.
+    '''
+
+    def __init__(self, num_experts: int, in_features: int, hidden_features: int, out_features: int, use_gate: bool = False, dropout: float = 0.1) -> None:
+        '''
+        Initializes the ParallelExpert module.
+
+        Args:
+            num_experts (int): The number of experts.
+            in_features (int): Size of each input sample.
+            hidden_features (int): Size of the hidden layer.
+            out_features (int): Size of each output sample.
+            use_gate (bool): If True, uses SiLU activation with gating; otherwise, uses GELU.
+            dropout (float): Dropout probability.
+        '''
+        super().__init__()
+        self.use_gate = use_gate
+        self.num_experts = num_experts
+        self.dropout_p = dropout
+        
+        inter_dim = hidden_features * 2 if use_gate else hidden_features
+
+        # [Experts, In, Hidden]
+        self.weight1 = nn.Parameter(torch.empty(num_experts, in_features, inter_dim))
+        self.weight2 = nn.Parameter(torch.empty(num_experts, hidden_features, out_features))
+        
+        self.act = nn.SiLU() if use_gate else nn.GELU()
+        self.dropout = nn.Dropout(dropout)
+        
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        '''
+        Resets the parameters of the experts using Kaiming Uniform initialization.
+        '''
+        for i in range(self.num_experts):
+            nn.init.kaiming_uniform_(self.weight1[i], a=math.sqrt(5))
+            nn.init.kaiming_uniform_(self.weight2[i], a=math.sqrt(5))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        '''
+        Performs the forward pass for all experts in parallel.
+
+        Args:
+            x (torch.Tensor): Input tensor with shape (num_experts, capacity, in_features).
+
+        Returns:
+            torch.Tensor: Output tensor with shape (num_experts, capacity, out_features).
+        '''
+        # x shape: [Num_Experts, Capacity, In_Features]
+        # Weight1: [Num_Experts, In_Features, Inter_Dim]
+        
+        # [Num_Experts, Capacity, Inter_Dim]
+        h = torch.bmm(x, self.weight1)
+        
+        if self.use_gate:
+            gate, val = h.chunk(2, dim=-1)
+            h = self.act(gate) * val
+        else:
+            h = self.act(h)
+            
+        h = self.dropout(h)
+        
+        # Weight2: [Num_Experts, Hidden_Features, Out_Features]
+        out = torch.bmm(h, self.weight2)
+        
+        return out
+
+class ParallelMoE(BasicModel):
+    '''
+    A Parallel Mixture-of-Experts (MoE) model.
+
+    This model routes tokens to the top-k experts and computes their outputs in parallel. It also supports shared experts and auxiliary loss for load balancing.
+
+    Attributes:
+        model_dim (int): The dimension of the model.
+        top_k (int): The number of experts to route each token to.
+        num_experts (int): The total number of experts.
+        num_shared_experts (int): The number of shared experts that process all tokens.
+        use_aux_loss (bool): Whether to use auxiliary loss for load balancing.
+        capacity_factor (float): Factor to determine the capacity of each expert.
+        use_gate (bool): Whether to use Gated Linear Unit (GLU) variants in experts.
+        gate (nn.Linear): The gating network to route tokens to experts.
+        parallel_experts (ParallelExpert): The parallel experts module.
+        shared_experts (nn.ModuleList): List of shared experts, if any.
+    '''
+
+    def __init__(
+        self,
+        model_dim: int,
+        top_k: int,
+        num_experts: int,
+        num_shared_experts: int = 0,
+        use_aux_loss: bool = False,
+        use_gate: bool = False,
+        capacity_factor: float = 1.25,
+    ) -> None:
+        '''
+        Initializes the ParallelMoE model.
+
+        Args:
+            model_dim (int): The dimension of the model.
+            top_k (int): The number of experts to route each token to.
+            num_experts (int): The total number of experts.
+            num_shared_experts (int): The number of shared experts that process all tokens.
+            use_aux_loss (bool): Whether to use auxiliary loss for load balancing.
+            use_gate (bool): Whether to use Gated Linear Unit (GLU) variants in experts.
+            capacity_factor (float): Factor to determine the capacity of each expert.
+        '''
+        super().__init__()
+        self.model_dim = model_dim
+        self.top_k = top_k
+        self.num_experts = num_experts
+        self.num_shared_experts = num_shared_experts
+        self.use_aux_loss = use_aux_loss
+        self.capacity_factor = capacity_factor
+        self.use_gate = use_gate
+
+        hidden_dim = model_dim * 4
+
+        self.gate = nn.Linear(model_dim, num_experts, bias=False)
+
+        self.parallel_experts = ParallelExpert(
+            num_experts, model_dim, hidden_dim, model_dim, use_gate=use_gate
+        )
+
+        self.shared_experts = None
+        if num_shared_experts > 0:
+            act_layer = "silu" if use_gate else "gelu"
+            self.shared_experts = nn.ModuleList([
+                MLP(
+                    in_features=model_dim,
+                    hidden_features=hidden_dim,
+                    out_features=model_dim,
+                    use_gate=use_gate,
+                    act_layer=act_layer
+                ) for _ in range(num_shared_experts)
+            ])
+
+    def count_params(self, trainable_only: bool = False, active_only: bool = False) -> int:
+        '''
+        Counts the number of parameters in the model.
+
+        Args:
+            trainable_only (bool): If True, counts only trainable parameters.
+            active_only (bool): If True, counts only active parameters (parameters used during a single forward pass).
+
+        Returns:
+            int: The number of parameters.
+        '''
+        if not active_only:
+            return super().count_params(trainable_only, active_only)
+        
+        total = self.gate.weight.numel()
+        
+        if self.shared_experts:
+            total += sum(p.numel() for split in self.shared_experts for p in split.parameters())
+            
+        parallel_params = sum(p.numel() for p in self.parallel_experts.parameters())
+        single_expert_params = parallel_params // self.num_experts
+        total += single_expert_params * self.top_k
+            
+        return total
+
+    @property
+    def info(self) -> MoEInfo:
+        '''
+        Returns information about the MoE model's parameters.
+
+        Returns:
+            MoEInfo: An object containing total and active parameter counts.
+        '''
+        total = self.count_params(active_only=False)
+        active = self.count_params(active_only=True)
+        return MoEInfo(total_count=total, active_count=active)
+
+    def forward(self, x: torch.Tensor) -> MoEOutput:
+        '''
+        Performs the forward pass of the ParallelMoE model.
+
+        This method routes tokens to experts, computes expert outputs in parallel, adds shared expert outputs, and returns the combined result.
+
+        Args:
+            x (torch.Tensor): Input tensor with shape (batch_size, seq_len, model_dim).
+
+        Returns:
+            MoEOutput: The output of the MoE model, containing the final output tensor and auxiliary loss.
+        '''
+        original_shape = x.shape
+        batch, seq_len, dim = original_shape
+        num_tokens = batch * seq_len
+        
+        x_flat = x.reshape(-1, dim)
+
+        # Shared Experts
+        shared_output = torch.zeros_like(x_flat)
+        if self.shared_experts is not None:
+            for expert in self.shared_experts:
+                shared_output = shared_output + expert(x_flat)
+
+        # Gating
+        router_logits = self.gate(x_flat) # [Tokens, Experts]
+        routing_probs = F.softmax(router_logits, dim=-1)
+        
+        # weights: [Tokens, TopK], indices: [Tokens, TopK]
+        topk_weights, topk_indices = torch.topk(routing_probs, self.top_k, dim=-1)
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+
+        # Aux Loss
+        aux_loss = None
+        if self.use_aux_loss and self.training:
+            mask = torch.zeros_like(routing_probs).scatter_(1, topk_indices, 1.0)
+            density = mask.mean(dim=0)
+            density_proxy = routing_probs.mean(dim=0)
+            aux_loss = (self.num_experts * (density * density_proxy).sum())
+
+        capacity = int(num_tokens * self.top_k / self.num_experts * self.capacity_factor)
+        capacity = max(capacity, 4) 
+
+        # [Tokens * TopK]
+        flat_topk_indices = topk_indices.view(-1)
+        
+        sort_vals, sort_indices = flat_topk_indices.sort()
+        
+        # x_flat: [Tokens, Dim] -> [Tokens, TopK, Dim] -> [Tokens*TopK, Dim]
+        x_expanded = x_flat.index_select(0, sort_indices // self.top_k)
+        
+        # expert_counts: [Num_Experts]
+        expert_counts = torch.histc(
+            flat_topk_indices.float(), 
+            bins=self.num_experts, 
+            min=0, 
+            max=self.num_experts - 1
+        ).int()
+
+        parallel_inputs = torch.zeros(
+            self.num_experts, capacity, dim, 
+            dtype=x.dtype, device=x.device
+        )
+        
+        cumsum_counts = torch.cat([torch.tensor([0], device=x.device), expert_counts.cumsum(0)])
+        expert_starts = cumsum_counts[sort_vals] 
+        range_indices = torch.arange(sort_vals.size(0), device=x.device)
+        indices_in_expert = range_indices - expert_starts
+
+        mask = indices_in_expert < capacity
+        
+        valid_indices = indices_in_expert[mask]    # [Valid_Count]
+        valid_experts = sort_vals[mask]            # [Valid_Count]
+        valid_inputs  = x_expanded[mask]           # [Valid_Count, Dim]
+
+        # index: (Expert_ID, Capacity_ID)
+        parallel_inputs[valid_experts, valid_indices] = valid_inputs
+
+        parallel_outputs = self.parallel_experts(parallel_inputs)
+        # [Num_Experts, Capacity, Dim]
+
+        # [Tokens * TopK, Dim]
+        combined_output = torch.zeros(
+            num_tokens * self.top_k, dim, 
+            dtype=x.dtype, device=x.device
+        )
+        
+        # parallel_outputs[valid_experts, valid_indices]
+        valid_outputs = parallel_outputs[valid_experts, valid_indices]
+        
+        original_positions = sort_indices[mask] # [Valid_Count]
+
+        token_ids = original_positions.div(self.top_k, rounding_mode='floor')
+        
+        # [Tokens * TopK]
+        flat_weights = topk_weights.view(-1)
+        valid_weights = flat_weights[original_positions].unsqueeze(-1) # W
+        
+        weighted_output = valid_outputs * valid_weights
+        
+        final_output = torch.zeros_like(x_flat)
+        final_output.index_add_(0, token_ids, weighted_output)
+
+        final_output = final_output + shared_output
+
+        # [Batch, Seq, Dim]
+        return MoEOutput(
+            output=final_output.reshape(original_shape), 
+            aux_loss=aux_loss
+        )

codon/model/motif/__init__.py

@@ -0,0 +1 @@
+from .motif_a1 import MotifA1, MotifA1Output

codon/model/motif/motif_a1.py

@@ -0,0 +1,121 @@
+from codon.base import *
+
+from codon.block.transformer import TransformerMoEDecoder
+from codon.block.embedding   import RotaryEmbedding
+
+from typing import Optional, List, Tuple
+from dataclasses import dataclass
+
+@dataclass
+class MotifA1Output:
+    logits: torch.Tensor
+    past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None
+    aux_loss: Optional[torch.Tensor] = None
+    attentions: Optional[List[torch.Tensor]] = None
+
+
+class MotifA1(BasicModel):
+    def __init__(
+        self,
+        vocab_size: int = 32000,
+        model_dim: int = 768,
+        num_layers: int = 8,
+        num_heads: int = 8,
+        num_kv_heads: int = 2,
+        dropout: float = 0.1,
+        tie_weights: bool = False
+    ):
+        super().__init__()
+
+        self.token_emb = nn.Embedding(vocab_size, model_dim)
+        self.position_emb = RotaryEmbedding(model_dim // num_heads)
+        self.dropout = nn.Dropout(dropout)
+
+        self.decoder = nn.ModuleList([
+            TransformerMoEDecoder(
+                model_dim=model_dim,
+                num_heads=num_heads,
+                num_kv_heads=num_kv_heads,
+                top_k=1,
+                num_experts=3,
+                num_shared_experts=1,
+                use_expert_gate=True,
+                use_qk_norm=True,
+                use_attn_gate=True,
+                dropout=dropout,
+                idx=str(idx)
+            )
+            for idx in range(num_layers)
+        ])
+
+        self.norm = nn.RMSNorm(model_dim)
+        self.proj_out = nn.Linear(model_dim, vocab_size, bias=False)
+
+        if tie_weights:
+            self.proj_out.weight = self.token_emb.weight
+
+        self.apply(self._init_weights)
+    
+    def _init_weights(self, module):
+        std = 0.02
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+            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=std)
+            if module.padding_idx is not None:
+                torch.nn.init.zeros_(module.weight[module.padding_idx])
+    
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        mask: torch.Tensor = None,
+        start_pos: int = 0,
+        past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
+        use_cache: bool = False,
+        output_attentions: bool = False
+    ) -> MotifA1Output:
+        x = self.token_emb(input_ids)
+        x = self.dropout(x)
+
+        new_kv_cache = [] if use_cache else None
+        all_attentions = [] if output_attentions else None
+        aux_loss = None
+
+        for i, layer in enumerate(self.decoder):
+            layer_past = past_key_values[i] if past_key_values is not None else None
+            
+            out = layer(
+                hidden_states=x,
+                attention_mask=mask,
+                output_attentions=output_attentions,
+                position_emb=self.position_emb,
+                embedding_start=start_pos,
+                past_key_value=layer_past,
+                use_cache=use_cache
+            )
+            
+            x = out.output
+            
+            if use_cache:
+                new_kv_cache.append(out.past_key_value)
+            
+            if output_attentions:
+                all_attentions.append(out.attention_weights)
+            
+            if out.aux_loss is not None:
+                if aux_loss is None:
+                    aux_loss = out.aux_loss
+                else:
+                    aux_loss += out.aux_loss
+
+        x = self.norm(x)
+        logits = self.proj_out(x)
+
+        return MotifA1Output(
+            logits=logits,
+            past_key_values=new_kv_cache,
+            aux_loss=aux_loss,
+            attentions=all_attentions
+        )