Upload TaoNet model to HuggingFace Hub

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

bitlinear.py

@@ -0,0 +1,83 @@
+"""
+BitLinear - Simplified for training stability.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization."""
+    
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        
+    def forward(self, x):
+        rms = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+        return (x / rms)
+
+
+class TernaryQuantize(torch.autograd.Function):
+    """Ternary quantization with straight-through estimator."""
+    
+    @staticmethod
+    def forward(ctx, w):
+        scale = 1.0 / w.abs().mean().clamp_(min=1e-5)
+        u = (w * scale).round().clamp_(-1, 1) / scale
+        return u
+    
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output
+
+
+class ActivationQuantize(torch.autograd.Function):
+    """INT8 activation quantization."""
+    
+    @staticmethod
+    def forward(ctx, x):
+        scale = 127.0 / x.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5)
+        y = (x * scale).round().clamp_(-128, 127) / scale
+        return y
+    
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output
+
+
+class BitLinear(nn.Linear):
+    """
+    Linear layer with ternary weight quantization.
+    
+    No internal normalization - caller handles it (Pre-Norm architecture).
+    """
+    
+    def __init__(self, in_features, out_features, bias=True):
+        super().__init__(in_features, out_features)
+        
+        # Gentler initialization for ternary stability
+        nn.init.normal_(self.weight, mean=0.0, std=0.02)
+        self.rmsnorm = RMSNorm(in_features)
+    
+    def forward(self, x):
+        w = self.weight # a weight tensor with shape [d, k]
+        x_norm = self.rmsnorm(x)
+        # A trick for implementing Straight−Through−Estimator (STE) using detach()
+        x_quant = x_norm + (ActivationQuantize.apply(x_norm) - x_norm).detach()
+        w_quant = w + (TernaryQuantize.apply(w) - w).detach()
+        y = F.linear(x_quant, w_quant)
+
+        return self.rmsnorm(y)
+    
+    def get_inference_params(self):
+        """Export for FPGA deployment."""
+        with torch.no_grad():
+            scale = self.weight.abs().mean(dim=-1, keepdim=True).clamp(min=1e-5)
+            w_ternary = (self.weight / scale).round().clamp(-1, 1).to(torch.int8)
+            
+            return {
+                'weight_ternary': w_ternary,
+                'weight_scale':  scale.squeeze()
+            }

config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "TaoNetForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_taonet.TaoNetConfig",
+    "AutoModelForCausalLM": "modeling_taonet.TaoNetForCausalLM"
+  },
+  "block_arrangement": "layered",
+  "bos_token_id": 1,
+  "d_embed_rank": 384,
+  "d_ff": 512,
+  "d_kv_comp": 384,
+  "d_model": 512,
+  "d_rope": 64,
+  "d_state": 512,
+  "dropout": 0.0,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "layered_mla_num": 0,
+  "max_seq_len": 256,
+  "model_type": "taonet",
+  "n_heads": 4,
+  "n_layers": 8,
+  "pad_token_id": 3,
+  "ssm_per_mla": 3,
+  "transformers_version": "4.57.6",
+  "unk_token_id": 0,
+  "vocab_size": 50257
+}

configuration_taonet.py

@@ -0,0 +1,56 @@
+"""
+Configuration class for TaoNet model.
+"""
+
+from transformers import PretrainedConfig
+
+
+class TaoNetConfig(PretrainedConfig):
+    """Configuration for TaoNet model."""
+    
+    model_type = "taonet"
+    
+    def __init__(
+        self,
+        vocab_size: int = 25000,
+        d_model: int = 512,
+        d_embed_rank: int = 384,
+        d_state: int = 512,
+        d_ff: int = 512,
+        n_heads: int = 4,
+        d_kv_comp: int = 384,
+        d_rope: int = 64,
+        n_layers: int = 8,
+        max_seq_len: int = 256,
+        dropout: float = 0.02,
+        block_arrangement: str = "layered",
+        ssm_per_mla: int = 3,
+        layered_mla_num: int = 0,
+        pad_token_id: int = 3,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        unk_token_id: int = 0,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            unk_token_id=unk_token_id,
+            **kwargs,
+        )
+        
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.d_embed_rank = d_embed_rank
+        self.d_state = d_state
+        self.d_ff = d_ff
+        self.n_heads = n_heads
+        self.d_kv_comp = d_kv_comp
+        self.d_rope = d_rope
+        self.n_layers = n_layers
+        self.max_seq_len = max_seq_len
+        self.dropout = dropout
+        self.block_arrangement = block_arrangement
+        self.ssm_per_mla = ssm_per_mla
+        self.layered_mla_num = layered_mla_num

factorized_embedding.py

