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)