diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml index 67b98cbd230..08e8b699d48 100644 --- a/.github/workflows/example_tests.yml +++ b/.github/workflows/example_tests.yml @@ -70,6 +70,7 @@ jobs: - "!**/*.{md,rst,png,ipynb}" vllm: - "{examples,tests/examples}/vllm_serve/**" + - examples/hf_ptq/convert_ggml_to_int8.py # tests/examples/vllm_serve serves its output - "!**/*.{md,rst,png,ipynb}" trtllm: - "{examples,tests/examples}/{gpt-oss,hf_ptq,llm_eval}/**" diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f20a33f1c17..e7abe4ac733 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -11,6 +11,7 @@ Changelog - On Linux x86-64 with Python 3.11 and later, the ONNX extra now selects ONNX Runtime 1.30.x and CUDA 13 packages by default. GPU calibration requires a compatible CUDA 13 runtime and driver; Python 3.10 and other platforms retain their existing dependency selections. - Add imatrix calibration (``algorithm: {method: imatrix}``) for the GGML formats with a weighted search: each weight is encoded once, its search weighted by the square root of its inputs' per-column ``E[x^2]`` (llama.cpp's imatrix statistic), and HF fused-MoE experts get one statistic per expert. A weight that already holds final codes is left alone, so ``[{method: gptq}, {method: imatrix}]`` runs GPTQ on the dense linears and imatrix on the experts GPTQ skips; the built-in IQ1 and IQ2 recipes now do this, and new ``_experts_only`` recipes apply imatrix to the routed MoE experts alone. - Add a checkpoint-specific NVIDIA Nemotron 3.5 Super VL AutoQuantize recipe with a 4.9-effective-bit NVFP4-MSE/FP8/BF16 search and fixed NVFP4 MTP experts. Select it under ``models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/auto_quantize/``. +- Add ``examples/hf_ptq/convert_ggml_to_int8.py`` to expand a checkpoint exported with a GGML recipe (IQ formats, ``q8_0``) into an INT8 weight-only compressed-tensors checkpoint. Stock vLLM serves it with its integer kernels at about half the BF16 weight memory, reproducing the GGML weights up to rounding their scales to the model dtype. *Speculative Decoding* diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 02aae3c07eb..58394dc7a84 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -707,6 +707,20 @@ llm_fp8 = LLM(model="", quantization="modelopt") print(llm_fp8.generate(["What's the age of the earth? "])) ``` +#### vLLM with GGML recipes (IQ, Q8_0) + +vLLM cannot load the GGML checkpoints ModelOpt exports, which are Hugging Face checkpoints holding packed GGML blocks. To serve one on vLLM's integer kernels, expand it into an INT8 weight-only compressed-tensors checkpoint with [convert_ggml_to_int8.py](./convert_ggml_to_int8.py): + +```bash +python hf_ptq.py --pyt_ckpt_path --recipe general/ptq/iq1_s --export_path +python convert_ggml_to_int8.py +vllm serve +``` + +Each weight keeps its GGML integer level, and the step of its 16- or 32-weight sub-block becomes the INT8 group scale. The INT8 checkpoint therefore holds the same weights up to rounding the scales to the model dtype, at about half the BF16 weight memory. Its config names each converted layer explicitly, so every other layer loads unquantized. The GGML checkpoint is left unchanged. + +vLLM serves group-32 layers (IQ1_S, IQ2_XXS, Q8_0) with Marlin and group-16 layers (IQ1_M, IQ2_XS, IQ2_S) with Humming. On Blackwell GPUs, Marlin is much slower at large batch sizes, so set `VLLM_DISABLED_KERNELS=MarlinLinearKernel` to use Humming for group 32 as well. + #### SGLang ```python diff --git a/examples/hf_ptq/convert_ggml_to_int8.py b/examples/hf_ptq/convert_ggml_to_int8.py new file mode 100644 index 00000000000..5ca496aa5fd --- /dev/null +++ b/examples/hf_ptq/convert_ggml_to_int8.py @@ -0,0 +1,255 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Expand a GGML checkpoint exported by hf_ptq.py into an INT8 checkpoint that vLLM serves. + +python hf_ptq.py --pyt_ckpt_path --recipe general/ptq/iq1_s --export_path +python convert_ggml_to_int8.py +vllm serve +""" + +import argparse +import json +import os +import shutil +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY, GGMLFormat +from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( + copy_non_model_files, + copy_off_index_safetensors, + indexed_weight_map, +) + +_INDEX = "model.safetensors.index.json" + + +def convert_ggml_export_to_int8( + ggml_dir: str | os.PathLike, + output_dir: str | os.PathLike, + max_shard_size: int = 10 * 2**30, +) -> None: + """Write a GGML export as an INT8 weight-only compressed-tensors checkpoint, which vLLM serves. + + Every GGML weight is an integer level of at most 8 bits times the step of its 16- or + 32-weight sub-block, so storing those levels with the steps as group scales reproduces the + GGML weights up to rounding the steps to the model dtype. Unquantized tensors and the other + files are copied as they are; the GGML export is left unchanged. + + Args: + ggml_dir: A Hugging Face export whose quantized layers are all GGML formats, as written + by hf_ptq.py with a GGML recipe. + output_dir: Where to write the INT8 checkpoint. Must be empty or not exist yet. + max_shard_size: Target bytes per output safetensors shard; a larger tensor gets a shard + of its own. Conversion holds one output shard in memory at a time. + """ + ggml_dir, output_dir = Path(ggml_dir), Path(output_dir) + config = json.loads((ggml_dir / "config.json").read_text()) + quantization = json.loads((ggml_dir / "hf_quant_config.json").read_text())["quantization"] + scale_dtype = getattr(torch, config.get("dtype") or config["torch_dtype"]) + weight_map = indexed_weight_map(ggml_dir) + layer_formats = _ggml_layer_formats(ggml_dir, weight_map, quantization) + if output_dir.exists() and any(output_dir.iterdir()): + raise ValueError(f"{output_dir} is not empty.") + output_dir.mkdir(parents=True, exist_ok=True) + + try: + shards = sorted(set(weight_map.values())) + new_weight_map = _write_int8_shards( + ggml_dir, shards, output_dir, layer_formats, scale_dtype, max_shard_size + ) + total_size = sum((output_dir / s).stat().st_size for s in set(new_weight_map.values())) + index = {"metadata": {"total_size": total_size}, "weight_map": new_weight_map} + (output_dir / _INDEX).write_text(json.dumps(index, indent=2)) + config["quantization_config"] = _compressed_tensors_config(layer_formats) + (output_dir / "config.json").write_text(json.dumps(config, indent=4)) + # Both skip files already written, and copy_non_model_files skips hf_quant_config.json, + # which describes the GGML payloads this replaced. + copy_off_index_safetensors(ggml_dir, output_dir) + copy_non_model_files(ggml_dir, output_dir) + except BaseException: + # Empty the directory again, so a rerun is not refused over a partial checkpoint. + for path in output_dir.iterdir(): + if path.is_dir(): + shutil.rmtree(path) + else: + path.unlink() + raise + + +def _ggml_layer_formats( + ggml_dir: Path, weight_map: dict[str, str], quantization: dict +) -> dict[str, GGMLFormat]: + """The GGML format of every layer holding a packed payload, refusing anything else.""" + if quantization.get("kv_cache_quant_algo") is not None: + raise ValueError("INT8 conversion of a GGML export does not carry KV cache quantization.") + mixed = quantization.get("quant_algo") == "MIXED_PRECISION" + layer_configs = quantization.get("quantized_layers", {}) if mixed else {} + algos = ( + {cfg.get("quant_algo") for cfg in layer_configs.values()} + if mixed + else {quantization.get("quant_algo")} + ) + unsupported = sorted(str(a) for a in algos if str(a).lower() not in GGML_FORMAT_REGISTRY) + if unsupported: + raise ValueError(f"Only GGML layers convert to INT8, but the export has {unsupported}.") + + layer_formats = {} + for shard in sorted(set(weight_map.values())): + with safe_open(ggml_dir / shard, framework="pt") as f: + # GGML payloads are the only uint8 weights an export writes. + payloads = { + name: f.get_slice(name).get_shape() + for name in f.keys() # noqa: SIM118 -- a safe_open handle is not a dict + if name.endswith(".weight") and f.get_slice(name).get_dtype() == "U8" + } + for name, shape in payloads.items(): + layer = name.removesuffix(".weight") + # A 2-D weight's payload is [rows, blocks, block bytes]; weight_shape and the int32 + # packing assume that layout. + if len(shape) != 3: + raise ValueError(f"{layer} is not a 2-D weight; only 2-D weights convert to INT8.") + if mixed and layer not in layer_configs: + raise ValueError(f"{layer} holds a uint8 weight but is not in quantized_layers.") + algo = layer_configs[layer]["quant_algo"] if mixed else quantization["quant_algo"] + layer_formats[layer] = GGML_FORMAT_REGISTRY[algo.lower()] + if not layer_formats: + raise ValueError(f"{ggml_dir} holds no GGML weights to convert.") + return layer_formats + + +def _write_int8_shards( + ggml_dir: Path, + shards: list[str], + output_dir: Path, + layer_formats: dict[str, GGMLFormat], + scale_dtype: torch.dtype, + max_shard_size: int, +) -> dict[str, str]: + """Rewrite every shard's tensors into size-bounded output shards; return the weight map.""" + device = "cuda" if torch.cuda.is_available() else "cpu" + pending: dict[str, torch.Tensor] = {} + written: list[list[str]] = [] + + def flush(): + save_file(pending, output_dir / f"part-{len(written):05d}.safetensors", {"format": "pt"}) + written.append(list(pending)) + pending.clear() + + for shard in shards: + with safe_open(ggml_dir / shard, framework="pt") as f: + for name in f.keys(): # noqa: SIM118 -- a safe_open handle is not a dict + layer = name.removesuffix(".weight") + ggml_format = layer_formats.get(layer) if name.endswith(".weight") else None + if ggml_format is None: + tensors = {name: f.get_tensor(name)} + else: + tensors = _int8_tensors( + layer, f.get_tensor(name), ggml_format, scale_dtype, device + ) + if pending and _nbytes(pending) + _nbytes(tensors) > max_shard_size: + flush() + pending.update(tensors) + flush() + + weight_map = {} + for i, names in enumerate(written): + shard_name = f"model-{i + 1:05d}-of-{len(written):05d}.safetensors" + os.replace(output_dir / f"part-{i:05d}.safetensors", output_dir / shard_name) + weight_map.update(dict.fromkeys(names, shard_name)) + return dict(sorted(weight_map.items())) + + +def _int8_tensors( + layer: str, + payload: torch.Tensor, + ggml_format: GGMLFormat, + scale_dtype: torch.dtype, + device: str, +) -> dict[str, torch.Tensor]: + """A layer's GGML payload as compressed-tensors ``weight_packed``, ``weight_scale`` and shape.""" + shape = torch.tensor((*payload.shape[:-2], payload.shape[-2] * ggml_format.block_size)) + levels, steps = ggml_format.decode_int8(payload.to(device), shape) + return { + # compressed-tensors packs 8-bit values offset by 128, four little-endian per int32. + f"{layer}.weight_packed": (levels.view(torch.uint8) ^ 0x80).view(torch.int32).cpu(), + f"{layer}.weight_scale": steps.to(scale_dtype).cpu(), + f"{layer}.weight_shape": shape, + } + + +def _nbytes(tensors: dict[str, torch.Tensor]) -> int: + return sum(t.numel() * t.element_size() for t in tensors.values()) + + +def _compressed_tensors_config(layer_formats: dict[str, GGMLFormat]) -> dict: + """One config group per INT8 group size, targeting its layers by name. + + vLLM maps explicit layer names onto its own modules, but not ``re:`` patterns. + """ + layers_by_group_size: dict[int, list[str]] = {} + for layer, ggml_format in sorted(layer_formats.items()): + layers_by_group_size.setdefault(ggml_format.int8_group_size, []).append(layer) + config_groups = { + f"group_{i}": { + "targets": layers, + "weights": { + "num_bits": 8, + "type": "int", + "symmetric": True, + "strategy": "group", + "group_size": group_size, + "dynamic": False, + }, + "input_activations": None, + "output_activations": None, + "format": "pack-quantized", + } + for i, (group_size, layers) in enumerate(sorted(layers_by_group_size.items())) + } + return { + "quant_method": "compressed-tensors", + "format": "pack-quantized", + "quantization_status": "compressed", + "config_groups": config_groups, + "ignore": [], + } + + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("ggml_dir", help="Checkpoint exported with a GGML recipe (IQ, q8_0).") + parser.add_argument("output_dir", help="Where to write the INT8 checkpoint; must be empty.") + parser.add_argument( + "--max_shard_size_gb", + type=float, + default=10, + help="Target size of each output shard in GiB; a larger tensor gets a shard of its own.", + ) + args = parser.parse_args() + convert_ggml_export_to_int8( + args.ggml_dir, args.output_dir, max_shard_size=int(args.max_shard_size_gb * 2**30) + ) + print(f"INT8 checkpoint written to {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py index 83dfc2f00d5..59a24cbdb14 100644 --- a/modelopt/torch/quantization/ggml/common.py +++ b/modelopt/torch/quantization/ggml/common.py @@ -247,6 +247,11 @@ class GGMLFormat: # its search temporaries, decoding runs every forward and is bounded by kernel launches. block_chunk_size: int decode_chunk_size: int + # Every decoded weight is an int8 level times a float32 step shared by ``int8_group_size`` + # weights. ``unpack_int8`` returns both for a chunk of blocks, and the torch decoder + # multiplies them out, so the integer view and the decoding cannot disagree. + int8_group_size: int + unpack_int8: Callable[[torch.Tensor], tuple[torch.Tensor, torch.Tensor]] # Whether ``quantize`` takes a per-input-column ``importance`` that weights its search. weighted: bool = False @@ -255,6 +260,36 @@ def effective_bits(self) -> float: """Packed storage cost per weight.""" return self.block_bytes * 8 / self.block_size + def decode_int8( + self, packed_weights: torch.Tensor, weight_shape: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Decode a payload into int8 levels and the float32 step of each ``int8_group_size``. + + ``levels * steps`` (each step repeated over its group) equals ``dequantize`` in float32 + exactly, so a weight-only integer kernel with these group scales reproduces the format. + + Returns: + ``levels`` shaped like the weight and ``steps`` shaped + ``[*weight_shape[:-1], weight_shape[-1] // int8_group_size]``. + """ + shape = validate_packed_weights( + packed_weights, + weight_shape, + block_bytes=self.block_bytes, + block_size=self.block_size, + format_name=self.name.upper(), + ) + blocks = packed_weights.contiguous().reshape(-1, self.block_bytes) + groups = self.block_size // self.int8_group_size + levels = torch.empty((len(blocks), self.block_size), dtype=torch.int8, device=blocks.device) + steps = torch.empty((len(blocks), groups), dtype=torch.float32, device=blocks.device) + for start in range(0, len(blocks), self.decode_chunk_size): + stop = start + self.decode_chunk_size + chunk_levels, chunk_steps = self.unpack_int8(blocks[start:stop]) + levels[start:stop] = chunk_levels + steps[start:stop] = chunk_steps + return levels.reshape(shape), steps.reshape(*shape[:-1], -1) + def fake_quant( self, inputs: torch.Tensor, @@ -354,6 +389,26 @@ def decoding(storage_dtype): return None +def dequantize_blocks( + blocks: torch.Tensor, + unpack_int8: Callable[[torch.Tensor], tuple[torch.Tensor, torch.Tensor]], + *, + block_size: int, + group_size: int, + dtype: torch.dtype, + block_chunk_size: int, +) -> torch.Tensor: + """Decode ``[blocks, block_bytes]`` payload rows a chunk at a time as ``levels * steps``.""" + decoded = torch.empty((len(blocks), block_size), dtype=dtype, device=blocks.device) + for start in range(0, len(blocks), block_chunk_size): + stop = start + block_chunk_size + levels, steps = unpack_int8(blocks[start:stop]) + decoded[start:stop] = ( + levels.reshape(len(levels), -1, group_size) * steps.unsqueeze(-1) + ).reshape(-1, block_size) + return decoded + + def narrow_to_float32(blocks: torch.Tensor) -> torch.Tensor: """Narrow ``blocks`` to float32 the way the CUDA ``load_float`` helper does. diff --git a/modelopt/torch/quantization/ggml/iq1_m.py b/modelopt/torch/quantization/ggml/iq1_m.py index e1e5907aaa4..b0cd994b771 100644 --- a/modelopt/torch/quantization/ggml/iq1_m.py +++ b/modelopt/torch/quantization/ggml/iq1_m.py @@ -45,6 +45,7 @@ GGML_BLOCK_SIZE, GGMLFormat, chunk_importance, + dequantize_blocks, importance_blocks, narrow_to_float32, validate_block_chunk_size, @@ -67,6 +68,7 @@ IQ1_M_BLOCK_BYTES = 56 IQ1_M_EFFECTIVE_BITS = IQ1_M_BLOCK_BYTES * 8 / IQ1_M_BLOCK_SIZE _IQ1_M_DELTA = 0.125 +_IQ1_M_INT8_GROUP_SIZE = 16 # Largest representable magnitude: (1 + delta) at local scale 7 -> 15 * 1.125. _IQ1_M_NATIVE_MAX = 16.875 # IQ1_M anchors differently from IQ1_S: the ratio rises with a block's peak-to-RMS @@ -213,45 +215,46 @@ def dequantize_iq1_m( if extension is not None: grid = iq1_s_grid(blocks.device) return extension.iq1_m_unpack(blocks, grid, dtype).reshape(shape) - grid = iq1_s_grid(blocks.device) + return dequantize_blocks( + blocks, + _unpack_int8, + block_size=IQ1_M_BLOCK_SIZE, + group_size=_IQ1_M_INT8_GROUP_SIZE, + dtype=dtype, + block_chunk_size=block_chunk_size, + ).reshape(shape) + + +def _unpack_int8(blocks: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Levels in eighths of a grid step (±1, ±7, ±9) and the step of each 16-value half.""" + count = blocks.shape[0] scale_shifts = torch.tensor([0, 3, 6, 9], dtype=torch.int64, device=blocks.device) - decoded = torch.empty((blocks.shape[0], IQ1_M_BLOCK_SIZE), dtype=dtype, device=blocks.device) - for start in range(0, blocks.shape[0], block_chunk_size): - stop = min(start + block_chunk_size, blocks.shape[0]) - block_chunk = blocks[start:stop] - count = block_chunk.shape[0] - low = block_chunk[:, :32].to(torch.int64).reshape(count, _IQ1_M_SUBBLOCKS, 4) - qh = block_chunk[:, 32:48].to(torch.int64).reshape(count, _IQ1_M_SUBBLOCKS, 2) - words = block_chunk[:, 48:56:2].to(torch.int64) | ( - block_chunk[:, 49:56:2].to(torch.int64) << 8 - ) - - # The FP16 super-block scale is the four words' top nibbles, low word first. - d_bits = ( - (words[:, 0] >> 12) - | ((words[:, 1] >> 8) & 0x00F0) - | ((words[:, 2] >> 4) & 0x0F00) - | (words[:, 3] & 0xF000) - ) - d = d_bits.to(torch.int16).view(torch.float16).float() + low = blocks[:, :32].to(torch.int64).reshape(count, _IQ1_M_SUBBLOCKS, 4) + qh = blocks[:, 32:48].to(torch.int64).reshape(count, _IQ1_M_SUBBLOCKS, 2) + words = blocks[:, 48:56:2].to(torch.int64) | (blocks[:, 49:56:2].to(torch.int64) << 8) + + # The FP16 super-block scale is the four words' top nibbles, low word first. + d_bits = ( + (words[:, 0] >> 12) + | ((words[:, 1] >> 8) & 0x00F0) + | ((words[:, 2] >> 4) & 0x0F00) + | (words[:, 3] & 0xF000) + ) + d = d_bits.to(torch.int16).view(torch.float16).float() - # Each qh byte carries two groups, low nibble first: three high index bits and a delta sign. - nibbles = torch.stack((qh & 0xF, qh >> 4), dim=-1).reshape(count, _IQ1_M_SUBBLOCKS, 4) - entries = low | ((nibbles & 0x7) << 8) - deltas = torch.where((nibbles & 0x8).bool(), -_IQ1_M_DELTA, _IQ1_M_DELTA) + # Each qh byte carries two groups, low nibble first: three high index bits and a delta sign. + nibbles = torch.stack((qh & 0xF, qh >> 4), dim=-1).reshape(count, _IQ1_M_SUBBLOCKS, 4) + entries = low | ((nibbles & 0x7) << 8) + delta_sign = torch.where((nibbles & 0x8).bool(), -1.0, 1.0).unsqueeze(-1) - # Each word packs four 3-bit local scales below its nibble of d: two sub-blocks, two - # halves each, so word w, slot k is sub-block 2 * w + k // 2, half k % 2. - local = ((words.unsqueeze(-1) >> scale_shifts) & 0x7).reshape(count, _IQ1_M_SUBBLOCKS, 2) - scales = d.unsqueeze(-1).unsqueeze(-1) * (2 * local + 1).float() + # Each word packs four 3-bit local scales below its nibble of d: two sub-blocks, two + # halves each, so word w, slot k is sub-block 2 * w + k // 2, half k % 2. Groups 0,1 take + # the first half's scale and groups 2,3 the second, which is the order the steps flatten in. + local = ((words.unsqueeze(-1) >> scale_shifts) & 0x7).reshape(count, _IQ1_M_SUBBLOCKS * 2) + steps = d.unsqueeze(-1) * (2 * local + 1).float() * _IQ1_M_DELTA - values = grid[entries] + deltas.unsqueeze(-1) - # Groups 0,1 take the first local scale and groups 2,3 the second, so the repeat - # runs along the half axis: [h0, h0, h1, h1], not [h0, h1, h0, h1]. - group_scale = scales.repeat_interleave(2, dim=2) - chunk_decoded = values * group_scale.unsqueeze(-1) - decoded[start:stop] = chunk_decoded.reshape(-1, IQ1_M_BLOCK_SIZE) - return decoded.reshape(shape) + levels = iq1_s_grid(blocks.device)[entries] / _IQ1_M_DELTA + delta_sign + return levels.reshape(count, IQ1_M_BLOCK_SIZE), steps IQ1_M_FORMAT = GGMLFormat( @@ -262,6 +265,8 @@ def dequantize_iq1_m( dequantize=dequantize_iq1_m, block_chunk_size=_DEFAULT_BLOCK_CHUNK_SIZE, decode_chunk_size=_DEFAULT_DECODE_CHUNK_SIZE, + int8_group_size=_IQ1_M_INT8_GROUP_SIZE, + unpack_int8=_unpack_int8, weighted=True, ) diff --git a/modelopt/torch/quantization/ggml/iq1_s.py b/modelopt/torch/quantization/ggml/iq1_s.py index 57ee47cf7cb..f00d777c0d2 100644 --- a/modelopt/torch/quantization/ggml/iq1_s.py +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -41,6 +41,7 @@ GGML_BLOCK_SIZE, GGMLFormat, chunk_importance, + dequantize_blocks, importance_blocks, narrow_to_float32, validate_block_chunk_size, @@ -62,6 +63,7 @@ IQ1_S_BLOCK_BYTES = 50 IQ1_S_EFFECTIVE_BITS = IQ1_S_BLOCK_BYTES * 8 / IQ1_S_BLOCK_SIZE _IQ1_S_DELTA = 0.125 +_IQ1_S_INT8_GROUP_SIZE = 32 _IQ1_S_NATIVE_MAX = 16.875 _IQ1_S_SCALE_ANCHOR = 0.61 # Bounds the torch encode fallback, whose codebook search holds the large temporaries: at @@ -273,26 +275,29 @@ def dequantize_iq1_s( if extension is not None: grid = iq1_s_grid(blocks.device) return extension.iq1_s_unpack(blocks, grid, dtype).reshape(shape) + return dequantize_blocks( + blocks, + _unpack_int8, + block_size=IQ1_S_BLOCK_SIZE, + group_size=_IQ1_S_INT8_GROUP_SIZE, + dtype=dtype, + block_chunk_size=block_chunk_size, + ).reshape(shape) + + +def _unpack_int8(blocks: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Levels in eighths of a grid step (±1, ±7, ±9) and the step of each 32-value sub-block.""" shifts = torch.tensor([0, 3, 6, 9], dtype=torch.int64, device=blocks.device) - grid = iq1_s_grid(blocks.device) - decoded = torch.empty((blocks.shape[0], IQ1_S_BLOCK_SIZE), dtype=dtype, device=blocks.device) - for start in range(0, blocks.shape[0], block_chunk_size): - stop = min(start + block_chunk_size, blocks.shape[0]) - block_chunk = blocks[start:stop] - d = block_chunk[:, :2].contiguous().view(torch.float16).reshape(-1).float() - low = block_chunk[:, 2:34].to(torch.int64).reshape(-1, 8, 4) - qh = block_chunk[:, 34:50:2].to(torch.int64) | ( - block_chunk[:, 35:50:2].to(torch.int64) << 8 - ) - high = (qh.unsqueeze(-1) >> shifts) & 0x7 - entries = low | (high << 8) - local = (qh >> 12) & 0x7 - delta = torch.where((qh & 0x8000).bool(), -_IQ1_S_DELTA, _IQ1_S_DELTA) - values = grid[entries] + delta.unsqueeze(-1).unsqueeze(-1) - scales = d.unsqueeze(-1) * (2 * local + 1).float() - chunk_decoded = values * scales.unsqueeze(-1).unsqueeze(-1) - decoded[start:stop] = chunk_decoded.reshape(-1, IQ1_S_BLOCK_SIZE) - return decoded.reshape(shape) + d = blocks[:, :2].contiguous().view(torch.float16).reshape(-1).float() + low = blocks[:, 2:34].to(torch.int64).reshape(-1, 8, 4) + qh = blocks[:, 34:50:2].to(torch.int64) | (blocks[:, 35:50:2].to(torch.int64) << 8) + high = (qh.unsqueeze(-1) >> shifts) & 0x7 + entries = low | (high << 8) + local = (qh >> 12) & 0x7 + delta_sign = torch.where((qh & 0x8000).bool(), -1.0, 1.0).unsqueeze(-1).unsqueeze(-1) + levels = iq1_s_grid(blocks.device)[entries] / _IQ1_S_DELTA + delta_sign + steps = d.unsqueeze(-1) * (2 * local + 1).float() * _IQ1_S_DELTA + return levels.reshape(-1, IQ1_S_BLOCK_SIZE), steps IQ1_S_FORMAT = GGMLFormat( @@ -303,6 +308,8 @@ def dequantize_iq1_s( dequantize=dequantize_iq1_s, block_chunk_size=_DEFAULT_BLOCK_CHUNK_SIZE, decode_chunk_size=_DEFAULT_DECODE_CHUNK_SIZE, + int8_group_size=_IQ1_S_INT8_GROUP_SIZE, + unpack_int8=_unpack_int8, weighted=True, ) diff --git a/modelopt/torch/quantization/ggml/iq2_s.py b/modelopt/torch/quantization/ggml/iq2_s.py index c7c24c826ba..c64ae46fdd7 100644 --- a/modelopt/torch/quantization/ggml/iq2_s.py +++ b/modelopt/torch/quantization/ggml/iq2_s.py @@ -43,6 +43,7 @@ GGML_BLOCK_SIZE, GGMLFormat, chunk_importance, + dequantize_blocks, importance_blocks, iq2_tile_terms, narrow_to_float32, @@ -68,6 +69,7 @@ _IQ2_S_LOCAL_SCALES = 16 _IQ2_S_GROUPS = 32 _IQ2_S_SUBBLOCKS = 8 +_IQ2_S_INT8_GROUP_SIZE = 16 _IQ2_S_NATIVE_MAX = 43 * 31 / 8 _IQ2_S_SCALE_ANCHOR_MIN = 0.65 _IQ2_S_SCALE_ANCHOR_MAX = 0.92 @@ -227,30 +229,35 @@ def dequantize_iq2_s( if extension is not None: grid = iq2_s_grid(blocks.device) return extension.iq2_s_unpack(blocks, grid, dtype).reshape(shape) + return dequantize_blocks( + blocks, + _unpack_int8, + block_size=IQ2_S_BLOCK_SIZE, + group_size=_IQ2_S_INT8_GROUP_SIZE, + dtype=dtype, + block_chunk_size=block_chunk_size, + ).reshape(shape) + + +def _unpack_int8(blocks: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Signed grid levels (±8, ±25, ±43) and the step of each 16-value half sub-block.""" + count = blocks.shape[0] bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device) high_shifts = torch.tensor([0, 2, 4, 6], dtype=torch.int64, device=blocks.device) - grid = iq2_s_grid(blocks.device) - decoded = torch.empty((blocks.shape[0], IQ2_S_BLOCK_SIZE), dtype=dtype, device=blocks.device) - for start in range(0, blocks.shape[0], block_chunk_size): - stop = min(start + block_chunk_size, blocks.shape[0]) - block_chunk = blocks[start:stop] - count = block_chunk.shape[0] - d = block_chunk[:, :2].contiguous().view(torch.float16).reshape(-1).float() - low = block_chunk[:, 2:34].to(torch.int64).reshape(count, _IQ2_S_SUBBLOCKS, 4) - sign_mask = block_chunk[:, 34:66].to(torch.int64).reshape(count, _IQ2_S_SUBBLOCKS, 4) - qh = block_chunk[:, 66:74].to(torch.int64).reshape(count, _IQ2_S_SUBBLOCKS) - scale_bytes = block_chunk[:, 74:].to(torch.int64) - - entries = low | (((qh.unsqueeze(-1) >> high_shifts) & 0x3) << 8) - local = torch.empty((count, 16), dtype=torch.int64, device=blocks.device) - local[:, 0::2] = scale_bytes & 0x0F - local[:, 1::2] = scale_bytes >> 4 - scales = d.unsqueeze(-1) * (0.5 + local.float()) * 0.25 - signs = 1.0 - 2.0 * ((sign_mask.unsqueeze(-1) >> bit_positions) & 1).float() - values = grid[entries] * signs - chunk_decoded = values.reshape(count, 16, 2, 8) * scales.unsqueeze(-1).unsqueeze(-1) - decoded[start:stop] = chunk_decoded.reshape(-1, IQ2_S_BLOCK_SIZE) - return decoded.reshape(shape) + d = blocks[:, :2].contiguous().view(torch.float16).reshape(-1).float() + low = blocks[:, 2:34].to(torch.int64).reshape(count, _IQ2_S_SUBBLOCKS, 4) + sign_mask = blocks[:, 34:66].to(torch.int64).reshape(count, _IQ2_S_SUBBLOCKS, 4) + qh = blocks[:, 66:74].to(torch.int64).reshape(count, _IQ2_S_SUBBLOCKS) + scale_bytes = blocks[:, 74:].to(torch.int64) + + entries = low | (((qh.unsqueeze(-1) >> high_shifts) & 0x3) << 8) + local = torch.empty((count, 16), dtype=torch.int64, device=blocks.device) + local[:, 0::2] = scale_bytes & 0x0F + local[:, 1::2] = scale_bytes >> 4 + steps = d.unsqueeze(-1) * (0.5 + local.float()) * 0.25 + signs = 1.0 - 2.0 * ((sign_mask.unsqueeze(-1) >> bit_positions) & 1).float() + levels = iq2_s_grid(blocks.device)[entries] * signs + return levels.reshape(count, IQ2_S_BLOCK_SIZE), steps IQ2_S_FORMAT = GGMLFormat( @@ -261,6 +268,8 @@ def dequantize_iq2_s( dequantize=dequantize_iq2_s, block_chunk_size=_DEFAULT_BLOCK_CHUNK_SIZE, decode_chunk_size=_DEFAULT_DECODE_CHUNK_SIZE, + int8_group_size=_IQ2_S_INT8_GROUP_SIZE, + unpack_int8=_unpack_int8, weighted=True, ) diff --git a/modelopt/torch/quantization/ggml/iq2_xs.py b/modelopt/torch/quantization/ggml/iq2_xs.py index 27e26a8fbe3..9f73c1f81b7 100644 --- a/modelopt/torch/quantization/ggml/iq2_xs.py +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -38,6 +38,7 @@ GGML_BLOCK_SIZE, GGMLFormat, chunk_importance, + dequantize_blocks, importance_blocks, iq2_tile_terms, narrow_to_float32, @@ -59,6 +60,7 @@ IQ2_XS_BLOCK_SIZE = GGML_BLOCK_SIZE IQ2_XS_BLOCK_BYTES = 74 IQ2_XS_EFFECTIVE_BITS = IQ2_XS_BLOCK_BYTES * 8 / IQ2_XS_BLOCK_SIZE +_IQ2_XS_INT8_GROUP_SIZE = 16 _IQ2_XS_NATIVE_MAX = 43 * 31 / 8 _IQ2_XS_SCALE_ANCHOR_MIN = 0.65 _IQ2_XS_SCALE_ANCHOR_MAX = 0.92 @@ -239,37 +241,40 @@ def dequantize_iq2_xs( if extension is not None: grid = iq2_xs_grid(blocks.device) return extension.iq2_xs_unpack(blocks, grid, dtype).reshape(shape) + return dequantize_blocks( + blocks, + _unpack_int8, + block_size=IQ2_XS_BLOCK_SIZE, + group_size=_IQ2_XS_INT8_GROUP_SIZE, + dtype=dtype, + block_chunk_size=block_chunk_size, + ).reshape(shape) + + +def _unpack_int8(blocks: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Signed grid levels (±8, ±25, ±43) and the step of each 16-value half sub-block.""" bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device) - grid = iq2_xs_grid(blocks.device) - decoded = torch.empty((blocks.shape[0], IQ2_XS_BLOCK_SIZE), dtype=dtype, device=blocks.device) - for start in range(0, blocks.shape[0], block_chunk_size): - stop = min(start + block_chunk_size, blocks.shape[0]) - block_chunk = blocks[start:stop] - d = block_chunk[:, :2].contiguous().view(torch.float16).reshape(-1).float() - codes = block_chunk[:, 2:66:2].to(torch.int64) | ( - block_chunk[:, 3:66:2].to(torch.int64) << 8 - ) - entries = codes & 0x1FF - sign_index = codes >> 9 - # XOR-fold the seven payload bits down to bit 0 to recover the eighth sign bit. - # The loop this replaces cost seven elementwise passes per chunk, and unlike packing - # the decode is not cached -- it runs again on every forward. - folded = sign_index ^ (sign_index >> 4) - folded ^= folded >> 2 - folded ^= folded >> 1 - sign_mask = sign_index | ((folded & 1) << 7) - signs = 1.0 - 2.0 * ((sign_mask.unsqueeze(-1) >> bit_positions) & 1).float() - - scale_bytes = block_chunk[:, 66:].to(torch.int64) - local = torch.empty((block_chunk.shape[0], 16), dtype=torch.int64, device=blocks.device) - local[:, 0::2] = scale_bytes & 0x0F - local[:, 1::2] = scale_bytes >> 4 - # Pinned format rule: d * (0.5 + local) * 0.25 == d * (2 * local + 1) / 8. - scales = d.unsqueeze(-1) * (2 * local + 1).float() / 8.0 - values = grid[entries] * signs - chunk_decoded = values * scales.repeat_interleave(2, dim=1).unsqueeze(-1) - decoded[start:stop] = chunk_decoded.reshape(-1, IQ2_XS_BLOCK_SIZE) - return decoded.reshape(shape) + d = blocks[:, :2].contiguous().view(torch.float16).reshape(-1).float() + codes = blocks[:, 2:66:2].to(torch.int64) | (blocks[:, 3:66:2].to(torch.int64) << 8) + entries = codes & 0x1FF + sign_index = codes >> 9 + # XOR-fold the seven payload bits down to bit 0 to recover the eighth sign bit. + # The loop this replaces cost seven elementwise passes per chunk, and unlike packing + # the decode is not cached -- it runs again on every forward. + folded = sign_index ^ (sign_index >> 4) + folded ^= folded >> 2 + folded ^= folded >> 1 + sign_mask = sign_index | ((folded & 1) << 7) + signs = 1.0 - 2.0 * ((sign_mask.unsqueeze(-1) >> bit_positions) & 1).float() + + scale_bytes = blocks[:, 66:].to(torch.int64) + local = torch.empty((blocks.shape[0], 16), dtype=torch.int64, device=blocks.device) + local[:, 0::2] = scale_bytes & 0x0F + local[:, 1::2] = scale_bytes >> 4 + # Pinned format rule: d * (0.5 + local) * 0.25 == d * (2 * local + 1) / 8. + steps = d.unsqueeze(-1) * (2 * local + 1).float() / 8.0 + levels = iq2_xs_grid(blocks.device)[entries] * signs + return levels.reshape(-1, IQ2_XS_BLOCK_SIZE), steps IQ2_XS_FORMAT = GGMLFormat( @@ -280,6 +285,8 @@ def dequantize_iq2_xs( dequantize=dequantize_iq2_xs, block_chunk_size=_DEFAULT_BLOCK_CHUNK_SIZE, decode_chunk_size=_DEFAULT_DECODE_CHUNK_SIZE, + int8_group_size=_IQ2_XS_INT8_GROUP_SIZE, + unpack_int8=_unpack_int8, weighted=True, ) diff --git a/modelopt/torch/quantization/ggml/iq2_xxs.py b/modelopt/torch/quantization/ggml/iq2_xxs.py index 7ea803a1358..50ddf8c3b97 100644 --- a/modelopt/torch/quantization/ggml/iq2_xxs.py +++ b/modelopt/torch/quantization/ggml/iq2_xxs.py @@ -43,6 +43,7 @@ GGML_BLOCK_SIZE, GGMLFormat, chunk_importance, + dequantize_blocks, importance_blocks, iq2_tile_terms, narrow_to_float32, @@ -69,6 +70,7 @@ _IQ2_XXS_GROUPS = 32 _IQ2_XXS_GROUPS_PER_SUBBLOCK = 4 _IQ2_XXS_SUBBLOCKS = _IQ2_XXS_GROUPS // _IQ2_XXS_GROUPS_PER_SUBBLOCK +_IQ2_XXS_INT8_GROUP_SIZE = IQ2_XXS_BLOCK_SIZE // _IQ2_XXS_SUBBLOCKS # Largest representable magnitude: grid entry 43 at local scale 15 -> (0.5 + 15) * 0.25. _IQ2_XXS_NATIVE_MAX = 43 * 31 / 8 _IQ2_XXS_SCALE_ANCHOR_MIN = 0.65 @@ -258,32 +260,37 @@ def dequantize_iq2_xxs( if extension is not None: grid = iq2_xxs_grid(blocks.device) return extension.iq2_xxs_unpack(blocks, grid, dtype).reshape(shape) + return dequantize_blocks( + blocks, + _unpack_int8, + block_size=IQ2_XXS_BLOCK_SIZE, + group_size=_IQ2_XXS_INT8_GROUP_SIZE, + dtype=dtype, + block_chunk_size=block_chunk_size, + ).reshape(shape) + + +def _unpack_int8(blocks: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Signed grid levels (±8, ±25, ±43) and the step of each 32-value sub-block.""" + count = blocks.shape[0] bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device) sign_shifts = 7 * torch.arange( _IQ2_XXS_GROUPS_PER_SUBBLOCK, dtype=torch.int64, device=blocks.device ) - grid = iq2_xxs_grid(blocks.device) - decoded = torch.empty((blocks.shape[0], IQ2_XXS_BLOCK_SIZE), dtype=dtype, device=blocks.device) - for start in range(0, blocks.shape[0], block_chunk_size): - stop = min(start + block_chunk_size, blocks.shape[0]) - block_chunk = blocks[start:stop] - count = block_chunk.shape[0] - d = block_chunk[:, :2].contiguous().view(torch.float16).reshape(-1).float() - body = block_chunk[:, 2:].reshape(count, _IQ2_XXS_SUBBLOCKS, 8).to(torch.int64) - entries = body[:, :, 0:4] - aux = body[:, :, 4] | (body[:, :, 5] << 8) | (body[:, :, 6] << 16) | (body[:, :, 7] << 24) - # Top nibble is the sub-block scale; the low 28 bits are four 7-bit sign indices. - scales = d.unsqueeze(-1) * (0.5 + ((aux >> 28) & 0xF).float()) * 0.25 - sign_index = (aux.unsqueeze(-1) >> sign_shifts) & 0x7F - folded = sign_index ^ (sign_index >> 4) - folded ^= folded >> 2 - folded ^= folded >> 1 - sign_mask = sign_index | ((folded & 1) << 7) - signs = 1.0 - 2.0 * ((sign_mask.unsqueeze(-1) >> bit_positions) & 1).float() - values = grid[entries] * signs - chunk_decoded = values * scales.unsqueeze(-1).unsqueeze(-1) - decoded[start:stop] = chunk_decoded.reshape(-1, IQ2_XXS_BLOCK_SIZE) - return decoded.reshape(shape) + d = blocks[:, :2].contiguous().view(torch.float16).reshape(-1).float() + body = blocks[:, 2:].reshape(count, _IQ2_XXS_SUBBLOCKS, 8).to(torch.int64) + entries = body[:, :, 0:4] + aux = body[:, :, 4] | (body[:, :, 5] << 8) | (body[:, :, 6] << 16) | (body[:, :, 7] << 24) + # Top nibble is the sub-block scale; the low 28 bits are four 7-bit sign indices. + steps = d.unsqueeze(-1) * (0.5 + ((aux >> 28) & 0xF).float()) * 0.25 + sign_index = (aux.unsqueeze(-1) >> sign_shifts) & 0x7F + folded = sign_index ^ (sign_index >> 4) + folded ^= folded >> 2 + folded ^= folded >> 1 + sign_mask = sign_index | ((folded & 1) << 7) + signs = 1.0 - 2.0 * ((sign_mask.unsqueeze(-1) >> bit_positions) & 1).float() + levels = iq2_xxs_grid(blocks.device)[entries] * signs + return levels.reshape(count, IQ2_XXS_BLOCK_SIZE), steps IQ2_XXS_FORMAT = GGMLFormat( @@ -294,6 +301,8 @@ def dequantize_iq2_xxs( dequantize=dequantize_iq2_xxs, block_chunk_size=_DEFAULT_BLOCK_CHUNK_SIZE, decode_chunk_size=_DEFAULT_DECODE_CHUNK_SIZE, + int8_group_size=_IQ2_XXS_INT8_GROUP_SIZE, + unpack_int8=_unpack_int8, weighted=True, ) diff --git a/modelopt/torch/quantization/ggml/q8_0.py b/modelopt/torch/quantization/ggml/q8_0.py index 242a823835d..8b3700e9d5d 100644 --- a/modelopt/torch/quantization/ggml/q8_0.py +++ b/modelopt/torch/quantization/ggml/q8_0.py @@ -24,6 +24,7 @@ from ..extensions import get_cuda_ext_ggml from .common import ( GGMLFormat, + dequantize_blocks, narrow_to_float32, validate_block_chunk_size, validate_packed_weights, @@ -127,14 +128,20 @@ def dequantize_q8_0( validate_block_chunk_size(block_chunk_size) blocks = packed_weights.contiguous().reshape(-1, Q8_0_BLOCK_BYTES) - decoded = torch.empty((blocks.shape[0], Q8_0_BLOCK_SIZE), dtype=dtype, device=blocks.device) - for start in range(0, blocks.shape[0], block_chunk_size): - stop = min(start + block_chunk_size, blocks.shape[0]) - block_chunk = blocks[start:stop] - d = block_chunk[:, :2].contiguous().view(torch.float16).reshape(-1).float() - quants = block_chunk[:, 2:].contiguous().view(torch.int8).float() - decoded[start:stop] = (d.unsqueeze(1) * quants).to(dtype) - return decoded.reshape(shape) + return dequantize_blocks( + blocks, + _unpack_int8, + block_size=Q8_0_BLOCK_SIZE, + group_size=Q8_0_BLOCK_SIZE, + dtype=dtype, + block_chunk_size=block_chunk_size, + ).reshape(shape) + + +def _unpack_int8(blocks: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """The stored quants and each block's scale.""" + d = blocks[:, :2].contiguous().view(torch.float16).float() + return blocks[:, 2:].contiguous().view(torch.int8).float(), d Q8_0_FORMAT = GGMLFormat( @@ -145,6 +152,8 @@ def dequantize_q8_0( dequantize=dequantize_q8_0, block_chunk_size=_DEFAULT_BLOCK_CHUNK_SIZE, decode_chunk_size=_DEFAULT_DECODE_CHUNK_SIZE, + int8_group_size=Q8_0_BLOCK_SIZE, + unpack_int8=_unpack_int8, ) # Keep the public per-format entry point on the same record used by backend dispatch. diff --git a/tests/_test_utils/examples/ggml_export.py b/tests/_test_utils/examples/ggml_export.py new file mode 100644 index 00000000000..cc4ef9bda5e --- /dev/null +++ b/tests/_test_utils/examples/ggml_export.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""A tiny LLaMA with GGML-quantized MLPs, for the ``convert_ggml_to_int8.py`` example tests.""" + +import torch +from _test_utils.torch.transformers_models import create_tiny_llama_dir +from transformers import AutoModelForCausalLM + +import modelopt.torch.quantization as mtq + +IQ1_S = {"num_bits": "iq1_s", "block_sizes": {-1: 256}, "backend": "ggml"} +IQ2_XS = {"num_bits": "iq2_xs", "block_sizes": {-1: 256}, "backend": "ggml"} + + +def quantize_tiny_llama(tmp_path, layer_cfgs=(IQ1_S, IQ2_XS)): + """Return a two-layer tiny LLaMA whose layer ``i`` MLP takes ``layer_cfgs[i]``, and its tokens. + + The default formats have 32- and 16-weight sub-blocks, so the INT8 checkpoint needs a config + group for each. The shapes suit vLLM's kernels, and fake quant runs on CPU. + """ + source = create_tiny_llama_dir( + tmp_path, + hidden_size=256, + intermediate_size=512, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=64, + max_position_embeddings=64, + vocab_size=128, + initializer_range=0.05, + ) + model = AutoModelForCausalLM.from_pretrained(source, dtype=torch.bfloat16) + quant_cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + *( + {"quantizer_name": f"*layers.{i}.mlp*weight_quantizer", "cfg": cfg} + for i, cfg in enumerate(layer_cfgs) + ), + ], + "algorithm": "max", + } + token_ids = torch.randint(0, 128, (1, 48), generator=torch.Generator().manual_seed(0)) + mtq.quantize(model, quant_cfg, lambda m: m(token_ids)) + return model, token_ids diff --git a/tests/_test_utils/examples/hf_ptq_example_utils.py b/tests/_test_utils/examples/hf_ptq_example_utils.py index 0a1b70503ec..21adae3227b 100644 --- a/tests/_test_utils/examples/hf_ptq_example_utils.py +++ b/tests/_test_utils/examples/hf_ptq_example_utils.py @@ -25,7 +25,8 @@ if str(_HF_PTQ_DIR) not in sys.path: sys.path.insert(0, str(_HF_PTQ_DIR)) +import convert_ggml_to_int8 import example_utils import nemotron_vl_calib -__all__ = ["example_utils", "nemotron_vl_calib"] +__all__ = ["convert_ggml_to_int8", "example_utils", "nemotron_vl_calib"] diff --git a/tests/examples/hf_ptq/test_convert_ggml_to_int8.py b/tests/examples/hf_ptq/test_convert_ggml_to_int8.py new file mode 100644 index 00000000000..9e1d9a5e4c8 --- /dev/null +++ b/tests/examples/hf_ptq/test_convert_ggml_to_int8.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import subprocess +from pathlib import Path + +import pytest +import torch +from _test_utils.examples.ggml_export import IQ1_S, quantize_tiny_llama +from _test_utils.examples.hf_ptq_example_utils import convert_ggml_to_int8 +from _test_utils.examples.run_command import run_example_command +from safetensors.torch import load_file, save_file + +from modelopt.torch.export import export_hf_checkpoint +from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY + +# quantize_tiny_llama's layer 0 is IQ1_S (group 32) and layer 1 IQ2_XS (group 16). +LAYER_FORMATS = {0: "iq1_s", 1: "iq2_xs"} + + +def _convert(ggml_dir, int8_dir, *args): + run_example_command( + ["python", "convert_ggml_to_int8.py", str(ggml_dir), str(int8_dir), *args], "hf_ptq" + ) + + +def _files(directory): + return {p.relative_to(directory): p.read_bytes() for p in directory.rglob("*") if p.is_file()} + + +@pytest.fixture(scope="module") +def ggml_export(tmp_path_factory): + tmp_path = tmp_path_factory.mktemp("ggml") + model, _ = quantize_tiny_llama(tmp_path) + export_dir = tmp_path / "export" + export_hf_checkpoint(model, export_dir=export_dir) + # Files a checkpoint can ship beside its indexed weights: a nested sidecar and an off-index + # weight file such as an MTP head. + (export_dir / "chat_templates").mkdir() + (export_dir / "chat_templates" / "tool_use.jinja").write_text("{{ messages }}") + save_file({"mtp.fc.weight": torch.ones(2, 2)}, export_dir / "mtp.safetensors") + return export_dir + + +def test_ggml_export_converts_to_int8_levels_and_steps(ggml_export, tmp_path): + ggml_files = _files(ggml_export) + before = load_file(ggml_export / "model.safetensors") + int8_dir = tmp_path / "int8" + + # Shards of about 200 KB split the output, so the index has to stitch them together. + _convert(ggml_export, int8_dir, "--max_shard_size_gb", "0.0002") + + assert _files(ggml_export) == ggml_files + assert not (int8_dir / "hf_quant_config.json").exists() + for sidecar in (Path("chat_templates/tool_use.jinja"), Path("mtp.safetensors")): + assert (int8_dir / sidecar).read_bytes() == ggml_files[sidecar] + config = json.loads((int8_dir / "config.json").read_text())["quantization_config"] + assert config["quant_method"] == "compressed-tensors" + layers = { + i: [f"model.layers.{i}.mlp.{proj}" for proj in ("down_proj", "gate_proj", "up_proj")] + for i in LAYER_FORMATS + } + assert { + group["weights"]["group_size"]: group["targets"] + for group in config["config_groups"].values() + } == {32: layers[0], 16: layers[1]} + + weight_map = json.loads((int8_dir / "model.safetensors.index.json").read_text())["weight_map"] + assert len(set(weight_map.values())) > 1 + after = {} + for shard in set(weight_map.values()): + after.update(load_file(int8_dir / shard)) + assert after.keys() == weight_map.keys() + for i, name in LAYER_FORMATS.items(): + for layer in layers[i]: + payload = before.pop(f"{layer}.weight") + shape = after.pop(f"{layer}.weight_shape") + levels, steps = GGML_FORMAT_REGISTRY[name].decode_int8(payload, shape) + # compressed-tensors' layout: offset by 128, four values per int32, lowest byte first. + packed = after.pop(f"{layer}.weight_packed") + unpacked = torch.stack([(packed >> (8 * k)) & 0xFF for k in range(4)], -1).flatten(1) + assert torch.equal(unpacked - 128, levels.long()) + assert torch.equal(after.pop(f"{layer}.weight_scale"), steps.to(torch.bfloat16)) + assert shape.tolist() == list(levels.shape) + assert after.keys() == before.keys() + for name, tensor in before.items(): + assert torch.equal(after[name], tensor) + + +def test_a_failed_conversion_leaves_the_output_directory_empty(ggml_export, tmp_path, monkeypatch): + def fail(*args, **kwargs): + raise OSError("disk full") + + # Called in-process to fail the last step: by then the shards, index and config are written. + monkeypatch.setattr(convert_ggml_to_int8, "copy_non_model_files", fail) + + with pytest.raises(OSError, match="disk full"): + convert_ggml_to_int8.convert_ggml_export_to_int8(ggml_export, tmp_path / "int8") + + assert not any((tmp_path / "int8").iterdir()) + + +def test_export_with_a_non_ggml_layer_is_refused_before_writing(tmp_path): + model, _ = quantize_tiny_llama(tmp_path, (IQ1_S, {"num_bits": (4, 3), "axis": None})) + export_hf_checkpoint(model, export_dir=tmp_path / "ggml") + + with pytest.raises(subprocess.CalledProcessError): + _convert(tmp_path / "ggml", tmp_path / "int8") + + assert not (tmp_path / "int8").exists() diff --git a/tests/examples/vllm_serve/test_convert_ggml_to_int8.py b/tests/examples/vllm_serve/test_convert_ggml_to_int8.py new file mode 100644 index 00000000000..252350e788d --- /dev/null +++ b/tests/examples/vllm_serve/test_convert_ggml_to_int8.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc +import sys + +import pytest +import torch +from _test_utils.examples.ggml_export import quantize_tiny_llama +from _test_utils.examples.run_command import run_example_command +from vllm import LLM, SamplingParams +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.inputs import TokensPrompt + +from modelopt.torch.export import export_hf_checkpoint + + +@pytest.fixture(scope="module") +def int8_export_in_vllm(tmp_path_factory): + """vLLM serving the INT8 conversion of a GGML export, the model's own logprobs and its tokens.""" + tmp_path = tmp_path_factory.mktemp("ggml_int8") + # Group 32 (IQ1_S) and group 16 (IQ2_XS) layers take different vLLM kernels. + model, token_ids = quantize_tiny_llama(tmp_path) + with torch.no_grad(): + logits = model(token_ids).logits[0, :-1].float() + expected = logits.log_softmax(-1).gather(-1, token_ids[0, 1:, None])[:, 0] + + export_hf_checkpoint(model, export_dir=tmp_path / "ggml") + # The vLLM image has no `python` binary. + run_example_command( + [sys.executable, "convert_ggml_to_int8.py", str(tmp_path / "ggml"), str(tmp_path / "int8")], + "hf_ptq", + ) + + llm = LLM( + model=str(tmp_path / "int8"), + enforce_eager=True, + gpu_memory_utilization=0.2, + max_model_len=64, + max_num_seqs=1, + dtype="bfloat16", + skip_tokenizer_init=True, + ) + yield llm, token_ids[0], expected + del llm + gc.collect() + cleanup_dist_env_and_memory(shutdown_ray=False) + + +def test_vllm_serves_the_int8_export_like_the_fake_quantized_model(int8_export_in_vllm): + llm, token_ids, expected = int8_export_in_vllm + (output,) = llm.generate( + TokensPrompt(prompt_token_ids=token_ids.tolist()), + SamplingParams(max_tokens=1, prompt_logprobs=0, detokenize=False), + ) + served = torch.tensor( + [lp[t].logprob for t, lp in zip(token_ids[1:].tolist(), output.prompt_logprobs[1:])] + ) + + # vLLM's kernels round differently from the fake-quant reference by up to ~0.02; giving each + # group its neighbour's scale moves these logprobs by over 1. + torch.testing.assert_close(served, expected, atol=0.05, rtol=0) diff --git a/tests/unit/torch/quantization/test_iq_formats.py b/tests/unit/torch/quantization/test_iq_formats.py index 992ec981019..5402b78c9b9 100644 --- a/tests/unit/torch/quantization/test_iq_formats.py +++ b/tests/unit/torch/quantization/test_iq_formats.py @@ -52,6 +52,9 @@ DISPATCHED = sorted(IQ_FORMAT_REGISTRY) # IQ1 grids are ternary; IQ2 grids hold the magnitudes 8, 25 and 43. TERNARY = {"iq1_s", "iq1_m"} +# Weights sharing one sub-block scale: the formats with an extra 0.19-0.25 bits per weight spend +# them on scales twice as fine. +INT8_GROUP_SIZES = {"iq1_s": 32, "iq1_m": 16, "iq2_xxs": 32, "iq2_xs": 16, "iq2_s": 16} # name -> (largest representable magnitude, scale anchor at peak-to-RMS 1, 4, 8 and 16), where the # anchor is d * native max / amax. IQ1_S is flat; IQ1_M rises with peak-to-RMS and the IQ2 formats # fall with it, each clamped at both ends. Written out rather than read from the modules, so a @@ -295,6 +298,24 @@ def test_decoder_matches_llama_cpp_on_captured_blocks(name): assert np.array_equal(decoded.numpy(), expected_values(name)) +@pytest.mark.parametrize("name", NAMES) +def test_int8_view_matches_llama_cpp_on_captured_blocks(name): + """Integer levels times their group's step are llama.cpp's decoding, exactly.""" + record = IQ_FORMAT_REGISTRY[name] + blocks = torch.from_numpy(packed_blocks(name)) + count, group_size = blocks.shape[0], INT8_GROUP_SIZES[name] + + levels, steps = record.decode_int8(blocks.reshape(count, 1, -1), torch.tensor([count, 256])) + + assert record.int8_group_size == group_size + assert levels.dtype == torch.int8 + assert set(levels.abs().unique().tolist()) <= ( + {8, 25, 43} if name not in TERNARY else {1, 7, 9} + ) + decoded = levels.float().reshape(count, -1, group_size) * steps.unsqueeze(-1) + assert np.array_equal(decoded.reshape(count, 256).numpy(), expected_values(name)) + + def test_every_format_has_conformance_vectors(): """A new format must arrive with blocks captured from a real llama.cpp checkpoint.""" assert sorted(formats()) == NAMES