@@ -0,0 +1,44 @@
+"""
+Low-Rank Factorized Embedding. 
+
+IMPORTANT: Uses standard nn.Linear for projection, NOT BitLinear. 
+Embeddings need full precision for good token representations.
+"""
+
+import torch
+import torch.nn as nn
+
+class FactorizedEmbedding(nn.Module):
+    """
+    Low-Rank Factorized Embedding:  vocab → d_embed_rank → d_model
+    
+    Uses standard Linear (not BitLinear) for the projection.
+    Embeddings are memory lookups - they benefit from full precision.
+    """
+    
+    def __init__(self, vocab_size, d_model, d_embed_rank=96):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self. d_model = d_model
+        self.d_embed_rank = d_embed_rank
+        
+        # Embedding table:  vocab → compressed
+        self.embed = nn.Embedding(vocab_size, d_embed_rank)
+        
+        # Projection: compressed → full (standard Linear, NOT BitLinear)
+        self.proj = nn.Linear(d_embed_rank, d_model, bias=False)
+        
+        # Initialize
+        nn.init.normal_(self.embed.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.proj.weight, mean=0.0, std=0.02)
+        
+        print(f"FactorizedEmbedding:  {vocab_size} × {d_embed_rank} → {d_model}")
+        print(f"  Params: {self.get_num_params()/1e6:.2f}M (vs {vocab_size * d_model/1e6:.2f}M dense)")
+    
+    def forward(self, input_ids):
+        x = self.embed(input_ids)  # [B, S, d_embed_rank]
+        x = self.proj(x)           # [B, S, d_model]
+        return x
+    
+    def get_num_params(self):
+        return self.vocab_size * self.d_embed_rank + self.d_embed_rank * self.d_model

mla.py

