initial commit

__init__.py

@@ -0,0 +1,2 @@
+from .configuration_ministu import MiniSTUConfig
+from .modeling_ministu import MiniSTU

added_tokens.json

@@ -0,0 +1,3 @@
+{
+    "<|endofprompt|>": 200018
+}

attn.py

@@ -0,0 +1,125 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import nearest_power_of_two
+
+try:
+    from flash_attn import flash_attn_func as fa2
+except ImportError as e:
+    print(
+        f"Unable to import Triton-based flash attention: {e}. No alternative currently available."
+    )
+    # TODO: Add FlexAttention + local attention mask when it's in stable release
+
+class Attention(nn.Module):
+    def __init__(self, config):
+        super(Attention, self).__init__()
+        if isinstance(config.torch_dtype, str):
+            torch_dtype = getattr(torch, config.torch_dtype)
+        else:
+            torch_dtype = config.torch_dtype
+        assert torch.cuda.is_available(), "CUDA is required."
+        assert config.n_embd % config.n_heads == 0
+        self.n_heads = config.n_heads
+
+        self.device = torch.device("cuda")
+        self.bsz = config.bsz
+        self.c_attn = nn.Linear(
+            config.n_embd, 3 * config.n_embd, bias=config.bias, dtype=torch_dtype
+        )
+        self.c_proj = nn.Linear(
+            config.n_embd, config.n_embd, bias=config.bias, dtype=torch_dtype
+        )
+        self.c_proj.SCALE_INIT = 1
+        self.dropout = config.dropout
+        self.resid_dropout = nn.Dropout(self.dropout)
+        self.alibi_slopes = self._get_alibi_slopes(self.n_heads)
+        self.window_size = config.window_size
+        self.softcap = config.softcap
+
+    def _generate_slopes(self, n: int):
+        start = 2 ** (-(2 ** -(math.log2(n) - 3)))
+        return [start * (start**i) for i in range(n)]
+
+    def _get_alibi_slopes(self, n_heads: int, interpolation_factor: float = 0.25):
+        # If n_heads is a power of 2, generate slopes directly
+        if math.log2(n_heads).is_integer():
+            slopes = self._generate_slopes(n_heads)
+        else:
+            # Get slopes for the nearest power of two
+            n = nearest_power_of_two(n_heads, round_up=False)
+            slopes_power_of_two = self._generate_slopes(n)
+    
+            # Generate extra slopes
+            extra_slopes = self._generate_slopes(2 * n)
+            extra_slopes_trunc = extra_slopes[0::2][: n_heads - n]
+            slopes = slopes_power_of_two + extra_slopes_trunc
+        slopes = torch.tensor(slopes, device=self.device)
+        slopes = slopes * interpolation_factor  # https://arxiv.org/pdf/2310.13017
+        return slopes.to(torch.float32)  # Ensure slopes are in float32
+
+
+    def forward(self, x):
+        bsz, seq_len, d_in = x.size()
+
+        qkv = self.c_attn(x)
+        q, k, v = torch.chunk(qkv, 3, dim=2)
+
+        q = q.view(bsz, seq_len, self.n_heads, d_in // self.n_heads)
+        k = k.view(bsz, seq_len, self.n_heads, d_in // self.n_heads)
+        v = v.view(bsz, seq_len, self.n_heads, d_in // self.n_heads)
+        y = fa2(  # https://arxiv.org/pdf/2307.08691
+            q,
+            k,
+            v,
+            dropout_p=self.dropout if self.training else 0.0,
+            causal=True,
+            window_size=(self.window_size, 0),
+            alibi_slopes=self.alibi_slopes,  # https://arxiv.org/pdf/2108.12409
+            softcap=self.softcap,  # https://arxiv.org/pdf/2408.00118
+        )
+        y = y.contiguous().view(bsz, seq_len, d_in)
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+class AttentionSDPA(nn.Module):
+    def __init__(self, config):
+        super(Attention, self).__init__()
+        if isinstance(config.torch_dtype, str):
+            torch_dtype = getattr(torch, config.torch_dtype)
+        else:
+            torch_dtype = config.torch_dtype
+        assert torch.cuda.is_available(), "CUDA is required."
+        assert config.n_embd % config.n_heads == 0
+        self.n_heads = config.n_heads
+
+        self.device = torch.device("cuda") # Technically don't need CUDA for SDPA
+        self.bsz = config.bsz
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias, dtype=torch_dtype)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias, dtype=torch_dtype)
+        self.dropout = config.dropout
+        self.resid_dropout = nn.Dropout(self.dropout)
+
+    def forward(self, x):
+        bsz, seq_len, d_in = x.size()
+
+        qkv = self.c_attn(x)
+        q, k, v = torch.chunk(qkv, 3, dim=2)
+
+        q = q.view(bsz, seq_len, self.n_heads, d_in // self.n_heads).transpose(1, 2)
+        k = k.view(bsz, seq_len, self.n_heads, d_in // self.n_heads).transpose(1, 2)
+        v = v.view(bsz, seq_len, self.n_heads, d_in // self.n_heads).transpose(1, 2)
+
+        y = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=True,
+            dropout_p=self.dropout if self.training else 0.0
+        )
+
+        y = y.transpose(1, 2).contiguous().view(bsz, seq_len, d_in)
+
+        y = self.resid_dropout(self.c_proj(y))
+        return y

config.json

@@ -0,0 +1,55 @@
+{
+  "model_type": "ministu",
+  "_name_or_path": "STU-500M",
+  "architectures": ["MiniSTU"],
+  "n_embd": 768,
+  "n_heads": 8,
+  "n_layers": 12,
+  "seq_len": 8192,
+  "window_size": 1024,
+  "vocab_size": 200064,
+  "mlp_scale": 12,
+  "bias": false,
+  "dropout": 0.0,
+  "num_eigh": 24,
+  "use_hankel_L": false,
+  "num_epochs": 1,
+  "global_bsz": 524288,
+  "bsz": 1,
+  "warmup_steps": 1907,
+  "eval_period": 25,
+  "save_period": 4500,
+  "max_lr": 3.0e-3,
+  "min_lr": 3.0e-5,
+  "max_norm": 1.0,
+  "dilation": 1,
+  "fsdp": true,
+  "ddp": false,
+  "mixed_precision": true,
+  "use_cpu_offload": false,
+  "sharding_strategy": "full_shard",
+  "state_dict_type": "full",
+  "auto_wrap_policy": "partial",
+  "backward_prefetch": "backward_pre",
+  "forward_prefetch": false,
+  "sync_module_states": true,
+  "use_orig_params": true,
+  "device_id": null,
+  "precision": {
+    "param": "bfloat16",
+    "reduce": "bfloat16",
+    "buffer": "bfloat16"
+  },
+  "fsdp_modules": [
+    "STU",
+    "Attention",
+    "MLP"
+  ],
+  "use_activation_checkpointing": true,
+  "use_flash_fft": false,
+  "use_approx": true,
+  "use_attn": true,
+  "softcap": 50.0,
+  "torch_compile": false
+}
+

configuration_ministu.py

@@ -0,0 +1,50 @@
+import torch
+from transformers import PretrainedConfig, AutoConfig
+
+class MiniSTUConfig(PretrainedConfig):
+    model_type = "ministu"
+
+    def __init__(
+        self,
+        bsz: int = 1,
+        n_embd: int = 768,
+        n_heads: int = 8,
+        n_layers: int = 12,
+        seq_len: int = 8192,
+        window_size: int = 1024,
+        vocab_size: int = 200064,
+        mlp_scale: int = 12,
+        bias: bool = False,
+        dropout: float = 0.0,
+        num_eigh: int = 24,
+        use_hankel_L: bool = False,
+        use_flash_fft: bool = True,
+        use_approx: bool = True,
+        use_attn: bool = True,
+        softcap: float = 50.0,
+        torch_dtype = torch.bfloat16,
+        device: str = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.bsz = bsz
+        self.n_embd = n_embd
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.seq_len = seq_len
+        self.window_size = window_size
+        self.vocab_size = vocab_size
+        self.hidden_size = n_embd
+        self.intermediate_size = n_embd * mlp_scale
+        self.mlp_scale = mlp_scale
+        self.hidden_act = "swish"
+        self.bias = bias
+        self.dropout = dropout
+        self.num_eigh = num_eigh
+        self.use_hankel_L = use_hankel_L
+        self.use_flash_fft = use_flash_fft
+        self.use_approx = use_approx
+        self.use_attn = use_attn
+        self.softcap = softcap
+        self.torch_dtype = torch_dtype  # Store as string
+        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')  # Store as string

convolve.py

@@ -0,0 +1,84 @@
+import torch
+import torch.nn.functional as F
+
+from .utils import nearest_power_of_two
+from flashfftconv import FlashFFTConv
+
+
+def convolve(u: torch.Tensor, v: torch.Tensor, n: int, use_approx: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
+    bsz, seq_len, d_in = u.shape
+
+    sgn = torch.full((1, seq_len, 1), 1, device=u.device, dtype=torch.float32)
+    sgn[:, 1::2] *= -1
+
+    # Cast u and v to float32 for FFT
+    u = u.to(torch.float32)
+    v = v.to(torch.float32)
+
+    if use_approx:
+        _, d_out = v.shape
+        v = v.view(1, -1, d_out, 1)
+    else:
+        _, K = v.shape
+        sgn = sgn.unsqueeze(-1)
+        v = v.view(1, -1, K, 1, 1)
+        u = u.view(bsz, -1, 1, d_in).expand(bsz, -1, K, d_in)
+
+    v = torch.fft.rfft(v, n=n, dim=1)
+    U = torch.stack([u, u * sgn], dim=-1)
+    U = torch.fft.rfft(U, n=n, dim=1)
+    U_conv = torch.fft.irfft(v * U, n=n, dim=1)[:, :seq_len]
+    U_plus, U_minus = torch.unbind(U_conv, dim=-1)
+    U_minus = U_minus * sgn
+
+    # Convert back to original dtype
+    U_plus = U_plus.to(u.dtype)
+    U_minus = U_minus.to(u.dtype)
+
+    return U_plus, U_minus
+
+def flash_convolve(
+    u: torch.Tensor, v: torch.Tensor, flash_fft: FlashFFTConv, use_approx: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    dtype = u.dtype  # Store the original dtype
+    u = u.to(torch.float32)
+    v = v.to(torch.float32)
+
+    bsz, seq_len, d_in = u.shape
+    _, K = v.shape
+
+    padded_len = nearest_power_of_two(seq_len, round_up=True)
+    pad_len = padded_len - seq_len
+
+    sgn = torch.full((1, 1, padded_len), 1, device=u.device, dtype=torch.float32)
+    sgn[:, :, 1::2] = -1
+
+    if use_approx:
+        u_padded = F.pad(u.transpose(1, 2), (0, pad_len)).contiguous()
+        v_padded = F.pad(v.transpose(0, 1), (0, pad_len)).contiguous()
+        u_conv = torch.stack([u_padded, u_padded * sgn], dim=0).reshape(2 * bsz, d_in, padded_len)
+    else:
+        u_k_padded = F.pad(u.transpose(1, 2), (0, pad_len)).repeat_interleave(K, dim=1).contiguous()
+        v_padded = F.pad(v.transpose(0, 1), (0, pad_len)).repeat(d_in, 1).contiguous()
+        u_conv = torch.stack([u_k_padded, u_k_padded * sgn], dim=0).reshape(2 * bsz, K * d_in, padded_len)
+
+    U_conv = flash_fft(u_conv, v_padded)
+
+    # Trim the output back to the original sequence length
+    U_conv = U_conv[..., :seq_len]
+
+    u_plus, u_minus = torch.chunk(U_conv, 2, dim=0)
+
+    if use_approx:
+        u_minus = u_minus * sgn[:, :, :seq_len]
+        U_plus, U_minus = u_plus.transpose(1, 2), u_minus.transpose(1, 2)
+    else:
+        sgn = sgn[:, :, :seq_len].unsqueeze(-1).transpose(1, 2)
+        U_plus = u_plus.view(bsz, d_in, K, seq_len).permute(0, 3, 2, 1).contiguous()
+        U_minus = u_minus.view(bsz, d_in, K, seq_len).permute(0, 3, 2, 1).contiguous() * sgn
+
+    # Convert back to original dtype
+    U_plus = U_plus.to(dtype)
+    U_minus = U_minus.to(dtype)
+
+    return U_plus, U_minus

filters.py

@@ -0,0 +1,106 @@
+import math
+
+import numpy as np
+import torch
+
+from .utils import logger
+from .utils import get_hankel
+
+def get_spectral_filters(
+    seq_len: int,
+    K: int,
+    use_hankel_L: bool = False,
+    device: torch.device = None,
+    dtype: torch.dtype = torch.bfloat16,
+) -> torch.Tensor:
+    # Generate the Hankel matrix using PyTorch
+    Z = get_hankel(seq_len, use_hankel_L, device=device, dtype=dtype)
+
+    # Cast Z to torch.float32 for the eigenvalue decomposition
+    Z_float32 = Z.to(torch.float32)
+
+    # Perform eigen decomposition using torch.float32
+    sigma, phi = torch.linalg.eigh(Z_float32)
+
+    # Cast the results back to the original dtype (torch.bfloat16)
+    sigma = sigma.to(dtype=dtype)
+    phi = phi.to(dtype=dtype)
+
+    # Select the top K eigenvalues and eigenvectors
+    sigma_k, phi_k = sigma[-K:], phi[:, -K:]
+
+    # Compute the spectral filters
+    phi_k = phi_k * sigma_k ** 0.25
+
+    # Ensure the filters are in the correct dtype and device
+    filters = phi_k.to(device=device, dtype=dtype)
+
+    return filters
+
+
+def compute_dimensions(n: int) -> tuple[int, int, int]:
+    if n <= 2:
+        raise ValueError("n must be greater than 2")
+
+    T_prime = (math.ceil(math.sqrt(n - 2)))**2 + 2
+    sqrt_T_prime = math.ceil(math.sqrt(T_prime - 2))
+    k_max = sqrt_T_prime
+    return T_prime, sqrt_T_prime, k_max
+
+def get_tensorized_spectral_filters_explicit(n: int, k: int, device: torch.device) -> torch.Tensor:
+    T_prime, sqrt_T_prime, k_max = compute_dimensions(n)
+    k = min(k, k_max)
+
+    Z = get_hankel(sqrt_T_prime).to(device)
+    sigma, phi = torch.linalg.eigh(Z)
+    sigma_k = sigma[-k:]
+    phi_k = phi[:, -k:]
+
+    result = torch.zeros(sqrt_T_prime * sqrt_T_prime, device=device)
+    
+    for i in range(k):
+        for j in range(k):
+            phi_i = phi_k[:, i] * (sigma_k[i] ** 0.25)
+            phi_j = phi_k[:, j] * (sigma_k[j] ** 0.25)
+            kron = torch.kron(phi_i, phi_j)
+            result += kron
+            
+    return result
+
+
+def get_tensorized_spectral_filters(
+    n: int = 8192,
+    k: int = 24,
+    use_hankel_L: bool = False,
+    device: torch.device = None,
+    dtype: torch.dtype = torch.bfloat16,
+) -> torch.Tensor:
+    """
+    Compute tensorized spectral filters for given sequence length and filter count.
+
+    Args:
+        n: Sequence length
+        k: Number of filters
+        use_hankel_L: Hankel_main ⊗ Hankel_L? Default is Hankel_main ⊗ Hankel_main.
+        device: Computation device
+        dtype: Computation dtype
+    """
+    assert torch.cuda.is_available(), "CUDA is required."
+
+    T_prime, sqrt_T_prime, k_max = compute_dimensions(n)
+    k = min(k, k_max)
+
+    Z = get_hankel(sqrt_T_prime)
+    sigma, phi = torch.linalg.eigh(Z)
+    phi_i = phi[:, -k:] * sigma[-k:] ** 0.25
+
+    if use_hankel_L: # TODO: We may want to use Hankel_L above too if use_hankel_L is true, make another variable for this (mix != use_hankel_L)
+        logger.info("Mixing Hankel_L with Hankel_main to generate tensorized filters.")
+        Z_L = get_hankel(sqrt_T_prime, True)
+        sigma_L, phi_L = torch.linalg.eigh(Z_L)
+        phi_j = phi_L[:, -k:] * sigma_L[-k:] ** 0.25
+    else:
+        phi_j = phi_i
+
+    filters = torch.kron(phi_i, phi_j)
+    return filters.to(device=device, dtype=dtype)

layers.py

@@ -0,0 +1,102 @@
+import torch
+import torch.nn as nn
+
+from .modules import STU
+from .modules import MLP
+from .modules import Attention
+try:
+    from liger_kernel.transformers.swiglu import LigerSwiGLUMLP as TritonMLP
+    triton_mlp = True
+except ImportError as e:
+    print(
+        f"Unable to import Triton-based MLP: {e}. Falling back to vanilla SwiGLU MLP instead."
+    )
+    triton_mlp = False
+
+try:
+    from liger_kernel.transformers.rms_norm import LigerRMSNorm as TritonNorm
+    triton_norm = True
+except ImportError as e:
+    print(
+        f"Unable to import Triton-based RMSNorm: {e}. Falling back to PyTorch implementation."
+    )
+    from torch.nn import RMSNorm
+    triton_norm = False
+
+
+class STULayer(nn.Module):
+    def __init__(self, config, phi, n):
+        super(STULayer, self).__init__()
+        if isinstance(config.torch_dtype, str):
+            torch_dtype = getattr(torch, config.torch_dtype)
+        else:
+            torch_dtype = config.torch_dtype
+        self.stu_norm = (
+            TritonNorm(config.n_embd)
+            if triton_norm
+            else RMSNorm(config.n_embd, dtype=torch_dtype)
+        )
+        self.stu = STU(config, phi, n)
+        self.stu = self.stu.to(dtype=torch_dtype)
+        self.mlp_norm = (
+            TritonNorm(config.n_embd)
+            if triton_norm
+            else RMSNorm(config.n_embd, dtype=torch_dtype)
+        )
+        self.mlp = (
+            TritonMLP(config) if triton_mlp else MLP(config, dtype=torch_dtype)
+        )
+
+        # TODO: Write Issue in Liger-Kernel repo to support user-defined dtype for MLP
+        self.stu_norm = self.stu_norm.to(dtype=torch_dtype)
+        self.mlp = self.mlp.to(dtype=torch_dtype)
+        self.mlp_norm = self.mlp_norm.to(dtype=torch_dtype)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Debug dtype
+    
+        # Normalize and apply STU
+        x_normed = self.stu_norm(x).to(dtype=self.stu.M_inputs.dtype)  # Match dtype for STU
+        x_stu = self.stu(x_normed).to(dtype=x.dtype)  # Ensure output matches `x`'s dtype
+        x = x + x_stu
+    
+        # Normalize and apply MLP
+        x_normed_mlp = self.mlp_norm(x).to(dtype=self.mlp.gate_proj.weight.dtype)  # Match dtype for MLP
+        x_mlp = self.mlp(x_normed_mlp).to(dtype=x.dtype)  # Ensure output matches `x`'s dtype
+        x = x + x_mlp
+    
+        return x
+
+class AttentionLayer(nn.Module):
+    def __init__(self, config) -> None:
+        super(AttentionLayer, self).__init__()
+        if isinstance(config.torch_dtype, str):
+            torch_dtype = getattr(torch, config.torch_dtype)
+        else:
+            torch_dtype = config.torch_dtype
+        self.attn_norm = (
+            TritonNorm(config.n_embd)
+            if triton_norm
+            else RMSNorm(config.n_embd, dtype=torch_dtype)
+        )
+        self.attn = Attention(config)
+        self.attn = self.attn.to(dtype=torch_dtype)
+        self.mlp_norm = (
+            TritonNorm(config.n_embd)
+            if triton_norm
+            else RMSNorm(config.n_embd, dtype=torch_dtype)
+        )
+        self.mlp = (
+            TritonMLP(config) if triton_mlp else MLP(config, dtype=torch_dtype)
+        )
+        self.mlp = self.mlp.to(dtype=torch_dtype)
+
+        # TODO: Write Issue in Liger-Kernel repo to support user-defined dtype for MLP
+        self.attn_norm = self.attn_norm.to(dtype=torch_dtype)
+        self.mlp = self.mlp.to(dtype=torch_dtype)
+        self.mlp_norm = self.mlp_norm.to(dtype=torch_dtype)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.attn_norm(x))
+        x = x + self.mlp(self.mlp_norm(x))
+        return x

mlp.py

@@ -0,0 +1,27 @@
+import torch.nn as nn
+from torch.nn import functional as F
+
+class MLP(nn.Module):
+    def __init__(self, config, dtype=None):
+        # https://arxiv.org/pdf/2002.05202
+        super().__init__()
+        torch_dtype = getattr(torch, config.torch_dtype, torch.float32)  # Use config dtype
+        dtype = dtype if dtype is not None else torch_dtype
+        self.hidden_size = config.n_embd
+        self.intermediate_size = config.n_embd * config.mlp_scale
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.bias, dtype=torch.bfloat16)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.bias, dtype=torch.bfloat16)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.bias, dtype=torch.bfloat16)
+        self.dropout = nn.Dropout(config.dropout)
+    
+    def forward(self, x):
+        dtype = self.gate_proj.weight.dtype  # Match the dtype of projection layers
+        x = x.to(dtype=dtype)  # Convert input to the same dtype
+        x = x.to(self.gate_proj.weight.dtype)
+        gate = self.gate_proj(x)
+        gate = F.gelu(gate, approximate="tanh").to(dtype=dtype)
+        up = self.up_proj(x).to(dtype=dtype)
+        fuse = gate * up
+        outputs = self.down_proj(fuse).to(dtype=dtype)
+        outputs = self.dropout(outputs)
+        return outputs

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:101763e8119c492ada9816aa38006cbf6ba8bbc0530224510d62b2c7e20a8bfd
+size 1140654808

modeling_ministu.py

@@ -0,0 +1,248 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutput
+from .modules import STU, Attention, MLP
+from .utils import nearest_power_of_two
+from .layers import STULayer, AttentionLayer
+from .configuration_ministu import MiniSTUConfig
+from .filters import get_spectral_filters
+
+try:
+    from liger_kernel.transformers.rms_norm import LigerRMSNorm as TritonNorm
+    triton_norm = True
+except ImportError as e:
+    print(
+        f"Unable to import Triton-based RMSNorm: {e}. Falling back to PyTorch implementation."
+    )
+    from torch.nn import RMSNorm
+    triton_norm = False
+# Load the tokenizer
+#from transformers import AutoModelForCausalLM, AutoTokenizer
+#model_name = "Hazan-Lab/STU-426M"
+#tokenizer = AutoTokenizer.from_pretrained(
+#    model_name,
+#    trust_remote_code=True
+#)
+
+class MiniSTU(PreTrainedModel):
+    config_class = MiniSTUConfig
+
+    def __init__(self, config) -> None:
+        super(MiniSTU, self).__init__(config)
+        self.n_layers = config.n_layers
+        self.n = nearest_power_of_two(config.seq_len * 2 - 1, round_up=True)
+        
+        if isinstance(config.torch_dtype, torch.dtype):
+            torch_dtype = config.torch_dtype
+        else:
+            torch_dtype = getattr(torch, config.torch_dtype)
+
+        device = torch.device(config.device)
+
+        self.phi = get_spectral_filters(
+            config.seq_len,
+            config.num_eigh,
+            config.use_hankel_L,
+            device=device,
+            dtype=torch_dtype,
+        )
+        
+        self.use_approx = config.use_approx
+        self.use_hankel_L = config.use_hankel_L
+
+        self.tok_emb = nn.Embedding(
+            config.vocab_size, config.n_embd, dtype=torch_dtype, device=device
+        )
+        self.dropout = nn.Dropout(config.dropout)
+
+        self.layers = nn.ModuleList()
+        for layer_idx in range(self.n_layers):
+            if layer_idx % 2 == 0:
+                self.layers.append(STULayer(config, self.phi, self.n))
+            else:
+                self.layers.append(
+                    AttentionLayer(config)
+                    if config.use_attn
+                    else STULayer(config, self.phi, self.n)
+                )
+
+        self.norm = TritonNorm(config.n_embd) if triton_norm else RMSNorm(config.n_embd)
+
+        self.lm_head = nn.Linear(
+            config.n_embd, config.vocab_size, bias=config.bias, dtype=torch_dtype, device=device
+        )
+        self.tok_emb.weight = self.lm_head.weight
+
+        self.std = (config.n_embd) ** -0.5
+        self.apply(self._init_weights)
+        print("Model Parameter Count: %.2fM\n" % (self._get_num_params() / 1e6,))
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: torch.Tensor = None,
+        **kwargs
+    ) -> CausalLMOutput:
+        # Compute embeddings
+        tok_emb = self.tok_emb(input_ids)
+        x = self.dropout(tok_emb)
+
+        # Pass through layers
+        for layer in self.layers:
+            x = layer(x)
+
+        # Normalize and project to vocabulary
+        x = self.norm(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens predict the next token
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1)
+            )
+
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits,
+        )
+
+    def _get_num_params(self):
+        n_params = sum(p.numel() for p in self.parameters())
+        if hasattr(self, "pos_emb") and self.pos_emb is not None:
+            n_params -= self.pos_emb.weight.numel()
+        if self.tok_emb.weight is not self.lm_head.weight:
+            n_params -= self.tok_emb.weight.numel()
+        return n_params
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            if hasattr(module, "SCALE_INIT"):
+                self.std *= (2 * self.n_layers) ** -0.5
+            torch.nn.init.normal_(module.weight, mean=0.0, std=self.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=self.std)
+        elif isinstance(module, STU):
+            if self.use_approx:
+                torch.nn.init.xavier_normal_(module.M_inputs)
+                torch.nn.init.xavier_normal_(module.M_filters)
+            else:
+                torch.nn.init.xavier_normal_(module.M_phi_plus)
+                if not self.use_hankel_L:
+                    torch.nn.init.xavier_normal_(module.M_phi_minus)
+        elif isinstance(module, Attention):
+            torch.nn.init.xavier_normal_(module.c_attn.weight)
+            torch.nn.init.xavier_normal_(module.c_proj.weight)
+            if module.c_attn.bias is not None:
+                torch.nn.init.zeros_(module.c_attn.bias)
+            if module.c_proj.bias is not None:
+                torch.nn.init.zeros_(module.c_proj.bias)
+    @staticmethod
+    def top_k_top_p_filtering(
+        logits: torch.Tensor,
+        top_k: int = 50,
+        top_p: float = 0.95,
+        filter_value: float = float("-inf"),
+    ):
+        """
+        Filters a distribution of logits using top-k and/or nucleus (top-p) filtering.
+        """
+        # top_k
+        if top_k > 0:
+            top_k = min(top_k, logits.size(-1))
+            # Remove all logits that are not in the top k
+            indices_to_remove = logits < torch.topk(logits, top_k, dim=-1).values[:, -1, None]
+            logits[indices_to_remove] = filter_value
+
+        # top_p (nucleus)
+        if 0 < top_p < 1.0:
+            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+            # Remove tokens with cumulative probability above the threshold
+            sorted_indices_to_remove = cumulative_probs > top_p
+            # Shift the indices to the right to keep also the first token above the threshold
+            sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
+            sorted_indices_to_remove[:, 0] = False
+
+            indices_to_remove = sorted_indices_to_remove.scatter(
+                dim=1, index=sorted_indices, src=sorted_indices_to_remove
+            )
+            logits[indices_to_remove] = filter_value
+
+        return logits
+
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        max_new_tokens: int = 50,
+        temperature: float = 0.5,
+        top_k: int = 50,
+        top_p: float = 0.95,
+        eos_token_id: int = None,
+        pad_token_id: int = 0,
+        **kwargs
+    ):
+        """
+        Naive token-by-token generation loop that uses top-k/top-p filtering and optional temperature.
+
+        Args:
+            input_ids (torch.LongTensor): shape (batch_size, sequence_length).
+            max_new_tokens (int): max number of tokens to generate (beyond input_ids length).
+            temperature (float): sampling temperature (>=0).
+            top_k (int): Top-K sampling cutoff.
+            top_p (float): Nucleus sampling cutoff.
+            eos_token_id (int): If set, stop generation when this token is produced.
+            pad_token_id (int): If set, can be used to pad sequences. (Not fully used here.)
+            kwargs: Unused arguments (like num_beams) for compatibility.
+
+        Returns:
+            torch.LongTensor: shape (batch_size, sequence_length + generated_tokens).
+        """
+        device = input_ids.device
+        #print("1=====================")
+        #print(tokenizer.decode(input_ids[0], skip_special_tokens=True))
+        #print("1=====================")
+
+        # We'll accumulate new tokens into generated_ids
+        generated_ids = input_ids.clone()
+
+        for _ in range(max_new_tokens):
+            # Forward pass to get logits for the last token
+            outputs = self.forward(generated_ids)
+            logits = outputs.logits[:, -1, :]  # shape: (batch_size, vocab_size)
+
+            # Scale logits by temperature
+            if temperature != 1.0:
+                logits = logits / temperature
+
+            # Filter logits using top-k and/or top-p
+            logits = self.top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+
+            # Convert to probabilities
+            probabilities = F.softmax(logits, dim=-1)
+
+            # Sample from the distribution
+            next_token = torch.multinomial(probabilities, num_samples=1)  # (batch_size, 1)
+
+            # Append next token
+            generated_ids = torch.cat([generated_ids, next_token], dim=1)
+
+            # If eos_token_id is set and any sample produced it, we optionally could break early
+            if eos_token_id is not None:
+                # Check if all sequences in the batch ended
+                # or if you want to do a more fine-grained approach
+                if (next_token == eos_token_id).all():
+                    break
+        #print("2=====================")
+        #print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))
+        #print("2=====================")
+        return generated_ids

