| import math
|
|
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| from rosa import rosa, rosa_batch
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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)
|
| 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)
|
| return codes @ self.B
|
|
|
| def logits(self, hidden):
|
| r = hidden @ self.B.t()
|
| return r @ self.A.t()
|
|
|
|
|
| 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)
|
|
|
|
|
| tok2d_oriented = tok2d.permute(0, 2, 1).contiguous()
|
|
|
| tok2d_flat = tok2d_oriented.view(B * Cg, T)
|
|
|
| idx_q_flat = rosa_batch_python(tok2d_flat)
|
|
|
|
|
| idx_q = idx_q_flat.view(B, Cg, T)
|
| idx_q = idx_q.transpose(1, 2).contiguous()
|
|
|
| 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)
|
|
|
| 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)]
|
| )
|
|
|
| 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)
|
| 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
|
| C = 256
|
| rank_emb = 48
|
| rank_rosa = 48
|
| num_rosa_layers = 6
|
|
|
| B, T = 16, 512
|
|
|
| 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)
|
|
|