@@ -0,0 +1,123 @@
+"""
+Basic Multi-headed Latent Attention (MLA).
+Simple implementation without KV cache.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+from .rope import RotaryEmbedding, apply_rotary
+
+
+class MemoryOptimizedMLA(nn.Module):
+    """
+    Basic MLA: Project to latent space, apply multi-head attention, project back.
+    Numerically stable implementation with proper normalization.
+    """
+    
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.n_heads = config.n_heads
+        self.d_head = config.d_kv_comp // config.n_heads
+        self.d_rope = config.d_rope
+        # Improved scaling: use sqrt(d_head) with a small epsilon for numerical stability
+        self.scale = 1.0 / math.sqrt(max(self.d_head, 1.0))
+        
+        # Layer normalization before projections for stability
+        self.norm_latent = nn.LayerNorm(config.d_model)
+        
+        # Projections
+        self.to_latent = nn.Linear(config.d_model, config.d_kv_comp, bias=False)
+        
+        # Q/K/V from latent
+        self.q_proj = nn.Linear(config.d_kv_comp, config.d_kv_comp, bias=False)
+        self.k_proj = nn.Linear(config.d_kv_comp, config.d_kv_comp, bias=False)
+        self.v_proj = nn.Linear(config.d_kv_comp, config.d_kv_comp, bias=False)
+        
+        # RoPE
+        self.rotary = RotaryEmbedding(config.d_rope)
+        
+        # Output
+        self.out_proj = nn.Linear(config.d_kv_comp, config.d_model, bias=False)
+        
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+    
+    def forward(self, x, mask=None):
+        """
+        Args:
+            x: (batch_size, seq_len, d_model)
+            mask: (batch_size, seq_len) or (batch_size, 1, seq_len, seq_len), optional
+        
+        Returns:
+            out: (batch_size, seq_len, d_model)
+        """
+        batch_size, seq_len, _ = x.shape
+
+        # Normalize input before projection to prevent activation explosion
+        x_norm = self.norm_latent(x)
+        
+        # Project to latent space
+        latent = self.to_latent(x_norm)
+
+        # Generate Q/K/V
+        q = self.q_proj(latent)
+        k = self.k_proj(latent)
+        v = self.v_proj(latent)
+        
+        # Reshape for multi-head attention: (batch_size, seq_len, d_kv_comp) -> (batch_size, n_heads, seq_len, d_head)
+        q = q.view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)
+        k = k.view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)
+        v = v.view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)
+        
+        # Normalize Q and K for stable attention (standard practice in modern attention mechanisms)
+        q = F.normalize(q, dim=-1, p=2)
+        k = F.normalize(k, dim=-1, p=2)
+        
+        # Apply RoPE
+        if self.d_rope > 0:
+            rotary_emb = self.rotary(seq_len, x.device)
+            cos = torch.cos(rotary_emb).unsqueeze(0).unsqueeze(0)
+            sin = torch.sin(rotary_emb).unsqueeze(0).unsqueeze(0)
+
+            q_rot = apply_rotary(q[..., :self.d_rope], cos, sin)
+            k_rot = apply_rotary(k[..., :self.d_rope], cos, sin)
+            
+            q = torch.cat([q_rot, q[..., self.d_rope:]], dim=-1)
+            k = torch.cat([k_rot, k[..., self.d_rope:]], dim=-1)
+        
+        # Attention computation with numerical stability
+        # Scale before matmul to prevent overflow
+        attn_scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+        
+        # Clamp attention scores to prevent inf/-inf in softmax
+        attn_scores = torch.clamp(attn_scores, min=-20.0, max=20.0)
+        
+        if mask is not None:
+            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))
+        
+        # Numerically stable softmax
+        attn_weights = F.softmax(attn_scores, dim=-1)
+        
+        # Check for NaN and print warning
+        if torch.isnan(attn_weights).any():
+            print(f"WARNING: NaN detected in attention weights! "
+                  f"attn_scores min={attn_scores.min():.4f}, max={attn_scores.max():.4f}, "
+                  f"attn_weights min={attn_weights.min():.4f}, max={attn_weights.max():.4f}")
+        
+        attn_weights = self.attn_dropout(attn_weights)
+        
+        # Apply attention to values
+        out = torch.matmul(attn_weights, v)
+        
+        # Reshape back: (batch_size, n_heads, seq_len, d_head) -> (batch_size, seq_len, d_kv_comp)
+        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
+        
+        # Project back to model dimension
+        out = self.out_proj(out)
+        out = self.resid_dropout(out)
+        
+        return out

model.py

@@ -0,0 +1,397 @@
+"""
+SimpleLLM - Mamba-style State-Space Model with ternary quantization.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn. functional as F
+
+from .ssm import SSMBlock
+from .bitlinear import BitLinear, RMSNorm, ActivationQuantize
+from .factorized_embedding import FactorizedEmbedding
+from .mla import MemoryOptimizedMLA
+
+class SSMBlockWrapper(nn.Module):
+    """
+    Pre-Norm SSM Block (Mamba-style) with nn.Sequential structure.
+    
+    Structure:
+        x → Norm → SSM → Add → Norm → FFN → Add → output
+    """
+    
+    def __init__(self, config):
+        super().__init__()
+        self.ssm = SSMBlock(config)
+        self.feed_forward = nn.Sequential(
+            BitLinear(config.d_model, config.d_ff, bias=False),
+            nn.ReLU(),
+            BitLinear(config.d_ff, config.d_model, bias=False),
+        )
+        self.dropout = nn.Dropout(config.dropout)
+    
+    def forward(self, x, mask=None):
+        # Pre-norm SSM with residual
+        x = x + self.dropout(self.ssm(x, mask))  # Normalize before SSM
+        # Pre-norm FFN with residual
+        x = x + self.dropout(self.feed_forward(x))  # Normalize before FFN
+        return x
+
+
+class MLABlockWrapper(nn.Module):
+    """
+    MLA Block with residual connection and FFN.
+    
+    Structure:
+        x → Norm → MLA → Add → Norm → FFN → Add → output
+    
+    Pre-norm structure stabilizes training and prevents gradient explosion.
+    """
+    
+    def __init__(self, config):
+        super().__init__()
+        self.mla = MemoryOptimizedMLA(config)
+        self.ffn = nn.Sequential(
+            nn.Linear(config.d_model, config.d_ff, bias=False),
+            nn.ReLU(),
+            nn.Linear(config.d_ff, config.d_model, bias=False),
+            nn.ReLU(),
+            nn.Linear(config.d_ff, config.d_model, bias=False),
+        )
+        self.dropout = nn.Dropout(config.dropout)
+    
+    def forward(self, x, mask=None):
+        # Pre-norm MLA with residual
+        x = x + self.dropout(self.mla(x, mask=mask))
+        # Pre-norm FFN with residual
+        x = x + self.dropout(self.ffn(x))
+        return x
+
+
+class SimpleLLM(nn.Module):
+    """
+    Language Model with Hybrid Mamba-style SSM + MLA blocks.
+    
+    Architecture: Token Embedding → (SSM Blocks + MLA Blocks) → Output Head
+    
+    Hybrid structure controlled by config.ssm_per_mla:
+    - ssm_per_mla = 2: SSM, SSM, MLA, SSM, SSM, MLA, ...
+    - ssm_per_mla = 3: SSM, SSM, SSM, MLA, SSM, SSM, SSM, MLA, ...
+    """
+    
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        
+        # Factorized embeddings
+        self.token_embedding = FactorizedEmbedding(
+            vocab_size=config.vocab_size,
+            d_model=config.d_model,
+            d_embed_rank=config.d_embed_rank
+        )
+        
+        self.dropout = nn.Dropout(config.dropout)
+        
+        # Build block architecture based on arrangement strategy
+        self.blocks = nn.ModuleList()
+        
+        if config.block_arrangement == "interleaving":
+            self._build_interleaving_blocks(config)
+        elif config.block_arrangement == "layered":
+            self._build_layered_blocks(config)
+        else:
+            raise ValueError(f"Unknown block_arrangement: {config.block_arrangement}")
+        
+        # =================================================================
+        # Two-stage output projection (mirrors factorized embedding)
+        # =================================================================
+        # Stage 1: d_model → d_embed_rank (reverse of embedding projection)
+        self.output_proj = nn.Linear(config.d_model, config.d_embed_rank, bias=False)
+        
+        # Stage 2: d_embed_rank → vocab_size (tied to embedding table)
+        self.lm_head = nn.Linear(config.d_embed_rank, config.vocab_size, bias=False)
+        
+        # Tie lm_head weights to embedding table
+        self.lm_head.weight = self.token_embedding.embed.weight
+        # =================================================================
+        
+        # Final layer norm before output head to stabilize predictions
+        self.pre_final_norm = nn.LayerNorm(config.d_model)
+        self.final_norm = nn.LayerNorm(config.d_embed_rank)
+        
+        self.apply(self._init_weights)
+        self.register_buffer("causal_mask_cache", None, persistent=False)
+        self._print_architecture()
+    
+    def _build_interleaving_blocks(self, config):
+        """
+        Build interleaving block arrangement: SSM blocks followed by MLA blocks in a pattern.
+        
+        Example with ssm_per_mla=3 and n_layers=16:
+        SSM, SSM, SSM, MLA, SSM, SSM, SSM, MLA, SSM, SSM, SSM, MLA, SSM, SSM, SSM, MLA
+        """
+        ssm_per_mla = config.ssm_per_mla
+        num_mla_blocks = max(1, config.n_layers // (ssm_per_mla + 1))
+        
+        block_idx = 0
+        for mla_idx in range(num_mla_blocks):
+            # Add SSM blocks before each MLA block
+            for _ in range(ssm_per_mla):
+                if block_idx < config.n_layers:
+                    self.blocks.append(SSMBlockWrapper(config))
+                    block_idx += 1
+            
+            # Add MLA block
+            if block_idx < config.n_layers:
+                self.blocks.append(MLABlockWrapper(config))
+                block_idx += 1
+        
+        # Add remaining SSM blocks (if n_layers is not evenly divisible)
+        while block_idx < config.n_layers:
+            self.blocks.append(SSMBlockWrapper(config))
+            block_idx += 1
+    
+    def _build_layered_blocks(self, config):
+        """
+        Build layered block arrangement: MLA blocks followed by SSM blocks.
+        
+        Example with layered_mla_num=4 and n_layers=16:
+        MLA, MLA, MLA, MLA, SSM, SSM, SSM, SSM, SSM, SSM, SSM, SSM, SSM, SSM, SSM, SSM
+        """
+        num_mla = config.layered_mla_num
+        
+        # Add MLA blocks first
+        for _ in range(min(num_mla, config.n_layers)):
+            self.blocks.append(MLABlockWrapper(config))
+        
+        # Add remaining SSM blocks
+        num_ssm = config.n_layers - len(self.blocks)
+        for _ in range(num_ssm):
+            self.blocks.append(SSMBlockWrapper(config))
+    
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear) and not isinstance(module, BitLinear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module. bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+    
+    def _print_architecture(self):
+        total_params = self.count_parameters()
+        embed_params = self.token_embedding.get_num_params()
+        output_proj_params = self.config.d_model * self.config.d_embed_rank
+        ssm_params = total_params - embed_params - output_proj_params
+        
+        # Count SSM and MLA blocks
+        num_ssm = sum(1 for b in self.blocks if isinstance(b, SSMBlockWrapper))
+        num_mla = sum(1 for b in self.blocks if isinstance(b, MLABlockWrapper))
+        
+        print(f"\n{'='*60}")
+        print("MODEL ARCHITECTURE - HYBRID SSM + MLA")
+        print(f"{'='*60}")
+        print(f"Embedding:     {embed_params/1e6:>6.2f}M params")
+        print(f"Hybrid Blocks: {num_ssm} SSM + {num_mla} MLA = {num_ssm + num_mla} total")
+        print(f"Output Proj:   {output_proj_params/1e6:>6.2f}M params")
+        print(f"Output Head:   tied to embedding (0 extra params)")
+        print(f"{'─'*60}")
+        print(f"Total:         {total_params/1e6:>6.2f}M params")
+        print(f"{'='*60}")
+        print(f"Config:  {self.config.n_layers} layers, {self.config.d_model} dim")
+        print(f"SSM:     d_state={self.config.d_state}")
+        print(f"MLA:     n_heads={self.config.n_heads}, d_kv_comp={self.config.d_kv_comp}")
+        
+        # Print arrangement-specific info
+        if self.config.block_arrangement == "interleaving":
+            print(f"Arrangement: INTERLEAVING (ssm_per_mla={self.config.ssm_per_mla})")
+        elif self.config.block_arrangement == "layered":
+            print(f"Arrangement: LAYERED (mla_blocks={self.config.layered_mla_num}, ssm_blocks={num_ssm})")
+        
+        print(f"{'='*60}\n")
+    
+    def _get_causal_mask(self, seq_len, device):
+        if self.causal_mask_cache is None or self.causal_mask_cache. size(-1) < seq_len:
+            mask = torch.tril(torch.ones(seq_len, seq_len, device=device))
+            mask = mask.unsqueeze(0).unsqueeze(0)
+            self.causal_mask_cache = mask
+        return self.causal_mask_cache[: , :, :seq_len, :seq_len]
+    
+    def forward(self, input_ids, attention_mask=None):
+        batch_size, seq_len = input_ids.shape
+        
+        # Causal mask
+        causal_mask = self._get_causal_mask(seq_len, input_ids.device)
+        if attention_mask is not None:
+            padding_mask = attention_mask.unsqueeze(1).unsqueeze(1)
+            causal_mask = causal_mask * padding_mask
+        
+        # Token embedding
+        x = self.token_embedding(input_ids)
+        x = self.dropout(x)
+        x = ActivationQuantize.apply(x)
+        
+        # Hybrid SSM + MLA blocks
+        for block in self.blocks:
+            x = block(x, causal_mask)
+        
+        # Two-stage output projection
+        x = self.pre_final_norm(x)
+        x = self.output_proj(x)      # d_model → d_embed_rank
+        x = self.final_norm(x)       # Normalize before output head
+        logits = self.lm_head(x)     # d_embed_rank → vocab_size
+        
+        return logits
+    
+    def init_ssm_states(self, batch_size, device, dtype):
+        """
+        Initialize SSM states for all SSM blocks (MLA blocks are stateless).
+        
+        Returns:
+            states: List of [batch, d_state] tensors for each SSM block
+        """
+        states = []
+        for block in self.blocks:
+            if isinstance(block, SSMBlockWrapper):
+                state = block.ssm.init_state(batch_size, device, dtype)
+                states.append(state)
+        return states
+    
+    def inference_step(self, input_id, states, return_hidden_states=False):
+        """
+        Single inference step for autoregressive generation (RNN-like).
+        
+        Args:
+            input_id: [batch, 1] or scalar token id
+            states: List of SSM states from previous step
+            return_hidden_states: If True, also return SSM hidden states for visualization
+        
+        Returns:
+            logits: [batch, vocab_size] - output logits for next token
+            new_states: List of updated SSM states for SSM blocks
+            hidden_states: (Optional) List of SSM hidden state values for each SSM layer
+        """
+        if isinstance(input_id, int):
+            input_id = torch.tensor([[input_id]], dtype=torch.long, device=next(self.parameters()).device)
+        elif input_id.dim() == 1:
+            input_id = input_id.unsqueeze(0)
+        
+        # Embed the token
+        x = self.token_embedding(input_id)  # [batch, 1, d_model]
+        x = x.squeeze(1)  # [batch, d_model]
+        x = ActivationQuantize.apply(x)
+
+        # Pass through hybrid blocks
+        new_states = []
+        hidden_states = [] if return_hidden_states else None
+        state_idx = 0  # Track position in states list (only for SSM blocks)
+        
+        for block in self.blocks:
+            if isinstance(block, SSMBlockWrapper):
+                # SSM block with state management
+                residual = x
+                ssm_out, new_state = block.ssm.step(x, states[state_idx])
+                
+                # Collect hidden state if requested
+                if return_hidden_states:
+                    hidden_states.append(new_state.clone().detach())
+                
+                x = residual + block.dropout(ssm_out)
+                
+                # FFN + residual
+                residual = x
+                ffn_out = block.feed_forward(x)
+                x = residual + block.dropout(ffn_out)
+                
+                new_states.append(new_state)
+                state_idx += 1
+            else:
+                # MLA block (stateless)
+                x = block(x.unsqueeze(1), mask=None).squeeze(1)
+        
+        # Output projection
+        x = self.pre_final_norm(x)
+        x = self.output_proj(x)
+        x = self.final_norm(x)
+        logits = self.lm_head(x)
+        
+        if return_hidden_states:
+            return logits, new_states, hidden_states
+        else:
+            return logits, new_states
+    
+    def count_parameters(self):
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+    
+    def count_non_embedding_parameters(self):
+        total = self.count_parameters()
+        embedding_params = self.token_embedding.get_num_params()
+        return total - embedding_params
+    
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids,
+        max_new_tokens=50,
+        temperature=1.0,
+        top_k=50,
+        top_p=0.9,
+        repetition_penalty=1.1,
+        do_sample=True
+    ):
+        """Generate tokens autoregressively."""
+        self.eval()
+        
+        for _ in range(max_new_tokens):
+            # Crop to max_seq_len
+            idx_cond = input_ids[:, -self.config.max_seq_len:]
+            
+            # Forward
+            logits = self(idx_cond)
+            logits = logits[:, -1, : ] / max(temperature, 1e-5)
+            
+            # Repetition penalty
+            if repetition_penalty != 1.0:
+                for i in range(input_ids.shape[0]):
+                    for token_id in set(input_ids[i].tolist()):
+                        if logits[i, token_id] > 0:
+                            logits[i, token_id] /= repetition_penalty
+                        else:
+                            logits[i, token_id] *= repetition_penalty
+            
+            # Top-k filtering
+            if top_k is not None and top_k > 0:
+                v, _ = torch.topk(logits, min(top_k, logits. size(-1)))
+                logits[logits < v[:, [-1]]] = float('-inf')
+            
+            # Top-p filtering
+            if top_p is not None and top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[: , 1:] = sorted_indices_to_remove[:, :-1].clone()
+                sorted_indices_to_remove[:, 0] = 0
+                
+                for i in range(logits.shape[0]):
+                    indices_to_remove = sorted_indices[i, sorted_indices_to_remove[i]]
+                    logits[i, indices_to_remove] = float('-inf')
+            
+            # Sample or greedy
+            probs = F.softmax(logits, dim=-1)
+            if do_sample:
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                next_token = torch.argmax(probs, dim=-1, keepdim=True)
+            
+            input_ids = torch. cat([input_ids, next_token], dim=1)
+            
+            # Stop on EOS
+            if self.config.eos_token_id is not None:
+                if (next_token == self.config. eos_token_id).all():
+                    break
+        
+        return input_ids
+    
+    def get_num_params(self, non_embedding=True):
+        if non_embedding:
+            return self.count_non_embedding_parameters()
+        return self.count_parameters()

