| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from config import * |
| import sys |
| import os |
|
|
| |
| 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 |
| ) |
|
|
| |
|
|
| 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) |
|
|
| return logits |
|
|
