Upload 11 files

config.json

@@ -0,0 +1,41 @@
+{
+  "model_type": "slm",
+  "model_name": "slm",
+  "architectures": [
+    "TinyGPTForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_tiny_gpt.TinyGPTConfig",
+    "AutoModel": "modeling_tiny_gpt.TinyGPTModel",
+    "AutoModelForCausalLM": "modeling_tiny_gpt.TinyGPTForCausalLM"
+  },
+  "vocab_size": 32000,
+  "ctx_len": 2048,
+  "max_position_embeddings": 2048,
+  "n_layer": 4,
+  "num_hidden_layers": 4,
+  "n_head": 4,
+  "num_attention_heads": 4,
+  "n_embd": 384,
+  "hidden_size": 384,
+  "dropout": 0.0,
+  "attention_backend": "torch",
+  "available_attention_backends": [
+    "sage",
+    "torch",
+    "flash2",
+    "flash3"
+  ],
+  "trained_attention_backend": "flash2",
+  "torch_fallback": true,
+  "positional_encoding": "rope",
+  "trained_positional_encoding": "rope",
+  "rope_base": 10000.0,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "custom",
+  "pad_token_id": 0,
+  "bos_token_id": 5,
+  "eos_token_id": 6,
+  "sep_token_id": 2,
+  "unk_token_id": 1
+}

configuration_tiny_gpt.py

@@ -0,0 +1,56 @@
+from transformers import PretrainedConfig
+
+
+class TinyGPTConfig(PretrainedConfig):
+    model_type = "slm"
+
+    def __init__(
+        self,
+        vocab_size=32768,
+        ctx_len=512,
+        n_layer=4,
+        n_head=4,
+        n_embd=384,
+        dropout=0.0,
+        attention_backend="torch",
+        torch_fallback=False,
+        rope_base=10000.0,
+        positional_encoding="rope",
+        pad_token_id=None,
+        bos_token_id=None,
+        eos_token_id=None,
+        sep_token_id=None,
+        unk_token_id=None,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            sep_token_id=sep_token_id,
+            unk_token_id=unk_token_id,
+            **kwargs,
+        )
+
+        if attention_backend not in ("sage", "torch", "flash2", "flash3"):
+            raise ValueError("attention_backend must be sage, torch, flash2 or flash3")
+
+        self.vocab_size = int(vocab_size)
+        self.ctx_len = int(ctx_len)
+        self.max_position_embeddings = int(ctx_len)
+
+        self.n_layer = int(n_layer)
+        self.n_head = int(n_head)
+        self.n_embd = int(n_embd)
+
+        self.num_hidden_layers = int(n_layer)
+        self.num_attention_heads = int(n_head)
+        self.hidden_size = int(n_embd)
+
+        self.dropout = float(dropout)
+        self.attention_backend = str(attention_backend)
+        self.available_attention_backends = ["sage", "torch", "flash2", "flash3"]
+        self.torch_fallback = bool(torch_fallback)
+
+        self.rope_base = float(rope_base)
+        self.positional_encoding = str(positional_encoding)

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "pad_token_id": 0,
+  "bos_token_id": 5,
+  "eos_token_id": 6,
+  "sep_token_id": 2
+}

model.safetensors

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

modeling_tiny_gpt.py