modules.py

@@ -0,0 +1,6 @@
+from .attn import Attention
+from .attn import AttentionSDPA
+from .mlp import MLP
+from .stu import STU
+
+

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

stu.py

@@ -0,0 +1,87 @@
+import torch
+import torch.nn as nn
+
+from .convolve import convolve, flash_convolve
+
+try:
+    from flashfftconv import FlashFFTConv
+
+    flash_fft_available = True
+except ImportError as e:
+    print(
+        f"Unable to import FlashFFTConv: {e}. Falling back to PyTorch implementation."
+    )
+    flash_fft_available = False
+
+
+class STU(nn.Module):
+    def __init__(self, config, phi, n) -> None:
+        super(STU, self).__init__()
+        self.config = config
+        if isinstance(config.torch_dtype, str):
+            torch_dtype = getattr(torch, config.torch_dtype)
+        else:
+            torch_dtype = config.torch_dtype
+        self.phi = phi.to(device=config.device, dtype=torch_dtype)
+        self.n = n
+        self.K = config.num_eigh
+        self.d_in = config.n_embd
+        self.d_out = config.n_embd
+        self.use_hankel_L = config.use_hankel_L
+        self.use_approx = config.use_approx
+        self.flash_fft = (
+            FlashFFTConv(self.n, dtype=torch.bfloat16)
+            if config.use_flash_fft and flash_fft_available
+            else None
+        )
+        if self.use_approx:
+            self.M_inputs = nn.Parameter(
+                torch.empty(self.d_in, self.d_out, dtype=torch_dtype)
+            )
+            self.M_filters = nn.Parameter(
+                torch.empty(self.K, self.d_in, dtype=torch_dtype)
+            )
+        else:
+            self.M_phi_plus = nn.Parameter(
+                torch.empty(self.K, self.d_in, self.d_out, dtype=torch_dtype)
+            )
+            if not self.use_hankel_L:
+                self.M_phi_minus = nn.Parameter(
+                    torch.empty(self.K, self.d_in, self.d_out, dtype=torch_dtype)
+                )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = self.M_inputs.dtype
+        x = x.to(dtype=dtype)
+        if self.use_approx:
+            # Contract inputs and filters over the K and d_in dimensions, then convolve
+            x_proj = x @ self.M_inputs
+            phi_proj = self.phi @ self.M_filters
+            x_proj = x_proj.to(dtype=dtype)
+            phi_proj = phi_proj.to(dtype=dtype)
+            if self.flash_fft:
+                spectral_plus, spectral_minus = flash_convolve(
+                    x_proj, phi_proj, self.flash_fft, self.use_approx
+                )
+            else:
+                spectral_plus, spectral_minus = convolve(
+                    x_proj, phi_proj, self.n, self.use_approx
+                )
+        else:
+            # Convolve inputs and filters,
+            if self.flash_fft:
+                U_plus, U_minus = flash_convolve(
+                    x, self.phi, self.flash_fft, self.use_approx
+                )
+            else:
+                U_plus, U_minus = convolve(x, self.phi, self.n, self.use_approx)
+            # Then, contract over the K and d_in dimensions
+            spectral_plus = torch.tensordot(
+                U_plus, self.M_phi_plus, dims=([2, 3], [0, 1])
+            )
+            if not self.use_hankel_L:
+                spectral_minus = torch.tensordot(
+                    U_minus, self.M_phi_minus, dims=([2, 3], [0, 1])
+                )
+
+        return spectral_plus if self.use_hankel_L else spectral_plus + spectral_minus

