File size: 4,777 Bytes
dbd41fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from config import *
import sys
import os
# Dynamically link the blazing fast Qsbits C++ Kernel
ext_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'Qsbits', 'ssm_extension')
if ext_path not in sys.path:
sys.path.append(ext_path)
import qsbits_ssm_extension
class RMSNorm(nn.Module):
def __init__(self, d_model, eps=1e-5):
super().__init__()
self.weight = nn.Parameter(torch.ones(d_model))
self.eps = eps
def forward(self, x):
variance = x.pow(2).mean(-1, keepdim=True)
return x * torch.rsqrt(variance + self.eps) * self.weight
class TernaryQuantize(torch.autograd.Function):
@staticmethod
def forward(ctx, input): return torch.round(torch.clamp(input, min=-1.0, max=1.0))
@staticmethod
def backward(ctx, grad_output): return grad_output
def ternary_quantize(x): return TernaryQuantize.apply(x)
class BinaryQuantize(torch.autograd.Function):
@staticmethod
def forward(ctx, input): return torch.sign(input + 1e-6)
@staticmethod
def backward(ctx, grad_output): return grad_output
def binary_quantize(x): return BinaryQuantize.apply(x)
class QsbitsTernarySSM(nn.Module):
def __init__(self, d_model, d_state):
super().__init__()
self.d_model = d_model
self.d_state = d_state
self.B_proj = nn.Linear(d_model, d_state, bias=False)
self.C_proj = nn.Linear(d_state, d_model, bias=False)
self.D_proj = nn.Linear(d_model, d_model, bias=False)
self.A = nn.Parameter(torch.randn(d_state))
def forward(self, x, hidden_state=None):
batch_size, seq_len, _ = x.shape
device = x.device
if hidden_state is None:
hidden_state = torch.zeros(batch_size, self.d_state, device=device)
b_weight_1bit = binary_quantize(self.B_proj.weight) / math.sqrt(self.d_model)
c_weight_1bit = binary_quantize(self.C_proj.weight) / math.sqrt(self.d_state)
d_weight_1bit = binary_quantize(self.D_proj.weight) / math.sqrt(self.d_model)
a_ternary = ternary_quantize(self.A)
delta = torch.ones(self.d_state, device=device)
y_t, hidden_state = qsbits_ssm_extension.forward(
x.contiguous(), hidden_state.contiguous(), delta.contiguous(),
a_ternary.contiguous(), b_weight_1bit.t().contiguous(),
c_weight_1bit.t().contiguous(), True
)
outputs = y_t + F.linear(x, d_weight_1bit)
return outputs, hidden_state
class QsbitsA2BBlock(nn.Module):
def __init__(self, d_model, d_state, d_ffn):
super().__init__()
self.norm1 = RMSNorm(d_model)
self.ssm = QsbitsTernarySSM(d_model, d_state)
self.norm2 = RMSNorm(d_model)
self.ffn_up = nn.Linear(d_model, d_ffn, bias=False)
self.ffn_down = nn.Linear(d_ffn, d_model, bias=False)
self.d_model = d_model
self.d_ffn = d_ffn
def forward(self, x, ssm_state):
normalized_x = self.norm1(x)
ssm_out, new_ssm_state = self.ssm(normalized_x, ssm_state)
normalized_ssm_out = self.norm2(ssm_out)
ffn_up_1bit = binary_quantize(self.ffn_up.weight) / math.sqrt(self.d_model)
ffn_down_1bit = binary_quantize(self.ffn_down.weight) / math.sqrt(self.d_ffn)
ffn_out = F.gelu(F.linear(normalized_ssm_out, ffn_up_1bit))
ffn_out = F.linear(ffn_out, ffn_down_1bit)
final_output = x + ssm_out + ffn_out
return final_output, new_ssm_state
class MiniTransformer(nn.Module):
def __init__(self):
super().__init__()
self.token_emb = nn.Embedding(
VOCAB_SIZE,
N_EMBD
)
# Removed self.pos_emb (Mamba/SSM handles time naturally without Positional Embeddings!)
# Dynamically scaled to match your old config variables!
self.blocks = nn.ModuleList(
[QsbitsA2BBlock(N_EMBD, N_EMBD // 2, 4 * N_EMBD) for _ in range(N_LAYER)]
)
self.ln = RMSNorm(N_EMBD)
self.head = nn.Linear(
N_EMBD,
VOCAB_SIZE,
bias=False
)
def forward(self, idx, ssm_states=None):
B, T = idx.shape
if ssm_states is None:
ssm_states = [None] * len(self.blocks)
new_ssm_states = []
x = self.token_emb(idx)
for i, block in enumerate(self.blocks):
x, new_state = block(x, ssm_states[i])
new_ssm_states.append(new_state)
x = self.ln(x)
logits = self.head(x)
# Returned EXACTLY like your old model so your train.py doesn't crash!
return logits
|