@@ -0,0 +1,504 @@
+import importlib
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+from .configuration_tiny_gpt import TinyGPTConfig
+
+
+_FLASH2_KERNEL = None
+_FLASH3_KERNEL = None
+
+
+def _get_flash2_kernel():
+    global _FLASH2_KERNEL
+
+    if _FLASH2_KERNEL is None:
+        kernels = importlib.import_module("kernels")
+        _FLASH2_KERNEL = kernels.get_kernel("kernels-community/flash-attn2", version=1)
+
+    return _FLASH2_KERNEL
+
+
+def _get_flash3_kernel():
+    global _FLASH3_KERNEL
+
+    if _FLASH3_KERNEL is None:
+        kernels = importlib.import_module("kernels")
+        _FLASH3_KERNEL = kernels.get_kernel("kernels-community/flash-attn3", version=1)
+
+    return _FLASH3_KERNEL
+
+
+def _get_sageattn():
+    module = importlib.import_module("sageattention")
+    return module.sageattn
+
+
+def rotate_half(x):
+    x_even = x[..., ::2]
+    x_odd = x[..., 1::2]
+    x_rot = torch.stack((-x_odd, x_even), dim=-1)
+    return x_rot.flatten(start_dim=-2)
+
+
+def apply_rope(x, cos, sin):
+    return (x * cos) + (rotate_half(x) * sin)
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings, base=10000.0):
+        super().__init__()
+
+        if dim % 2 != 0:
+            raise ValueError(f"RoPE dim must be even, got {dim}")
+
+        self.dim = int(dim)
+        self.max_position_embeddings = int(max_position_embeddings)
+        self.base = float(base)
+
+        inv_freq = 1.0 / (
+            self.base
+            ** (
+                torch.arange(
+                    0,
+                    self.dim,
+                    2,
+                    dtype=torch.float32,
+                )
+                / self.dim
+            )
+        )
+
+        self.register_buffer(
+            "inv_freq",
+            inv_freq,
+            persistent=False,
+        )
+
+        self._cos_cached = None
+        self._sin_cached = None
+        self._seq_len_cached = 0
+        self._device_cached = None
+        self._dtype_cached = None
+
+    def _build_cache(self, seq_len, device, dtype):
+        t = torch.arange(
+            seq_len,
+            device=device,
+            dtype=torch.float32,
+        )
+
+        freqs = torch.einsum(
+            "i,j->ij",
+            t,
+            self.inv_freq.to(device=device, dtype=torch.float32),
+        )
+
+        emb = torch.repeat_interleave(freqs, repeats=2, dim=-1)
+
+        cos = emb.cos().to(dtype=dtype).view(1, 1, seq_len, self.dim)
+        sin = emb.sin().to(dtype=dtype).view(1, 1, seq_len, self.dim)
+
+        self._cos_cached = cos
+        self._sin_cached = sin
+        self._seq_len_cached = int(seq_len)
+        self._device_cached = device
+        self._dtype_cached = dtype
+
+    def forward(self, seq_len, device, dtype):
+        if (
+            self._cos_cached is None
+            or self._sin_cached is None
+            or self._seq_len_cached < seq_len
+            or self._device_cached != device
+            or self._dtype_cached != dtype
+        ):
+            self._build_cache(
+                seq_len=seq_len,
+                device=device,
+                dtype=dtype,
+            )
+
+        return (
+            self._cos_cached[:, :, :seq_len, :],
+            self._sin_cached[:, :, :seq_len, :],
+        )
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config: TinyGPTConfig):
+        super().__init__()
+
+        if config.n_embd % config.n_head != 0:
+            raise ValueError("n_embd must be divisible by n_head")
+
+        self.n_head = int(config.n_head)
+        self.head_dim = int(config.n_embd // config.n_head)
+        self.attention_backend = str(getattr(config, "attention_backend", "torch"))
+        self.torch_fallback = bool(getattr(config, "torch_fallback", False))
+        self.dropout_p = float(config.dropout)
+
+        if self.head_dim % 2 != 0:
+            raise ValueError(f"RoPE requires even head_dim, got {self.head_dim}")
+
+        if self.attention_backend not in ("sage", "torch", "flash2", "flash3"):
+            raise ValueError("attention_backend must be sage, torch, flash2 or flash3")
+
+        if self.attention_backend == "sage" and self.head_dim not in (64, 96, 128):
+            raise ValueError(f"SageAttention requires head_dim in [64, 96, 128], got {self.head_dim}")
+
+        if self.attention_backend == "sage" and self.dropout_p != 0.0:
+            raise ValueError("SageAttention requires dropout=0.0")
+
+        if self.attention_backend == "flash3" and self.dropout_p != 0.0:
+            raise ValueError("FlashAttention3 requires dropout=0.0")
+
+        if self.attention_backend in ("flash2", "flash3") and self.head_dim % 8 != 0:
+            raise ValueError(f"FlashAttention requires head_dim multiple of 8, got {self.head_dim}")
+
+        self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
+        self.proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
+        self.dropout = nn.Dropout(config.dropout)
+
+        self.rope = RotaryEmbedding(
+            dim=self.head_dim,
+            max_position_embeddings=config.ctx_len,
+            base=float(getattr(config, "rope_base", 10000.0)),
+        )
+
+        mask = torch.tril(torch.ones(config.ctx_len, config.ctx_len, dtype=torch.bool))
+        self.register_buffer(
+            "mask",
+            mask.view(1, 1, config.ctx_len, config.ctx_len),
+            persistent=False,
+        )
+
+        self.sageattn = None
+        self.flash_kernel = None
+
+        if self.attention_backend == "sage":
+            try:
+                self.sageattn = _get_sageattn()
+            except Exception:
+                if self.torch_fallback:
+                    self.attention_backend = "torch"
+                else:
+                    raise
+
+        if self.attention_backend == "flash2":
+            try:
+                self.flash_kernel = _get_flash2_kernel()
+            except Exception:
+                if self.torch_fallback:
+                    self.attention_backend = "torch"
+                else:
+                    raise
+
+        if self.attention_backend == "flash3":
+            try:
+                self.flash_kernel = _get_flash3_kernel()
+            except Exception:
+                if self.torch_fallback:
+                    self.attention_backend = "torch"
+                else:
+                    raise
+
+    def _torch_attention(self, q, k, v, t):
+        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        scores = scores.masked_fill(self.mask[:, :, :t, :t] == 0, float("-inf"))
+
+        att = F.softmax(scores.float(), dim=-1).to(q.dtype)
+        att = self.dropout(att)
+
+        return att @ v
+
+    def _sage_attention(self, q, k, v):
+        if self.sageattn is None or not q.is_cuda:
+            if self.torch_fallback:
+                return None
+            raise RuntimeError("SageAttention requires CUDA + sageattention")
+
+        return self.sageattn(
+            q.contiguous(),
+            k.contiguous(),
+            v.contiguous(),
+            tensor_layout="HND",
+            is_causal=True,
+        )
+
+    def _flash2_attention(self, q, k, v):
+        if self.flash_kernel is None or not q.is_cuda:
+            if self.torch_fallback:
+                return None
+            raise RuntimeError("FlashAttention2 requires CUDA + kernels")
+
+        q = q.transpose(1, 2).contiguous()
+        k = k.transpose(1, 2).contiguous()
+        v = v.transpose(1, 2).contiguous()
+
+        dropout_p = self.dropout_p if self.training else 0.0
+
+        y = self.flash_kernel.flash_attn_func(
+            q,
+            k,
+            v,
+            dropout_p=dropout_p,
+            causal=True,
+        )
+
+        return y.transpose(1, 2).contiguous()
+
+    def _flash3_attention(self, q, k, v):
+        if self.flash_kernel is None or not q.is_cuda:
+            if self.torch_fallback:
+                return None
+            raise RuntimeError("FlashAttention3 requires CUDA + kernels")
+
+        q = q.transpose(1, 2).contiguous()
+        k = k.transpose(1, 2).contiguous()
+        v = v.transpose(1, 2).contiguous()
+
+        y = self.flash_kernel.flash_attn_func(
+            q,
+            k,
+            v,
+            causal=True,
+        )
+
+        return y.transpose(1, 2).contiguous()
+
+    def forward(self, x):
+        b, t, c = x.shape
+
+        qkv = self.qkv(x)
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        q = q.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()
+        k = k.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()
+        v = v.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()
+
+        cos, sin = self.rope(
+            seq_len=t,
+            device=q.device,
+            dtype=q.dtype,
+        )
+
+        q = apply_rope(q, cos, sin)
+        k = apply_rope(k, cos, sin)
+
+        if self.attention_backend == "sage":
+            y = self._sage_attention(q, k, v)
+            if y is None:
+                y = self._torch_attention(q, k, v, t)
+        elif self.attention_backend == "flash2":
+            y = self._flash2_attention(q, k, v)
+            if y is None:
+                y = self._torch_attention(q, k, v, t)
+        elif self.attention_backend == "flash3":
+            y = self._flash3_attention(q, k, v)
+            if y is None:
+                y = self._torch_attention(q, k, v, t)
+        else:
+            y = self._torch_attention(q, k, v, t)
+
+        y = y.transpose(1, 2).contiguous().view(b, t, c)
+
+        return self.proj(y)
+
+
+class MLP(nn.Module):
+    def __init__(self, config: TinyGPTConfig):
+        super().__init__()
+
+        self.fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        x = self.fc(x)
+        x = F.gelu(x)
+        x = self.proj(x)
+        x = self.dropout(x)
+
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, config: TinyGPTConfig):
+        super().__init__()
+
+        self.ln1 = nn.LayerNorm(config.n_embd)
+        self.attn = CausalSelfAttention(config)
+        self.ln2 = nn.LayerNorm(config.n_embd)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+
+        return x
+
+
+class TinyGPTPreTrainedModel(PreTrainedModel):
+    config_class = TinyGPTConfig
+    base_model_prefix = "tiny_gpt"
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            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)
+
+
+class TinyGPTModel(TinyGPTPreTrainedModel):
+    _tied_weights_keys = ["head.weight"]
+
+    def __init__(self, config: TinyGPTConfig):
+        super().__init__(config)
+
+        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
+        self.drop = nn.Dropout(config.dropout)
+
+        self.blocks = nn.ModuleList(
+            [Block(config) for _ in range(config.n_layer)]
+        )
+
+        self.ln_f = nn.LayerNorm(config.n_embd)
+        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        self.post_init()
+        self.tie_weights()
+
+    def get_input_embeddings(self):
+        return self.tok_emb
+
+    def set_input_embeddings(self, value):
+        self.tok_emb = value
+        self.tie_weights()
+
+    def get_output_embeddings(self):
+        return self.head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.head = new_embeddings
+
+    def tie_weights(self):
+        self._tie_or_clone_weights(self.head, self.tok_emb)
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        return_dict=True,
+        return_logits=False,
+        **kwargs,
+    ):
+        b, t = input_ids.shape
+
+        if t > self.config.ctx_len:
+            raise ValueError(
+                f"Input length {t} > ctx_len {self.config.ctx_len}. "
+                "Truncate before calling the model."
+            )
+
+        x = self.tok_emb(input_ids)
+        x = self.drop(x)
+
+        for block in self.blocks:
+            x = block(x)
+
+        hidden = self.ln_f(x)
+        logits = self.head(hidden) if return_logits else None
+
+        if not return_dict:
+            return (hidden, logits) if return_logits else (hidden,)
+
+        if return_logits:
+            return hidden, logits
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+
+class TinyGPTForCausalLM(TinyGPTPreTrainedModel):
+    _tied_weights_keys = ["tiny_gpt.head.weight"]
+
+    def __init__(self, config: TinyGPTConfig):
+        super().__init__(config)
+
+        self.tiny_gpt = TinyGPTModel(config)
+
+        self.post_init()
+        self.tie_weights()
+
+    def get_input_embeddings(self):
+        return self.tiny_gpt.tok_emb
+
+    def set_input_embeddings(self, value):
+        self.tiny_gpt.tok_emb = value
+        self.tie_weights()
+
+    def get_output_embeddings(self):
+        return self.tiny_gpt.head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.tiny_gpt.head = new_embeddings
+
+    def tie_weights(self):
+        self._tie_or_clone_weights(
+            self.tiny_gpt.head,
+            self.tiny_gpt.tok_emb,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        labels=None,
+        return_dict=True,
+        **kwargs,
+    ):
+        hidden, logits = self.tiny_gpt(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            return_dict=True,
+            return_logits=True,
+        )
+
+        loss = None
+
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)).float(),
+                shift_labels.view(-1),
+            )
+
+        if not return_dict:
+            return ((loss, logits) if loss is not None else (logits,))
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )

requirements.txt

@@ -0,0 +1,5 @@
+torch
+transformers
+safetensors
+tokenizers
+kernels

special_tokens_map.json

@@ -0,0 +1,9 @@
+{
+  "unk_token": "[UNK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "bos_token": "[BOS]",
+  "eos_token": "[EOS]"
+}

tokenizer_config.json

@@ -0,0 +1,14 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "clean_up_tokenization_spaces": false,
+  "padding_side": "right",
+  "truncation_side": "right",
+  "unk_token": "[UNK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "bos_token": "[BOS]",
+  "eos_token": "[EOS]",
+  "model_max_length": 2048
+}