modeling_taonet.py

@@ -0,0 +1,181 @@
+"""
+Modeling class for TaoNet model.
+"""
+
+import torch
+import torch.nn as nn
+from dataclasses import dataclass
+from transformers import PreTrainedModel
+
+from .model import SimpleLLM
+from .configuration_taonet import TaoNetConfig
+
+
+@dataclass
+class InternalModelConfig:
+    """Internal config for SimpleLLM."""
+    vocab_size: int = 25000
+    d_model: int = 512
+    d_embed_rank: int = 384
+    d_state: int = 512
+    d_ff: int = 512
+    n_heads: int = 4
+    d_kv_comp: int = 384
+    d_rope: int = 64
+    n_layers: int = 8
+    max_seq_len: int = 256
+    dropout: float = 0.02
+    block_arrangement: str = "layered"
+    ssm_per_mla: int = 3
+    layered_mla_num: int = 0
+    pad_token_id: int = 3
+    bos_token_id: int = 1
+    eos_token_id: int = 2
+    unk_token_id: int = 0
+
+
+class TaoNetForCausalLM(PreTrainedModel):
+    """TaoNet model for causal language modeling."""
+    
+    config_class = TaoNetConfig
+    base_model_prefix = "taonet"
+    
+    def __init__(self, config: TaoNetConfig):
+        super().__init__(config)
+        
+        # Convert HF config to internal config
+        internal_config = InternalModelConfig(
+            vocab_size=config.vocab_size,
+            d_model=config.d_model,
+            d_embed_rank=config.d_embed_rank,
+            d_state=config.d_state,
+            d_ff=config.d_ff,
+            n_heads=config.n_heads,
+            d_kv_comp=config.d_kv_comp,
+            d_rope=config.d_rope,
+            n_layers=config.n_layers,
+            max_seq_len=config.max_seq_len,
+            dropout=config.dropout,
+            block_arrangement=config.block_arrangement,
+            ssm_per_mla=config.ssm_per_mla,
+            layered_mla_num=config.layered_mla_num,
+            pad_token_id=config.pad_token_id,
+            bos_token_id=config.bos_token_id,
+            eos_token_id=config.eos_token_id,
+            unk_token_id=config.unk_token_id,
+        )
+        
+        self.taonet = SimpleLLM(internal_config)
+        
+        # Tie the lm_head weights to the token embedding weights
+        self._tie_weights()
+    
+    def _tie_weights(self):
+        """Tie the lm_head weight to the token embedding weight."""
+        if hasattr(self.taonet, 'token_embedding') and hasattr(self.taonet, 'lm_head'):
+            # Tie the weights - make lm_head.weight reference the same tensor as token_embedding.embed.weight
+            self.taonet.lm_head.weight = self.taonet.token_embedding.embed.weight
+    
+    def _init_weights(self, module):
+        """Initialize weights (override to maintain tied weights)."""
+        # Let the parent handle initialization, then retie weights
+        super()._init_weights(module) if hasattr(super(), '_init_weights') else None
+        self._tie_weights()
+    
+    @property
+    def all_tied_weights_keys(self):
+        """Return the tied weights keys to satisfy transformers requirements."""
+        # Return as a dict with tied_weight -> main_weight mapping
+        return {"taonet.lm_head.weight": "taonet.token_embedding.embed.weight"}
+    
+    def mark_tied_weights_as_initialized(self):
+        """Mark tied weights as initialized by actually tying them together."""
+        # Tie the weights so they reference the same tensor
+        self._tie_weights()
+    
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask=None,
+        labels=None,
+        **kwargs,
+    ):
+        """Forward pass."""
+        logits = self.taonet(input_ids)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            loss = loss_fct(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1)
+            )
+        
+        return {
+            "loss": loss,
+            "logits": logits,
+        }
+
+    def init_ssm_states(self, batch_size: int, device: torch.device, dtype: torch.dtype):
+        """Initialize SSM states for all SSM blocks."""
+        return self.taonet.init_ssm_states(batch_size, device, dtype)
+    
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        max_length: int = 100,
+        temperature: float = 1.0,
+        top_k=None,
+        top_p=None,
+        **kwargs,
+    ):
+        """Generate text using RNN-style inference with state management."""
+        self.taonet.eval()
+        batch_size = input_ids.shape[0]
+        device = input_ids.device
+        dtype = next(self.taonet.parameters()).dtype
+        
+        # Initialize SSM states for all SSM blocks
+        states = self.taonet.init_ssm_states(batch_size, device, dtype)
+        
+        current_ids = input_ids.clone()
+        
+        # Process initial tokens to prime the states
+        with torch.no_grad():
+            for i in range(input_ids.shape[1]):
+                token_id = input_ids[:, i:i+1]
+                _, states = self.taonet.inference_step(token_id, states)
+        
+        # Generate new tokens
+        for _ in range(max_length - input_ids.shape[1]):
+            with torch.no_grad():
+                # Get logits for next token using inference_step
+                next_token_id = current_ids[:, -1:]
+                logits, states = self.taonet.inference_step(next_token_id, states)
+            
+            next_logits = logits / temperature
+            
+            if top_k is not None:
+                top_k_logits, top_k_indices = torch.topk(next_logits, min(top_k, next_logits.size(-1)), dim=-1)
+                indices_to_remove = next_logits < top_k_logits[..., -1, None]
+                next_logits[indices_to_remove] = float('-inf')
+            
+            if top_p is not None:
+                sorted_logits, sorted_indices = torch.sort(next_logits, descending=True, dim=-1)
+                cumsum_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+                sorted_indices_to_remove = cumsum_probs > top_p
+                sorted_indices_to_remove[..., 0] = False
+                indices_to_remove = sorted_indices[sorted_indices_to_remove]
+                next_logits[..., indices_to_remove] = float('-inf')
+            
+            probs = torch.softmax(next_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            current_ids = torch.cat([current_ids, next_token], dim=1)
+            
+            if (next_token == self.config.eos_token_id).any():
+                break
+        
+        return current_ids

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:68d530216dc370cf96f066f1897a26896b98e12cab647dba61a4314364c7d05e
+size 112420148

rope.py

@@ -0,0 +1,47 @@
+"""Rotary Position Embedding (RoPE) implementation."""
+
+import torch
+import torch.nn as nn
+import math
+
+
+class RotaryEmbedding(nn.Module):
+    """Rotary position embeddings."""
+    
+    def __init__(self, dim, scale=40):
+        super().__init__()
+        assert dim % 2 == 0, "Dimension must be even for rotary embeddings"
+        self.dim = dim
+        self.scale = scale
+        
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+    
+    def forward(self, seq_len, device):
+        """Generate rotary embeddings for sequence."""
+        t = torch.arange(seq_len, device=device).type_as(self.inv_freq) / self.scale
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        return torch.cat((freqs, freqs), dim=-1)
+
+
+def rotate_half(x):
+    """Rotate half the hidden dims of the input."""
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary(x, cos, sin):
+    """Apply rotary embeddings to input tensor."""
+    # Handle case where cos/sin may be shorter than x
+    cos = cos[..., :x.shape[-1]]
+    sin = sin[..., :x.shape[-1]]
+    
+    # Split x based on cos dimensions
+    x_rot = x[..., :cos.shape[-1]]
+    x_base = x[..., cos.shape[-1]:]
+    
+    # Apply rotation
+    x_rot = (x_rot * cos) + (rotate_half(x_rot) * sin)
+    
+    # Concatenate rotated and base parts
+    return torch.cat([x_rot, x_base], dim=-1) if x_base.shape[-1] > 0 else x_rot

ssm.py

@@ -0,0 +1,147 @@
+"""
+Ternary Quantized Diagonal State-Space Model (Parallel)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .bitlinear import BitLinear, RMSNorm
+
+class Q88Quantize(torch.autograd.Function):
+    """Q8.8 fixed-point quantization with straight-through estimator."""
+    
+    @staticmethod
+    def forward(ctx, x):
+        """
+        Quantize to Q8.8 format (8 integer bits, 8 fractional bits)
+        Range: [-128, 127.99609375]
+        """
+        scale = 2**8  # 256
+        # Quantize: scale -> round -> clamp to int16 range -> dequantize
+        x_scaled = x * scale
+        x_int = torch.clamp(torch.round(x_scaled), -32768, 32767)
+        x_quant = x_int / scale
+        return x_quant
+    
+    @staticmethod
+    def backward(ctx, grad_output):
+        # Straight-through estimator: pass gradients unchanged
+        return grad_output
+
+class SSMBlock(nn.Module):
+    """
+    Diagonal (convolutional) SSM Block with ternary BitLinear projections.
+
+    Architecture:
+        Input → B projection → diagonal SSM convolution → C projection → Output
+
+    State dynamics (training, parallel):
+        s_t = sum_{k=0}^t (B x_k)
+
+    State dynamics (inference, step-wise):
+        s_t = s_{t-1} + B x_t
+        y_t = C s_t
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.d_model = config.d_model
+        self.d_state = config.d_state
+
+        # =====================================================================
+        # Stationary ternary projections
+        # =====================================================================
+        self.b_proj = BitLinear(self.d_model, self.d_state, bias=False)
+        self.c_proj = BitLinear(self.d_state, self.d_model, bias=False)
+        # A matrix: identity matrix scaled by a single scalar decay factor
+        self.register_buffer("a_log", torch.log(torch.tensor(0.9)))
+
+        self.dropout = nn.Dropout(config.dropout)
+
+    # ---------------------------------------------------------------------
+    # Training / parallel forward
+    # ---------------------------------------------------------------------
+    def forward(self, x, mask=None):
+        """
+        Args:
+            x: [batch, seq_len, d_model]
+            mask: unused (SSM is causal by construction)
+
+        Returns:
+            y: [batch, seq_len, d_model]
+        """
+        B, L, _ = x.shape
+
+        # Input projection
+        u = self.b_proj(x)
+
+        # Compute decay with Q8.8 quantization
+        decay = torch.exp(self.a_log)  # scalar
+        decay_quant = decay + (Q88Quantize.apply(decay) - decay).detach()
+
+        L = u.size(1)
+        device = u.device
+        dtype = u.dtype
+
+        # Decay powers with quantization: [L, 1] (broadcasts across d_state)
+        t = torch.arange(L, device=device, dtype=dtype).unsqueeze(1)  # [L, 1]
+        decay_pows = decay_quant ** t  # [L, 1]
+        #decay_pows = decay_pows + (Q88Quantize.apply(decay_pows) - decay_pows).detach()
+        
+        inv_decay_pows = decay_pows.reciprocal()
+        #inv_decay_pows = inv_decay_pows + (Q88Quantize.apply(inv_decay_pows) - inv_decay_pows).detach()
+
+        # Reweight, cumsum, reweight back
+        s = torch.cumsum(u * inv_decay_pows.unsqueeze(0), dim=1)  # [B, L, d_state]
+        s = s * decay_pows.unsqueeze(0)
+
+        # Output projection
+        y = self.c_proj(s)
+        y = self.dropout(y)
+
+        return y
+
+    # ---------------------------------------------------------------------
+    # Autoregressive single-step inference
+    # ---------------------------------------------------------------------
+    def step(self, x, state):
+        """
+        Single timestep SSM update (for autoregressive decoding).
+
+        Args:
+            x: [batch, d_model]
+            state: [batch, d_state]
+
+        Returns:
+            output: [batch, d_model]
+            new_state: [batch, d_state]
+        """
+        decay = torch.exp(self.a_log)  # scalar
+        new_state = decay * state + self.b_proj(x)  # [batch, d_state]
+        output = self.c_proj(new_state)
+        return output, new_state
+
+    # ---------------------------------------------------------------------
+    # State utilities
+    # ---------------------------------------------------------------------
+    def init_state(self, batch_size, device, dtype):
+        """Initialize hidden state."""
+        return torch.zeros(batch_size, self.d_state, device=device, dtype=dtype)
+
+    # ---------------------------------------------------------------------
+    # Export parameters for inference / FPGA
+    # ---------------------------------------------------------------------
+    def get_inference_params(self):
+        """
+        Export parameters for deployment.
+
+        Returns:
+            dict with quantized projections and diagonal A
+        """
+        with torch.no_grad():
+            return {
+                "b_proj": self.b_proj.get_inference_params(),
+                "c_proj": self.c_proj.get_inference_params(),
+            }