| """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 <metal_stdlib> |
| 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] |
| |
| 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) |
|
|
