| """WeMM-Embedding INT8: asymmetric int8 with a zero point and one scale per 64-weight block.
|
|
|
| Two builds are served by this one file, selected by config.json:
|
|
|
| quant_upgraded_int4 = false attention and down_proj are int8; the MLP gate/up, ViT and
|
| merger keep their original uint8 group-16 int4
|
| quant_upgraded_int4 = true every quantized matrix in the model is int8, so the int4
|
| unpacking path is unused
|
|
|
| Layout comes from `quant_scheme` and `quant_group_size` in config.json, where a group size
|
| of 0 means one scale per output channel.
|
|
|
| Why a zero point and why 64-weight blocks: SLQ (arXiv 2605.02404) proves symmetric grids
|
| inflate output variance by gamma^2 with gamma = 2M/R over the dynamic range, and its
|
| allocation counts zero-point overhead against the budget. Measured on the shapes in this
|
| checkpoint, one scalar per matrix becomes 64-weight blocks with a zero point and output
|
| error drops from 3.5e-04 to 1.5e-05 in 1-cosine terms.
|
|
|
| model = AutoModel.from_pretrained(repo, trust_remote_code=True)
|
| emb = model.embedding(input_ids=..., attention_mask=...)
|
| """
|
|
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| from transformers import Qwen3_5ForConditionalGeneration
|
|
|
| DEFAULT_GROUP = 64
|
|
|
|
|
| def _dequant_int8(q, s, z, group):
|
| """(q - z) * s, with s and z shared across `group` input features."""
|
| shape = q.shape
|
| in_features = shape[-1]
|
| g = in_features if not group else group
|
| v = q.to(torch.float32).view(*shape[:-1], in_features // g, g)
|
| sc = s.to(torch.float32).unsqueeze(-1)
|
| zz = z.to(torch.float32).unsqueeze(-1)
|
| return ((v - zz) * sc).view(shape)
|
|
|
|
|
| class Int8AsymLinear(nn.Module):
|
| def __init__(self, in_features, out_features, bias=False, group=DEFAULT_GROUP):
|
| super().__init__()
|
| self.in_features = in_features
|
| self.out_features = out_features
|
| self.group = in_features if not group else group
|
| self.register_buffer("weight", torch.zeros((out_features, in_features), dtype=torch.int8))
|
| self.register_buffer("weight_scale", torch.zeros((out_features, in_features // self.group), dtype=torch.bfloat16))
|
| self.register_buffer("weight_zero", torch.zeros((out_features, in_features // self.group), dtype=torch.int8))
|
| if bias:
|
| self.bias = nn.Parameter(torch.zeros(out_features, dtype=torch.bfloat16))
|
| else:
|
| self.register_parameter("bias", None)
|
|
|
| def forward(self, x):
|
| w = _dequant_int8(self.weight, self.weight_scale, self.weight_zero, self.group).to(x.dtype)
|
| return F.linear(x, w, self.bias.to(x.dtype) if self.bias is not None else None)
|
|
|
|
|
| class Int8AsymRowEmbedding(nn.Module):
|
| def __init__(self, num_embeddings, embedding_dim, group=DEFAULT_GROUP):
|
| super().__init__()
|
| self.num_embeddings = num_embeddings
|
| self.embedding_dim = embedding_dim
|
| self.group = embedding_dim if not group else group
|
| self.register_buffer("weight", torch.zeros((num_embeddings, embedding_dim), dtype=torch.int8))
|
| self.register_buffer("weight_scale", torch.zeros((num_embeddings, embedding_dim // self.group), dtype=torch.bfloat16))
|
| self.register_buffer("weight_zero", torch.zeros((num_embeddings, embedding_dim // self.group), dtype=torch.int8))
|
|
|
| def forward(self, input_ids):
|
|
|
|
|
| flat = input_ids.reshape(-1)
|
| w = torch.index_select(self.weight, 0, flat).view(*input_ids.shape, self.embedding_dim)
|
| s = torch.index_select(self.weight_scale, 0, flat).view(*input_ids.shape, -1)
|
| z = torch.index_select(self.weight_zero, 0, flat).view(*input_ids.shape, -1)
|
| return _dequant_int8(w, s, z, self.group).to(torch.bfloat16)
|
|
|
|
|
| class FlatQuantW4A8Linear(nn.Module):
|
| """uint8 packed int4, group 16, two nibbles per byte in even/odd column order.
|
|
|
| Copied from the FP8 checkpoint so those tensors load byte for byte, and kept only for
|
| builds where the int4 modules were not upgraded to int8.
|
| """
|
|
|
| def __init__(self, in_features, out_features, bias=False, group_size=16):
|
| super().__init__()
|
| self.in_features = in_features
|
| self.out_features = out_features
|
| self.group_size = group_size
|
| self.register_buffer("weight", torch.zeros((out_features, in_features // 2), dtype=torch.uint8))
|
| self.register_buffer("weight_scale", torch.zeros((out_features, in_features // group_size), dtype=torch.bfloat16))
|
| if bias:
|
| self.bias = nn.Parameter(torch.zeros(out_features, dtype=torch.bfloat16))
|
| else:
|
| self.register_parameter("bias", None)
|
|
|
| def forward(self, x):
|
| low = (self.weight & 0x0F).to(torch.int8)
|
| high = (self.weight >> 4).to(torch.int8)
|
| low = torch.where(low >= 8, low - 16, low)
|
| high = torch.where(high >= 8, high - 16, high)
|
| unpacked = torch.empty(self.out_features, self.in_features, device=x.device, dtype=x.dtype)
|
| unpacked[:, 0::2] = low.to(x.dtype)
|
| unpacked[:, 1::2] = high.to(x.dtype)
|
| scales = self.weight_scale.to(x.dtype).repeat_interleave(self.group_size, dim=1)
|
| return F.linear(x, unpacked * scales, self.bias.to(x.dtype) if self.bias is not None else None)
|
|
|
|
|
| class WeMMEmbedding(Qwen3_5ForConditionalGeneration):
|
| def __init__(self, config):
|
| super().__init__(config)
|
| group = int(getattr(config, "quant_group_size", DEFAULT_GROUP))
|
| upgraded = bool(getattr(config, "quant_upgraded_int4", False))
|
| for name, mod in list(self.model.named_modules()):
|
| if name.endswith("embed_tokens") and isinstance(mod, nn.Embedding):
|
| parent = self.model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else self.model
|
| setattr(parent, name.rsplit(".", 1)[-1],
|
| Int8AsymRowEmbedding(mod.num_embeddings, mod.embedding_dim, group=group))
|
| elif isinstance(mod, nn.Linear):
|
| parent = self.model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else self.model
|
| is_attn = ("self_attn" in name) or ("linear_attn" in name)
|
| if is_attn or "down_proj" in name or upgraded:
|
| new = Int8AsymLinear(mod.in_features, mod.out_features,
|
| bias=mod.bias is not None, group=group)
|
| else:
|
| new = FlatQuantW4A8Linear(mod.in_features, mod.out_features,
|
| bias=mod.bias is not None, group_size=16)
|
| setattr(parent, name.rsplit(".", 1)[-1], new)
|
|
|
| def embedding(self, input_ids=None, attention_mask=None, **kwargs):
|
| self.model.rope_deltas = None
|
| outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
|
| last_hidden_state = outputs.last_hidden_state
|
| if attention_mask is not None:
|
| eos_positions = attention_mask.sum(dim=1) - 1
|
| else:
|
| eos_positions = torch.full((last_hidden_state.shape[0],), last_hidden_state.shape[1] - 1, device=last_hidden_state.device)
|
| eos_positions = eos_positions.clamp(min=0)
|
| batch_indices = torch.arange(last_hidden_state.size(0), device=last_hidden_state.device)
|
| embeddings = last_hidden_state[batch_indices, eos_positions]
|
| return F.normalize(embeddings, dim=-1)
|
|
|
| def tie_weights(self, *args, **kwargs):
|
| pass
|
|
|
