Duplicate from Martico2432/srlm-1m

Co-authored-by: Martí <Martico2432@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,36 @@
+---
+datasets:
+- HuggingFaceTB/smollm-corpus
+language:
+- en
+pipeline_tag: text-generation
+tags:
+- SLM
+- ROSA
+license: mit
+---
+
+# SRLM-1M
+
+A **S**mall **R**OSA based **L**anguage **M**odel, of 900k parameters.
+
+## Training inforation
+
+Trained on smollm-corpus fineweb-edu-dedup subset, over 16M tokens.
+
+## Evaluations
+
+- Wikitext v2 byte_perplexity: 4.4553
+- Arc_easy acc_norm: 28.66%
+- Arc_easy acc: 26.39%
+- Blimp acc: 53.1%
+- Hellaswag acc: 26.65%
+- Hellaswag acc_norm: 27.07%
+
+## How to run it
+
+To run this model, you have to:
+1. Download the files
+2. Compile rosa.pyx using `python setup.py build_ext --inplace`
+3. Load the model with the correct configuration
+4. Call the model

config.json

@@ -0,0 +1,11 @@
+{
+    "architectures": ["SRLMForCausalLM"],
+    "model_type": "srlm",
+    "vocab_size": 5000,
+    "d_model": 256,
+    "rank_emb": 48,
+    "rank_rosa": 48,
+    "num_rosa_layers": 6,
+    "max_position_embeddings": 512,
+    "torch_dtype": "float32"
+}

model.py

