"""Export and load ModernBERT encoder weights packed for 4-bit inference.""" import argparse import json import shutil from functools import cache from pathlib import Path from tempfile import TemporaryDirectory import torch from safetensors import safe_open from safetensors.torch import load_file, save_file from torch import nn from torch.nn import functional as F from transformers import AutoConfig, AutoModelForSequenceClassification from transformers.modeling_utils import no_init_weights from .io import read_json WEIGHT_FORMAT = "int4-group64" GROUP_SIZE = 64 @cache def unpack_shader(): return torch.mps.compile_shader(""" #include using namespace metal; kernel void unpack_int4(device float *out, device const uchar *packed, device const float2 *params, constant uint &inputs, constant uint &outputs, constant uint &group_size, uint i [[thread_position_in_grid]]) { uint row = i / (inputs / 2); uint column = (i % (inputs / 2)) * 2; uchar byte = packed[i]; float2 group = params[column / group_size * outputs + row]; out[i * 2] = (int(byte >> 4) - 8) * group.x + group.y; out[i * 2 + 1] = (int(byte & 15) - 8) * group.x + group.y; } """) class Int4Linear(nn.Module): def __init__(self, linear, unpack): super().__init__() self.unpack = unpack self.out_features = linear.out_features self.in_features = linear.in_features self.register_buffer("packed_weight", torch.empty( (self.out_features, self.in_features // 2), dtype=torch.uint8)) self.register_buffer("scales_and_zeros", torch.empty( (self.in_features // GROUP_SIZE, self.out_features, 2), dtype=torch.float32)) def forward(self, inputs): weight = self.unpack(self.packed_weight, self.scales_and_zeros, self.in_features, GROUP_SIZE) return F.linear(inputs, weight) def unpack_mps(packed, params, inputs, group_size): outputs = packed.shape[0] weight = params.new_empty((outputs, inputs)) unpack_shader().unpack_int4(weight, packed, params, inputs, outputs, group_size, threads=packed.numel()) return weight def unpack_cuda(packed, params, inputs, group_size): outputs = packed.shape[0] values = torch.stack((packed >> 4, packed & 15), dim=-1) weights = values.reshape(outputs, -1, group_size).float() scale, midpoint = params.transpose(0, 1).unsqueeze(-2).unbind(-1) return weights.sub_(8).mul_(scale).add_(midpoint).reshape(outputs, inputs) def pack_weight(weight): outputs, inputs = weight.shape if outputs % 8 or inputs % GROUP_SIZE: raise ValueError(f"Unsupported linear weight shape for 4-bit MPS: {weight.shape}") groups = weight.reshape(outputs, -1, GROUP_SIZE) minimum, maximum = groups.aminmax(dim=-1, keepdim=True) scale = ((maximum - minimum) / 15).clamp_min(1e-8) values = ((groups - minimum) / scale).round_().clamp_(0, 15).to(torch.uint8) values = values.reshape_as(weight) packed = (values[:, ::2] << 4) | values[:, 1::2] # The unpacker interprets zero as the real-valued midpoint at code 8. scales_and_zeros = torch.cat((scale, minimum + 8 * scale), dim=-1) return packed.contiguous(), scales_and_zeros.transpose(0, 1).contiguous() def export_int4(source, output): config = read_json(source / "config.json") if (config.get("model_type") != "modernbert" or config.get("decision_model_version") != 1 or config.get("decision_weight_format") or config.get("dtype") != "float32" or config.get("attention_bias") or config.get("mlp_bias")): raise ValueError("Expected an FP32 ModernBERT decision export with bias-free encoder projections.") if output.exists(): raise FileExistsError(output) print(f"Packing encoder weights from {source}", flush=True) tensors = {} with safe_open(source / "model.safetensors", framework="pt") as original: for name in original.keys(): weight = original.get_tensor(name) if name.startswith("model.layers.") and weight.ndim == 2: prefix = name.removesuffix(".weight") tensors[f"{prefix}.packed_weight"], tensors[f"{prefix}.scales_and_zeros"] = pack_weight(weight) else: tensors[name] = weight config["decision_weight_format"] = WEIGHT_FORMAT output.parent.mkdir(parents=True, exist_ok=True) with TemporaryDirectory(prefix=".quantize-", dir=output.parent) as temporary: stage = Path(temporary) / output.name stage.mkdir() save_file(tensors, stage / "model.safetensors", metadata={"format": "pt"}) (stage / "config.json").write_text(json.dumps(config, indent=2) + "\n") for name in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"): shutil.copyfile(source / name, stage / name) calibration = source / "calibration.json" if calibration.is_file(): shutil.copyfile(calibration, stage / calibration.name) stage.rename(output) print(f"Saved 4-bit export to {output}", flush=True) def load_int4(path, device): backend = torch.device(device).type if backend == "mps" and torch.backends.mps.is_available(): unpack_shader() unpack = unpack_mps elif backend == "cuda" and torch.cuda.is_available(): unpack = unpack_cuda else: raise ValueError("This 4-bit export requires an MPS or CUDA GPU.") config = AutoConfig.from_pretrained(path, local_files_only=True, trust_remote_code=False, reference_compile=False) if (config.model_type != "modernbert" or config.decision_model_version != 1 or config.decision_weight_format != WEIGHT_FORMAT): raise ValueError("Unsupported 4-bit decision checkpoint format.") with no_init_weights(): model = AutoModelForSequenceClassification.from_config(config, attn_implementation="sdpa") for name, layer in tuple(model.model.layers.named_modules()): if isinstance(layer, nn.Linear): model.model.layers.set_submodule(name, Int4Linear(layer, unpack)) model.load_state_dict(load_file(path / "model.safetensors"), strict=True, assign=True) return model.to(device).eval() if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("source", type=Path) parser.add_argument("output", type=Path) args = parser.parse_args() torch.set_num_threads(min(4, torch.get_num_threads())) export_int4(args.source, args.output)