tokenizer_config.json

@@ -0,0 +1,27 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "199999": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "200018": {
+      "content": "<|endofprompt|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 128000,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

utils.py

@@ -0,0 +1,105 @@
+import math
+import numpy as np
+import torch
+
+import logging
+import os
+import sys
+from colorama import Fore, Style, init
+from dotenv import load_dotenv
+
+load_dotenv()
+init(autoreset=True)
+
+def nearest_power_of_two(x: int, round_up: bool = False) -> int:
+    return (
+        1 << math.floor(math.log2(x)) if not round_up else 1 << math.ceil(math.log2(x))
+    )
+
+def get_hankel(seq_len: int, use_hankel_L: bool = False, device: torch.device = None, dtype: torch.dtype = torch.float32) -> torch.Tensor:
+    entries = torch.arange(1, seq_len + 1, dtype=dtype, device=device)
+    i_plus_j = entries[:, None] + entries[None, :]
+
+    if use_hankel_L:
+        sgn = (-1.0) ** (i_plus_j - 2.0) + 1.0
+        denom = (i_plus_j + 3.0) * (i_plus_j - 1.0) * (i_plus_j + 1.0)
+        Z = sgn * (8.0 / denom)
+    elif not use_hankel_L:
+        Z = 2.0 / (i_plus_j**3 - i_plus_j)
+    else:
+        raise ValueError("use_hankel_L must be a boolean")
+
+    return Z
+
+
+class ColorFormatter(logging.Formatter):
+    """
+    A custom log formatter that applies color based on the log level using the Colorama library.
+    
+    Attributes:
+        LOG_COLORS (dict): A dictionary mapping log levels to their corresponding color codes.
+    """
+
+    # Colors for each log level
+    LOG_COLORS = {
+        logging.DEBUG: Fore.LIGHTMAGENTA_EX + Style.BRIGHT,
+        logging.INFO: Fore.CYAN,
+        logging.WARNING: Fore.YELLOW + Style.BRIGHT,
+        logging.ERROR: Fore.RED + Style.BRIGHT,
+        logging.CRITICAL: Fore.RED + Style.BRIGHT + Style.NORMAL,
+    }
+
+    # Colors for other parts of the log message
+    TIME_COLOR = Fore.GREEN
+    FILE_COLOR = Fore.BLUE
+    LEVEL_COLOR = Style.BRIGHT
+
+    def __init__(self, fmt=None):
+        super().__init__(fmt or "%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s", "%Y-%m-%d %H:%M:%S")
+
+    def format(self, record):
+        """
+        Formats a log record with the appropriate color based on the log level.
+        
+        Args:
+            record (logging.LogRecord): The log record to format.
+
+        Returns:
+            str: The formatted log message with colors applied.
+        """
+        # Apply color based on the log level
+        level_color = self.LOG_COLORS.get(record.levelno, Fore.WHITE)
+        time_str = f"{self.TIME_COLOR}{self.formatTime(record)}{Style.RESET_ALL}"
+        levelname_str = f"{level_color}{record.levelname}{Style.RESET_ALL}"
+        file_info_str = f"{self.FILE_COLOR}{record.filename}:{record.lineno}{Style.RESET_ALL}"
+
+        # Format the log message with color
+        log_msg = f"{time_str} - {levelname_str} - {file_info_str} - {record.msg}"
+        return log_msg
+
+def setup_logger():
+    """
+    Sets up a logger with a custom color formatter that logs to standard output (stdout).
+    
+    The logger is configured with the ColorFormatter to format log messages with color based on the log level.
+    The log level is set to INFO by default, but this can be changed to show more or less detailed messages.
+
+    Returns:
+        logging.Logger: A logger instance that logs formatted messages to stdout.
+    """
+    handler = logging.StreamHandler(sys.stdout)
+
+    # Set custom formatter
+    formatter = ColorFormatter()
+    handler.setFormatter(formatter)
+    logger = logging.getLogger(__name__)
+    
+    # Set to DEBUG to capture all logging levels
+    DEBUG = os.environ.get("DEBUG", "False").lower() in ("true", "1", "t")
+    logger.setLevel(logging.DEBUG) if DEBUG else logger.setLevel(logging.INFO)
+    logger.addHandler(handler)
+    logger.propagate = False  # Prevents multiple logging if re-initialized
+
+    return logger
+
+logger = setup_logger()  # Initialize once to prevent multiple loggers