| """Weight-only int8 / int4 checkpoints for Auto's ModernBERT classifiers, in plain PyTorch. |
| |
| Runs wherever PyTorch runs (CPU, Apple MPS, NVIDIA CUDA, AMD ROCm): weights stay quantized in memory and each layer |
| dequantizes its own weight just before the matrix multiply. Activations keep the device's float dtype, so no |
| quantized kernels are needed. Requires torch and transformers (5.x); nothing is compiled. |
| |
| from auto_quant import load, score |
| model, tokenizer = load("ProCreations/auto-200m-2-int4") # device and dtype chosen automatically |
| p_deny = score(model, tokenizer, text) # one request, any length up to 65,536 tokens |
| |
| By default attention is Auto's exact memory-linear attention (the same as the Auto runtime): no N x N score matrix or |
| mask is ever built, so a 65,536-token request fits in a few GB on CPU or Apple MPS. It scores one unpadded request |
| at a time through score(). Pass attention="sdpa" (or any Transformers attention implementation) to call |
| model(**inputs) on padded batches instead. |
| |
| Format (quant_config.json, "format": "auto-quant-v1"): symmetric integers, one float16 scale per output channel |
| (int8) or per group of `group_size` consecutive input weights (int4). Dequantized weight = float32(q) * |
| float32(scale), then cast to the compute dtype. int4 values are stored as q + 8 in nibbles, two per byte, low first. |
| """ |
|
|
| import json |
| import os |
| import types |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| FORMAT = "auto-quant-v1" |
| RANGES = {8: (-127, 127), 4: (-8, 7)} |
|
|
|
|
| |
|
|
|
|
| def pack_int4(q): |
| """int8 tensor in [-8, 7] with an even last dimension -> uint8 with two values per byte (low nibble first).""" |
| u = (q.to(torch.int16) + 8).to(torch.uint8) |
| return u[..., 0::2] | (u[..., 1::2] << 4) |
|
|
|
|
| def unpack_int4(packed): |
| low = (packed & 15).to(torch.int8) - 8 |
| high = (packed >> 4).to(torch.int8) - 8 |
| return torch.stack([low, high], dim=-1).flatten(-2) |
|
|
|
|
| def dequantize(qweight, scales, bits, group_size, dtype): |
| """Float weight from stored integers and scales (rows = output channels / embedding rows).""" |
| q = unpack_int4(qweight) if bits == 4 else qweight |
| if bits == 8: |
| w = q.float() * scales.float()[..., None] |
| else: |
| cols = q.shape[-1] |
| w = (q.float().reshape(*q.shape[:-1], cols // group_size, group_size) * scales.float()[..., None]).reshape(q.shape) |
| return w.to(dtype) |
|
|
|
|
| def quantize(weight, bits, group_size, scales): |
| """Integers for `weight` given its scales (rounded to float16): per row for int8, per row and group for int4.""" |
| lo, hi = RANGES[bits] |
| w = weight.detach().float() |
| s = scales.to(torch.float16).float() |
| if bits == 8: |
| q = torch.clamp(torch.round(w / s[:, None]), lo, hi) |
| else: |
| rows, cols = w.shape |
| q = torch.clamp(torch.round(w.reshape(rows, cols // group_size, group_size) / s[..., None]), lo, hi).reshape(rows, cols) |
| q = q.to(torch.int8) |
| return (pack_int4(q) if bits == 4 else q), scales.to(torch.float16) |
|
|
|
|
| class QuantLinear(nn.Module): |
| def __init__(self, in_features, out_features, bits, group_size, bias=False, device=None): |
| super().__init__() |
| self.in_features, self.out_features, self.bits, self.group_size = in_features, out_features, bits, group_size |
| cols = in_features // 2 if bits == 4 else in_features |
| groups = in_features // group_size if bits == 4 else None |
| self.register_buffer("qweight", torch.empty(out_features, cols, dtype=torch.uint8 if bits == 4 else torch.int8, device=device)) |
| self.register_buffer("scales", torch.empty((out_features, groups) if groups else (out_features,), dtype=torch.float16, device=device)) |
| self.register_buffer("bias", torch.empty(out_features, dtype=torch.float16, device=device) if bias else None) |
|
|
| def forward(self, x): |
| w = dequantize(self.qweight, self.scales, self.bits, self.group_size, x.dtype) |
| return F.linear(x, w, None if self.bias is None else self.bias.to(x.dtype)) |
|
|
|
|
| class QuantEmbedding(nn.Module): |
| def __init__(self, num_embeddings, embedding_dim, bits, group_size, dtype=torch.float32, device=None): |
| super().__init__() |
| self.num_embeddings, self.embedding_dim, self.bits, self.group_size = num_embeddings, embedding_dim, bits, group_size |
| self.dtype = dtype |
| cols = embedding_dim // 2 if bits == 4 else embedding_dim |
| groups = embedding_dim // group_size if bits == 4 else None |
| self.register_buffer("qweight", torch.empty(num_embeddings, cols, dtype=torch.uint8 if bits == 4 else torch.int8, device=device)) |
| self.register_buffer("scales", torch.empty((num_embeddings, groups) if groups else (num_embeddings,), dtype=torch.float16, device=device)) |
|
|
| def forward(self, input_ids): |
| |
| return dequantize(self.qweight[input_ids], self.scales[input_ids], self.bits, self.group_size, self.dtype) |
|
|
|
|
| def targets(model): |
| """Quantized modules: token embeddings, every encoder Linear and the head's dense layer. LayerNorm weights and |
| the final 2-way classifier stay in float16.""" |
| names = ["model.embeddings.tok_embeddings"] |
| for name, module in model.named_modules(): |
| if isinstance(module, nn.Linear) and (name.startswith("model.layers.") or name == "head.dense"): |
| names.append(name) |
| return names |
|
|
|
|
| def swap(model, config, dtype=None, device=None): |
| """Replace the target modules with (empty) quantized modules; then load with load_state_dict(strict=True).""" |
| bits, group = config["bits"], config.get("group_size") |
| dtype = dtype or next(model.parameters()).dtype |
| for name in config["modules"]: |
| parent_name, _, child = name.rpartition(".") |
| parent = model.get_submodule(parent_name) |
| old = getattr(parent, child) |
| dev = device if device is not None else old.weight.device |
| if isinstance(old, nn.Embedding): |
| new = QuantEmbedding(old.num_embeddings, old.embedding_dim, bits, group, dtype=dtype, device=dev) |
| else: |
| new = QuantLinear(old.in_features, old.out_features, bits, group, bias=old.bias is not None, device=dev) |
| setattr(parent, child, new) |
| return model |
|
|
|
|
| |
|
|
| LOCAL_BLOCK, MLP_CHUNK = 1024, 8192 |
| SCORE_BUDGET = {"cpu": 256 * 2**20} |
| GPU_SCORE_BUDGET = 1024 * 2**20 |
| try: |
| from torch.nn.attention import SDPBackend, sdpa_kernel |
|
|
| _FUSED = {"cuda": [b for b in ("FLASH_ATTENTION", "EFFICIENT_ATTENTION", "CUDNN_ATTENTION") if hasattr(SDPBackend, b)], |
| "cpu": ["FLASH_ATTENTION"]} |
| except ImportError: |
| sdpa_kernel, _FUSED = None, {} |
| _UNAVAILABLE = set() |
|
|
|
|
| def _sdpa(query, key, value, mask, scaling): |
| return F.scaled_dot_product_attention(query, key, value, attn_mask=mask, dropout_p=0.0, is_causal=False, scale=scaling) |
|
|
|
|
| def _rows(device, heads, keys, element_size): |
| return max(1, SCORE_BUDGET.get(device.type, GPU_SCORE_BUDGET) // max(1, heads * keys * max(4, element_size))) |
|
|
|
|
| def _global_attention(query, key, value, scaling): |
| kind = query.device.type |
| if sdpa_kernel is not None and kind in _FUSED and (kind, query.dtype) not in _UNAVAILABLE: |
| try: |
| with sdpa_kernel([getattr(SDPBackend, b) for b in _FUSED[kind]]): |
| return _sdpa(query, key, value, None, scaling) |
| except RuntimeError: |
| _UNAVAILABLE.add((kind, query.dtype)) |
| n = query.shape[-2] |
| block = _rows(query.device, query.shape[1], n, query.element_size()) |
| if block >= n: |
| return _sdpa(query, key, value, None, scaling) |
| out = torch.empty_like(query) |
| for s in range(0, n, block): |
| out[:, :, s : s + block] = _sdpa(query[:, :, s : s + block], key, value, None, scaling) |
| return out |
|
|
|
|
| def _local_attention(query, key, value, radius, scaling): |
| """Each query i attends to keys j with |i - j| <= radius, one block of queries at a time.""" |
| n = query.shape[-2] |
| block = max(1, min(LOCAL_BLOCK, _rows(query.device, query.shape[1], LOCAL_BLOCK + 2 * radius, query.element_size()))) |
| out, cached = torch.empty_like(query), None |
| for start in range(0, n, block): |
| end = min(n, start + block) |
| low, high = max(0, start - radius), min(n, end + radius) |
| interior = low == start - radius and high == end + radius and end - start == block |
| if interior and cached is not None: |
| mask = cached |
| else: |
| mask = (torch.arange(start, end, device=query.device)[:, None] - torch.arange(low, high, device=query.device)[None, :]).abs() <= radius |
| if interior: |
| cached = mask |
| out[:, :, start:end] = _sdpa(query[:, :, start:end], key[:, :, low:high], value[:, :, low:high], mask, scaling) |
| return out |
|
|
|
|
| def linear_attention(module, query, key, value, attention_mask=None, scaling=None, dropout=0.0, sliding_window=None, **kwargs): |
| """Transformers attention interface (registered as "auto_linear"); one unpadded example, no mask.""" |
| if query.shape[0] != 1 or attention_mask is not None: |
| raise ValueError('attention="linear" scores one unpadded request at a time: use score(), or load with attention="sdpa"') |
| out = _global_attention(query, key, value, scaling) if sliding_window is None else _local_attention(query, key, value, sliding_window - 1, scaling) |
| return out.transpose(1, 2).contiguous(), None |
|
|
|
|
| def _chunked_mlp(self, hidden): |
| if hidden.shape[-2] <= MLP_CHUNK or torch.is_grad_enabled(): |
| return self._full_forward(hidden) |
| |
| return torch.cat([self._full_forward(hidden[..., s : s + MLP_CHUNK, :]) for s in range(0, hidden.shape[-2], MLP_CHUNK)], dim=-2) |
|
|
|
|
| |
|
|
|
|
| def _resolve(path_or_repo, revision): |
| if os.path.isdir(path_or_repo): |
| return path_or_repo |
| from huggingface_hub import snapshot_download |
|
|
| return snapshot_download(path_or_repo, revision=revision) |
|
|
|
|
| def load(path_or_repo, device="auto", dtype=None, revision=None, attention="linear", **model_kwargs): |
| """Load a quantized Auto checkpoint -> (model, tokenizer), in eval mode on `device` ("auto": CUDA/ROCm, then |
| Apple MPS, then CPU). dtype defaults to bfloat16 on GPUs that support it and float32 elsewhere. No full-precision |
| copy of the weights is ever materialized.""" |
| from safetensors.torch import load_file |
| from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer |
| from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS |
| from transformers.models.modernbert.modeling_modernbert import ModernBertRotaryEmbedding |
|
|
| path = _resolve(path_or_repo, revision) |
| with open(os.path.join(path, "quant_config.json"), encoding="utf-8") as f: |
| qconfig = json.load(f) |
| if qconfig.get("format") != FORMAT: |
| raise ValueError(f"unsupported quantization format {qconfig.get('format')!r}") |
| if device == "auto": |
| device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu") |
| device = torch.device(device) |
| if dtype is None: |
| dtype = torch.bfloat16 if device.type == "cuda" and torch.cuda.is_bf16_supported() else torch.float32 |
| implementation = attention |
| if attention == "linear": |
| ALL_ATTENTION_FUNCTIONS.register("auto_linear", linear_attention) |
| implementation = "auto_linear" |
| config = AutoConfig.from_pretrained(path) |
| with torch.device("meta"): |
| model = AutoModelForSequenceClassification.from_config(config, attn_implementation=implementation, dtype=dtype, **model_kwargs) |
| swap(model, qconfig, dtype=dtype) |
| model.to_empty(device=device) |
| model.load_state_dict(load_file(os.path.join(path, qconfig["weights"]), device=str(device)), strict=True) |
| with torch.device(device): |
| model.model.rotary_emb = ModernBertRotaryEmbedding(config=config) |
| if attention == "linear": |
| for layer in model.model.layers: |
| layer.mlp._full_forward = layer.mlp.forward |
| layer.mlp.forward = types.MethodType(_chunked_mlp, layer.mlp) |
| model.auto_attention = attention |
| return model.eval(), AutoTokenizer.from_pretrained(path) |
|
|
|
|
| @torch.inference_mode() |
| def logits(model, input_ids): |
| """Classifier logits for one unpadded request (input_ids of shape [1, n]) on the model's device.""" |
| input_ids = input_ids.to(model.device) |
| if getattr(model, "auto_attention", None) != "linear": |
| return model(input_ids=input_ids).logits |
| hidden = model.model(input_ids=input_ids, attention_mask={"full_attention": None, "sliding_attention": None})[0] |
| pooled = hidden[:, 0] if model.config.classifier_pooling == "cls" else hidden.mean(dim=1) |
| return model.classifier(model.drop(model.head(pooled))) |
|
|
|
|
| def score(model, tokenizer, text): |
| """P(deny) for one serialized request (see build_input); deny if >= 0.5.""" |
| ids = tokenizer(text, return_tensors="pt", truncation=False)["input_ids"] |
| if ids.shape[-1] > model.config.max_position_embeddings: |
| raise ValueError(f"{ids.shape[-1]} tokens exceed the {model.config.max_position_embeddings}-token context") |
| return logits(model, ids).float().softmax(-1)[0, 1].item() |
|
|
|
|
| def build_input(user_request, history, call): |
| """The serialization every Auto model was trained on: proposed call, user request, agent history.""" |
| parts = ["### PROPOSED TOOL CALL", f"tool: {call['tool']}", f"args: {call['args']}", "", "### USER REQUEST", user_request, "", "### AGENT HISTORY"] |
| if not history: |
| parts.append("(no prior actions)") |
| else: |
| for i, h in enumerate(history): |
| parts.append(f"[{i + 1}] {h['tool']}({h['args']})\n-> {h.get('result', '')}") |
| return "\n".join(parts) |
|
|
