| """Trit sentence embedder: 3 post-norm BERT layers, width 384, BitNet b1.58 weights, int8 activations. |
| |
| Reference inference code for the packed release format (config "format": "trit-packed-v1"): |
| - every linear layer inside the encoder is stored as 2-bit ternary codes plus one bf16 scale per output row; |
| - the word and position embedding tables are int8 with one fp32 scale per row; |
| - LayerNorms, biases, the token-type embedding and the 384x384 output projection are fp32. |
| |
| from trit_embed import TritEmbed |
| model = TritEmbed.from_pretrained(".") |
| vectors = model.encode(["how do index funds work", "Index funds track a market index."]) # (2, 384), unit norm |
| |
| PyTorch has no ternary matmul, so this decodes the layers to fp32 once at load and emulates the int8 activation |
| rounding in float. The NPU encoder (phlegm, open_kernels/designs/trit_embed) runs the 2-bit codes natively. |
| """ |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| GROUP = 256 |
|
|
|
|
| def ste(x: torch.Tensor, q: torch.Tensor) -> torch.Tensor: |
| return x + (q - x).detach() |
|
|
|
|
| def quant_act(x: torch.Tensor) -> torch.Tensor: |
| """int8 per token (per 256-channel block when the width divides), power-of-two scale.""" |
| shape = x.shape |
| if shape[-1] > GROUP and shape[-1] % GROUP == 0: |
| x = x.reshape(*shape[:-1], shape[-1] // GROUP, GROUP) |
| m = x.detach().float().abs().amax(-1, keepdim=True).clamp(min=1e-5) |
| f, k = torch.frexp(m) |
| s = torch.ldexp(torch.ones_like(m), k - 7 + (2 * f > 127 / 64).to(k.dtype)) |
| q = ((x.detach().float() / s).round().clamp(-128, 127) * s).to(x.dtype) |
| return ste(x, q).reshape(shape) |
|
|
|
|
| def pack_codes(codes: torch.Tensor) -> torch.Tensor: |
| """Ternary codes [out, in] -> uint8 [out, in/4]: code + 1 in {0, 1, 2}, column 4b + j at bits 2j..2j+1.""" |
| c = (codes.to(torch.int16) + 1).to(torch.uint8).reshape(codes.shape[0], -1, 4) |
| return c[..., 0] | (c[..., 1] << 2) | (c[..., 2] << 4) | (c[..., 3] << 6) |
|
|
|
|
| def unpack_codes(packed: torch.Tensor) -> torch.Tensor: |
| """uint8 [out, in/4] -> int8 ternary codes [out, in].""" |
| shifts = torch.tensor([0, 2, 4, 6], dtype=torch.uint8) |
| c = (packed.unsqueeze(-1) >> shifts) & 3 |
| if (c == 3).any(): |
| raise ValueError("Reserved 2-bit ternary code 3") |
| return (c.to(torch.int8) - 1).reshape(packed.shape[0], -1) |
|
|
|
|
| class TernLinear(nn.Module): |
| """y = int8(x) @ (codes * scale).T + bias. codes in {-1, 0, 1}; scale is one bf16 value per output row.""" |
|
|
| def __init__(self, n_in: int, n_out: int): |
| super().__init__() |
| self.register_buffer("codes", torch.zeros(n_out, n_in, dtype=torch.int8)) |
| self.register_buffer("scale", torch.zeros(n_out, dtype=torch.bfloat16)) |
| self.register_buffer("weight", torch.zeros(n_out, n_in), persistent=False) |
| self.bias = nn.Parameter(torch.zeros(n_out)) |
|
|
| def decode(self) -> None: |
| self.weight = self.codes.float() * self.scale.float()[:, None] |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return F.linear(quant_act(x), self.weight, self.bias) |
|
|
|
|
| class Int8Embedding(nn.Module): |
| """Lookup table stored as int8 rows with one fp32 scale per row.""" |
|
|
| def __init__(self, rows: int, dim: int): |
| super().__init__() |
| self.register_buffer("qweight", torch.zeros(rows, dim, dtype=torch.int8)) |
| self.register_buffer("scale", torch.zeros(rows)) |
|
|
| def forward(self, ids: torch.Tensor) -> torch.Tensor: |
| return self.qweight[ids].float() * self.scale[ids].unsqueeze(-1) |
|
|
|
|
| class Layer(nn.Module): |
| def __init__(self, d: int, heads: int, ffn: int, eps: float, gelu: str): |
| super().__init__() |
| self.heads, self.gelu = heads, gelu |
| self.q, self.k, self.v, self.o = (TernLinear(d, d) for _ in range(4)) |
| self.ln1 = nn.LayerNorm(d, eps=eps) |
| self.fc1, self.fc2 = TernLinear(d, ffn), TernLinear(ffn, d) |
| self.ln2 = nn.LayerNorm(d, eps=eps) |
|
|
| def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: |
| b, t, d = x.shape |
| q, k, v = (m(x).view(b, t, self.heads, -1).transpose(1, 2) for m in (self.q, self.k, self.v)) |
| attention = F.scaled_dot_product_attention(q, k, v, attn_mask=mask[:, None, None, :]) |
| x = self.ln1(x + self.o(attention.transpose(1, 2).reshape(b, t, d))) |
| return self.ln2(x + self.fc2(F.gelu(self.fc1(x), approximate=self.gelu))) |
|
|
|
|
| class TritEmbed(nn.Module): |
| def __init__(self, vocab, emb, d, heads, ffn, layers, max_pos, eps, out_dim, loops=1, pool="mean", gelu="tanh"): |
| super().__init__() |
| assert emb == d and loops == 1 and pool == "mean", "this release is a single-pass, mean-pooled model" |
| self.max_pos = max_pos |
| self.word = Int8Embedding(vocab, emb) |
| self.pos = Int8Embedding(max_pos, emb) |
| self.tok_type = nn.Embedding(2, emb) |
| self.ln_e = nn.LayerNorm(emb, eps=eps) |
| self.layers = nn.ModuleList(Layer(d, heads, ffn, eps, gelu) for _ in range(layers)) |
| self.out = nn.Linear(d, out_dim, bias=False) |
| self.tokenizer = None |
|
|
| def embed(self, ids: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: |
| """Input states: word + position + token-type embeddings, then LayerNorm.""" |
| return self.ln_e(self.word(ids) + self.pos(positions) + self.tok_type.weight[0]) |
|
|
| def forward(self, ids: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: |
| x = self.embed(ids, torch.arange(ids.size(1), device=ids.device)[None]) |
| for layer in self.layers: |
| x = layer(x, mask) |
| m = mask.unsqueeze(-1).to(x.dtype) |
| return F.normalize(self.out((x * m).sum(1) / m.sum(1).clamp(min=1)), dim=-1) |
|
|
| def ternary_layers(self): |
| """(name, module) for every ternary linear, in encoder order.""" |
| return [(n, m) for n, m in self.named_modules() if isinstance(m, TernLinear)] |
|
|
| @classmethod |
| def from_pretrained(cls, path) -> "TritEmbed": |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer |
| path = Path(path) |
| config = json.loads((path / "config.json").read_text()) |
| if config.get("format") != "trit-packed-v1": |
| raise ValueError(f"Unknown weight format {config.get('format')!r}") |
| model = cls(**config["dims"]) |
| state = load_file(path / "model.safetensors") |
| for name, _ in model.ternary_layers(): |
| state[f"{name}.codes"] = unpack_codes(state[f"{name}.codes"]) |
| model.load_state_dict(state) |
| for _, m in model.ternary_layers(): |
| m.decode() |
| model.tokenizer = AutoTokenizer.from_pretrained(path) |
| return model.eval() |
|
|
| @torch.inference_mode() |
| def encode(self, texts: list[str], batch: int = 128) -> np.ndarray: |
| """Unit-norm float32 vectors; inputs are truncated to 128 tokens. Queries need no prefix. |
| Runs on the model's device: model.to("cuda") for a GPU.""" |
| device = self.out.weight.device |
| tokens = self.tokenizer(texts, truncation=True, max_length=self.max_pos)["input_ids"] |
| order = sorted(range(len(tokens)), key=lambda i: len(tokens[i])) |
| result = np.empty((len(tokens), self.out.out_features), dtype=np.float32) |
| for start in range(0, len(order), batch): |
| indexes = order[start:start + batch] |
| enc = self.tokenizer.pad({"input_ids": [tokens[i] for i in indexes]}, padding=True, return_tensors="pt") |
| vectors = self(enc["input_ids"].to(device), enc["attention_mask"].bool().to(device)) |
| result[indexes] = vectors.float().cpu().numpy() |
| return result |
|
|