@@ -0,0 +1,232 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from rosa import rosa, rosa_batch
+
+# def rosa(x):
+#     n = len(x)
+#     y = [-1] * n
+#     s = 2 * n + 1
+#     b = [None] * s
+#     c = [-1] * s
+#     d = [0] * s
+#     e = [-1] * s
+#     b[0] = {}
+#     g = 0
+#     z = 1
+#     for i, t in enumerate(x):
+#         r = z
+#         z += 1
+#         b[r] = {}
+#         d[r] = d[g] + 1
+#         p = g
+#         while p != -1 and t not in b[p]:
+#             b[p][t] = r
+#             p = c[p]
+#         if p == -1:
+#             c[r] = 0
+#         else:
+#             q = b[p][t]
+#             if d[p] + 1 == d[q]:
+#                 c[r] = q
+#             else:
+#                 u = z
+#                 z += 1
+#                 b[u] = b[q].copy()
+#                 d[u] = d[p] + 1
+#                 c[u] = c[q]
+#                 e[u] = e[q]
+#                 while p != -1 and b[p][t] == q:
+#                     b[p][t] = u
+#                     p = c[p]
+#                 c[q] = c[r] = u
+#         v = g = r
+#         a = -1
+#         while v != -1:
+#             if d[v] > 0 and e[v] >= 0:
+#                 a = x[e[v] + 1]
+#                 break
+#             v = c[v]
+#         y[i] = a
+#         v = g
+#         while v != -1 and e[v] < i:
+#             e[v] = i
+#             v = c[v]
+#     return y
+
+
+def rosa_batch_python_orig(z: torch.Tensor, alphabet: int) -> torch.Tensor:
+    assert z.dtype == torch.long and z.ndim == 2
+    zc = z.detach().contiguous().cpu().numpy()
+    out = rosa_batch(zc, alphabet)
+    return torch.from_numpy(out).to(z.device)
+
+
+def rosa_batch_python(z: torch.Tensor) -> torch.Tensor:
+    assert z.dtype == torch.uint8 and z.ndim == 2
+    zc = z.detach().contiguous().cpu().to(torch.int64).numpy()
+    out = rosa_batch(zc, 16)  # 4-bit alphabet
+    out = out.clip(min=0).astype("uint8")
+    return torch.from_numpy(out).to(z.device)
+
+
+class FactorizedTiedEmbedding(nn.Module):
+    def __init__(self, vocab_size, d_model, rank):
+        super().__init__()
+        self.A = nn.Parameter(torch.randn(vocab_size, rank) * 0.02)
+        self.B = nn.Parameter(torch.randn(rank, d_model) * (1.0 / math.sqrt(rank)))
+
+    def embed(self, ids):
+        codes = F.embedding(ids, self.A)  # (B, T, r) gather, not (V, d) matmul
+        return codes @ self.B  # (B, T, d)
+
+    def logits(self, hidden):
+        r = hidden @ self.B.t()  # (B, T, r)
+        return r @ self.A.t()  # (B, T, vocab)
+
+
+class rosa_emb_layer(nn.Module):
+    def __init__(self, V, C, rank):
+        super().__init__()
+        self.emb = FactorizedTiedEmbedding(V, C, rank)
+        self.V = V
+
+    def forward(self, idx):
+        idx = rosa_batch_python_orig(idx, self.V)
+        out = self.emb.embed(idx.clamp_min(0))
+        return out.masked_fill(idx.eq(-1).unsqueeze(-1), 0.0)
+
+
+class rosa_4bit_layer(nn.Module):
+    def __init__(self, C: int, eps: float = 1e-5):
+        super().__init__()
+        assert C % 4 == 0
+        self.emb0 = nn.Parameter(torch.full((1, 1, C), -eps))
+        self.emb1 = nn.Parameter(torch.full((1, 1, C), eps))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+        Cg = C // 4
+
+        b = (x.reshape(B, T, Cg, 4) > 0).to(torch.uint8)
+        tok2d = b[..., 0] | (b[..., 1] << 1) | (b[..., 2] << 2) | (b[..., 3] << 3)
+
+        # Orient to (B, Cg, T)
+        tok2d_oriented = tok2d.permute(0, 2, 1).contiguous()
+
+        tok2d_flat = tok2d_oriented.view(B * Cg, T)
+
+        idx_q_flat = rosa_batch_python(tok2d_flat)
+
+        # Reshape back to the 3D track orientation
+        idx_q = idx_q_flat.view(B, Cg, T)
+        idx_q = idx_q.transpose(1, 2).contiguous()  # (B, T, Cg)
+
+        bit0 = (idx_q & 1).bool()
+        bit1 = ((idx_q >> 1) & 1).bool()
+        bit2 = ((idx_q >> 2) & 1).bool()
+        bit3 = ((idx_q >> 3) & 1).bool()
+        bits = torch.stack([bit0, bit1, bit2, bit3], dim=-1)
+
+        e0 = self.emb0.view(1, 1, Cg, 4).expand(B, T, -1, -1)
+        e1 = self.emb1.view(1, 1, Cg, 4).expand(B, T, -1, -1)
+
+        return torch.where(bits, e1, e0).reshape(B, T, C)
+
+
+class Model(nn.Module):
+    def __init__(self, V, C, rank_emb, rank_rosa, num_rosa_layers):
+        super().__init__()
+        self.embedding = FactorizedTiedEmbedding(V, C, rank_emb)
+        self.emb_rosa = rosa_emb_layer(V, C, rank_rosa)
+        # Now a list of Rosa embeddings
+        self.emb_rosa_list = nn.ModuleList(
+            [rosa_4bit_layer(C) for _ in range(num_rosa_layers)]
+        )
+        self.num_rosa_layers = num_rosa_layers
+        self.linear_list = nn.ModuleList(
+            [nn.Linear(C, C) for _ in range(num_rosa_layers)]
+        )
+        self.norm_list = nn.ModuleList(
+            [nn.RMSNorm(C) for _ in range(num_rosa_layers)]
+        )  # me save params, me repeat
+
+    def forward(self, x):
+        x = self.embedding.embed(x) + self.emb_rosa(x)
+        for i in range(self.num_rosa_layers):
+            x = self.norm_list[i](x)
+            x = x + self.emb_rosa_list[i](x)  # Really want to add RMSNorm here
+            x = x + self.linear_list[i](x)
+        return self.embedding.logits(x)
+
+
+if __name__ == "__main__":
+    import time
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device.upper()}")
+
+    V = 5000  # Vocab size
+    C = 256  # Hidden Dimension
+    rank_emb = 48  # Factorization Rank
+    rank_rosa = 48
+    num_rosa_layers = 6  # Deeper structural depth
+
+    B, T = 16, 512  # Batch size and Context length for benchmark loop
+
+    print(f"Initializing model on {device.upper()}...")
+    model = Model(V, C, rank_emb, rank_rosa, num_rosa_layers).to(device)
+
+    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print("\n" + "=" * 60)
+    print(f"  TOTAL TRAINABLE FOOTPRINT: {total_params:,} parameters")
+    print("=" * 60)
+
+    def get_batch():
+        x = torch.randint(0, V, (B, T), device=device)
+        y = torch.roll(x, shifts=-1, dims=1)
+        y[:, -1] = 0
+        return x, y
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
+    criterion = nn.CrossEntropyLoss()
+
+    print("\nStarting Benchmarking iterations with Loss tracking...")
+    print(
+        f"Config: Batch={B}, SeqLen={T}, Vocab={V}, Channels={C}, Layers={num_rosa_layers}"
+    )
+    print("-" * 60)
+
+    model.train()
+    total_time = 0.0
+    steps = 5
+
+    for step in range(1, steps + 1):
+        x, y = get_batch()
+
+        torch.cuda.synchronize() if device == "cuda" else None
+        start_time = time.perf_counter()
+
+        logits = model(x)
+        loss = criterion(logits.view(-1, V), y.view(-1))
+
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        optimizer.step()
+
+        torch.cuda.synchronize() if device == "cuda" else None
+        end_time = time.perf_counter()
+
+        step_time = end_time - start_time
+        total_time += step_time
+
+        tokens_per_sec = (B * T) / step_time
+        print(
+            f"Step {step}/{steps} | Loss: {loss.item():.4f} | Time: {step_time * 1000:.2f}ms | Throughput: {tokens_per_sec:.2f} tok/sec"
+        )
+
+    print("-" * 60)
+    print(f"Average Benchmark Step Velocity: {(total_time / steps) * 1000:.2f} ms")
+    print("=" * 60)

