model.py

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)