model.safetensors

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

pytorch_model.bin

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

rosa.pyx

@@ -0,0 +1,154 @@
+# cython: language_level=3
+# cython: boundscheck=False
+# cython: wraparound=False
+# cython: initializedcheck=False
+# cython: nonecheck=False
+# cython: cdivision=True
+
+from libc.stdlib cimport malloc, free
+from libc.string cimport memset, memcpy
+from cython.parallel cimport prange
+cimport cython
+import numpy as np
+cimport numpy as cnp
+
+cnp.import_array()
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+@cython.cdivision(True)
+cdef void _rosa_row(const cnp.int64_t* x, int n, int alphabet, cnp.int64_t* y_out) noexcept nogil:
+    cdef int s = 2 * n + 1
+    cdef int *trans = <int*> malloc(<size_t>s * alphabet * sizeof(int))
+    cdef int *c = <int*> malloc(<size_t>s * sizeof(int))
+    cdef int *d = <int*> malloc(<size_t>s * sizeof(int))
+    cdef cnp.int64_t *e = <cnp.int64_t*> malloc(<size_t>s * sizeof(cnp.int64_t))
+
+    memset(trans, 0xFF, <size_t>s * alphabet * sizeof(int))  # all -1
+
+    cdef int g = 0, z = 1, i, t, r, p, q, u, v
+    cdef cnp.int64_t a
+
+    d[0] = 0
+    c[0] = -1
+    e[0] = -1
+
+    for i in range(n):
+        t = <int> x[i]
+        r = z
+        z += 1
+        d[r] = d[g] + 1
+        e[r] = -1
+        p = g
+        while p != -1 and trans[p * alphabet + t] == -1:
+            trans[p * alphabet + t] = r
+            p = c[p]
+        if p == -1:
+            c[r] = 0
+        else:
+            q = trans[p * alphabet + t]
+            if d[p] + 1 == d[q]:
+                c[r] = q
+            else:
+                u = z
+                z += 1
+                memcpy(trans + <size_t>u * alphabet, trans + <size_t>q * alphabet,
+                       <size_t>alphabet * sizeof(int))
+                d[u] = d[p] + 1
+                c[u] = c[q]
+                e[u] = e[q]
+                while p != -1 and trans[p * alphabet + t] == q:
+                    trans[p * alphabet + t] = u
+                    p = c[p]
+                c[q] = u
+                c[r] = u
+        v = g = r
+        a = -1
+        while v != -1:
+            if d[v] > 0 and e[v] >= 0:
+                a = x[e[v] + 1]
+                break
+            v = c[v]
+        y_out[i] = a
+        v = g
+        while v != -1 and e[v] < i:
+            e[v] = i
+            v = c[v]
+
+    free(trans)
+    free(c)
+    free(d)
+    free(e)
+
+
+def rosa_batch(cnp.int64_t[:, :] x not None, int alphabet):
+    """
+    x: (num_rows, n) int64, values in [0, alphabet)
+    returns: (num_rows, n) int64 numpy array, -1 where no next-distinct-symbol exists
+    """
+    cdef Py_ssize_t num_rows = x.shape[0]
+    cdef Py_ssize_t n = x.shape[1]
+    y_np = np.empty((num_rows, n), dtype=np.int64)
+    cdef cnp.int64_t[:, :] y = y_np
+    cdef Py_ssize_t i
+
+    for i in prange(num_rows, nogil=True, schedule='static'):
+        _rosa_row(&x[i, 0], <int>n, alphabet, &y[i, 0])
+
+    return y_np
+
+
+cpdef list rosa(list x):
+    cdef int n = len(x)
+    cdef list y = [-1] * n
+    cdef int s = 2 * n + 1
+    cdef list b = [None] * s
+    cdef list c = [-1] * s
+    cdef list d = [0] * s
+    cdef list e = [-1] * s
+    b[0] = {}
+    cdef int g = 0
+    cdef int z = 1
+
+    cdef int i, t
+    cdef int r, p, q, u, v, a
+    for i in range(n):
+        t = x[i]
+        r = z
+        z += 1
+        b[r] = {}
+        d[r] = d[g] + 1
+        p = g
+        while p != -1 and t not in b[p]:
+            b[p][t] = r
+            p = c[p]
+        if p == -1:
+            c[r] = 0
+        else:
+            q = b[p][t]
+            if d[p] + 1 == d[q]:
+                c[r] = q
+            else:
+                u = z
+                z += 1
+                b[u] = b[q].copy()
+                d[u] = d[p] + 1
+                c[u] = c[q]
+                e[u] = e[q]
+                while p != -1 and b[p][t] == q:
+                    b[p][t] = u
+                    p = c[p]
+                c[q] = c[r] = u
+        v = g = r
+        a = -1
+        while v != -1:
+            if d[v] > 0 and e[v] >= 0:
+                a = x[e[v] + 1]
+                break
+            v = c[v]
+        y[i] = a
+        v = g
+        while v != -1 and e[v] < i:
+            e[v] = i
+            v = c[v]
+    return y

setup.py

@@ -0,0 +1,12 @@
+import numpy
+from Cython.Build import cythonize
+from setuptools import Extension, setup
+
+ext = Extension(
+    "rosa",
+    ["rosa.pyx"],
+    include_dirs=[numpy.get_include()],
+    extra_compile_args=["/O3", "/openmp"],
+)
+
+setup(ext_modules=cythonize([ext], language_level=3))

tiny_lm_tokenizer/tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "[BOS]",
+  "eos_token": "[EOS]",
+  "model_max_length": 512,
+  "pad_token": "[PAD]",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "[UNK]"
+}