diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1ab4d7e7b85..48e973c6dee 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -103,6 +103,8 @@ repos: exclude: > (?x)^( modelopt/torch/quantization/utils/calib_utils.py| + modelopt/torch/quantization/ggml/iq1_s.py| + modelopt/torch/quantization/ggml/iq2_xs.py| modelopt/onnx/quantization/operators.py| modelopt/onnx/quantization/ort_patching.py| modelopt/torch/_deploy/utils/onnx_utils.py| diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 59e999ea2ee..fa492f04dc1 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,6 +8,7 @@ Changelog *Quantization* +- Add IQ1_S and IQ2_XS weight quantization and unified checkpoint export from Hugging Face and TP=1 Megatron models. Supported weights are encoded in GGML-compatible 256-value blocks and stored as packed ``uint8`` weight tensors in safetensors. A model-specific Nemotron 3.5 Lightning recipe combines IQ2_XS expert weights with dynamic weight-only NVFP4 Mamba projections. - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Calibration writes the layer shards; ``finalize()`` on the exporter left on the model adds the tail shard, the index and the config artifacts, and the checkpoint does not load until it runs. ``examples/hf_ptq`` does this for you. Supports FP8 and NVFP4 on single-process models, resident or offloaded, including multimodal models and models with MTP layers; other formats and placements raise ``NotImplementedError`` before calibration starts. *Misc* diff --git a/LICENSE b/LICENSE index c58bddda878..57c1104ae67 100644 --- a/LICENSE +++ b/LICENSE @@ -250,6 +250,7 @@ the following copyright holders, licensed under the MIT License: Copyright (c) 2023 DeepSeek Copyright (c) 2025 sgl-project Copyright (c) 2026 The DeepSpec Authors + Copyright (c) 2023-2026 The ggml authors Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/docs/source/deployment/3_unified_hf.rst b/docs/source/deployment/3_unified_hf.rst index ccef639d00e..0fae422e6db 100644 --- a/docs/source/deployment/3_unified_hf.rst +++ b/docs/source/deployment/3_unified_hf.rst @@ -50,6 +50,43 @@ The unified HF export API supports the following quantization formats: 4. NVFP4_AWQ - NVIDIA 4-bit floating point with AWQ optimization 5. INT4_AWQ - 4-bit integer with AWQ optimization 6. W4A8_AWQ - 4-bit weights and 8-bit activations with AWQ optimization +7. IQ1_S - 1-bit importance-aware quantization using the GGML block layout +8. IQ2_XS - 2-bit importance-aware quantization using the GGML block layout + +.. note:: + GGML has no equivalent for ModelOpt's per-tensor FP8 weight-and-activation format. In particular, + GGML does not define a first-class FP8 tensor type with the corresponding per-tensor weight and + activation scale semantics. Converting a ModelOpt FP8 checkpoint to GGUF therefore requires + conversion to another GGML-supported tensor type rather than a lossless FP8 encoding. + +IQ weight representation +~~~~~~~~~~~~~~~~~~~~~~~~ + +For IQ1_S and IQ2_XS, unified export replaces each floating-point ``.weight`` with a +``uint8`` tensor containing byte-exact GGML blocks. Its shape is +``[*logical_shape[:-1], ceil(logical_shape[-1] / 256), payload_bytes]``, where ``payload_bytes`` +is 50 for IQ1_S and 74 for IQ2_XS. Export right-pads each logical row with zeros to a multiple of +256 before packing, so blocks never cross row boundaries. Two ``int64`` sidecars preserve the +shape contract: + +* ``.weight_logical_shape`` records the original tensor shape; and +* ``.weight_padded_shape`` records the per-row padded shape represented by the payload. + +Loaders use the padded shape for block addressing and the logical shape for the matrix operation. +The padding does not add logical weights and is discarded by dequantization. + +Each 74-byte IQ2_XS block represents 256 logical weights: + +* bytes 0--1 are the little-endian FP16 super-block scale ``d``; +* bytes 2--65 are 32 little-endian ``uint16`` codes, one per group of eight weights. Each code + contains a 9-bit codebook index and seven stored sign bits; the eighth sign bit is derived from + parity; and +* bytes 66--73 contain sixteen 4-bit local-scale codes, packed two per byte. Each local scale is + shared by two adjacent eight-weight groups. + +The canonical 512-by-8 IQ2_XS codebook is part of the implementation rather than the checkpoint. +The complete block costs ``74 * 8 / 256 = 2.3125`` bits per represented weight before any final +row-padding overhead. Minimum Framework Versions -------------------------- diff --git a/examples/megatron_bridge/materialize_mixed_iq_gguf.py b/examples/megatron_bridge/materialize_mixed_iq_gguf.py new file mode 100644 index 00000000000..1fba0950291 --- /dev/null +++ b/examples/megatron_bridge/materialize_mixed_iq_gguf.py @@ -0,0 +1,326 @@ +# 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. + +"""Materialize a mixed ModelOpt checkpoint as GGUF without changing llama.cpp. + +The pinned stock converter owns model metadata, tensor naming, BF16 conversion, +and NVFP4 repacking. This validation bridge consumes only the already-packed +IQ2_XS tensors, writes their canonical bytes through stock ``gguf-py``, and then +verifies the resulting GGUF payloads byte for byte. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import re +import runpy +import subprocess +import sys +from collections import defaultdict +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +_IQ2_XS_BLOCK_BYTES = 74 +_IQ2_XS_BLOCK_SIZE = 256 +_ROUTED_EXPERT = re.compile( + r"^backbone\.layers\.(\d+)\.mixer\.experts\.(\d+)\.(up_proj|down_proj)\.weight$" +) + + +def packed_rows(tensor: torch.Tensor) -> tuple[np.ndarray, list[int]]: + """Collapse the explicit block/payload axes into GGUF's byte-row layout.""" + + if tensor.dtype != torch.uint8: + raise ValueError(f"IQ2_XS payload must use uint8, got {tensor.dtype}") + if tensor.ndim < 3 or tensor.shape[-1] != _IQ2_XS_BLOCK_BYTES: + raise ValueError( + f"IQ2_XS payload must have shape [..., blocks, {_IQ2_XS_BLOCK_BYTES}], " + f"got {tuple(tensor.shape)}" + ) + raw_shape = (*tensor.shape[:-2], tensor.shape[-2] * _IQ2_XS_BLOCK_BYTES) + raw = tensor.contiguous().reshape(raw_shape).cpu().numpy() + logical_shape = [*raw_shape[:-1], tensor.shape[-2] * _IQ2_XS_BLOCK_SIZE] + return raw, logical_shape + + +def _shape_sidecar_names(weight_name: str) -> tuple[str, str]: + base = weight_name.removesuffix(".weight") + return base + ".weight_logical_shape", base + ".weight_padded_shape" + + +def _git_commit(source: Path) -> str: + return subprocess.run( + ["git", "-C", str(source), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + + +def _normalize_layer_name(name: str) -> str: + if name.startswith("model.layers.") and ".mixer." in name: + return name.replace("model.layers.", "backbone.layers.", 1) + return name + + +def _load_iq2_layers(checkpoint: Path) -> set[str]: + config_path = checkpoint / "hf_quant_config.json" + config = json.loads(config_path.read_text()) + layers = config["quantization"]["quantized_layers"] + return { + _normalize_layer_name(name) + for name, layer_config in layers.items() + if layer_config.get("quant_algo") == "IQ2_XS" + } + + +def _matches_iq2_layer(name: str, layers: set[str]) -> bool: + return any( + name == f"{layer}.weight" or (name.startswith(f"{layer}.") and name.endswith(".weight")) + for layer in layers + ) + + +def _target_record( + raw: np.ndarray, + source_tensors: list[str], + logical_shape: list[int], + padded_shape: list[int], +) -> dict: + payload = raw.tobytes() + return { + "source_tensors": source_tensors, + "logical_shape": logical_shape, + "padded_shape": padded_shape, + "payload_bytes": len(payload), + "source_sha256": hashlib.sha256(payload).hexdigest(), + } + + +def _verify_gguf_payloads( + gguf: Any, path: Path, expected: dict[str, dict] +) -> tuple[list[dict], list[str]]: + tensors = {tensor.name: tensor for tensor in gguf.GGUFReader(path).tensors} + results = [] + errors = [] + for name, record in sorted(expected.items()): + tensor = tensors.get(name) + if tensor is None: + errors.append(f"Missing IQ2_XS tensor in GGUF: {name}") + continue + payload = tensor.data.tobytes() + actual_sha256 = hashlib.sha256(payload).hexdigest() + result = { + "gguf_tensor": name, + **record, + "tensor_type": tensor.tensor_type.name, + "gguf_sha256": actual_sha256, + "byte_identical": actual_sha256 == record["source_sha256"], + } + results.append(result) + if tensor.tensor_type.name != "IQ2_XS": + errors.append(f"Tensor {name} has type {tensor.tensor_type.name}, expected IQ2_XS") + if not result["byte_identical"]: + errors.append(f"IQ2_XS payload differs for {name}") + return results, errors + + +def materialize(checkpoint: Path, llama_source: Path, outfile: Path) -> tuple[dict[str, dict], str]: + """Run the stock converter with an in-process IQ2_XS payload adapter.""" + + checkpoint = checkpoint.resolve() + llama_source = llama_source.resolve() + sys.path.insert(0, str(llama_source)) + sys.path.insert(1, str(llama_source / "gguf-py")) + + import gguf + from conversion.base import LazyTorchTensor, ModelBase + from conversion.nemotron import NemotronHModel + + expected: dict[str, dict] = {} + + @ModelBase.register("NemotronHForCausalLM") + class ModelOptIQNemotronHModel(NemotronHModel): + model_arch = NemotronHModel.model_arch + + def _pop_iq2_shapes(self, source_name: str) -> tuple[list[int], list[int]]: + logical_name, padded_name = _shape_sidecar_names(source_name) + try: + logical = LazyTorchTensor.to_eager(self.model_tensors.pop(logical_name)()) + padded = LazyTorchTensor.to_eager(self.model_tensors.pop(padded_name)()) + except KeyError as error: + raise ValueError(f"Missing IQ2_XS shape sidecar for {source_name}") from error + return logical.tolist(), padded.tolist() + + def _write_iq2_tensor(self, source_name: str, target_name: str) -> None: + generator = self.model_tensors.pop(source_name) + tensor = LazyTorchTensor.to_eager(generator()) + raw, inferred_padded_shape = packed_rows(tensor) + logical_shape, padded_shape = self._pop_iq2_shapes(source_name) + if padded_shape != inferred_padded_shape: + raise ValueError( + f"IQ2_XS padded shape mismatch for {source_name}: metadata " + f"{padded_shape}, payload {inferred_padded_shape}" + ) + self.gguf_writer.add_tensor( + target_name, raw, raw_dtype=gguf.GGMLQuantizationType.IQ2_XS + ) + expected[target_name] = _target_record(raw, [source_name], logical_shape, padded_shape) + + def _write_routed_iq2_tensors(self, names: list[str]) -> set[str]: + grouped: dict[tuple[int, str], list[tuple[int, str]]] = defaultdict(list) + for name in names: + match = _ROUTED_EXPERT.match(name) + if match: + layer, expert, projection = match.groups() + grouped[(int(layer), projection)].append((int(expert), name)) + + consumed = set() + expected_experts = int(self.hparams["n_routed_experts"]) + for (layer, projection), experts in sorted(grouped.items()): + experts.sort() + expert_ids = [expert for expert, _ in experts] + if expert_ids != list(range(expected_experts)): + raise ValueError( + f"Layer {layer} {projection} has expert IDs {expert_ids}, " + f"expected 0..{expected_experts - 1}" + ) + source_names = [name for _, name in experts] + tensors = [ + LazyTorchTensor.to_eager(self.model_tensors.pop(name)()) + for name in source_names + ] + raw_parts = [packed_rows(tensor)[0] for tensor in tensors] + source_shapes = [self._pop_iq2_shapes(name) for name in source_names] + if any(shapes != source_shapes[0] for shapes in source_shapes[1:]): + raise ValueError( + f"Layer {layer} {projection} has inconsistent expert shape metadata" + ) + raw = np.stack(raw_parts, axis=0) + logical_shape = [ + len(tensors), + *source_shapes[0][0], + ] + padded_shape = [len(tensors), *source_shapes[0][1]] + inferred_padded_shape = [len(tensors), *packed_rows(tensors[0])[1]] + if padded_shape != inferred_padded_shape: + raise ValueError( + f"Layer {layer} {projection} padded shape metadata {padded_shape} " + f"does not match payload {inferred_padded_shape}" + ) + merged_name = f"model.layers.{layer}.mlp.experts.{projection}.weight" + target_name = self.map_tensor_name(merged_name) + self.gguf_writer.add_tensor( + target_name, raw, raw_dtype=gguf.GGMLQuantizationType.IQ2_XS + ) + expected[target_name] = _target_record( + raw, source_names, logical_shape, padded_shape + ) + consumed.update(source_names) + return consumed + + def _write_iq2_tensors(self) -> None: + layers = _load_iq2_layers(checkpoint) + names = sorted(name for name in self.model_tensors if _matches_iq2_layer(name, layers)) + if not names: + raise ValueError("No IQ2_XS tensors were found in the checkpoint") + + consumed = self._write_routed_iq2_tensors(names) + for name in names: + if name in consumed: + continue + target_name = self.map_tensor_name(name) + self._write_iq2_tensor(name, target_name) + + def prepare_tensors(self) -> None: + self._write_iq2_tensors() + super().prepare_tensors() + + original_argv = sys.argv + try: + sys.argv = [ + str(llama_source / "convert_hf_to_gguf.py"), + str(checkpoint), + "--outfile", + str(outfile.resolve()), + "--outtype", + "bf16", + "--use-temp-file", + "--no-mtp", + ] + runpy.run_path(str(llama_source / "convert_hf_to_gguf.py"), run_name="__main__") + finally: + sys.argv = original_argv + return expected, _git_commit(llama_source) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--llama-source", type=Path, required=True) + parser.add_argument("--outfile", type=Path, required=True) + parser.add_argument("--report", type=Path, required=True) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + expected, commit = materialize(args.checkpoint, args.llama_source, args.outfile) + + sys.path.insert(0, str(args.llama_source.resolve() / "gguf-py")) + import gguf + + payloads, errors = _verify_gguf_payloads(gguf, args.outfile, expected) + source_tensors = sum(len(payload["source_tensors"]) for payload in payloads) + summary = { + "iq2_xs_source_tensors": source_tensors, + "iq2_xs_gguf_tensors": len(payloads), + "iq2_xs_payload_bytes": sum(payload["payload_bytes"] for payload in payloads), + "payload_mismatches": sum(not payload["byte_identical"] for payload in payloads), + } + report: dict[str, Any] = { + "schema_version": 2, + "created_at": datetime.now(UTC).isoformat(), + "status": "passed" if not errors else "failed", + "checkpoint": str(args.checkpoint.resolve()), + "gguf": str(args.outfile.resolve()), + "llama_cpp": { + "source": str(args.llama_source.resolve()), + "commit": commit, + "source_modified": subprocess.run( + ["git", "-C", str(args.llama_source), "diff", "--quiet"], check=False + ).returncode + != 0, + }, + "summary": summary, + "iq2_xs_payloads": payloads, + "errors": errors, + } + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + print(json.dumps({"status": report["status"], **summary}, indent=2)) + print(f"Validation report: {args.report}") + if errors: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/megatron_bridge/validate_ggml_row_alignment.py b/examples/megatron_bridge/validate_ggml_row_alignment.py new file mode 100644 index 00000000000..85c5a2123ce --- /dev/null +++ b/examples/megatron_bridge/validate_ggml_row_alignment.py @@ -0,0 +1,181 @@ +# 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. + +"""Validate GGML block alignment from safetensors metadata. + +GGML block-quantized tensors require every logical row to contain an integral +number of format blocks. This preflight reads only safetensors headers and can +therefore reject an incompatible tensor policy before loading a model or +allocating a GPU. +""" + +from __future__ import annotations + +import argparse +import fnmatch +import json +import struct +from collections import Counter +from pathlib import Path +from typing import Any + +_INDEX_NAME = "model.safetensors.index.json" + + +def _read_safetensors_header(path: Path) -> dict[str, Any]: + with path.open("rb") as file: + header_length_raw = file.read(8) + if len(header_length_raw) != 8: + raise ValueError(f"Invalid safetensors header in {path}") + header_length = struct.unpack(" dict[str, dict[str, Any]]: + index_path = checkpoint / _INDEX_NAME + if index_path.is_file(): + weight_map = json.loads(index_path.read_text())["weight_map"] + headers: dict[str, dict[str, Any]] = {} + tensors = {} + for name, shard_name in weight_map.items(): + if shard_name not in headers: + headers[shard_name] = _read_safetensors_header(checkpoint / shard_name) + tensors[name] = headers[shard_name][name] + return tensors + + tensors = {} + for shard_path in sorted(checkpoint.glob("*.safetensors")): + header = _read_safetensors_header(shard_path) + for name, metadata in header.items(): + if name == "__metadata__": + continue + if name in tensors: + raise ValueError(f"Duplicate tensor {name!r} without {_INDEX_NAME}") + tensors[name] = metadata + if not tensors: + raise FileNotFoundError(f"No safetensors files found in {checkpoint}") + return tensors + + +def _matches(name: str, includes: list[str], excludes: list[str]) -> bool: + return any(fnmatch.fnmatchcase(name, pattern) for pattern in includes) and not any( + fnmatch.fnmatchcase(name, pattern) for pattern in excludes + ) + + +def validate_row_alignment( + checkpoint: Path, + *, + block_size: int, + includes: list[str], + excludes: list[str], +) -> dict[str, Any]: + """Return a machine-readable GGML row-alignment report.""" + if block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}") + if not includes: + raise ValueError("At least one include pattern is required") + + selected = [] + shape_groups: Counter[tuple[tuple[int, ...], int, bool]] = Counter() + for name, metadata in _tensor_metadata(checkpoint).items(): + if not _matches(name, includes, excludes): + continue + shape = tuple(int(value) for value in metadata["shape"]) + if not shape: + raise ValueError(f"Selected tensor {name!r} has scalar shape") + remainder = shape[-1] % block_size + compatible = remainder == 0 + selected.append( + { + "name": name, + "shape": list(shape), + "row_width": shape[-1], + "remainder": remainder, + "compatible": compatible, + } + ) + shape_groups[(shape, remainder, compatible)] += 1 + + if not selected: + raise ValueError(f"No tensors matched include patterns: {includes}") + + incompatible = [tensor for tensor in selected if not tensor["compatible"]] + groups = [ + { + "shape": list(shape), + "row_width": shape[-1], + "remainder": remainder, + "compatible": compatible, + "tensor_count": count, + } + for (shape, remainder, compatible), count in sorted(shape_groups.items()) + ] + return { + "schema_version": 1, + "checkpoint": str(checkpoint), + "block_size": block_size, + "include_patterns": includes, + "exclude_patterns": excludes, + "selected_tensors": len(selected), + "compatible_tensors": len(selected) - len(incompatible), + "incompatible_tensors": len(incompatible), + "shape_groups": groups, + "incompatible": incompatible, + "status": "passed" if not incompatible else "failed", + } + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--block-size", type=int, required=True) + parser.add_argument( + "--include", + action="append", + required=True, + help="fnmatch pattern selecting tensor names; may be repeated", + ) + parser.add_argument( + "--exclude", + action="append", + default=[], + help="fnmatch pattern excluding tensor names; may be repeated", + ) + parser.add_argument("--report", type=Path) + return parser.parse_args() + + +def main() -> int: + args = _parse_args() + report = validate_row_alignment( + args.checkpoint, + block_size=args.block_size, + includes=args.include, + excludes=args.exclude, + ) + rendered = json.dumps(report, indent=2, sort_keys=True) + "\n" + if args.report: + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text(rendered) + print(rendered, end="") + return 0 if report["status"] == "passed" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/megatron_bridge/validate_iq2_xs_reconstruction.py b/examples/megatron_bridge/validate_iq2_xs_reconstruction.py new file mode 100644 index 00000000000..ba293c9e34e --- /dev/null +++ b/examples/megatron_bridge/validate_iq2_xs_reconstruction.py @@ -0,0 +1,371 @@ +# 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. + +"""Compare packed IQ2_XS rows with their source checkpoint tensors. + +This validator closes a gap left by structural and decoder-oracle checks: a byte-valid block can +still belong to the wrong source row. It samples logical rows from a packed unified Hugging Face +checkpoint, reconstructs them, and compares them with the same rows in the floating-point source +checkpoint. Optionally, it also requantizes every sampled source row and requires byte-identical +payloads. The report separates complete 256-value blocks from the final partial block so row-padding +regressions are visible. +""" + +from __future__ import annotations + +import argparse +import fnmatch +import json +import math +from contextlib import ExitStack +from dataclasses import asdict, dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +import torch +from safetensors import safe_open +from validate_mixed_quantized_hf import TensorRecord, _read_int64_vector, read_checkpoint_index + +from modelopt.torch.quantization.ggml import dequantize_iq2_xs, quantize_iq2_xs + +_BLOCK_SIZE = 256 +_PAYLOAD_BYTES = 74 + + +@dataclass(frozen=True) +class ErrorMetrics: + """Numerical reconstruction metrics for one vector region.""" + + values: int + sum_squared_error: float + normalized_squared_error: float | None + cosine_similarity: float | None + max_absolute_error: float + + +def _metrics(reference: torch.Tensor, reconstructed: torch.Tensor) -> ErrorMetrics | None: + if reference.numel() == 0: + return None + reference = reference.float() + reconstructed = reconstructed.float() + error = reconstructed - reference + squared_error = error.square().sum().item() + signal_energy = reference.square().sum().item() + reference_norm = reference.norm().item() + reconstructed_norm = reconstructed.norm().item() + return ErrorMetrics( + values=reference.numel(), + sum_squared_error=squared_error, + normalized_squared_error=squared_error / signal_energy if signal_energy else None, + cosine_similarity=( + torch.dot(reference, reconstructed).item() / (reference_norm * reconstructed_norm) + if reference_norm and reconstructed_norm + else None + ), + max_absolute_error=error.abs().max().item(), + ) + + +def _row_coordinates(shape: tuple[int, ...], flat_row: int) -> tuple[int, ...]: + coordinates = [] + for dimension in reversed(shape[:-1]): + coordinates.append(flat_row % dimension) + flat_row //= dimension + return tuple(reversed(coordinates)) + + +def _sample_indices(count: int, samples: int) -> list[int]: + if count <= 0 or samples <= 0: + return [] + if samples >= count: + return list(range(count)) + if samples == 1: + return [0] + return sorted({round(index * (count - 1) / (samples - 1)) for index in range(samples)}) + + +def _select_tensors(names: list[str], patterns: tuple[str, ...], maximum: int | None) -> list[str]: + selected = [ + name + for name in sorted(names) + if not patterns or any(fnmatch.fnmatchcase(name, pattern) for pattern in patterns) + ] + if maximum is None or maximum >= len(selected): + return selected + return [selected[index] for index in _sample_indices(len(selected), maximum)] + + +class _CheckpointReader: + """Keep only the required safetensors shards open while reading row slices.""" + + def __init__(self, root: Path, records: dict[str, TensorRecord], stack: ExitStack): + self.root = root + self.records = records + self.stack = stack + self.handles: dict[str, Any] = {} + + def row(self, name: str, flat_row: int, *, trailing_dimensions: int = 1) -> torch.Tensor: + record = self.records[name] + if record.shard not in self.handles: + self.handles[record.shard] = self.stack.enter_context( + safe_open(str(self.root / record.shard), framework="pt", device="cpu") + ) + logical_row_shape = (*record.shape[:-trailing_dimensions], 1) + coordinates = _row_coordinates(logical_row_shape, flat_row) + trailing_slices = (slice(None),) * trailing_dimensions + return self.handles[record.shard].get_slice(name)[(*coordinates, *trailing_slices)] + + +def _iq2_tensor_names( + checkpoint: Path, records: dict[str, TensorRecord] +) -> dict[str, tuple[tuple[int, ...], tuple[int, ...]]]: + shapes = {} + for name, record in records.items(): + if record.dtype != "U8" or not name.endswith(".weight"): + continue + if len(record.shape) < 2 or record.shape[-1] != _PAYLOAD_BYTES: + continue + base = name.removesuffix(".weight") + logical_record = records.get(base + ".weight_logical_shape") + padded_record = records.get(base + ".weight_padded_shape") + if logical_record is None or padded_record is None: + continue + logical_shape = _read_int64_vector(checkpoint, logical_record) + padded_shape = _read_int64_vector(checkpoint, padded_record) + expected_padded_width = math.ceil(logical_shape[-1] / _BLOCK_SIZE) * _BLOCK_SIZE + expected_packed_shape = (*logical_shape[:-1], expected_padded_width // _BLOCK_SIZE, 74) + if padded_shape != (*logical_shape[:-1], expected_padded_width): + raise ValueError(f"{name} has invalid padded shape {padded_shape}") + if record.shape != expected_packed_shape: + raise ValueError(f"{name} has invalid packed shape {record.shape}") + shapes[name] = (logical_shape, padded_shape) + if not shapes: + raise ValueError(f"No packed IQ2_XS tensors found in {checkpoint}") + return shapes + + +def _aggregate_metrics(samples: list[dict], region: str) -> dict: + metrics = [sample[region] for sample in samples if sample[region] is not None] + values = sum(metric["values"] for metric in metrics) + squared_error = sum(metric["sum_squared_error"] for metric in metrics) + signal_energy = sum(metric["signal_energy"] for metric in metrics) + reconstructed_energy = sum(metric["reconstructed_energy"] for metric in metrics) + dot_product = sum(metric["dot_product"] for metric in metrics) + return { + "values": values, + "sum_squared_error": squared_error, + "normalized_squared_error": squared_error / signal_energy if signal_energy else None, + "cosine_similarity": ( + dot_product / math.sqrt(signal_energy * reconstructed_energy) + if signal_energy and reconstructed_energy + else None + ), + "max_absolute_error": max( + (metric["max_absolute_error"] for metric in metrics), default=0.0 + ), + } + + +def _metric_report(reference: torch.Tensor, reconstructed: torch.Tensor) -> dict | None: + metrics = _metrics(reference, reconstructed) + if metrics is None: + return None + return { + **asdict(metrics), + "signal_energy": reference.float().square().sum().item(), + "reconstructed_energy": reconstructed.float().square().sum().item(), + "dot_product": torch.dot(reference.float(), reconstructed.float()).item(), + } + + +@torch.no_grad() +def validate_reconstruction( + checkpoint: Path, + reference_checkpoint: Path, + *, + rows_per_tensor: int = 1, + maximum_tensors: int | None = 32, + tensor_patterns: tuple[str, ...] = (), + require_repack_match: bool = False, + device: str = "cpu", + maximum_normalized_error: float | None = None, +) -> dict: + """Sample packed rows, compare reconstruction, and return a JSON-serializable report.""" + + checkpoint = checkpoint.resolve() + reference_checkpoint = reference_checkpoint.resolve() + records = read_checkpoint_index(checkpoint) + reference_records = read_checkpoint_index(reference_checkpoint) + shapes = _iq2_tensor_names(checkpoint, records) + tensor_names = _select_tensors(list(shapes), tensor_patterns, maximum_tensors) + if not tensor_names: + raise ValueError(f"No IQ2_XS tensor names match {list(tensor_patterns)}") + resolved_device = torch.device(device) + if resolved_device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA validation was requested, but CUDA is not available") + + errors: list[str] = [] + samples: list[dict] = [] + with ExitStack() as stack: + packed_reader = _CheckpointReader(checkpoint, records, stack) + reference_reader = _CheckpointReader(reference_checkpoint, reference_records, stack) + for name in tensor_names: + logical_shape, padded_shape = shapes[name] + reference_record = reference_records.get(name) + if reference_record is None: + errors.append(f"{name} is absent from the reference checkpoint") + continue + if reference_record.shape != logical_shape: + errors.append( + f"{name} source shape {reference_record.shape} does not match {logical_shape}" + ) + continue + + row_count = math.prod(logical_shape[:-1]) + for flat_row in _sample_indices(row_count, rows_per_tensor): + coordinates = _row_coordinates(logical_shape, flat_row) + packed = packed_reader.row(name, flat_row, trailing_dimensions=2).to( + device=resolved_device + ) + reference = reference_reader.row(name, flat_row).to( + device=resolved_device, dtype=torch.float32 + ) + row_shape = torch.tensor([logical_shape[-1]], device=resolved_device) + reconstructed = dequantize_iq2_xs(packed, row_shape, dtype=torch.float32) + complete_values = logical_shape[-1] // _BLOCK_SIZE * _BLOCK_SIZE + + repack_match = None + mismatched_bytes = None + if require_repack_match: + repacked, _ = quantize_iq2_xs(reference) + difference = repacked != packed + mismatched_bytes = int(difference.sum().item()) + repack_match = not bool(mismatched_bytes) + if not repack_match: + errors.append( + f"{name} row {coordinates} differs from direct packing in " + f"{mismatched_bytes} bytes" + ) + + sample: dict[str, Any] = { + "tensor": name, + "row": list(coordinates), + "logical_width": logical_shape[-1], + "padded_width": padded_shape[-1], + "repack_match": repack_match, + "mismatched_bytes": mismatched_bytes, + "all_values": _metric_report(reference, reconstructed), + "complete_blocks": _metric_report( + reference[:complete_values], reconstructed[:complete_values] + ), + "partial_tail": _metric_report( + reference[complete_values:], reconstructed[complete_values:] + ), + } + samples.append(sample) + normalized_error = sample["all_values"]["normalized_squared_error"] + if ( + maximum_normalized_error is not None + and normalized_error is not None + and normalized_error > maximum_normalized_error + ): + errors.append( + f"{name} row {coordinates} normalized error {normalized_error:.6g} " + f"exceeds {maximum_normalized_error:.6g}" + ) + + summary = { + "available_iq2_xs_tensors": len(shapes), + "sampled_tensors": len({sample["tensor"] for sample in samples}), + "sampled_rows": len(samples), + "repack_matches": sum(sample["repack_match"] is True for sample in samples), + "repack_mismatches": sum(sample["repack_match"] is False for sample in samples), + "all_values": _aggregate_metrics(samples, "all_values"), + "complete_blocks": _aggregate_metrics(samples, "complete_blocks"), + "partial_tails": _aggregate_metrics(samples, "partial_tail"), + } + return { + "schema_version": 1, + "created_at": datetime.now(UTC).isoformat(), + "status": "passed" if not errors else "failed", + "checkpoint": str(checkpoint), + "reference_checkpoint": str(reference_checkpoint), + "device": str(resolved_device), + "settings": { + "rows_per_tensor": rows_per_tensor, + "maximum_tensors": maximum_tensors, + "tensor_patterns": list(tensor_patterns), + "require_repack_match": require_repack_match, + "maximum_normalized_error": maximum_normalized_error, + }, + "summary": summary, + "samples": samples, + "errors": errors, + } + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--reference-checkpoint", type=Path, required=True) + parser.add_argument("--report", type=Path, required=True) + parser.add_argument("--rows-per-tensor", type=int, default=1) + parser.add_argument( + "--max-tensors", + type=int, + default=32, + help="Sample tensors evenly from the sorted matching names; use 0 for all tensors.", + ) + parser.add_argument( + "--tensor-pattern", + action="append", + default=[], + help="Optional shell-style tensor-name pattern. May be specified more than once.", + ) + parser.add_argument("--require-repack-match", action="store_true") + parser.add_argument("--device", default="cpu") + parser.add_argument("--max-normalized-error", type=float) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + if args.rows_per_tensor <= 0: + raise ValueError("--rows-per-tensor must be positive") + if args.max_tensors < 0: + raise ValueError("--max-tensors cannot be negative") + report = validate_reconstruction( + args.checkpoint, + args.reference_checkpoint, + rows_per_tensor=args.rows_per_tensor, + maximum_tensors=args.max_tensors or None, + tensor_patterns=tuple(args.tensor_pattern), + require_repack_match=args.require_repack_match, + device=args.device, + maximum_normalized_error=args.max_normalized_error, + ) + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + print(json.dumps({"status": report["status"], **report["summary"]}, indent=2)) + print(f"Validation report: {args.report}") + for error in report["errors"]: + print(f"ERROR: {error}") + if report["errors"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/megatron_bridge/validate_iq2_xs_stock_ggml.py b/examples/megatron_bridge/validate_iq2_xs_stock_ggml.py new file mode 100644 index 00000000000..42d20fc2443 --- /dev/null +++ b/examples/megatron_bridge/validate_iq2_xs_stock_ggml.py @@ -0,0 +1,296 @@ +# 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. + +"""Compare exported IQ2_XS blocks with an unmodified stock GGML decoder. + +The script samples complete 74-byte blocks from every IQ2_XS tensor in a unified +Hugging Face checkpoint. It sends those exact bytes to GGML's exported +``dequantize_row_iq2_xs`` function and to ModelOpt's decoder, then requires +bit-identical FP32 reconstruction. The report records the pinned llama.cpp commit, +sample locations, encoded-field digests, and mismatch counts. +""" + +from __future__ import annotations + +import argparse +import ctypes +import hashlib +import json +import subprocess +import sys +from collections import defaultdict +from datetime import UTC, datetime +from pathlib import Path + +import numpy as np +import torch +from validate_mixed_quantized_hf import ( + _load_quantization_config, + _matching_quantized_layer, + read_checkpoint_index, +) + +from modelopt.torch.quantization.ggml import dequantize_iq2_xs + +_BLOCK_SIZE = 256 +_BLOCK_BYTES = 74 + + +def _llama_commit(source: Path) -> str: + return subprocess.run( + ["git", "-C", str(source), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + + +def _iq2_tensor_names(checkpoint: Path) -> list[str]: + records = read_checkpoint_index(checkpoint) + quantization, _ = _load_quantization_config(checkpoint) + quantized_layers = quantization.get("quantized_layers", {}) + names = [] + for name, record in records.items(): + layer = _matching_quantized_layer(name, quantized_layers) + if layer is None or quantized_layers[layer].get("quant_algo") != "IQ2_XS": + continue + if record.dtype != "U8" or record.payload_bytes % _BLOCK_BYTES: + raise ValueError(f"Invalid IQ2_XS payload for {name}: {record}") + names.append(name) + if not names: + raise ValueError(f"No IQ2_XS tensors found in {checkpoint}") + return sorted(names) + + +def _sample_indices(num_blocks: int, blocks_per_tensor: int) -> list[int]: + if num_blocks <= 0: + raise ValueError("An IQ2_XS tensor must contain at least one block") + if blocks_per_tensor <= 0: + raise ValueError("blocks_per_tensor must be positive") + if blocks_per_tensor == 1: + return [0] + if num_blocks <= blocks_per_tensor: + return list(range(num_blocks)) + return sorted( + {round(i * (num_blocks - 1) / (blocks_per_tensor - 1)) for i in range(blocks_per_tensor)} + ) + + +def _read_samples(checkpoint: Path, blocks_per_tensor: int) -> tuple[bytes, list[dict]]: + records = read_checkpoint_index(checkpoint) + chunks: list[bytes] = [] + samples: list[dict] = [] + open_shard: str | None = None + file = None + try: + for name in _iq2_tensor_names(checkpoint): + record = records[name] + num_blocks = record.payload_bytes // _BLOCK_BYTES + indices = _sample_indices(num_blocks, blocks_per_tensor) + if record.shard != open_shard: + if file is not None: + file.close() + file = (checkpoint / record.shard).open("rb") + open_shard = record.shard + assert file is not None + for block_index in indices: + file.seek(record.payload_offset + block_index * _BLOCK_BYTES) + block = file.read(_BLOCK_BYTES) + if len(block) != _BLOCK_BYTES: + raise EOFError(f"Truncated IQ2_XS block {block_index} in {name}") + chunks.append(block) + samples.append( + { + "tensor": name, + "shard": record.shard, + "block_index": block_index, + "sha256": hashlib.sha256(block).hexdigest(), + } + ) + finally: + if file is not None: + file.close() + return b"".join(chunks), samples + + +def _stock_decode(library: Path, packed: bytes) -> np.ndarray: + if len(packed) % _BLOCK_BYTES: + raise ValueError(f"Packed byte count {len(packed)} is not divisible by {_BLOCK_BYTES}") + num_values = len(packed) // _BLOCK_BYTES * _BLOCK_SIZE + ggml = ctypes.CDLL(str(library.resolve())) + decode = ggml.dequantize_row_iq2_xs + decode.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_float), ctypes.c_int64] + decode.restype = None + source = (ctypes.c_uint8 * len(packed)).from_buffer_copy(packed) + output = np.empty(num_values, dtype=np.float32) + decode(source, output.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), num_values) + return output + + +def _write_test_gguf( + output_path: Path, llama_source: Path, packed_bytes: bytes, samples: list[dict] +) -> dict: + """Round-trip sampled payloads through the stock GGUF writer and reader.""" + + gguf_package = str(llama_source.resolve() / "gguf-py") + if gguf_package not in sys.path: + sys.path.insert(0, gguf_package) + import gguf + + source_payloads: dict[str, bytearray] = defaultdict(bytearray) + for offset, sample in enumerate(samples): + start = offset * _BLOCK_BYTES + source_payloads[sample["tensor"]].extend(packed_bytes[start : start + _BLOCK_BYTES]) + + output_path.parent.mkdir(parents=True, exist_ok=True) + writer = gguf.GGUFWriter(output_path, "llama") + writer.add_name("ModelOpt IQ2_XS compatibility sample") + tensor_sources = {} + expected_by_name = {} + for index, (source_name, payload) in enumerate(sorted(source_payloads.items())): + gguf_name = f"iq2_xs_sample_{index:05d}" + raw = np.frombuffer(payload, dtype=np.uint8).reshape(1, -1) + writer.add_tensor(gguf_name, raw, raw_dtype=gguf.GGMLQuantizationType.IQ2_XS) + tensor_sources[gguf_name] = source_name + expected_by_name[gguf_name] = bytes(payload) + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + payloads = [] + for tensor in gguf.GGUFReader(output_path).tensors: + expected = expected_by_name[tensor.name] + actual = tensor.data.tobytes() + payloads.append( + { + "gguf_tensor": tensor.name, + "source_tensor": tensor_sources[tensor.name], + "tensor_type": tensor.tensor_type.name, + "payload_bytes": len(actual), + "source_sha256": hashlib.sha256(expected).hexdigest(), + "gguf_sha256": hashlib.sha256(actual).hexdigest(), + "byte_identical": actual == expected, + } + ) + mismatches = sum(not payload["byte_identical"] for payload in payloads) + return { + "path": str(output_path.resolve()), + "payloads": payloads, + "payload_mismatches": mismatches, + } + + +def compare_checkpoint( + checkpoint: Path, + ggml_library: Path, + llama_source: Path, + *, + blocks_per_tensor: int = 3, + test_gguf: Path | None = None, +) -> dict: + """Return the stock-GGML compatibility report for one checkpoint.""" + + checkpoint = checkpoint.resolve() + packed_bytes, samples = _read_samples(checkpoint, blocks_per_tensor) + num_blocks = len(packed_bytes) // _BLOCK_BYTES + + packed = torch.frombuffer(bytearray(packed_bytes), dtype=torch.uint8).reshape( + num_blocks, 1, _BLOCK_BYTES + ) + logical_shape = torch.tensor([num_blocks, _BLOCK_SIZE], dtype=torch.int64) + modelopt = dequantize_iq2_xs(packed, logical_shape, dtype=torch.float32).numpy().reshape(-1) + stock = _stock_decode(ggml_library, packed_bytes) + modelopt_bits = modelopt.view(np.uint32) + stock_bits = stock.view(np.uint32) + different = modelopt_bits != stock_bits + + fields = np.frombuffer(packed_bytes, dtype=np.uint8).reshape(num_blocks, _BLOCK_BYTES) + scale_bytes = fields[:, :2] + vector_words = fields[:, 2:66] + local_scale_bytes = fields[:, 66:74] + max_abs_difference = float(np.max(np.abs(modelopt - stock))) if modelopt.size else 0.0 + gguf_report = ( + _write_test_gguf(test_gguf, llama_source, packed_bytes, samples) + if test_gguf is not None + else None + ) + passed = not np.any(different) and ( + gguf_report is None or gguf_report["payload_mismatches"] == 0 + ) + report = { + "schema_version": 1, + "created_at": datetime.now(UTC).isoformat(), + "status": "passed" if passed else "failed", + "checkpoint": str(checkpoint), + "llama_cpp": { + "source": str(llama_source.resolve()), + "commit": _llama_commit(llama_source), + "library": str(ggml_library.resolve()), + "symbol": "dequantize_row_iq2_xs", + }, + "summary": { + "iq2_xs_tensors": len({sample["tensor"] for sample in samples}), + "sampled_blocks": num_blocks, + "decoded_values": int(modelopt.size), + "bitwise_differences": int(np.count_nonzero(different)), + "max_abs_difference": max_abs_difference, + }, + "encoded_fields": { + "global_fp16_scale_bytes_sha256": hashlib.sha256(scale_bytes.tobytes()).hexdigest(), + "vector_words_sha256": hashlib.sha256(vector_words.tobytes()).hexdigest(), + "local_scale_bytes_sha256": hashlib.sha256(local_scale_bytes.tobytes()).hexdigest(), + "sample_payload_sha256": hashlib.sha256(packed_bytes).hexdigest(), + }, + "test_gguf": gguf_report, + "samples": samples, + } + return report + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--ggml-library", type=Path, required=True) + parser.add_argument("--llama-source", type=Path, required=True) + parser.add_argument("--report", type=Path, required=True) + parser.add_argument("--blocks-per-tensor", type=int, default=3) + parser.add_argument( + "--test-gguf", + type=Path, + help="Write sampled blocks with stock gguf-py and verify every payload digest.", + ) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + report = compare_checkpoint( + args.checkpoint, + args.ggml_library, + args.llama_source, + blocks_per_tensor=args.blocks_per_tensor, + test_gguf=args.test_gguf, + ) + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + print(json.dumps({"status": report["status"], **report["summary"]}, indent=2)) + print(f"Validation report: {args.report}") + if report["status"] != "passed": + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/megatron_bridge/validate_iq_cuda_parity.py b/examples/megatron_bridge/validate_iq_cuda_parity.py new file mode 100644 index 00000000000..2e9c6f7f334 --- /dev/null +++ b/examples/megatron_bridge/validate_iq_cuda_parity.py @@ -0,0 +1,98 @@ +# 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. + +"""Require byte-identical CPU and CUDA GGML IQ packing on deterministic inputs.""" + +from __future__ import annotations + +import argparse +import json +from datetime import UTC, datetime +from pathlib import Path + +import torch + +from modelopt.torch.quantization.extensions import get_cuda_ext_iq1_s, get_cuda_ext_iq2_xs +from modelopt.torch.quantization.ggml.iq1_s import quantize_iq1_s +from modelopt.torch.quantization.ggml.iq2_xs import quantize_iq2_xs + + +def _compare_format(name: str, weight: torch.Tensor, quantize, get_extension) -> dict: + extension = get_extension() + if extension is None: + raise RuntimeError(f"The CUDA extension for {name} is unavailable") + + packed_cpu, shape_cpu = quantize(weight) + packed_cuda, shape_cuda = quantize(weight.cuda()) + packed_cuda_cpu = packed_cuda.cpu() + differing_bytes = int(torch.count_nonzero(packed_cpu != packed_cuda_cpu).item()) + result = { + "format": name, + "logical_shape": list(weight.shape), + "padded_shape": [*weight.shape[:-1], packed_cpu.shape[-2] * 256], + "packed_shape": list(packed_cpu.shape), + "payload_bytes": packed_cpu.numel(), + "differing_payload_bytes": differing_bytes, + "logical_shape_matches": torch.equal(shape_cpu, shape_cuda.cpu()), + "cuda_extension_loaded": True, + } + if differing_bytes or not result["logical_shape_matches"]: + raise AssertionError(f"CPU/CUDA parity failed for {name}: {result}") + return result + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--report", type=Path, required=True) + parser.add_argument("--rows", type=int, default=64) + parser.add_argument("--row-width", type=int, default=257) + parser.add_argument("--seed", type=int, default=5918) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if args.rows <= 0: + raise ValueError("--rows must be positive") + if args.row_width <= 0: + raise ValueError("--row-width must be positive") + + generator = torch.Generator().manual_seed(args.seed) + weight = torch.randn((args.rows, args.row_width), generator=generator, dtype=torch.bfloat16) + results = [ + _compare_format("IQ1_S", weight, quantize_iq1_s, get_cuda_ext_iq1_s), + _compare_format("IQ2_XS", weight, quantize_iq2_xs, get_cuda_ext_iq2_xs), + ] + report = { + "schema_version": 1, + "created_at": datetime.now(UTC).isoformat(), + "status": "passed", + "seed": args.seed, + "rows": args.rows, + "row_width": args.row_width, + "device": torch.cuda.get_device_name(), + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "results": results, + } + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/examples/megatron_bridge/validate_mixed_quantized_hf.py b/examples/megatron_bridge/validate_mixed_quantized_hf.py new file mode 100644 index 00000000000..78f91b94943 --- /dev/null +++ b/examples/megatron_bridge/validate_mixed_quantized_hf.py @@ -0,0 +1,466 @@ +# 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. + +"""Validate a mixed-precision unified Hugging Face checkpoint without loading its tensors. + +The validator reads safetensors headers and streams packed payload bytes directly from each +shard. For the Nemotron 3.5 Lightning policy it checks that expert weights use canonical +IQ2_XS blocks, Mamba projections use NVFP4, and all remaining weights stay unquantized. +It also writes stable per-tensor and aggregate digests that can be compared with payloads +extracted from another GGML-compatible container. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import struct +from dataclasses import asdict, dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import BinaryIO + +_IQ2_XS_BLOCK_SIZE = 256 +_IQ2_XS_PAYLOAD_BYTES = 74 +_HASH_CHUNK_BYTES = 8 * 1024 * 1024 + + +@dataclass(frozen=True) +class TensorRecord: + """Location and metadata for one tensor payload in a safetensors shard.""" + + name: str + dtype: str + shape: tuple[int, ...] + shard: str + payload_offset: int + payload_bytes: int + + +def _read_shard_header(checkpoint: Path, shard_name: str) -> dict[str, TensorRecord]: + shard_path = checkpoint / shard_name + with shard_path.open("rb") as file: + header_size_bytes = file.read(8) + if len(header_size_bytes) != 8: + raise ValueError(f"{shard_path} has a truncated safetensors header") + header_size = struct.unpack(" dict[str, TensorRecord]: + """Return all tensor records in a sharded or single-file checkpoint.""" + + index_path = checkpoint / "model.safetensors.index.json" + expected_shards: dict[str, str] | None = None + if index_path.exists(): + expected_shards = json.loads(index_path.read_text())["weight_map"] + shard_names = sorted(set(expected_shards.values())) + else: + shard_names = sorted(path.name for path in checkpoint.glob("*.safetensors")) + if not shard_names: + raise FileNotFoundError(f"No safetensors shards found in {checkpoint}") + + records: dict[str, TensorRecord] = {} + for shard_name in shard_names: + for name, record in _read_shard_header(checkpoint, shard_name).items(): + # Distributed exporters may repeat shared or tied tensors in more than one shard. + # The index is authoritative: ignore a duplicate copy unless this is the shard named + # by weight_map. Unique unindexed tensors remain in ``records`` and are rejected by + # the index-consistency check below. + if ( + expected_shards is not None + and name in expected_shards + and expected_shards[name] != shard_name + ): + continue + if name in records: + raise ValueError(f"Tensor {name!r} appears in more than one shard") + records[name] = record + + if expected_shards is not None: + if set(records) != set(expected_shards): + missing = sorted(set(expected_shards) - set(records)) + extra = sorted(set(records) - set(expected_shards)) + raise ValueError(f"Safetensors index mismatch: missing={missing}, extra={extra}") + for name, shard_name in expected_shards.items(): + if records[name].shard != shard_name: + raise ValueError( + f"Safetensors index maps {name!r} to {shard_name}, " + f"but its header is in {records[name].shard}" + ) + return records + + +def _load_quantization_config(checkpoint: Path) -> tuple[dict, str]: + hf_quant_config = checkpoint / "hf_quant_config.json" + if hf_quant_config.exists(): + config = json.loads(hf_quant_config.read_text()) + return config["quantization"], hf_quant_config.name + + config_path = checkpoint / "config.json" + config = json.loads(config_path.read_text()) + if "quantization_config" not in config: + raise ValueError("Neither hf_quant_config.json nor config.json contains quantization data") + return config["quantization_config"], config_path.name + + +def _matching_quantized_layer(name: str, quantized_layers: dict[str, dict]) -> str | None: + matches = [ + layer + for layer in quantized_layers + if name == f"{layer}.weight" or (name.startswith(f"{layer}.") and name.endswith(".weight")) + ] + return max(matches, key=len) if matches else None + + +def _is_expert_weight(name: str) -> bool: + return ( + not _is_mtp_tensor(name) + and name.endswith(".weight") + and (".mixer.experts." in name or ".mixer.shared_experts." in name) + ) + + +def _is_mamba_projection_weight(name: str) -> bool: + return not _is_mtp_tensor(name) and name.endswith( + (".mixer.in_proj.weight", ".mixer.out_proj.weight") + ) + + +def _is_mtp_tensor(name: str) -> bool: + return name.startswith("mtp.") or ".mtp." in name + + +def _read_int64_vector(checkpoint: Path, record: TensorRecord) -> tuple[int, ...]: + """Read one small int64 shape sidecar directly from its safetensors payload.""" + if record.dtype != "I64" or len(record.shape) != 1: + raise ValueError( + f"Shape sidecar {record.name} must be a one-dimensional I64 tensor, got " + f"{record.dtype} {record.shape}" + ) + with (checkpoint / record.shard).open("rb") as file: + file.seek(record.payload_offset) + payload = file.read(record.payload_bytes) + if len(payload) != record.shape[0] * 8: + raise ValueError(f"Shape sidecar {record.name} has an invalid byte count") + return struct.unpack(f"<{record.shape[0]}q", payload) + + +def _digest_iq2_tensors( + checkpoint: Path, records: dict[str, TensorRecord], names: list[str] +) -> tuple[dict[str, str], str]: + """Hash each packed tensor and the ordered concatenation of all packed payloads.""" + + per_tensor = {} + aggregate = hashlib.sha256() + open_shard: str | None = None + file: BinaryIO | None = None + try: + for name in sorted(names): + record = records[name] + if record.shard != open_shard: + if file is not None: + file.close() + file = (checkpoint / record.shard).open("rb") + open_shard = record.shard + assert file is not None + digest = hashlib.sha256() + file.seek(record.payload_offset) + remaining = record.payload_bytes + while remaining: + chunk = file.read(min(remaining, _HASH_CHUNK_BYTES)) + if not chunk: + raise EOFError( + f"Unexpected end of {record.shard} while hashing tensor {record.name}" + ) + digest.update(chunk) + aggregate.update(chunk) + remaining -= len(chunk) + per_tensor[name] = digest.hexdigest() + finally: + if file is not None: + file.close() + return per_tensor, aggregate.hexdigest() + + +def validate_checkpoint( + checkpoint: Path, + reference_checkpoint: Path, + *, + compute_digests: bool = True, +) -> dict: + """Validate the Nemotron mixed-format policy and return a JSON-serializable report.""" + + checkpoint = checkpoint.resolve() + reference_checkpoint = reference_checkpoint.resolve() + records = read_checkpoint_index(checkpoint) + reference_records = read_checkpoint_index(reference_checkpoint) + quantization, quant_config_file = _load_quantization_config(checkpoint) + quantized_layers = quantization.get("quantized_layers", {}) + errors: list[str] = [] + + if quantization.get("quant_algo") != "MIXED_PRECISION": + errors.append( + f"Expected quant_algo=MIXED_PRECISION, got {quantization.get('quant_algo')!r}" + ) + if not quantized_layers: + errors.append("The quantization config does not declare quantized_layers") + + iq2_layers = { + name: cfg for name, cfg in quantized_layers.items() if cfg.get("quant_algo") == "IQ2_XS" + } + nvfp4_layers = { + name + for name, cfg in quantized_layers.items() + if cfg.get("quant_algo") in {"NVFP4", "W4A16_NVFP4"} + } + other_layers = { + name: cfg.get("quant_algo") + for name, cfg in quantized_layers.items() + if name not in iq2_layers and name not in nvfp4_layers + } + if not iq2_layers: + errors.append("No IQ2_XS layers were declared") + if not nvfp4_layers: + errors.append("No NVFP4 layers were declared") + if other_layers: + errors.append(f"Unexpected quantized formats: {other_layers}") + + for layer, cfg in iq2_layers.items(): + if not (".mixer.experts" in layer or ".mixer.shared_experts" in layer): + errors.append(f"IQ2_XS is applied outside expert weights: {layer}") + expected = { + "quant_algo": "IQ2_XS", + "group_size": _IQ2_XS_BLOCK_SIZE, + "block_payload_bytes": _IQ2_XS_PAYLOAD_BYTES, + "packing": "ggml", + "row_padding": "right", + "logical_shape_key": "weight_logical_shape", + "padded_shape_key": "weight_padded_shape", + } + if cfg != expected: + errors.append(f"Unexpected IQ2_XS metadata for {layer}: {cfg}") + errors.extend( + f"NVFP4 is applied outside Mamba in_proj/out_proj: {layer}" + for layer in nvfp4_layers + if not layer.endswith((".mixer.in_proj", ".mixer.out_proj")) + ) + if any("mtp" in name.lower() for name in quantized_layers): + errors.append("MTP must remain unquantized") + + iq2_tensor_names: list[str] = [] + iq2_shapes: dict[str, tuple[tuple[int, ...], tuple[int, ...]]] = {} + nvfp4_tensor_names: list[str] = [] + unquantized_weight_names: list[str] = [] + for name, record in records.items(): + if not name.endswith(".weight"): + continue + layer = _matching_quantized_layer(name, quantized_layers) + algo = quantized_layers[layer]["quant_algo"] if layer else None + reference = reference_records.get(name) + + if algo == "IQ2_XS": + iq2_tensor_names.append(name) + base = name.removesuffix(".weight") + logical_record = records.get(base + ".weight_logical_shape") + padded_record = records.get(base + ".weight_padded_shape") + if record.dtype != "U8": + errors.append(f"IQ2_XS tensor {name} has dtype {record.dtype}, expected U8") + if len(record.shape) < 2 or record.shape[-1] != _IQ2_XS_PAYLOAD_BYTES: + errors.append( + f"IQ2_XS tensor {name} has shape {record.shape}, expected [..., blocks, 74]" + ) + if record.payload_bytes != math.prod(record.shape): + errors.append( + f"IQ2_XS tensor {name} stores {record.payload_bytes} bytes for shape " + f"{record.shape}" + ) + if reference is None: + errors.append(f"IQ2_XS tensor {name} is absent from the BF16 reference") + if logical_record is None or padded_record is None: + errors.append(f"IQ2_XS tensor {name} is missing logical or padded shape metadata") + else: + try: + logical_shape = _read_int64_vector(checkpoint, logical_record) + padded_shape = _read_int64_vector(checkpoint, padded_record) + except ValueError as error: + errors.append(str(error)) + else: + iq2_shapes[name] = (logical_shape, padded_shape) + expected_padded_shape = ( + *logical_shape[:-1], + math.ceil(logical_shape[-1] / _IQ2_XS_BLOCK_SIZE) * _IQ2_XS_BLOCK_SIZE, + ) + expected_packed_shape = ( + *padded_shape[:-1], + padded_shape[-1] // _IQ2_XS_BLOCK_SIZE, + _IQ2_XS_PAYLOAD_BYTES, + ) + if reference is not None and logical_shape != reference.shape: + errors.append( + f"IQ2_XS tensor {name} has logical shape {logical_shape}, " + f"but the BF16 reference shape is {reference.shape}" + ) + if padded_shape != expected_padded_shape: + errors.append( + f"IQ2_XS tensor {name} has padded shape {padded_shape}, expected " + f"{expected_padded_shape}" + ) + if record.shape != expected_packed_shape: + errors.append( + f"IQ2_XS tensor {name} has packed shape {record.shape}, expected " + f"{expected_packed_shape}" + ) + elif algo in {"NVFP4", "W4A16_NVFP4"}: + nvfp4_tensor_names.append(name) + if record.dtype != "U8": + errors.append(f"NVFP4 tensor {name} has dtype {record.dtype}, expected U8") + base = name.removesuffix(".weight") + errors.extend( + f"NVFP4 tensor {name} is missing {base}{suffix}" + for suffix in (".weight_scale", ".weight_scale_2") + if f"{base}{suffix}" not in records + ) + if reference is None: + errors.append(f"NVFP4 tensor {name} is absent from the BF16 reference") + elif record.shape != (*reference.shape[:-1], reference.shape[-1] // 2): + errors.append( + f"NVFP4 tensor {name} has packed shape {record.shape}, expected " + f"{(*reference.shape[:-1], reference.shape[-1] // 2)}" + ) + else: + unquantized_weight_names.append(name) + if record.dtype == "U8": + errors.append(f"Unaccounted packed U8 weight tensor: {name}") + + reference_experts = sorted(name for name in reference_records if _is_expert_weight(name)) + reference_mamba = sorted( + name for name in reference_records if _is_mamba_projection_weight(name) + ) + missing_experts = sorted(set(reference_experts) - set(iq2_tensor_names)) + unexpected_iq2 = sorted(set(iq2_tensor_names) - set(reference_experts)) + missing_mamba = sorted(set(reference_mamba) - set(nvfp4_tensor_names)) + unexpected_nvfp4 = sorted(set(nvfp4_tensor_names) - set(reference_mamba)) + if missing_experts: + errors.append(f"Expert weights not exported as IQ2_XS: {missing_experts}") + if unexpected_iq2: + errors.append(f"IQ2_XS tensors are not expert weights: {unexpected_iq2}") + if missing_mamba: + errors.append(f"Mamba projections not exported as NVFP4: {missing_mamba}") + if unexpected_nvfp4: + errors.append(f"NVFP4 tensors are not Mamba projections: {unexpected_nvfp4}") + + tensor_digests: dict[str, str] = {} + aggregate_digest: str | None = None + if compute_digests and iq2_tensor_names: + tensor_digests, aggregate_digest = _digest_iq2_tensors( + checkpoint, records, iq2_tensor_names + ) + + iq2_blocks = sum(math.prod(records[name].shape[:-1]) for name in iq2_tensor_names) + iq2_logical_weights = sum(math.prod(shapes[0]) for shapes in iq2_shapes.values()) + iq2_padded_weights = sum(math.prod(shapes[1]) for shapes in iq2_shapes.values()) + report = { + "schema_version": 2, + "created_at": datetime.now(UTC).isoformat(), + "status": "passed" if not errors else "failed", + "checkpoint": str(checkpoint), + "reference_checkpoint": str(reference_checkpoint), + "quantization_config_file": quant_config_file, + "quant_algo": quantization.get("quant_algo"), + "summary": { + "total_tensors": len(records), + "total_weight_tensors": sum(name.endswith(".weight") for name in records), + "iq2_xs_layers": len(iq2_layers), + "iq2_xs_tensors": len(iq2_tensor_names), + "iq2_xs_blocks": iq2_blocks, + "iq2_xs_logical_weights": iq2_logical_weights, + "iq2_xs_padded_weights": iq2_padded_weights, + "iq2_xs_payload_bytes": iq2_blocks * _IQ2_XS_PAYLOAD_BYTES, + "nvfp4_layers": len(nvfp4_layers), + "nvfp4_tensors": len(nvfp4_tensor_names), + "unquantized_weight_tensors": len(unquantized_weight_names), + }, + "iq2_xs": { + "block_size": _IQ2_XS_BLOCK_SIZE, + "block_payload_bytes": _IQ2_XS_PAYLOAD_BYTES, + "aggregate_sha256": aggregate_digest, + "tensors": [ + { + **asdict(records[name]), + "shape": list(records[name].shape), + "logical_shape": list(iq2_shapes[name][0]) if name in iq2_shapes else None, + "padded_shape": list(iq2_shapes[name][1]) if name in iq2_shapes else None, + "sha256": tensor_digests.get(name), + } + for name in sorted(iq2_tensor_names) + ], + }, + "nvfp4_tensors": sorted(nvfp4_tensor_names), + "errors": errors, + } + return report + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--reference-checkpoint", type=Path, required=True) + parser.add_argument("--report", type=Path, required=True) + parser.add_argument( + "--skip-digests", + action="store_true", + help="Validate layouts and policy without hashing packed IQ2_XS payloads.", + ) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + report = validate_checkpoint( + args.checkpoint, + args.reference_checkpoint, + compute_digests=not args.skip_digests, + ) + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + print(json.dumps({"status": report["status"], **report["summary"]}, indent=2)) + print(f"Validation report: {args.report}") + if report["errors"]: + for error in report["errors"]: + print(f"ERROR: {error}") + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py index 45fa0c30f3b..eb49ad700ae 100644 --- a/modelopt/torch/export/convert_hf_config.py +++ b/modelopt/torch/export/convert_hf_config.py @@ -117,6 +117,22 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) }, "weights": {"dynamic": False, "num_bits": 8, "type": "float", "group_size": gs}, } + elif quant_algo in ("IQ1_S", "IQ2_XS"): + effective_bits, payload_bytes = (1.5625, 50) if quant_algo == "IQ1_S" else (2.3125, 74) + return { + "weights": { + "dynamic": False, + "num_bits": 1 if quant_algo == "IQ1_S" else 2, + "effective_bits": effective_bits, + "type": "int", + "group_size": 256, + "packing": "ggml", + "block_payload_bytes": payload_bytes, + "row_padding": "right", + "logical_shape_key": "weight_logical_shape", + "padded_shape_key": "weight_padded_shape", + } + } else: warnings.warn( f"Unsupported quantization algorithm '{quant_algo}' in " @@ -209,6 +225,10 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An "targets": ["Linear"], } new_config["config_groups"] = {"group_0": config_group_details} + elif quant_algo_value in ("IQ1_S", "IQ2_XS"): + config_group_details = _quant_algo_to_group_config(quant_algo_value, 256) + config_group_details["targets"] = ["Linear"] + new_config["config_groups"] = {"group_0": config_group_details} elif quant_algo_value == "NVFP4_SVD": # NVFP4 + SVDQuant: NVFP4 weights/activations plus an AWQ-style # pre_quant_scale and a low-rank residual (svdquant_lora_a/b) stored as diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 734302690f1..a75bc516862 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -218,8 +218,17 @@ def _export_fused_experts( _export_quantized_weight(wrapper, dtype) proj = nn.Module() - proj.weight = wrapper.weight - for attr in ("weight_scale", "weight_scale_2", "input_scale"): + if isinstance(wrapper.weight, nn.Parameter): + proj.weight = wrapper.weight + else: + proj.register_buffer("weight", wrapper.weight) + for attr in ( + "weight_scale", + "weight_scale_2", + "input_scale", + "weight_logical_shape", + "weight_padded_shape", + ): if hasattr(wrapper, attr): proj.register_buffer(attr, getattr(wrapper, attr)) diff --git a/modelopt/torch/export/quant_format.py b/modelopt/torch/export/quant_format.py index b270aec1a68..b4475fa1f7c 100644 --- a/modelopt/torch/export/quant_format.py +++ b/modelopt/torch/export/quant_format.py @@ -36,11 +36,21 @@ QUANTIZATION_FP8_PB_REAL = "fp8_pb_real" QUANTIZATION_FP8_PB_WO = "fp8_pb_wo" QUANTIZATION_FP8_PC_PT = "fp8_pc_pt" +QUANTIZATION_IQ1_S = "iq1_s" +QUANTIZATION_IQ2_XS = "iq2_xs" # Formats whose scales are purely per-module, so export never merges them across the q/k/v # and gate/up groups that share an input. Every other format unifies input_amax (and, for # NVFP4, weight_scale_2) across such a group, which only a whole-model forward can discover. -FUSION_FREE_FORMATS = frozenset({QUANTIZATION_FP8, QUANTIZATION_NONE, QUANTIZATION_FP8_PB_REAL}) +FUSION_FREE_FORMATS = frozenset( + { + QUANTIZATION_FP8, + QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_XS, + QUANTIZATION_NONE, + QUANTIZATION_FP8_PB_REAL, + } +) KV_CACHE_FP8 = "FP8" KV_CACHE_INT8 = "INT8" diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 5320f012a94..aa179e8da61 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -61,6 +61,8 @@ QUANTIZATION_INT4_AWQ, QUANTIZATION_INT8_SQ, QUANTIZATION_INT8_WO, + QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_XS, QUANTIZATION_MXFP4, QUANTIZATION_MXFP8, QUANTIZATION_NONE, @@ -434,6 +436,11 @@ def _get_quantization_from_layer(layer, quantizer_attr_names: QuantizerAttrNames return QUANTIZATION_W4A8_AWQ # Handle individual num_bits cases + if weight_quantizer.num_bits in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if weight_quantizer.backend != "ggml": + raise ValueError("IQ formats require the built-in 'ggml' quantization backend") + return weight_quantizer.num_bits + if weight_quantizer.num_bits == 4: assert len(weight_quantizer.block_sizes) > 0 and weight_quantizer.block_sizes[-1] > 0, ( "Invalid block_sizes for INT4 quantizer" @@ -682,6 +689,17 @@ def process_layer_quant_config(layer_config_dict): "quant_algo": "MXFP8", "group_size": block_size_value, } + elif v in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + payload_bytes = 50 if v == QUANTIZATION_IQ1_S else 74 + layer_config = { + "quant_algo": v.upper(), + "group_size": 256, + "block_payload_bytes": payload_bytes, + "packing": "ggml", + "row_padding": "right", + "logical_shape_key": "weight_logical_shape", + "padded_shape_key": "weight_padded_shape", + } else: layer_config = {"quant_algo": v} @@ -1075,6 +1093,9 @@ def _export_key(key: str) -> str: # (pre_quant_scale is the AWQ / NVFP4_AWQ / SVDQuant companion, renamed in the KV-cache pass.) weight_suffixes = ( "weight", + "weight_shape", + "weight_logical_shape", + "weight_padded_shape", "weight_scale", "weight_scale_2", "input_scale", diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 354dda65f40..9014761de6a 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -59,6 +59,7 @@ from modelopt.torch.opt.conversion import ModeloptStateManager, modelopt_state from modelopt.torch.opt.plugins.huggingface import _MODELOPT_STATE_SAVE_NAME from modelopt.torch.quantization import set_quantizer_by_cfg_context +from modelopt.torch.quantization.ggml import quantize_iq1_s, quantize_iq2_xs from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper @@ -100,6 +101,8 @@ QUANTIZATION_FP8, QUANTIZATION_FP8_PB_REAL, QUANTIZATION_FP8_PC_PT, + QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_XS, QUANTIZATION_MXFP8, QUANTIZATION_NONE, QUANTIZATION_NVFP4, @@ -622,6 +625,25 @@ def _export_quantized_weight( "which dispatches to the streaming writer that materialises weights layer-by-layer." ) + if quantization_format in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + if weight_name != "weight": + raise NotImplementedError( + "IQ unified export currently supports modules with a standard 'weight' " + f"attribute, got {weight_name!r} on {type(sub_module).__name__}" + ) + quantize_iq = ( + quantize_iq1_s if quantization_format == QUANTIZATION_IQ1_S else quantize_iq2_xs + ) + packed_weight, logical_shape = quantize_iq(weight.to(dtype)) + padded_shape = logical_shape.clone() + padded_shape[-1] = packed_weight.shape[-2] * 256 + delattr(sub_module, weight_name) + sub_module.register_buffer("weight", packed_weight) + sub_module.register_buffer("weight_logical_shape", logical_shape) + sub_module.register_buffer("weight_padded_shape", padded_shape) + maybe_clear_cuda_cache() + return + weight_quantizer: TensorQuantizer | SequentialQuantizer = getattr( sub_module, quantizer_attrs.weight_quantizer ) diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 37c896fd609..42609256036 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -35,6 +35,7 @@ from safetensors.torch import save_file from modelopt import __version__ +from modelopt.torch.quantization.ggml import quantize_iq1_s, quantize_iq2_xs from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer from modelopt.torch.utils import import_plugin, warn_rank_0 @@ -61,6 +62,8 @@ QUANTIZATION_FP8, QUANTIZATION_FP8_PB_REAL, QUANTIZATION_FP8_PB_WO, + QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_XS, QUANTIZATION_NONE, QUANTIZATION_NVFP4, QUANTIZATION_W4A16_NVFP4, @@ -94,6 +97,7 @@ get_pipeline_model_parallel_rank, get_pipeline_model_parallel_world_size, get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, ) from megatron.core.ssm.mamba_layer import MambaLayer from megatron.core.transformer.identity_op import IdentityOp @@ -309,10 +313,19 @@ def save_pretrained( is_last_stage_main_rank = pp_rank == pp_size - 1 and tp_rank == 0 is_writer_rank = self._is_sidecar_writer_rank(is_last_stage_main_rank) + quantization_format = self._get_quantization_format(self.model) + if ( + quantization_format in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) + and get_tensor_model_parallel_world_size() != 1 + ): + raise NotImplementedError( + "Megatron IQ1_S/IQ2_XS unified export currently requires tensor model " + "parallel size 1" + ) + # Main export process layer_state_dicts = self.layer_state_dicts - quantization_format = self._get_quantization_format(self.model) quantization = None if quantization_format in ( QUANTIZATION_FP8_PB_REAL, @@ -325,6 +338,8 @@ def save_pretrained( quantization = "NVFP4" elif quantization_format == QUANTIZATION_W4A16_NVFP4: quantization = "W4A16_NVFP4" + elif quantization_format in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + quantization = quantization_format.upper() if is_last_stage_main_rank: if is_writer_rank: @@ -1027,6 +1042,7 @@ def _get_weight_bias( module: torch.nn.Module, dtype: torch.dtype = torch.float16, name_to_value: dict[str, torch.Tensor] | None = None, + keep_weight_device: bool = False, ) -> dict[str, torch.Tensor]: """Get the weight and bias of the module. @@ -1035,6 +1051,7 @@ def _get_weight_bias( dtype: The data type of the weight and bias. name_to_value: The dictionary to store the weight and bias. A new dict is created if not provided. + keep_weight_device: Keep the weight on its current device instead of moving it to CPU. Returns: The dictionary containing the weight and bias. @@ -1045,7 +1062,9 @@ def _get_weight_bias( # layers whose weight is a placeholder) so callers can use "weight" in name_to_value # as a reliable guard without re-inspecting module.weight. if hasattr(module, "weight") and module.weight is not None and module.weight.numel() > 0: - weight = module.weight.to(dtype).cpu() + weight = module.weight.to(dtype) + if not keep_weight_device: + weight = weight.cpu() name_to_value["weight"] = weight if hasattr(module, "bias") and module.bias is not None and module.bias.numel() > 0: @@ -1082,13 +1101,21 @@ def _get_quantized_state( self._record_excluded_module(prefix) block_size = get_weight_block_size(module) - name_to_value = self._get_weight_bias(module, dtype, name_to_value) + is_iq = qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS) + name_to_value = self._get_weight_bias( + module, dtype, name_to_value, keep_weight_device=is_iq + ) if "weight" not in name_to_value: return name_to_value, qformat, block_size if qformat == QUANTIZATION_NONE: return name_to_value, qformat, block_size + # IQ formats derive all block metadata directly from the weight and do not use amax or + # separately exported scaling tensors. Keep the weight on-device until its final HF layout + # has been produced, so the CUDA packer can be used. + if is_iq: + return name_to_value, qformat, block_size # Getting the weight scales weight_scale = get_weight_scaling_factor(module) weight_scale_2 = get_weight_scaling_factor_2(module) @@ -1124,6 +1151,22 @@ def _get_weight_scales(self, quantized_state: dict[str, Any], qformat: str): return weight_scale, weight_scale_2 + @staticmethod + def _get_iq_weight_state( + weight_key: str, weight: torch.Tensor, qformat: str + ) -> dict[str, torch.Tensor]: + """Pack one final-layout weight into the IQ unified-checkpoint representation.""" + quantize_iq = quantize_iq1_s if qformat == QUANTIZATION_IQ1_S else quantize_iq2_xs + packed_weight, logical_shape = quantize_iq(weight) + padded_shape = logical_shape.clone() + padded_shape[-1] = packed_weight.shape[-2] * 256 + prefix = weight_key.removesuffix("weight") + return { + weight_key: packed_weight.detach().cpu(), + prefix + "weight_logical_shape": logical_shape.detach().cpu(), + prefix + "weight_padded_shape": padded_shape.detach().cpu(), + } + def _record_layer_quant_config(self, prefix: str, qformat: str | None, block_size: int | None): """Record per-HF-layer quantization metadata for mixed precision exports.""" if qformat in (None, QUANTIZATION_NONE): @@ -1188,7 +1231,9 @@ def _name_remapping( weight = weight + 1.0 weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) - if weight_scale is None: + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + self._state_dict.update(self._get_iq_weight_state(prefix + "weight", weight, qformat)) + elif weight_scale is None: self._state_dict[prefix + "weight"] = weight else: self._state_dict[prefix + "weight"] = to_quantized_weight( @@ -1233,7 +1278,14 @@ def _gated_mlp_slicing( gate_proj_weight = weight[:ffn_hidden_size, :] up_proj_weight = weight[ffn_hidden_size:, :] - if weight_scale is None: + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + self._state_dict.update( + self._get_iq_weight_state(gate_proj_prefix + "weight", gate_proj_weight, qformat) + ) + self._state_dict.update( + self._get_iq_weight_state(up_proj_prefix + "weight", up_proj_weight, qformat) + ) + elif weight_scale is None: self._state_dict[gate_proj_prefix + "weight"] = gate_proj_weight self._state_dict[up_proj_prefix + "weight"] = up_proj_weight else: @@ -1399,7 +1451,9 @@ def _grouped_mlp_slicing( name_to_value.pop("weight", None) seen_qformat, seen_block_size = qformat, block_size - weight = state_dict[weight_key].to(self.dtype).cpu() + weight = state_dict[weight_key].to(self.dtype) + if qformat not in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + weight = weight.cpu() weight_scale_cpu = ( weight_scale.detach().cpu().clone() if weight_scale is not None else None ) @@ -1430,7 +1484,13 @@ def _grouped_mlp_slicing( ] for shard_prefix, shard_weight, shard_scale in shards: - if shard_scale is None: + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + local_expert_state.update( + self._get_iq_weight_state( + shard_prefix + "weight", shard_weight, qformat + ) + ) + elif shard_scale is None: local_expert_state[shard_prefix + "weight"] = shard_weight else: local_expert_state[shard_prefix + "weight"] = to_quantized_weight( @@ -1593,7 +1653,10 @@ def _take(tensor, index, last_dim, with_gate=False): proj_weights = [_take(weight, s, hidden_size, g) for s, g in zip(slices, gated)] proj_keys = [p + "weight" for p in prefixes] - if weight_scale is None: + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + for key, weight in zip(proj_keys, proj_weights): + self._state_dict.update(self._get_iq_weight_state(key, weight, qformat)) + elif weight_scale is None: for key, weight in zip(proj_keys, proj_weights): self._state_dict[key] = weight else: @@ -1708,7 +1771,15 @@ def _gated_delta_net_slicing(self, module, prefix, is_mtp=False): proj_keys = [p + "weight" for p in proj_prefixes] weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) - if weight_scale is None: + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + for proj_prefix, proj_weight in zip(proj_prefixes, proj_weights): + if proj_prefix in keep_bf16: + self._state_dict[proj_prefix + "weight"] = proj_weight.cpu() + else: + self._state_dict.update( + self._get_iq_weight_state(proj_prefix + "weight", proj_weight, qformat) + ) + elif weight_scale is None: for key, proj_weight in zip(proj_keys, proj_weights): self._state_dict[key] = proj_weight else: @@ -1842,7 +1913,9 @@ def _pack_name_remapping(self, module, prefix, layer_type=None, is_mtp=False, tr merged_input_scale = None # Save the merged weights - if merged_weight_scale is None: + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + self._state_dict.update(self._get_iq_weight_state(prefix, merged_weight, qformat)) + elif merged_weight_scale is None: self._state_dict[prefix] = merged_weight else: self._state_dict[prefix] = to_quantized_weight( @@ -1953,7 +2026,9 @@ def _pack_name_remapping_gpt_oss(self, module, prefix, layer_type=None, is_mtp=F merged_input_scale = None # Save the merged weights - if merged_weight_scale is None: + if qformat in (QUANTIZATION_IQ1_S, QUANTIZATION_IQ2_XS): + self._state_dict.update(self._get_iq_weight_state(prefix, merged_weight, qformat)) + elif merged_weight_scale is None: # TODO: May need to modify the key name later. self._state_dict[prefix] = merged_weight else: diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp b/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp new file mode 100644 index 00000000000..bc8a8fa325d --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cpp @@ -0,0 +1,31 @@ +/* + * 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. + */ + +#include +#include + +at::Tensor iq1_s_pack_cuda(at::Tensor input, at::Tensor grid); + +at::Tensor iq1_s_pack(at::Tensor input, at::Tensor grid) { + TORCH_CHECK(input.is_cuda(), "IQ1_S packing requires a CUDA input"); + TORCH_CHECK(grid.is_cuda(), "IQ1_S packing requires a CUDA grid"); + return iq1_s_pack_cuda(input.contiguous(), grid.contiguous()); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("pack", &iq1_s_pack, "Pack a tensor into GGML IQ1_S blocks"); +} diff --git a/modelopt/torch/kernels/quantization/ggml/iq1_s.cu b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu new file mode 100644 index 00000000000..ddf1f762e16 --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq1_s.cu @@ -0,0 +1,270 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace { + +constexpr int kBlockSize = 256; +constexpr int kVectorSize = 8; +constexpr int kEntries = 2048; +constexpr int kGroups = 8; +constexpr int kLocalScales = 8; +constexpr int kChoices = 16; +constexpr int kPayloadBytes = 50; +constexpr float kDelta = 0.125f; +constexpr float kNativeMax = 16.875f; + +template __device__ __forceinline__ float load_float(const scalar_t *input) { + return static_cast(*input); +} + +__device__ __forceinline__ float quant_error(float xnorm, float xsum, const float *x, + const float *q, float scale, float delta) { + float dot = 0.0f; + float qnorm = 0.0f; + float qsum = 0.0f; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + dot = fmaf(x[j], q[j], dot); + qnorm = fmaf(q[j], q[j], qnorm); + qsum += q[j]; + } + const float shifted_dot = dot + delta * xsum; + const float shifted_norm = qnorm + 2.0f * delta * qsum + 8.0f * delta * delta; + return fmaxf(fmaf(scale * scale, shifted_norm, fmaf(-2.0f * scale, shifted_dot, xnorm)), 0.0f); +} + +template +__global__ void find_scale(const scalar_t *input, int64_t num_blocks, int64_t *scale_bits) { + const int64_t block = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (block >= num_blocks) + return; + + float amax = 0.0f; + const scalar_t *values = input + block * kBlockSize; +#pragma unroll 1 + for (int i = 0; i < kBlockSize; ++i) + amax = fmaxf(amax, fabsf(load_float(values + i))); + const __half scale = __float2half_rn(fminf((amax / kNativeMax) * 0.61f, 65504.0f)); + scale_bits[block] = static_cast(__half_as_ushort(scale)); +} + +template +__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid, + const int64_t *scale_bits, uint8_t *output) { + __shared__ float warp_best[8 * kChoices]; + __shared__ float group_error[kChoices]; + __shared__ unsigned long long warp_keys[8]; + __shared__ int selected_choice; + __shared__ uint16_t selected_entries[4]; + + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int64_t block = blockIdx.x; + if (block >= num_blocks) + return; + + const scalar_t *source = input + block * kBlockSize; + uint8_t *payload = output + block * kPayloadBytes; + const uint16_t d_bits = static_cast(scale_bits[block]); + const float d = __half2float(__ushort_as_half(d_bits)); + if (d_bits == 0) { + if (tid < kPayloadBytes) + payload[tid] = 0; + return; + } + if (tid == 0) { + payload[0] = static_cast(d_bits); + payload[1] = static_cast(d_bits >> 8); + } + +#pragma unroll 1 + for (int group = 0; group < kGroups; ++group) { + if (tid < kChoices) + group_error[tid] = 0.0f; + __syncthreads(); + +#pragma unroll + for (int vector = 0; vector < 4; ++vector) { + float x[kVectorSize]; + float xnorm = 0.0f; + float xsum = 0.0f; + const int offset = group * 32 + vector * 8; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + x[j] = load_float(source + offset + j); + xnorm = fmaf(x[j], x[j], xnorm); + xsum += x[j]; + } + float local_best[kChoices]; +#pragma unroll + for (int choice = 0; choice < kChoices; ++choice) + local_best[choice] = FLT_MAX; + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + const float *q = grid + entry * kVectorSize; + float dot = 0.0f; + float qnorm = 0.0f; + float qsum = 0.0f; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + dot = fmaf(x[j], q[j], dot); + qnorm = fmaf(q[j], q[j], qnorm); + qsum += q[j]; + } +#pragma unroll + for (int choice = 0; choice < kChoices; ++choice) { + const int local = choice & 7; + const float delta = choice < 8 ? kDelta : -kDelta; + const float scale = d * (2 * local + 1); + const float shifted_dot = dot + delta * xsum; + const float shifted_norm = qnorm + 2.0f * delta * qsum + 8.0f * delta * delta; + const float error = fmaxf( + fmaf(scale * scale, shifted_norm, fmaf(-2.0f * scale, shifted_dot, xnorm)), 0.0f); + local_best[choice] = fminf(local_best[choice], error); + } + } +#pragma unroll + for (int choice = 0; choice < kChoices; ++choice) { + float value = local_best[choice]; +#pragma unroll + for (int delta = 16; delta > 0; delta >>= 1) + value = fminf(value, __shfl_down_sync(0xffffffff, value, delta)); + if (lane == 0) + warp_best[warp * kChoices + choice] = value; + } + __syncthreads(); + if (tid < kChoices) { + float value = warp_best[tid]; +#pragma unroll + for (int w = 1; w < 8; ++w) + value = fminf(value, warp_best[w * kChoices + tid]); + group_error[tid] += value; + } + __syncthreads(); + } + + if (tid == 0) { + selected_choice = 0; + float best = group_error[0]; +#pragma unroll + for (int choice = 1; choice < kChoices; ++choice) { + if (group_error[choice] < best) { + best = group_error[choice]; + selected_choice = choice; + } + } + } + __syncthreads(); + const int selected_local = selected_choice & 7; + const float selected_delta = selected_choice < 8 ? kDelta : -kDelta; + const float selected_scale = d * (2 * selected_local + 1); + +#pragma unroll + for (int vector = 0; vector < 4; ++vector) { + float x[kVectorSize]; + float xnorm = 0.0f; + float xsum = 0.0f; + const int offset = group * 32 + vector * 8; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + x[j] = load_float(source + offset + j); + xnorm = fmaf(x[j], x[j], xnorm); + xsum += x[j]; + } + unsigned long long key = ~0ULL; + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + const float error = + quant_error(xnorm, xsum, x, grid + entry * kVectorSize, selected_scale, selected_delta); + const unsigned long long candidate = + (static_cast(__float_as_uint(error)) << 32) | + static_cast(entry); + key = candidate < key ? candidate : key; + } +#pragma unroll + for (int delta = 16; delta > 0; delta >>= 1) { + const auto other = __shfl_down_sync(0xffffffff, key, delta); + key = other < key ? other : key; + } + if (lane == 0) + warp_keys[warp] = key; + __syncthreads(); + if (tid == 0) { + key = warp_keys[0]; +#pragma unroll + for (int w = 1; w < 8; ++w) + key = warp_keys[w] < key ? warp_keys[w] : key; + const uint16_t entry = static_cast(key & 0x7ff); + selected_entries[vector] = entry; + payload[2 + group * 4 + vector] = static_cast(entry); + } + __syncthreads(); + } + + if (tid == 0) { + const uint16_t qh = static_cast( + ((selected_entries[0] >> 8) & 7) | (((selected_entries[1] >> 8) & 7) << 3) | + (((selected_entries[2] >> 8) & 7) << 6) | (((selected_entries[3] >> 8) & 7) << 9) | + (selected_local << 12) | ((selected_choice >> 3) << 15)); + payload[34 + 2 * group] = static_cast(qh); + payload[35 + 2 * group] = static_cast(qh >> 8); + } + __syncthreads(); + } +} + +} // namespace + +at::Tensor iq1_s_pack_cuda(at::Tensor input, at::Tensor grid) { + TORCH_CHECK(input.is_contiguous() && grid.is_contiguous(), "inputs must be contiguous"); + TORCH_CHECK(input.numel() > 0 && input.numel() % kBlockSize == 0, + "input size must be a positive multiple of 256"); + TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.numel() == kEntries * kVectorSize, + "grid must be float32 [2048, 8]"); + TORCH_CHECK(input.get_device() == grid.get_device(), "input and grid must share a device"); + c10::cuda::CUDAGuard guard(input.device()); + const int64_t num_blocks = input.numel() / kBlockSize; + TORCH_CHECK(num_blocks <= std::numeric_limits::max(), "IQ1_S CUDA grid is too large"); + auto scales = at::empty({num_blocks}, input.options().dtype(at::kLong)); + auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte)); + const auto stream = c10::cuda::getCurrentCUDAStream(); + const int scale_grid = static_cast((num_blocks + 255) / 256); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq1_s_pack", [&] { + find_scale<<>>(input.data_ptr(), num_blocks, + scales.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + encode<<(num_blocks), 256, 0, stream>>>( + input.data_ptr(), num_blocks, grid.data_ptr(), + scales.data_ptr(), output.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + return output; +} diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp new file mode 100644 index 00000000000..ee90cae2650 --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cpp @@ -0,0 +1,30 @@ +/* + * 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. + */ + +#include + +at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid); + +at::Tensor iq2_xs_pack(at::Tensor input, at::Tensor grid) { + TORCH_CHECK(input.is_cuda(), "IQ2_XS packing requires a CUDA input"); + TORCH_CHECK(grid.is_cuda(), "IQ2_XS packing requires a CUDA grid"); + return iq2_xs_pack_cuda(input.contiguous(), grid.contiguous()); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("pack", &iq2_xs_pack, "Pack a tensor into GGML IQ2_XS blocks"); +} diff --git a/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu new file mode 100644 index 00000000000..cc798903aff --- /dev/null +++ b/modelopt/torch/kernels/quantization/ggml/iq2_xs.cu @@ -0,0 +1,292 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace { + +constexpr int kBlockSize = 256; +constexpr int kVectorSize = 8; +constexpr int kEntries = 512; +constexpr int kGroups = 16; +constexpr int kLocalScales = 16; +constexpr int kPayloadBytes = 74; +constexpr float kNativeMax = 166.625f; + +template __device__ __forceinline__ float load_float(const scalar_t *input) { + return static_cast(*input); +} + +__device__ __forceinline__ float quant_error(float xnorm, float dot, float qnorm, float scale) { + return fmaxf(fmaf(scale * scale, qnorm, fmaf(-2.0f * scale, dot, xnorm)), 0.0f); +} + +__device__ __forceinline__ float even_parity_dot(const float *x, const float *q, bool odd_parity) { + float dot = 0.0f; + float weakest = FLT_MAX; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + const float term = fabsf(x[j]) * q[j]; + dot += term; + weakest = fminf(weakest, term); + } + return odd_parity ? dot - 2.0f * weakest : dot; +} + +template +__global__ void find_scale(const scalar_t *input, int64_t num_blocks, int64_t *scale_bits) { + const int64_t block = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (block >= num_blocks) + return; + + float amax = 0.0f; + float sumsq = 0.0f; + const scalar_t *values = input + block * kBlockSize; +#pragma unroll 1 + for (int i = 0; i < kBlockSize; ++i) { + const float value = load_float(values + i); + amax = fmaxf(amax, fabsf(value)); + sumsq = fmaf(value, value, sumsq); + } + if (amax == 0.0f) { + scale_bits[block] = 0; + return; + } + const float rms = sqrtf(sumsq / kBlockSize); + const float peak_to_rms = rms > 0.0f ? amax / rms : 0.0f; + const float anchor = fminf(0.92f, fmaxf(0.65f, 1.0f - 0.035f * peak_to_rms)); + const __half scale = __float2half_rn(fminf((amax / kNativeMax) * anchor, 65504.0f)); + scale_bits[block] = static_cast(__half_as_ushort(scale)); +} + +template +__global__ void encode(const scalar_t *input, int64_t num_blocks, const float *grid, + const int64_t *scale_bits, uint8_t *output) { + __shared__ float shared_grid[kEntries * kVectorSize]; + __shared__ float grid_norm[kEntries]; + __shared__ float warp_best[8 * kLocalScales]; + __shared__ float group_error[kLocalScales]; + __shared__ unsigned long long warp_keys[8]; + __shared__ int selected_local; + __shared__ uint8_t locals[kGroups]; + + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + for (int i = tid; i < kEntries * kVectorSize; i += blockDim.x) + shared_grid[i] = grid[i]; + __syncthreads(); + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + float norm = 0.0f; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + const float q = shared_grid[entry * kVectorSize + j]; + norm = fmaf(q, q, norm); + } + grid_norm[entry] = norm; + } + __syncthreads(); + + const int64_t block = blockIdx.x; + if (block >= num_blocks) + return; + const scalar_t *source = input + block * kBlockSize; + uint8_t *payload = output + block * kPayloadBytes; + const uint16_t d_bits = static_cast(scale_bits[block]); + const float d = __half2float(__ushort_as_half(d_bits)); + if (tid == 0) { + payload[0] = static_cast(d_bits); + payload[1] = static_cast(d_bits >> 8); + } + +#pragma unroll 1 + for (int group = 0; group < kGroups; ++group) { + if (tid < kLocalScales) + group_error[tid] = 0.0f; + __syncthreads(); + +#pragma unroll + for (int vector = 0; vector < 2; ++vector) { + float x[kVectorSize]; + float xnorm = 0.0f; + int negative_count = 0; + const int offset = group * 16 + vector * 8; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + x[j] = load_float(source + offset + j); + xnorm = fmaf(x[j], x[j], xnorm); + negative_count += x[j] < 0.0f; + } + const bool odd_parity = (negative_count & 1) != 0; + float local_best[kLocalScales]; +#pragma unroll + for (int local = 0; local < kLocalScales; ++local) + local_best[local] = FLT_MAX; + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + const float *q = shared_grid + entry * kVectorSize; + const float dot = even_parity_dot(x, q, odd_parity); +#pragma unroll + for (int local = 0; local < kLocalScales; ++local) { + const float scale = d * (2 * local + 1) * 0.125f; + local_best[local] = + fminf(local_best[local], quant_error(xnorm, dot, grid_norm[entry], scale)); + } + } +#pragma unroll + for (int local = 0; local < kLocalScales; ++local) { + float value = local_best[local]; +#pragma unroll + for (int delta = 16; delta > 0; delta >>= 1) + value = fminf(value, __shfl_down_sync(0xffffffff, value, delta)); + if (lane == 0) + warp_best[warp * kLocalScales + local] = value; + } + __syncthreads(); + if (tid < kLocalScales) { + float value = warp_best[tid]; +#pragma unroll + for (int w = 1; w < 8; ++w) + value = fminf(value, warp_best[w * kLocalScales + tid]); + group_error[tid] += value; + } + __syncthreads(); + } + + if (tid == 0) { + selected_local = 0; + float best = group_error[0]; +#pragma unroll + for (int local = 1; local < kLocalScales; ++local) { + if (group_error[local] < best) { + best = group_error[local]; + selected_local = local; + } + } + locals[group] = static_cast(selected_local); + } + __syncthreads(); + const float selected_scale = d * (2 * selected_local + 1) * 0.125f; + +#pragma unroll + for (int vector = 0; vector < 2; ++vector) { + float x[kVectorSize]; + float xnorm = 0.0f; + int negative_count = 0; + const int offset = group * 16 + vector * 8; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + x[j] = load_float(source + offset + j); + xnorm = fmaf(x[j], x[j], xnorm); + negative_count += x[j] < 0.0f; + } + const bool odd_parity = (negative_count & 1) != 0; + unsigned long long key = ~0ULL; + for (int entry = tid; entry < kEntries; entry += blockDim.x) { + const float error = + quant_error(xnorm, even_parity_dot(x, shared_grid + entry * kVectorSize, odd_parity), + grid_norm[entry], selected_scale); + const unsigned long long candidate = + (static_cast(__float_as_uint(error)) << 32) | + static_cast(entry); + key = candidate < key ? candidate : key; + } +#pragma unroll + for (int delta = 16; delta > 0; delta >>= 1) { + const auto other = __shfl_down_sync(0xffffffff, key, delta); + key = other < key ? other : key; + } + if (lane == 0) + warp_keys[warp] = key; + __syncthreads(); + if (tid == 0) { + key = warp_keys[0]; +#pragma unroll + for (int w = 1; w < 8; ++w) + key = warp_keys[w] < key ? warp_keys[w] : key; + const int entry = static_cast(key & 0x1ff); + const float *q = shared_grid + entry * kVectorSize; + int flip_index = 0; + float weakest = fabsf(x[0]) * q[0]; +#pragma unroll + for (int j = 1; j < kVectorSize; ++j) { + const float term = fabsf(x[j]) * q[j]; + if (term < weakest) { + weakest = term; + flip_index = j; + } + } + int sign_mask = 0; +#pragma unroll + for (int j = 0; j < kVectorSize; ++j) { + bool is_negative = x[j] < 0.0f; + if (odd_parity && j == flip_index) + is_negative = !is_negative; + sign_mask |= static_cast(is_negative) << j; + } + const uint16_t code = static_cast(entry | ((sign_mask & 0x7f) << 9)); + const int code_offset = 2 + 2 * (group * 2 + vector); + payload[code_offset] = static_cast(code); + payload[code_offset + 1] = static_cast(code >> 8); + } + __syncthreads(); + } + } + + if (tid < 8) + payload[66 + tid] = locals[2 * tid] | (locals[2 * tid + 1] << 4); +} + +} // namespace + +at::Tensor iq2_xs_pack_cuda(at::Tensor input, at::Tensor grid) { + TORCH_CHECK(input.is_contiguous() && grid.is_contiguous(), "inputs must be contiguous"); + TORCH_CHECK(input.numel() > 0 && input.numel() % kBlockSize == 0, + "input size must be a positive multiple of 256"); + TORCH_CHECK(grid.scalar_type() == at::kFloat && grid.numel() == kEntries * kVectorSize, + "grid must be float32 [512, 8]"); + TORCH_CHECK(input.get_device() == grid.get_device(), "input and grid must share a device"); + c10::cuda::CUDAGuard guard(input.device()); + const int64_t num_blocks = input.numel() / kBlockSize; + TORCH_CHECK(num_blocks <= std::numeric_limits::max(), "IQ2_XS CUDA grid is too large"); + auto scales = at::empty({num_blocks}, input.options().dtype(at::kLong)); + auto output = at::empty({num_blocks, kPayloadBytes}, input.options().dtype(at::kByte)); + const auto stream = c10::cuda::getCurrentCUDAStream(); + const int scale_grid = static_cast((num_blocks + 255) / 256); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "iq2_xs_pack", [&] { + find_scale<<>>(input.data_ptr(), num_blocks, + scales.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + encode<<(num_blocks), 256, 0, stream>>>( + input.data_ptr(), num_blocks, grid.data_ptr(), + scales.data_ptr(), output.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + return output; +} diff --git a/modelopt/torch/quantization/__init__.py b/modelopt/torch/quantization/__init__.py index 87dbf30bb57..80dc83e7cb8 100644 --- a/modelopt/torch/quantization/__init__.py +++ b/modelopt/torch/quantization/__init__.py @@ -22,6 +22,7 @@ from .compress import * from .config import * from .conversion import * +from .ggml import * from .model_quant import * from .nn.modules.quant_module import QuantModuleRegistry from .utils import update_quant_cfg_with_kv_cache_quant diff --git a/modelopt/torch/quantization/extensions.py b/modelopt/torch/quantization/extensions.py index a65396d64ff..367f863ba4f 100644 --- a/modelopt/torch/quantization/extensions.py +++ b/modelopt/torch/quantization/extensions.py @@ -19,10 +19,18 @@ from modelopt.torch.utils import load_cpp_extension -__all__ = ["get_cuda_ext", "get_cuda_ext_fp8", "get_cuda_ext_mx", "precompile"] +__all__ = [ + "get_cuda_ext", + "get_cuda_ext_fp8", + "get_cuda_ext_iq1_s", + "get_cuda_ext_iq2_xs", + "get_cuda_ext_mx", + "precompile", +] path = Path(__file__).parent kernels_gemm = path.parent / "kernels" / "quantization" / "gemm" +kernels_ggml = path.parent / "kernels" / "quantization" / "ggml" def get_cuda_ext(raise_if_failed: bool = False): @@ -72,6 +80,34 @@ def get_cuda_ext_mx(raise_if_failed: bool = False): return get_cuda_ext_mx.extension # type:ignore[attr-defined] +def get_cuda_ext_iq1_s(raise_if_failed: bool = False): + """Return the GGML-compatible IQ1_S packing extension.""" + if not hasattr(get_cuda_ext_iq1_s, "extension"): + get_cuda_ext_iq1_s.extension = load_cpp_extension( # type:ignore[attr-defined] + name="modelopt_cuda_ext_iq1_s", + sources=[kernels_ggml / "iq1_s.cpp", kernels_ggml / "iq1_s.cu"], + cuda_version_specifiers=">=11.8", + fail_msg="IQ1_S CUDA packing is unavailable; using the PyTorch reference encoder.", + extra_cuda_cflags=["-O3", "--use_fast_math"], + raise_if_failed=raise_if_failed, + ) + return get_cuda_ext_iq1_s.extension # type:ignore[attr-defined] + + +def get_cuda_ext_iq2_xs(raise_if_failed: bool = False): + """Return the GGML-compatible IQ2_XS packing extension.""" + if not hasattr(get_cuda_ext_iq2_xs, "extension"): + get_cuda_ext_iq2_xs.extension = load_cpp_extension( # type:ignore[attr-defined] + name="modelopt_cuda_ext_iq2_xs", + sources=[kernels_ggml / "iq2_xs.cpp", kernels_ggml / "iq2_xs.cu"], + cuda_version_specifiers=">=11.8", + fail_msg="IQ2_XS CUDA packing is unavailable; using the PyTorch reference encoder.", + extra_cuda_cflags=["-O3", "--use_fast_math"], + raise_if_failed=raise_if_failed, + ) + return get_cuda_ext_iq2_xs.extension # type:ignore[attr-defined] + + def __getattr__(name): if name == "cuda_ext": return get_cuda_ext() @@ -79,6 +115,10 @@ def __getattr__(name): return get_cuda_ext_fp8() elif name == "cuda_ext_mx": return get_cuda_ext_mx() + elif name == "cuda_ext_iq1_s": + return get_cuda_ext_iq1_s() + elif name == "cuda_ext_iq2_xs": + return get_cuda_ext_iq2_xs() else: raise AttributeError(f"module {__name__} has no attribute {name}") @@ -88,3 +128,5 @@ def precompile(): print(get_cuda_ext()) print(get_cuda_ext_fp8()) print(get_cuda_ext_mx()) + print(get_cuda_ext_iq1_s()) + print(get_cuda_ext_iq2_xs()) diff --git a/modelopt/torch/quantization/ggml/__init__.py b/modelopt/torch/quantization/ggml/__init__.py new file mode 100644 index 00000000000..97d8863ce3f --- /dev/null +++ b/modelopt/torch/quantization/ggml/__init__.py @@ -0,0 +1,21 @@ +# 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. + +"""GGML-compatible block quantization formats.""" + +# Importing the backend installs its TensorQuantizer dispatch entry. +from . import backend as _backend +from .iq1_s import * +from .iq2_xs import * diff --git a/modelopt/torch/quantization/ggml/backend.py b/modelopt/torch/quantization/ggml/backend.py new file mode 100644 index 00000000000..95547659bb5 --- /dev/null +++ b/modelopt/torch/quantization/ggml/backend.py @@ -0,0 +1,35 @@ +# 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. + +"""TensorQuantizer backend dispatch for GGML-compatible IQ formats.""" + +import torch + +from ..nn.modules.tensor_quantizer import register_quant_backend +from .iq1_s import iq1_s_fake_quant +from .iq2_xs import iq2_xs_fake_quant + + +def ggml_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: + """Dispatch an IQ quantizer to its format-specific implementation.""" + num_bits = getattr(quantizer, "num_bits", None) + if num_bits == "iq1_s": + return iq1_s_fake_quant(inputs, quantizer) + if num_bits == "iq2_xs": + return iq2_xs_fake_quant(inputs, quantizer) + raise ValueError("The ggml backend requires num_bits='iq1_s' or 'iq2_xs'") + + +register_quant_backend("ggml", ggml_fake_quant) diff --git a/modelopt/torch/quantization/ggml/common.py b/modelopt/torch/quantization/ggml/common.py new file mode 100644 index 00000000000..4d95672c08e --- /dev/null +++ b/modelopt/torch/quantization/ggml/common.py @@ -0,0 +1,72 @@ +# 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. + +"""Shared validation for GGML-compatible block quantizers.""" + +import torch + +GGML_BLOCK_SIZE = 256 + + +def padded_weight_shape(shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]: + """Return ``shape`` with its last dimension rounded up to one GGML block.""" + if not shape: + raise ValueError("GGML block quantization requires a tensor with at least one dimension") + return (*shape[:-1], ((shape[-1] + GGML_BLOCK_SIZE - 1) // GGML_BLOCK_SIZE) * GGML_BLOCK_SIZE) + + +def pad_weight_rows(weight: torch.Tensor) -> torch.Tensor: + """Right-pad every logical row to a complete GGML block.""" + padded_shape = padded_weight_shape(weight.shape) + padding = padded_shape[-1] - weight.shape[-1] + return torch.nn.functional.pad(weight, (0, padding)) if padding else weight.contiguous() + + +def validate_weight(weight: torch.Tensor, format_name: str) -> None: + """Validate a weight accepted by the current GGML block encoders.""" + if weight.numel() == 0: + raise ValueError(f"{format_name} requires a non-empty weight") + if weight.dim() == 0: + raise ValueError(f"{format_name} requires a tensor with at least one dimension") + if not weight.is_floating_point(): + raise TypeError(f"{format_name} requires a floating-point weight, got {weight.dtype}") + if not torch.isfinite(weight).all(): + raise ValueError(f"{format_name} requires finite weight values") + + +def validate_packed_weights( + packed_weights: torch.Tensor, + weight_shape: torch.Tensor, + *, + block_bytes: int, + format_name: str, +) -> tuple[int, ...]: + """Validate a packed payload and return its logical shape.""" + if packed_weights.dtype != torch.uint8 or packed_weights.shape[-1] != block_bytes: + raise ValueError( + f"packed_weights must be uint8 with last dimension {block_bytes}, " + f"got {packed_weights.dtype} {tuple(packed_weights.shape)}" + ) + shape = tuple(int(v) for v in weight_shape.detach().cpu().tolist()) + if not shape or any(v <= 0 for v in shape): + raise ValueError(f"invalid {format_name} logical weight shape: {shape}") + padded_shape = padded_weight_shape(shape) + expected_shape = (*shape[:-1], padded_shape[-1] // GGML_BLOCK_SIZE, block_bytes) + if tuple(packed_weights.shape) != expected_shape: + raise ValueError( + f"packed_weights shape does not match {format_name} logical weight shape: " + f"expected {expected_shape}, got {tuple(packed_weights.shape)}" + ) + return shape diff --git a/modelopt/torch/quantization/ggml/iq1_s.py b/modelopt/torch/quantization/ggml/iq1_s.py new file mode 100644 index 00000000000..ac746ee4b82 --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq1_s.py @@ -0,0 +1,314 @@ +# This file includes the IQ1_S codebook adapted from: +# https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h +# +# MIT License +# +# Copyright (c) 2023-2026 The ggml authors +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# 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. + +"""IQ1_S fake quantization and GGML-compatible block packing. + +The encoder follows the canonical GGML-compatible ``search_impl="auto"`` search. Every complete +or right-padded 256-value row segment becomes one 50-byte ``block_iq1_s`` payload: + +* bytes 0..1: little-endian FP16 super-block scale ``d`` +* bytes 2..33: low eight bits of 32 codebook indices +* bytes 34..49: eight little-endian uint16 metadata words + +Each metadata word describes four consecutive eight-value vectors. Bits 0..11 +hold the three high index bits, bits 12..14 select one of eight local scales, +and bit 15 selects the shared -0.125 rather than +0.125 delta. The canonical +2048 x 8 ternary grid below comes from llama.cpp ``ggml-common.h`` revision +9b05354ec6fb58b4e665e9a39ebc40285c015638. +""" + +import base64 +import zlib +from functools import cache + +import torch + +from .common import ( + GGML_BLOCK_SIZE, + pad_weight_rows, + padded_weight_shape, + validate_packed_weights, + validate_weight, +) + +__all__ = [ + "IQ1_S_BLOCK_BYTES", + "IQ1_S_BLOCK_SIZE", + "IQ1_S_EFFECTIVE_BITS", + "dequantize_iq1_s", + "iq1_s_fake_quant", + "iq1_s_grid", + "quantize_iq1_s", +] + +IQ1_S_BLOCK_SIZE = GGML_BLOCK_SIZE +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_NATIVE_MAX = 16.875 + +# zlib-compressed little-endian bytes of the canonical uint64_t table. The +# decoded int8 values are -1, 0, and 1. +_IQ1_S_GRID_ZLIB_B64 = ( + "eNp1W4tWJEsII///0ew6lQTC6J7rxdaeflRBCAG73z/QVuUPQFtd/H2egHMiaIsfKF2RHwRtFeJCOBcE7f/T4wY4N4Jv+Pnv" + "c7fi89Z6END+3PZzTPtz+eLH33nP/vz4c8wHr3oWfIH6XPFd8H2OxwDPb/A5sF8YtP9NLADOQoD25257YXAWCLQ/rweuN/hY" + "4PO9Y8RC4iwoaH8uw/uC14kFx1l40NbZANB+noLffv4H2nobgrMxJfvOfI//fo59C22YvuDvKzYStJ+PL0/8nPb5fdsj3yNX" + "bHjRQseo8GTQViEcBLSzFM9BqtJhQFt2HPC8Z6FjpGP9LBO4rr85GmirjuN9bvz8F+/67xi0nwehHyxH/fjBclgcx5VTcVF8" + "Eo5Dg/az6yA+gKsJPv1yeBzHB+0HF1Yg4AQEaN/aEQ9WoHzcDAx0EAcwLivkKdnSMSLAPtsDLj+4/OByv/uA148AxAlEfX2W" + "YQUmToCCthio9SLjPQ6fD1zHHbha+s/lVyAD+nlFYOMEOBToga7lYJQTHhRy6F+AqBsd4c0l76UXavcGULzqHavgp+u4atkH" + "9OygfTD0AGT/RkD1YKceXry7CbBqXwG0AraihY5RkYkeDDAQ3+fBz/FJmKEK8WQvXF72AdcZzHTgxUD0BJ8LfB5wvUG02QAL" + "2ud+8+YCXtA+N5ncSoxwbPUG1wXMAhcGeTMouE1c3OJiPAxZDzM7IIAvWugYPCbwv2XcCaDB50Rk+Hq2aD8PVozmB4bYO69E" + "UrJ4Vrlaj6hYlodoTfQFfQ/+nh4EeVIhEhSQHtazBFVct+15OAmsaKFjJjQgPVUJrmihY2TiE2T3bEEkRCgx2vMrEmXRQsfg" + "MTKRgvblw4kgQQdzc8PHjA1+Nfg97fOQ53+bwUFMblyJCZq/p3XkMpFDCR38uawS/Huxh8cr4RctZEkAAFoSgaokAD8GC7P0" + "jvWSjhFF4IdaoIAhEHKajpuTOGE+LKIB2gKtQu0TUYwjMC4WIUGPjxeIahvxCjxvIL9WCIvICBlFaIoWOgaPkYSn3oY8PgY+" + "LZ+bzwdeLwgRPfHxr/feEGfA8qnPBUF/WilLS1VIIgXaKhIr8Pjd6d2eHyZ0gRkAvO+zeLajChtiJg4hbOmoHsD8PMRNpICs" + "fEBYQdCbVQ7BwyF6YmOPLQ3xE5t5eXqIoNiGIbonR0xWVBYTpRviSJQ26nWgxmRKRTmjltEJRgMJ3cOG5RWYXTuEVG9FRDF2" + "a41xCCtonYKYuUEi++gkSB9fpt5LCtqXl8H6bAgvDvEFrVKdiDCxhPmYfrAIck1ID4moeXImA4ItwcpBhdmUQ6xfPgXzKP1g" + "n6Kcp70GgogrRcu1xGhEzEv2AYJX9qUXMJ2AS7AJ/LNF+xwDhOlhSiL6oH1wPJSBzJseR0zj+mIzLS1FsUAw86KFCoZHPRjq" + "ruAfLEwhgVNQqFYQd+hdNK1CQ+Tf1KYXaQHjfBUiOAUJTmGCU6CAlktnRqnChZjqvVbqkGvIw9+2DgNVoQPatx2M46VYgAXQ" + "q7+4f2A8r8IIp0DSK/SK2V0o6cvUj4wYLKDeYzGOV0GFU1hpid7HptDCKbhAa+2it+pkecwyliK+UqaqP2QpJr0p2FzFdsg/" + "1VHtTYmFI/MceedUOapmVMV8yzsVbN2yjSuKTllGcozkl5LsUimrSD5xRdK/yyhGRrGNZlbPbG+Z48oY5pKUJyw/CGmFWkc2" + "kDzwoGcVtCzLy2JqZXlsESDL11u2+oxTjrriEuKf8rNEPkQKkGWnt1ygL3DuDSqrbDzloksjQUNluVenvDM3FFdRblBM1inM" + "T5nlVNNZPnllTnnkAt5lUUc5U6d8cZkCF/pRlqj88CudcsLlQ2WZoDKARHbKAEudtwxImu9a4A/6LtpuJVYSWR263knPS/T8" + "0HDRbjODSlo9lblKAsjjgh5f+lukvVMyyxVEV0VTyRXQQR8NPZX0UDTQHn7pn0vzpG9ftM3JS2B86BYIBla4k1a5Vjs0aujT" + "oUtKoQqFQ48uLRIdshBEOiT6U6I9ojeXxhz6UtQ5TV+kmJCuFOnKq1pWLd2/04Cb7p3emYmUvo0oTJ83PZ4u01eatMiv2qyy" + "i6Luh0XOyrRlqOlMR3XSj4tWpQ2c9CASa0EsVGCpuJMGBP+9tenvNGAQrVAbR2CrI7TJaY4KaNiuhOc6cGwf7IRdu2onvF44" + "rQOfVsg64bEOHAruRgLrgKc6cCQYGvghzODAiqvdDrj4Cw6c4xXOorJ1BMYThq4+GF4KK7Pw/j086oTDyoDh3k51nW5q9+xg" + "O27CuLbq393PGbbSnYw9x23sA51uUGe762yvtxW5vXW201njCqkqJSq3o87y32Wvr+Um2mhZkctbdZaxcrnqLI9DovK167yO" + "k1XnY/txKx+rzu29S+fyduo6QnHlj/1p/cMRkvUDIDulOqErT8QRnvVBnAt0nQuBUh9SqK4jWOtGOAL2vXFXPgDOg0joppLI" + "Sp10g/u8hfBCPnhXvkAvrXIL5bU0IDHH/cLUqS2o/7UA4l69cj3oHlt4r4U9G3q1cDjCvBYSZ0G7zsJSuJfSUiPFhpBPyYh8" + "73sjwI1AIzYEtA3+nFYNgdo97rVhahTUaRRo416tPRvXu9m1NrIrNxSn0VCn4QCJ85Ubj84N76FgsfF9Nx7ZUZGCJUeQNFKn" + "gVG7qNiOUukYXekgOI7SG6wXZexKB8JpjOgLx7G60sFwHa3depJjheN1pQPqEtTv39FyTM+aHAeFeh/HQfGHo6oxU7Jd4cB6" + "hV7vvJVDIQ2FvEeL18gI6Pii4AoAnvpULjBOIiB6TS/BQyY3INw46orA0JJ3TJ1MwHRMi2CaF085s3hKfd/iU69Zk+nak2X+" + "EnhWUOuPQMTvAYkbmNsZtGirEgYDFDdAe3dh3G1x6a7ugboFHao6XPPWcllsTLiB3VvNXIEe6tgEunuKqOhcathEQNBLW9rA" + "0KuW2kDRlYCBAxwQltEX0AhAaWeeXRWPpKbSsaMadbXprcgqEa4SLa5GFTWzLhfAcICsKwEMB8g6WNhwzAry8A1wXQl0OIAH" + "OeVI0r8CIA4QQrYRwEjUcENTHQS3LsqTXfR+emV42Wi/dRqhksy1il0JtPlW8FP1Ad6n84CNfJAXTIsfsBLEMysarlqCbkQD" + "FrLQMYEd5R6PNF8PbyzAF/a66Y3DUE5DV08qEd0ids9MxE4MvXo4Fi9WA9ijfX7zbAyr+OlNqjGpTK7QJ8E4KRC81VDGH4mn" + "h2llIuo1k0CxEEtDlxKnBNVdkah6XWordmZw0A62Sm7wc+DnsHe4cwzZjW6chjdO41u9FzHEM97rxriGTW7C7BzDdeNcLaM6" + "DXScRrqaOmagvUTrGXulKKqprEnMEmfUYuwcV3XiRsUYqhN553jpd0JHNvQVQTfBd453ugWulpmYdOd4pnuG6n1Jq5EyO4QB" + "MSiAMzCgllydwQGcAQKcQQITjo5xP3UPTDzk8mcMzzEkCciEpEOtc5NDBMWidU1Pco2HwapLxbiXNazOMS0jl6Adh+gIUzvI" + "4Ez2dI4/aezJrSol2zqDEQILLY4Q0+NAYjUdYz4mVMixniFYOY7jSzprKAvomdmyIGoJdTwhrKjvHGNxid05luIetryHvcYh" + "eF0x0PFF+N7djQk4Ax9qGavC6xzf8Blq9tZuKq6WTN/xClWIPb03fm7EwzUZjR7NxMX0IqJtTtLYJY8pA8cNOscK/EYKQhHZ" + "zrEAH6jXbaKb7X2XWJ3td6+UvkSMO9vnXkm1yz3S19H2NlU39VEF3dGW/iLcne1jE29cAp5tXg+N/EXMO6ebTdQ1iFNnEEep" + "V6HWu8myCL2aGCL2ne1Pi4kVU51DBbVV4qidbUOT8sr2nif7OqfXYBKgpNQxtTXUk07Z+yKrcOhsd7mAUGSIcd0CorONBAq5" + "/lOJzjaQZ8xVaCji3J3vaIu419DZ1jDD62xbOBVWdF2nUOlsM/gdOtsGjngP6yAHpqTYdMr37hqpIOqU4c1ATf0bWSilHO7C" + "qVPGNoOVXN0pP1tc75SVzXg7VVtTTKWSTpnXDNlct7cqNwNhlbKo17RSfXIxc2RNQ1alTAkgJ09N6jpUCRwZEPijMOzVW92T" + "rOioQv0wt4DUABsf3INswnrICqtSvnIHWKydeqv/dkms0qxRrK43e5reptiQ2AwqZBlrAMMSesspltCcpVjZdMoh/pubTpnD" + "hW4nKnFIYArgThTxTJTkBkVzR9RNodwpF7jykldpt+uU+TiFtDlnluX2TRXDLqcryuFFDhEdeYGzQFPgpiDqr0I8C/BeMbEr" + "x86yzMxB5ZjKrc4yysxCZZPLo45yxxKp3rCzDDEzEUdQzHaWEWYuKhuWUDBDT2swUhWxhlcgmzR/Rllq//WD6bgprVKTBIhO" + "umxxx3T4DGDCg5gk/Y+OmjJ10sslaAQNHM8gJHTSOl9MXEvY10nLfIpomP+YUDSo9zTaHgAN+mGGqOkm044cvkcnvbDwctI0" + "Tpq24KK0K1dWF11p8gownWnJSscVZJQGVDu3OwBbRJ4Ik6gpGBemuVjvgEdLJC5CCIudcAaTnA74wYEfgjV44/mTgj7CkHKA" + "MYalGPfEX53hbkaPM2CLM2ircFT13Bk2RqBO93eFIHdXLu50W3/JTcXmNTJh92r8KmyJPWn7le21ragUvLo3mo8YhTMIjDMQ" + "jDMYrMeWQNZ5e68uzuCwPq7TO3/ss/TvHzM5DA8=" +) + +_GRID_CACHE: dict[torch.device, torch.Tensor] = {} + + +@cache +def _grid_bytes() -> bytes: + return zlib.decompress(base64.b64decode(_IQ1_S_GRID_ZLIB_B64)) + + +def iq1_s_grid(device: torch.device | str | None = None) -> torch.Tensor: + """Return the canonical IQ1_S ternary grid as float32.""" + resolved_device = torch.device(device or "cpu") + if resolved_device not in _GRID_CACHE: + raw = torch.tensor(list(_grid_bytes()), dtype=torch.uint8).view(torch.int8) + _GRID_CACHE[resolved_device] = raw.reshape(2048, 8).to( + device=resolved_device, dtype=torch.float32 + ) + return _GRID_CACHE[resolved_device] + + +def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: + """Encode a moderate-size batch of flattened 256-value blocks.""" + x = blocks.float() + block_count = x.shape[0] + vectors = x.reshape(block_count, 32, 8) + xnorm = vectors.square().sum(dim=-1) + xsum = vectors.sum(dim=-1) + + amax = x.abs().amax(dim=1) + d = ((amax / _IQ1_S_NATIVE_MAX) * 0.61).clamp(max=65504.0).to(torch.float16) + d_float = d.float() + + best_error = torch.full((block_count, 32, 16), torch.inf, device=x.device) + best_entry = torch.zeros((block_count, 32, 16), dtype=torch.int64, device=x.device) + grid_norm = grid.square().sum(dim=-1) + grid_sum = grid.sum(dim=-1) + + # Tile the 2048-entry codebook to bound temporary memory. A strict update + # retains the lowest codebook index when two candidates have equal error. + for entry_start in range(0, 2048, 128): + grid_tile = grid[entry_start : entry_start + 128] + dot = torch.matmul(vectors, grid_tile.T) + tile_norm = grid_norm[entry_start : entry_start + 128].reshape(1, 1, -1) + tile_sum = grid_sum[entry_start : entry_start + 128].reshape(1, 1, -1) + + for shift in range(2): + delta = -_IQ1_S_DELTA if shift else _IQ1_S_DELTA + shifted_dot = dot + delta * xsum.unsqueeze(-1) + shifted_norm = tile_norm + 2 * delta * tile_sum + 8 * delta * delta + for local in range(8): + choice = shift * 8 + local + scale = d_float.reshape(-1, 1, 1) * (2 * local + 1) + error = ( + xnorm.unsqueeze(-1) - 2 * scale * shifted_dot + scale.square() * shifted_norm + ) + tile_error, tile_index = error.min(dim=-1) + replace = tile_error < best_error[:, :, choice] + best_error[:, :, choice] = torch.where( + replace, tile_error, best_error[:, :, choice] + ) + best_entry[:, :, choice] = torch.where( + replace, tile_index + entry_start, best_entry[:, :, choice] + ) + + group_error = best_error.reshape(block_count, 8, 4, 16).sum(dim=2) + selected_choice = group_error.argmin(dim=-1) + vector_choice = selected_choice.repeat_interleave(4, dim=1) + selected_entry = best_entry.gather(2, vector_choice.unsqueeze(-1)).squeeze(-1) + selected_local = selected_choice & 0x7 + selected_shift = selected_choice >> 3 + + high = (selected_entry >> 8).reshape(block_count, 8, 4) + qh = ( + high[:, :, 0] + | (high[:, :, 1] << 3) + | (high[:, :, 2] << 6) + | (high[:, :, 3] << 9) + | (selected_local << 12) + | (selected_shift << 15) + ) + + packed = torch.empty((block_count, IQ1_S_BLOCK_BYTES), dtype=torch.uint8, device=x.device) + packed[:, :2] = d.contiguous().view(torch.uint8).reshape(block_count, 2) + packed[:, 2:34] = (selected_entry & 0xFF).to(torch.uint8) + packed[:, 34:50:2] = (qh & 0xFF).to(torch.uint8) + packed[:, 35:50:2] = (qh >> 8).to(torch.uint8) + packed[d_float == 0] = 0 + return packed + + +@torch.no_grad() +def quantize_iq1_s( + weight: torch.Tensor, *, block_chunk_size: int = 4 +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack a floating-point weight into GGML-compatible IQ1_S blocks. + + Each logical row is right-padded to a multiple of 256. Returned shapes are + ``[*weight.shape[:-1], ceil(weight.shape[-1] / 256), 50]`` and ``[weight.ndim]``. + Both tensors remain on the weight's device. + """ + validate_weight(weight, "IQ1_S") + if block_chunk_size <= 0: + raise ValueError(f"block_chunk_size must be positive, got {block_chunk_size}") + + logical_shape = torch.tensor(weight.shape, dtype=torch.int64, device=weight.device) + padded_weight = pad_weight_rows(weight) + blocks = padded_weight.reshape(-1, IQ1_S_BLOCK_SIZE) + blocks_per_row = padded_weight.shape[-1] // IQ1_S_BLOCK_SIZE + grid = iq1_s_grid(weight.device) + if weight.is_cuda: + from ..extensions import get_cuda_ext_iq1_s + + extension = get_cuda_ext_iq1_s() + if extension is not None: + packed = extension.pack(blocks, grid) + packed_shape = ( + *weight.shape[:-1], + blocks_per_row, + IQ1_S_BLOCK_BYTES, + ) + return packed.reshape(packed_shape), logical_shape + + chunks = [ + _encode_blocks(blocks[start : start + block_chunk_size], grid) + for start in range(0, blocks.shape[0], block_chunk_size) + ] + packed_shape = ( + *weight.shape[:-1], + blocks_per_row, + IQ1_S_BLOCK_BYTES, + ) + return torch.cat(chunks).reshape(packed_shape), logical_shape + + +@torch.no_grad() +def dequantize_iq1_s( + packed_weights: torch.Tensor, + weight_shape: torch.Tensor, + *, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Decode GGML-compatible IQ1_S payload bytes.""" + shape = validate_packed_weights( + packed_weights, weight_shape, block_bytes=IQ1_S_BLOCK_BYTES, format_name="IQ1_S" + ) + + blocks = packed_weights.contiguous().reshape(-1, IQ1_S_BLOCK_BYTES) + 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) + shifts = torch.tensor([0, 3, 6, 9], dtype=torch.int64, device=blocks.device) + 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 = iq1_s_grid(blocks.device)[entries] + delta.unsqueeze(-1).unsqueeze(-1) + scales = d.unsqueeze(-1) * (2 * local + 1).float() + decoded = values * scales.unsqueeze(-1).unsqueeze(-1) + padded_shape = padded_weight_shape(shape) + return decoded.reshape(padded_shape)[..., : shape[-1]].to(dtype) + + +def iq1_s_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: + """IQ1_S backend for TensorQuantizer, with pass-through backward.""" + if getattr(quantizer, "num_bits", None) != "iq1_s": + raise ValueError("The ggml IQ1_S backend requires num_bits='iq1_s'") + extra_args = getattr(quantizer, "backend_extra_args", None) or {} + search_impl = extra_args.get("search_impl", extra_args.get("iq_search_impl", "auto")) + if search_impl != "auto": + raise NotImplementedError("Only IQ1_S search_impl='auto' is currently supported") + packed, shape = quantize_iq1_s(inputs) + reconstructed = dequantize_iq1_s(packed, shape, dtype=inputs.dtype) + return inputs + (reconstructed - inputs).detach() diff --git a/modelopt/torch/quantization/ggml/iq2_xs.py b/modelopt/torch/quantization/ggml/iq2_xs.py new file mode 100644 index 00000000000..cd3ab39760e --- /dev/null +++ b/modelopt/torch/quantization/ggml/iq2_xs.py @@ -0,0 +1,311 @@ +# This file includes the IQ2_XS codebook adapted from: +# https://github.com/ggml-org/llama.cpp/blob/9b05354ec6fb58b4e665e9a39ebc40285c015638/ggml/src/ggml-common.h +# +# MIT License +# +# Copyright (c) 2023-2026 The ggml authors +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# 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. + +"""IQ2_XS fake quantization and GGML-compatible block packing. + +The encoder follows the canonical GGML-compatible search_impl="auto" search. Every complete or +right-padded 256-value row segment becomes one 74-byte block_iq2_xs payload: + +* bytes 0..1: little-endian FP16 super-block scale d +* bytes 2..65: 32 little-endian uint16 codes (9-bit grid + 7-bit sign) +* bytes 66..73: 16 four-bit local scales, two per byte + +The canonical 512 x 8 magnitude grid below comes from llama.cpp +ggml-common.h revision 9b05354ec6fb58b4e665e9a39ebc40285c015638. +""" + +import base64 +from functools import cache + +import torch + +from .common import ( + GGML_BLOCK_SIZE, + pad_weight_rows, + padded_weight_shape, + validate_packed_weights, + validate_weight, +) + +__all__ = [ + "IQ2_XS_BLOCK_BYTES", + "IQ2_XS_BLOCK_SIZE", + "IQ2_XS_EFFECTIVE_BITS", + "dequantize_iq2_xs", + "iq2_xs_fake_quant", + "iq2_xs_grid", + "quantize_iq2_xs", +] + +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 + +# Compact byte representation of the canonical [512, 8] grid. Values are only +# 8, 25, and 43. Keeping this as checkpoint-independent package data avoids +# adding a pickle-backed torch.save artifact to the wheel. +_IQ2_XS_GRID_B64 = ( + "CAgICAgICAgrCAgICAgICBkZCAgICAgICCsICAgICAgrKwgICAgICBkIGQgICAgICBkZCAgICAgrGRkICAgICBkrGQgICAgICAgr" + "CAgICAgrCCsICAgICBkZKwgICAgICCsrCAgICAgZCAgZCAgICAgZCBkICAgIKxkIGQgICAgZKwgZCAgICAgIGRkICAgIKwgZGQgI" + "CAgZGRkZCAgICAgrGRkICAgIGQgrGQgICAgIGSsZCAgICAgICCsICAgIKwgIKwgICAgZGQgrCAgICAgrCCsICAgIGQgZKwgICAgI" + "GRkrCAgICBkrGSsICAgICAgrKwgICAgZCAgIGQgICAgZCAgZCAgIKxkICBkICAgZKwgIGQgICAgIGQgZCAgIKwgZCBkICAgZGRkI" + "GQgICAgrGQgZCAgIKysZCBkICAgZCCsIGQgICAgZKwgZCAgICAgIGRkICAgrCAgZGQgICBkZCBkZCAgICCsIGRkICAgZCBkZGQgI" + "CAgZGRkZCAgICAgrGRkICAgIKysZGQgICBkICCsZCAgICBkIKxkICAgICBkrGQgICAgICAgrCAgIKwgICCsICAgZGQgIKwgICAgr" + "CAgrCAgIGQgZCCsICAgIGRkIKwgICAgIKwgrCAgIGQgIGSsICAgIGQgZKwgICAgIGRkrCAgIGRkZGSsICAgICAgrKwgICCsrCCsr" + "CAgIGQgICAgZCAgIGQgICBkICCsZCAgIGQgIGSsICAgZCAgICBkICBkICCsIGQgIGQgIGRkZCAgZCAgIKxkICBkICBkIKwgIGQgI" + "CBkrCAgZCAgICAgZCBkICCsICBkIGQgIGRkIGQgZCAgIKwgZCBkICBkIGRkIGQgICBkZGQgZCAgrGRkZCBkICAgIKxkIGQgIGQgI" + "KwgZCAgIGQgrCBkICAgIGSsIGQgICAgICBkZCAgrCAgIGRkICBkZCAgZGQgICCsICBkZCAgZCBkIGRkICAgZGQgZGQgICAgrCBkZ" + "CAgZCAgZGRkICAgZCBkZGQgICAgZGRkZCAgZCCsZGRkICAgICCsZGQgIGQgICCsZCAgIGQgIKxkICAgIGQgrGQgIKxkrCCsZCAgI" + "CAgZKxkICCsICBkrGQgICBkIKysZCAgICAgICCsICCsICAgIKwgIGRkICAgrCAgIKwgICCsICCsrCAgIKwgIGQgZCAgrCAgIGRkI" + "CCsICAgIKwgIKwgIGRkrCAgrCAgZCAgZCCsICAgZCBkIKwgICAgZGQgrCAgIKxkZCCsICAgICCsIKwgICAgrKwgrCAgrKysrCCsI" + "CBkICAgZKwgICBkICBkrCAgICBkIGSsICAgICBkZKwgIGQgIKxkrCAgZKwgrGSsICAgICAgrKwgICAgrCCsrCAgIKysIKysICCsZ" + "GSsrKwgICAgrKysrCAgZCAgICAgZCAgZCAgICBkIKxkICAgIGQgZKwgICAgZCAgIGQgICBkIKwgZCAgIGQgZGRkICAgZCAgrGQgI" + "CBkIGQgrCAgIGQgIGSsICAgZCAgICBkICBkIKwgIGQgIGQgZGQgZCAgZCAgrCBkICBkIGQgZGQgIGQgIGRkZCAgZCAgIKxkICBkI" + "KysrGQgIGQgZCAgrCAgZCAgZCCsICBkICAgZKwgIGQgICAgIGQgZCCsICAgZCBkIGRkICBkIGQgIKwgIGQgZCBkIGQgZCBkICBkZ" + "CBkIGQgICCsIGQgZCBkICBkZCBkICBkIGRkIGQgICBkZGQgZCAgICCsZCBkICBkZKxkIGQgrGRkrGQgZCBkICAgrCBkICBkICCsI" + "GQgrGQgIKwgZCAgIGQgrCBkICAgIGSsIGQgICCsZKwgZCAgICAgIGRkIKwgICAgZGQgZGQgICBkZCAgrCAgIGRkIGQgZCAgZGQgI" + "GRkICBkZCAgIKwgIGRkIGQgIGQgZGQgIGQgZCBkZCBkrCBkIGRkICAgZGQgZGQgIGSsZCBkZCAgICCsIGRkIGQgICBkZGQgIGQgI" + "GRkZCAgIGQgZGRkICAgIGRkZGQgICAgIKxkZCAgZGQgrGRkIGSsIGSsZGQgZCAgICCsZCAgZCAgIKxkICAgZCAgrGQgrCBkICCsZ" + "CAgICBkIKxkICBkZGQgrGQgrGQgrCCsZCAgICAgZKxkIGRkICBkrGQgrGSsZGSsZCBkIGRkrKxkIGSsrKysrGQgICAgICAgrCCsI" + "CAgICCsIGRkICAgIKwgIKwgICAgrCCsrCAgICCsIGQgZCAgIKwgIGRkICAgrCAgIKwgICCsIGQgIGQgIKwgIGQgZCAgrCAgIGRkI" + "CCsICAgIKwgIKwgICCsrCAgrCBkICAgZCCsICBkICBkIKwgICBkIGQgrCAgICBkZCCsICCsIGRkIKwgZGSsZGQgrCAgICAgrCCsI" + "KwgrCCsIKwgICAgrKwgrCAgrKysrCCsIGQgICAgZKwgIGQgICBkrCAgIGQgIGSsIGSsrCAgZKwgICAgZCBkrCAgICAgZGSsIGQgI" + "GRkZKwgrCBkZGRkrCBkrGSsZGSsIGQgICCsZKwgrKxkIKxkrCCsZKysrGSsICAgICAgrKwgIKwgICCsrCCsrCAgIKysICAgrCAgr" + "KwgZGRkZCCsrCAgrCCsIKysIKwgrKwgrKwgIKysZGSsrCAgIGSsZKysICCsICCsrKwgICCsIKysrCCsICCsrKysICCsIKysrKwgr" + "KwgrKysrCBkICAgICAgZCBkICAgICBkrGQgICAgIGRkrCAgICAgZCAgZCAgICBkrCBkICAgIGRkZGQgICAgZCCsZCAgICBkZCCsI" + "CAgIGQgZKwgICAgZCAgIGQgICBkrCAgZCAgIGRkZCBkICAgZCCsIGQgICBkrKwgZCAgIGRkIGRkICAgZCBkZGQgICBkICCsZCAgI" + "GRkZKxkICAgZGQgIKwgICBkIGQgrCAgIGQgIGSsICAgZCAgICBkICBkrCAgIGQgIGRkZCAgZCAgZCCsICBkICBkZCBkIGQgIGQgZ" + "GQgZCAgZCAgrCBkICBkZCAgZGQgIGQgZCBkZCAgZCAgZGRkICBkICAgrGQgIGRkZCCsZCAgZKwgrKxkICBkZCAgIKwgIGQgZCAgr" + "CAgZCAgZCCsICBkrCBkIKwgIGRkrKwgrCAgZCAgIGSsICBkICAgICBkIGSsICAgIGQgZGRkICAgZCBkIKwgICBkIGRkIGQgIGQgZ" + "CBkZCAgZCBkZKxkICBkIGQgIKwgIGQgZGQgIGQgZCBkIGQgZCBkIGQgIGRkIGQgZCAgIKwgZCBkIGRkrCBkIGRkICAgZGQgZCBkI" + "CBkZCBkICBkIGRkIGQgZKwgZGQgZCAgIGRkZCBkrKxkrGRkIGQgICAgrGQgZKysICCsZCBkIGQgZKxkIGQgIGRkrGQgZGQgICAgr" + "CBkIGQgICCsIGQgIGQgIKwgZCAgIGQgrCBkZGQgZCCsIGQgZGRkIKwgZKwgrGQgrCBkICAgIGSsIGRkIGQgZKwgZCBkIGRkrCBkI" + "CBkZGSsIGRkrKxkZKwgZCBkICCsrCBkICAgICAgZGSsICAgICBkZGRkICAgIGRkIKwgICAgZGRkIGQgICBkZCBkZCAgIGRkICCsI" + "CAgZGQgrKwgICBkZGQgIGQgIGRkIGQgZCAgZGQgIGRkICBkZCAgIKwgIGRkZCAgIGQgZGQgZCAgZCBkZCAgZCBkIGRkZGRkIGQgZ" + "GQgICBkZCBkZKwgIGRkIGRkICAgIKwgZGQgZCBkrCBkZKysrKysIGRkZCAgICBkZGQgZCAgIGRkZCAgZCAgZGRkZCCsICBkZGQgI" + "CBkIGRkZCAgrGQgZGRkZCAgrCBkZGRkIKysIGRkZCAgICBkZGRkIKwgIGRkZGQgICCsZGRkZCCsIKxkZGRkZCCsIKxkZGQgrKxkr" + "GRkZGQgrKysZGRkICAgICCsZGQgZGQgIKxkZGQgIGQgrGRkICBkZCCsZGRkrGSsIKxkZKysZCBkrGRkICAgZGSsZGSsICBkZKxkZ" + "GRkIKysrGRkZCAgICAgrGQgZCAgICCsZCAgZCAgIKxkICAgZCAgrGQgZGRkICCsZKwgrGQgIKxkrGQgrCAgrGRkrKysICCsZCAgI" + "CBkIKxkIGSsIKwgrGSsrCBkrCCsZKwgZKysIKxkICAgICBkrGSsZGQgIGSsZCAgZCBkZKxkICAgZGRkrGRkZCBkZGSsZCBkrKxkZ" + "KxkZCAgICCsrGSsrKxkIKysZGRkrCBkrKxkrGQgIKysrGQgZGRkrKysZKwgrGSsrKxkICAgICAgIKysICAgICAgrGRkICAgICCsI" + "KwgICAgIKxkIGQgICAgrCBkZCAgICCsICCsICAgIKysrKwgICAgrGQgIGQgICCsIGQgZCAgIKwgIGRkICAgrCAgIKwgICCsrCAgr" + "CAgIKwgrKysICAgrKysrKwgICCsZCAgIGQgIKwgZCAgZCAgrKxkICBkICCsICBkIGQgIKwgICBkZCAgrGQgZGRkICCsZKxkZGQgI" + "KwgICAgrCAgrCAgrCCsICCsICAgrKwgIKysICCsrCAgrCAgrKysICCsIKysrKwgIKxkICAgIGQgrCBkICAgZCCsICBkICBkIKysI" + "GQgIGQgrGRkZCAgZCCsICAgZCBkIKwgIKxkIGQgrGSsIKwgZCCsICAgIGRkIKwgZCBkZGQgrGRkrKxkZCCsIKxkIKxkIKysrKxkr" + "GQgrCAgICAgrCCsIKwgICCsIKxkZKwgIKwgrKysZGQgrCCsICAgrCCsIKysICCsIKwgrCCsrKwgrCCsrGQgIGSsIKysIKwgrKwgr" + "CAgIKysrCCsIKwgrKysIKysZGSsrKwgrCCsrKysrCCsZCAgICAgZKwgZCAgICBkrCAgZCAgIGSsICAgZCAgZKysZGRkICBkrCBkI" + "KwgIGSsICAgIGQgZKysIKwgZCBkrCBkrGRkIGSsrGRkZKwgZKxkrCCsrCBkrCAgICAgZGSsZGQgICBkZKwgZCBkIGRkrCAgZGQgZ" + "GSsIKxkZCBkZKxkrKwgZGRkrCAgZKxkZGSsrCBkrGRkZKxkICBkrGRkrGQgZGQgrGSsrGSsrCCsZKxkrCBkZKxkrGRkZCCsrGSsI" + "CCsZKysZKwgICAgICCsrKwgICAgIKysIKwgICAgrKysrCAgICCsrCAgrCAgIKysrKysICAgrKwgIKysICCsrGQgZGRkIKysZKxkZ" + "GQgrKysZKysZCCsrCAgICCsIKysrCAgIKwgrKwgrCAgrCCsrKysrCCsIKysICAgrKwgrKwgIKysrCCsrCAgIGQgZKysZGRkrCBkr" + "KxkZKxkrGSsrCCsZKysZKysrKwgICCsrKwgIKwgIKysrKwgrCAgrKysIKysICCsrKwgIKysIKysrCCsrKwgrKysIGQgIGSsrKwgZ" + "CCsZKysrKxkIKxkrKysIKysIKysrKysrKwgrKysrGQgZKysrKysrKysrKysrKw==" +) + +_GRID_CACHE: dict[torch.device, torch.Tensor] = {} + + +@cache +def _grid_bytes() -> bytes: + return base64.b64decode(_IQ2_XS_GRID_B64) + + +def iq2_xs_grid(device: torch.device | str | None = None) -> torch.Tensor: + """Return the canonical IQ2_XS magnitude grid as float32.""" + resolved_device = torch.device(device or "cpu") + if resolved_device not in _GRID_CACHE: + values = torch.tensor(list(_grid_bytes()), dtype=torch.float32) + _GRID_CACHE[resolved_device] = values.reshape(512, 8).to(device=resolved_device) + return _GRID_CACHE[resolved_device] + + +def _encode_blocks(blocks: torch.Tensor, grid: torch.Tensor) -> torch.Tensor: + """Encode a moderate-size batch of flattened 256-value blocks.""" + x = blocks.float() + block_count = x.shape[0] + vectors = x.reshape(block_count, 32, 8) + magnitudes = vectors.abs() + negative = vectors < 0 + odd_parity = negative.sum(dim=-1).remainder(2).bool() + + amax = x.abs().amax(dim=1) + rms = x.square().mean(dim=1).sqrt() + peak_to_rms = torch.where(rms > 0, amax / rms, torch.zeros_like(rms)) + anchor_ratio = (1.0 - 0.035 * peak_to_rms).clamp(0.65, 0.92) + d = ((amax / 166.625) * anchor_ratio).clamp(max=65504.0).to(torch.float16) + d_float = d.float() + + xnorm = vectors.square().sum(dim=-1) + qnorm = grid.square().sum(dim=-1) + best_error = torch.full((block_count, 32, 16), torch.inf, dtype=torch.float32, device=x.device) + best_entry = torch.zeros((block_count, 32, 16), dtype=torch.int64, device=x.device) + # Search the codebook in tiles to cap temporary memory. Strict comparison + # preserves the lowest grid index on equal error, matching the CUDA key. + for entry_start in range(0, 512, 64): + grid_tile = grid[entry_start : entry_start + 64] + products = magnitudes.unsqueeze(2) * grid_tile.reshape(1, 1, -1, 8) + dot = products.sum(dim=-1) + dot = torch.where(odd_parity.unsqueeze(-1), dot - 2.0 * products.amin(dim=-1), dot) + tile_qnorm = qnorm[entry_start : entry_start + 64].reshape(1, 1, -1) + + for local in range(16): + scale = d_float.reshape(-1, 1, 1) * ((2 * local + 1) / 8.0) + error = ( + xnorm.unsqueeze(-1) - 2.0 * scale * dot + scale.square() * tile_qnorm + ).clamp_min_(0) + tile_error, tile_index = error.min(dim=-1) + replace = tile_error < best_error[:, :, local] + best_error[:, :, local] = torch.where(replace, tile_error, best_error[:, :, local]) + best_entry[:, :, local] = torch.where( + replace, tile_index + entry_start, best_entry[:, :, local] + ) + + group_error = best_error.reshape(block_count, 16, 2, 16).sum(dim=2) + selected_local = group_error.argmin(dim=-1) + vector_local = selected_local.repeat_interleave(2, dim=1) + selected_entry = best_entry.gather(2, vector_local.unsqueeze(-1)).squeeze(-1) + + selected_grid = grid[selected_entry] + weakest_index = (magnitudes * selected_grid).argmin(dim=-1) + flip = torch.nn.functional.one_hot(weakest_index, num_classes=8).bool() + encoded_negative = negative ^ (flip & odd_parity.unsqueeze(-1)) + sign_bits = torch.arange(8, dtype=torch.int64, device=x.device) + sign_mask = (encoded_negative.to(torch.int64) << sign_bits).sum(dim=-1) + + codes = selected_entry | ((sign_mask & 0x7F) << 9) + packed = torch.empty((block_count, IQ2_XS_BLOCK_BYTES), dtype=torch.uint8, device=x.device) + packed[:, :2] = d.contiguous().view(torch.uint8).reshape(block_count, 2) + packed[:, 2:66:2] = (codes & 0xFF).to(torch.uint8) + packed[:, 3:66:2] = (codes >> 8).to(torch.uint8) + packed[:, 66:] = (selected_local[:, 0::2] | (selected_local[:, 1::2] << 4)).to(torch.uint8) + return packed + + +@torch.no_grad() +def quantize_iq2_xs( + weight: torch.Tensor, *, block_chunk_size: int = 64 +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack a floating-point weight into GGML-compatible IQ2_XS blocks. + + Each logical row is right-padded to a multiple of 256. Returned shapes are + ``[*weight.shape[:-1], ceil(weight.shape[-1] / 256), 74]`` and ``[weight.ndim]``. + Both tensors remain on the weight's device. + """ + validate_weight(weight, "IQ2_XS") + if block_chunk_size <= 0: + raise ValueError(f"block_chunk_size must be positive, got {block_chunk_size}") + + logical_shape = torch.tensor(weight.shape, dtype=torch.int64, device=weight.device) + padded_weight = pad_weight_rows(weight) + blocks = padded_weight.reshape(-1, IQ2_XS_BLOCK_SIZE) + blocks_per_row = padded_weight.shape[-1] // IQ2_XS_BLOCK_SIZE + grid = iq2_xs_grid(weight.device) + if weight.is_cuda: + from ..extensions import get_cuda_ext_iq2_xs + + extension = get_cuda_ext_iq2_xs() + if extension is not None: + packed = extension.pack(blocks, grid) + packed_shape = ( + *weight.shape[:-1], + blocks_per_row, + IQ2_XS_BLOCK_BYTES, + ) + return packed.reshape(packed_shape), logical_shape + + chunks = [ + _encode_blocks(blocks[start : start + block_chunk_size], grid) + for start in range(0, blocks.shape[0], block_chunk_size) + ] + packed_shape = ( + *weight.shape[:-1], + blocks_per_row, + IQ2_XS_BLOCK_BYTES, + ) + return torch.cat(chunks).reshape(packed_shape), logical_shape + + +@torch.no_grad() +def dequantize_iq2_xs( + packed_weights: torch.Tensor, + weight_shape: torch.Tensor, + *, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Decode GGML-compatible IQ2_XS payload bytes.""" + shape = validate_packed_weights( + packed_weights, weight_shape, block_bytes=IQ2_XS_BLOCK_BYTES, format_name="IQ2_XS" + ) + + blocks = packed_weights.contiguous().reshape(-1, IQ2_XS_BLOCK_BYTES) + 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 + + parity = torch.zeros_like(sign_index) + for bit in range(7): + parity ^= (sign_index >> bit) & 1 + sign_mask = sign_index | (parity << 7) + bit_positions = torch.arange(8, dtype=torch.int64, device=blocks.device) + 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 + scales = d.unsqueeze(-1) * (2 * local + 1).float() / 8.0 + values = iq2_xs_grid(blocks.device)[entries] * signs + decoded = values * scales.repeat_interleave(2, dim=1).unsqueeze(-1) + padded_shape = padded_weight_shape(shape) + return decoded.reshape(padded_shape)[..., : shape[-1]].to(dtype) + + +def iq2_xs_fake_quant(inputs: torch.Tensor, quantizer) -> torch.Tensor: + """IQ2_XS backend for TensorQuantizer, with pass-through backward.""" + if getattr(quantizer, "num_bits", None) != "iq2_xs": + raise ValueError("The ggml IQ2_XS backend requires num_bits='iq2_xs'") + extra_args = getattr(quantizer, "backend_extra_args", None) or {} + search_impl = extra_args.get("search_impl", extra_args.get("iq_search_impl", "auto")) + if search_impl != "auto": + raise NotImplementedError("Only IQ2_XS search_impl='auto' is currently supported") + packed, shape = quantize_iq2_xs(inputs) + reconstructed = dequantize_iq2_xs(packed, shape, dtype=inputs.dtype) + return inputs + (reconstructed - inputs).detach() diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 18b97ac2774..9b8a7880bdc 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -179,7 +179,7 @@ class TensorQuantizer(nn.Module): "ds_grads_remaining", "ds_id", "pre_bwd_fn", - # quantizer cache for custom backends, like luts + # Quantizer cache for registered format-specific backends. "_quantizer_cache", # Runtime-only set of storage attributes tied to shared state. The tied # aliases are rebuilt from calibration config and tensor state during restore. diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 2515910ec6e..58f8851402d 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -2117,6 +2117,14 @@ def _reconstruct_fused_moe_linear(model: nn.Module) -> None: torch.stack([getattr(experts[i], attr) for i in range(n)]), ) + for attr in ("weight_logical_shape", "weight_padded_shape"): + if not all(hasattr(experts[i], attr) for i in range(n)): + continue + expert_shape = getattr(experts[0], attr) + if not all(torch.equal(getattr(experts[i], attr), expert_shape) for i in range(1, n)): + raise ValueError(f"Cannot reconstruct fused experts with inconsistent {attr}") + module.register_buffer(attr, torch.cat([expert_shape.new_tensor([n]), expert_shape])) + # Remove expanded experts — the reconstructed 3D tensors replace them del module.experts diff --git a/modelopt_recipes/configs/numerics/iq1_s.yaml b/modelopt_recipes/configs/numerics/iq1_s.yaml new file mode 100644 index 00000000000..3aeb910a774 --- /dev/null +++ b/modelopt_recipes/configs/numerics/iq1_s.yaml @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# IQ1_S weight quantizer using the built-in GGML-compatible search. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig +num_bits: iq1_s +effective_bits: 1.5625 +block_sizes: + -1: 256 +backend: ggml +backend_extra_args: + search_impl: auto +pass_through_bwd: true diff --git a/modelopt_recipes/configs/numerics/iq2_xs.yaml b/modelopt_recipes/configs/numerics/iq2_xs.yaml new file mode 100644 index 00000000000..74bf9128abe --- /dev/null +++ b/modelopt_recipes/configs/numerics/iq2_xs.yaml @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# IQ2_XS weight quantizer using the built-in GGML-compatible search. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig +num_bits: iq2_xs +effective_bits: 2.3125 +block_sizes: + -1: 256 +backend: ggml +backend_extra_args: + search_impl: auto +pass_through_bwd: true diff --git a/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml b/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml new file mode 100644 index 00000000000..31fab0f7a44 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq1_s.yaml @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# QuantizeConfig preset for IQ1_S weight-only quantization. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + iq1_s: configs/numerics/iq1_s + +algorithm: +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: iq1_s + - quantizer_name: '*input_quantizer' + enable: false + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml b/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml new file mode 100644 index 00000000000..96938c79f09 --- /dev/null +++ b/modelopt_recipes/configs/ptq/presets/model/iq2_xs.yaml @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# QuantizeConfig preset for IQ2_XS weight-only quantization. + +# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + iq2_xs: configs/numerics/iq2_xs + +algorithm: +quant_cfg: + - $import: base_disable_all + - quantizer_name: '*weight_quantizer' + cfg: + $import: iq2_xs + - quantizer_name: '*input_quantizer' + enable: false + - $import: default_disabled_quantizers diff --git a/modelopt_recipes/general/ptq/iq1_s.yaml b/modelopt_recipes/general/ptq/iq1_s.yaml new file mode 100644 index 00000000000..06c328a8da8 --- /dev/null +++ b/modelopt_recipes/general/ptq/iq1_s.yaml @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# IQ1_S weight-only PTQ and unified-checkpoint export. + +imports: + preset: configs/ptq/presets/model/iq1_s + +metadata: + recipe_type: ptq + description: >- + Applies GGML-compatible IQ1_S weight-only quantization. No calibration data is required; + unified export writes the GGML block payload to weight for every quantized Linear weight. +quantize: + $import: preset diff --git a/modelopt_recipes/general/ptq/iq2_xs.yaml b/modelopt_recipes/general/ptq/iq2_xs.yaml new file mode 100644 index 00000000000..1131a5a29c7 --- /dev/null +++ b/modelopt_recipes/general/ptq/iq2_xs.yaml @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# IQ2_XS weight-only PTQ and unified-checkpoint export. + +imports: + preset: configs/ptq/presets/model/iq2_xs + +metadata: + recipe_type: ptq + description: >- + Applies GGML-compatible IQ2_XS weight-only quantization. No calibration data is required; + unified export writes the GGML block payload to weight for every quantized Linear weight. +quantize: + $import: preset diff --git a/modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/iq2_xs_experts-nvfp4_mamba.yaml b/modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/iq2_xs_experts-nvfp4_mamba.yaml new file mode 100644 index 00000000000..48d467b8a19 --- /dev/null +++ b/modelopt_recipes/models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/iq2_xs_experts-nvfp4_mamba.yaml @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Mixed weight-only quantization for NVIDIA Nemotron 3.5 Lightning: +# IQ2_XS routed/shared MoE experts, dynamic NVFP4 Mamba projections, and BF16 elsewhere. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + iq2_xs: configs/numerics/iq2_xs + nvfp4: configs/numerics/nvfp4 + +metadata: + recipe_type: ptq + description: >- + Calibration-free mixed weight-only quantization for NVIDIA Nemotron 3.5 Lightning: + IQ2_XS routed and shared MoE expert weights, dynamic NVFP4 Mamba in_proj/out_proj + weights, and BF16 attention, routers, latent projections, embeddings, lm_head, + activations, and MTP modules. + +quantize: + algorithm: max + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*mixer.experts*weight_quantizer' + cfg: + $import: iq2_xs + - quantizer_name: '*mixer.shared_experts*weight_quantizer' + cfg: + $import: iq2_xs + # Megatron Core names the same expert modules under ``mlp`` rather than ``mixer``. + - quantizer_name: '*mlp.experts*weight_quantizer' + cfg: + $import: iq2_xs + - quantizer_name: '*mlp.shared_experts*weight_quantizer' + cfg: + $import: iq2_xs + - quantizer_name: '*mixer.in_proj*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mixer.out_proj*weight_quantizer' + cfg: + $import: nvfp4 + # Keep MTP in BF16 if its module names overlap the selectors above. + - quantizer_name: '*mtp*' + enable: false diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 2da6a0ff2b0..88a2fe3d8f9 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -28,7 +28,7 @@ supported combinations. ### The shipped recipes
-All 25 general/ptq/ recipes (click to expand) +All 27 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -57,6 +57,8 @@ supported combinations. | `int4_blockwise_weight_only` | INT4 W4A16, block 128, weights only | none | max | | `nvfp4_mlp_weight_only` | NVFP4 W4A16 (block 32), MLP + MoE weights only | none | max | | `mxfp4_mlp_weight_only` | MXFP4 W4A16, MLP + MoE weights only | none | none (no calibration) | +| `iq1_s` | IQ1_S W1A16, all linears | none | GGML-compatible auto search (no calibration) | +| `iq2_xs` | IQ2_XS W2A16, all linears | none | GGML-compatible auto search (no calibration) |
@@ -135,6 +137,10 @@ activations and tensor-core math are what deliver the throughput. - **`mxfp4_mlp_weight_only`** — MXFP4 weights on MLP/MoE layers only, BF16 activations. Needs no calibration forward pass; the QAT starting point for the GPT-OSS family (see `examples/gpt-oss`). +- **`iq1_s` / `iq2_xs`** — GGML-compatible IQ1_S or IQ2_XS weights on all linear + layers, with BF16 activations. Unified export stores each logical weight as a packed + 50- or 74-byte-per-256-values payload plus its original shape. No calibration data is + required. --- @@ -410,6 +416,18 @@ checkpoint's** quant config verbatim: Four-over-Six NVFP4 W4A16 to routed experts, shared experts, and the language model head; Mamba `in/out_proj` weights and inputs plus the KV cache use FP8, while attention remains BF16. +- **`models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/iq2_xs_experts-nvfp4_mamba`** + applies a calibration-free mixed weight-only policy: routed and shared MoE expert + weights use IQ2_XS, Mamba `in_proj` and `out_proj` weights use dynamic NVFP4, and + attention, routers, embeddings, `lm_head`, activations, and MTP modules remain BF16. The + corresponding `tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/` + `mbridge_iq2_xs_export_validate.yaml` pipeline exports a unified-HF checkpoint and records + CPU/CUDA packing parity, tensor-policy, packed-layout, and payload-digest evidence. To verify + the packed bytes independently, build an unmodified pinned llama.cpp checkout as shared + libraries and run `examples/megatron_bridge/validate_iq2_xs_stock_ggml.py`. The validator + samples every IQ2_XS tensor, calls the stock `dequantize_row_iq2_xs` symbol, requires + bit-identical FP32 reconstruction, and can round-trip the samples through stock `gguf-py` to + verify payload SHA-256 digests. - **`models/nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16/ptq/nvfp4_w4a16`** mirrors the GGUF **Q4_K_M** bit allocation of the Nemotron-H hybrid, mapped onto NVFP4/FP8 **per layer**: Q4_K/Q5_0 linears → NVFP4 W4A4 (attention q/k/v/o kept uniform so export can diff --git a/tests/gpu/torch/quantization/test_iq1_s_cuda.py b/tests/gpu/torch/quantization/test_iq1_s_cuda.py new file mode 100644 index 00000000000..11cc10ff9c9 --- /dev/null +++ b/tests/gpu/torch/quantization/test_iq1_s_cuda.py @@ -0,0 +1,54 @@ +# 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 torch + +from modelopt.torch.quantization.ggml.iq1_s import dequantize_iq1_s, quantize_iq1_s + + +def test_iq1_s_cuda_pack_is_deterministic_and_decodable(): + generator = torch.Generator(device="cuda").manual_seed(1234) + weight = torch.randn((8, 256), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed, shape = quantize_iq1_s(weight) + packed_again, _ = quantize_iq1_s(weight) + reconstructed = dequantize_iq1_s(packed, shape) + + assert packed.shape == (8, 1, 50) + assert torch.equal(packed, packed_again) + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.25 + + +def test_iq1_s_cuda_zero_encoding_matches_ggml_block_layout(): + weight = torch.zeros((1, 256), device="cuda", dtype=torch.bfloat16) + packed, shape = quantize_iq1_s(weight) + + assert not packed.any() + assert torch.equal(dequantize_iq1_s(packed, shape), weight) + + +def test_iq1_s_cuda_pack_is_byte_identical_to_cpu(): + generator = torch.Generator().manual_seed(5918) + weight = torch.randn((8, 257), generator=generator, dtype=torch.bfloat16) + + packed_cpu, shape_cpu = quantize_iq1_s(weight) + packed_cuda, shape_cuda = quantize_iq1_s(weight.cuda()) + + assert torch.equal(shape_cuda.cpu(), shape_cpu) + assert torch.equal(packed_cuda.cpu(), packed_cpu) + assert packed_cpu.shape == (8, 2, 50) diff --git a/tests/gpu/torch/quantization/test_iq2_xs_cuda.py b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py new file mode 100644 index 00000000000..fcd6c6860ee --- /dev/null +++ b/tests/gpu/torch/quantization/test_iq2_xs_cuda.py @@ -0,0 +1,54 @@ +# 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 torch + +from modelopt.torch.quantization.ggml.iq2_xs import dequantize_iq2_xs, quantize_iq2_xs + + +def test_iq2_xs_cuda_pack_is_deterministic_and_decodable(): + generator = torch.Generator(device="cuda").manual_seed(1234) + weight = torch.randn((8, 512), generator=generator, device="cuda", dtype=torch.bfloat16) + + packed, shape = quantize_iq2_xs(weight) + packed_again, _ = quantize_iq2_xs(weight) + reconstructed = dequantize_iq2_xs(packed, shape) + + assert packed.shape == (8, 2, 74) + assert torch.equal(packed, packed_again) + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.1 + + +def test_iq2_xs_cuda_zero_encoding_matches_ggml_block_layout(): + weight = torch.zeros((1, 256), device="cuda", dtype=torch.bfloat16) + packed, shape = quantize_iq2_xs(weight) + + assert not packed.any() + assert torch.equal(dequantize_iq2_xs(packed, shape), weight) + + +def test_iq2_xs_cuda_pack_is_byte_identical_to_cpu(): + generator = torch.Generator().manual_seed(5918) + weight = torch.randn((8, 257), generator=generator, dtype=torch.bfloat16) + + packed_cpu, shape_cpu = quantize_iq2_xs(weight) + packed_cuda, shape_cuda = quantize_iq2_xs(weight.cuda()) + + assert torch.equal(shape_cuda.cpu(), shape_cpu) + assert torch.equal(packed_cuda.cpu(), packed_cpu) + assert packed_cpu.shape == (8, 2, 74) diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py index 2da25e90e15..6c37710bb0f 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py @@ -46,6 +46,9 @@ import modelopt.torch.speculative as mtsp from modelopt.torch.export import KV_CACHE_FP8, export_mcore_gpt_to_hf, import_mcore_gpt_from_hf from modelopt.torch.export.unified_export_megatron import GPTModelExporter +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.ggml import dequantize_iq1_s, dequantize_iq2_xs +from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.speculative.eagle.default_config import default_eagle_config from modelopt.torch.speculative.plugins.megatron_eagle import _DynamicEagleGPTModel from modelopt.torch.speculative.plugins.megatron_medusa import _DynamicMedusaGPTModel @@ -86,6 +89,73 @@ def _verify_model_quant_config( assert quant_config_dict["kv_cache_quant_algo"] == KV_CACHE_FP8 +@pytest.mark.parametrize( + ("qformat", "payload_bytes", "dequantize"), + [("iq1_s", 50, dequantize_iq1_s), ("iq2_xs", 74, dequantize_iq2_xs)], +) +def test_megatron_name_remapping_exports_iq_payload(qformat, payload_bytes, dequantize): + """Megatron export writes the same scale-free IQ representation as HF export.""" + linear = torch.nn.Linear(257, 2, bias=False, dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=qformat, + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + ) + exporter = object.__new__(GPTModelExporter) + exporter.dtype = torch.bfloat16 + exporter._state_dict = {} + exporter.exclude_modules = [] + exporter.layer_config_dict = {} + + exporter._name_remapping(linear, "model.layers.0.mlp.down_proj.") + + packed_key = "model.layers.0.mlp.down_proj.weight" + assert exporter._state_dict[packed_key].shape == (2, 2, payload_bytes) + assert exporter._state_dict[packed_key].dtype == torch.uint8 + logical_shape = exporter._state_dict["model.layers.0.mlp.down_proj.weight_logical_shape"] + padded_shape = exporter._state_dict["model.layers.0.mlp.down_proj.weight_padded_shape"] + assert logical_shape.tolist() == [2, 257] + assert padded_shape.tolist() == [2, 512] + reconstructed = dequantize( + exporter._state_dict[packed_key], + logical_shape, + dtype=torch.bfloat16, + ) + torch.testing.assert_close(reconstructed, linear.weight_quantizer(linear.weight)) + assert exporter.layer_config_dict == { + "model.layers.0.mlp.down_proj.quantization": qformat, + "model.layers.0.mlp.down_proj.awq_block_size": 256, + } + + +def test_megatron_iq_export_rejects_tensor_parallelism(): + """IQ packing is intentionally limited to complete TP=1 weights.""" + linear = torch.nn.Linear(256, 2, bias=False, dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + ) + exporter = object.__new__(GPTModelExporter) + exporter.model = torch.nn.Sequential(linear) + + with ( + patch.object(exporter, "_is_sidecar_writer_rank", return_value=False), + patch.object(uem, "get_pipeline_model_parallel_rank", return_value=0), + patch.object(uem, "get_pipeline_model_parallel_world_size", return_value=1), + patch.object(uem, "get_tensor_model_parallel_rank", return_value=0), + patch.object(uem, "get_tensor_model_parallel_world_size", return_value=2), + pytest.raises(NotImplementedError, match="tensor model parallel size 1"), + ): + exporter.save_pretrained("unused", "unused") + + def _test_unified_export_megatron( tmp_path, model_type, diff --git a/tests/unit/examples/test_materialize_mixed_iq_gguf.py b/tests/unit/examples/test_materialize_mixed_iq_gguf.py new file mode 100644 index 00000000000..4f1c74453f8 --- /dev/null +++ b/tests/unit/examples/test_materialize_mixed_iq_gguf.py @@ -0,0 +1,121 @@ +# 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. + +"""Tests for the mixed-IQ GGUF validation bridge.""" + +import importlib.util +import sys +from pathlib import Path +from types import ModuleType + +import pytest +import torch + +_SCRIPT = ( + Path(__file__).resolve().parents[3] + / "examples" + / "megatron_bridge" + / "materialize_mixed_iq_gguf.py" +) +_SPEC = importlib.util.spec_from_file_location("materialize_mixed_iq_gguf", _SCRIPT) +assert _SPEC is not None and _SPEC.loader is not None +bridge = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = bridge +_SPEC.loader.exec_module(bridge) + + +def test_packed_rows_preserves_bytes_and_recovers_logical_shape(): + tensor = torch.arange(2 * 3 * 74, dtype=torch.int64).remainder(256).to(torch.uint8) + tensor = tensor.reshape(2, 3, 74) + + raw, logical_shape = bridge.packed_rows(tensor) + + assert raw.shape == (2, 3 * 74) + assert raw.tobytes() == tensor.numpy().tobytes() + assert logical_shape == [2, 3 * 256] + + +@pytest.mark.parametrize( + "tensor", + [ + torch.zeros((2, 3, 74), dtype=torch.int8), + torch.zeros((2, 3, 73), dtype=torch.uint8), + torch.zeros((2, 74), dtype=torch.uint8), + ], +) +def test_packed_rows_rejects_noncanonical_payloads(tensor): + with pytest.raises(ValueError, match="IQ2_XS payload"): + bridge.packed_rows(tensor) + + +def test_iq2_layer_matching_normalizes_model_prefix(tmp_path): + (tmp_path / "hf_quant_config.json").write_text( + """{ + "quantization": { + "quantized_layers": { + "model.layers.1.mixer.experts": {"quant_algo": "IQ2_XS"}, + "model.layers.2.mixer.in_proj": {"quant_algo": "NVFP4"} + } + } + }""" + ) + + layers = bridge._load_iq2_layers(tmp_path) + + assert layers == {"backbone.layers.1.mixer.experts"} + assert bridge._matches_iq2_layer("backbone.layers.1.mixer.experts.7.up_proj.weight", layers) + assert not bridge._matches_iq2_layer("backbone.layers.2.mixer.in_proj.weight", layers) + + +def test_materialize_declares_model_arch_for_pinned_converter(monkeypatch, tmp_path): + registered = {} + + class FakeModelBase: + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + if "model_arch" not in cls.__dict__: + raise TypeError(f"Missing property 'model_arch' for {cls.__name__!r}") + + @classmethod + def register(cls, architecture): + def decorator(model_cls): + registered[architecture] = model_cls + return model_cls + + return decorator + + class FakeNemotronHModel(FakeModelBase): + model_arch = "nemotron_h" + + conversion = ModuleType("conversion") + conversion.__path__ = [] + base = ModuleType("conversion.base") + base.LazyTorchTensor = object + base.ModelBase = FakeModelBase + nemotron = ModuleType("conversion.nemotron") + nemotron.NemotronHModel = FakeNemotronHModel + + monkeypatch.setitem(sys.modules, "gguf", ModuleType("gguf")) + monkeypatch.setitem(sys.modules, "conversion", conversion) + monkeypatch.setitem(sys.modules, "conversion.base", base) + monkeypatch.setitem(sys.modules, "conversion.nemotron", nemotron) + monkeypatch.setattr(bridge.runpy, "run_path", lambda *args, **kwargs: None) + monkeypatch.setattr(bridge, "_git_commit", lambda source: "pinned-commit") + + expected, commit = bridge.materialize(tmp_path, tmp_path, tmp_path / "model.gguf") + + assert expected == {} + assert commit == "pinned-commit" + assert registered["NemotronHForCausalLM"].model_arch == "nemotron_h" diff --git a/tests/unit/examples/test_validate_ggml_row_alignment.py b/tests/unit/examples/test_validate_ggml_row_alignment.py new file mode 100644 index 00000000000..13284189beb --- /dev/null +++ b/tests/unit/examples/test_validate_ggml_row_alignment.py @@ -0,0 +1,109 @@ +# 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. + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file + +_SCRIPT = ( + Path(__file__).parents[3] / "examples" / "megatron_bridge" / "validate_ggml_row_alignment.py" +) +_SPEC = importlib.util.spec_from_file_location("validate_ggml_row_alignment", _SCRIPT) +assert _SPEC and _SPEC.loader +_MODULE = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_MODULE) + + +def _write_indexed_checkpoint(checkpoint: Path) -> None: + first = { + "model.layers.0.mixer.experts.0.up_proj.weight": torch.zeros((8, 256)), + "model.layers.0.mixer.experts.0.down_proj.weight": torch.zeros((8, 320)), + } + second = { + "model.layers.0.mixer.shared_experts.up_proj.weight": torch.zeros((16, 512)), + "mtp.layers.0.mixer.experts.0.up_proj.weight": torch.zeros((8, 384)), + } + save_file(first, checkpoint / "model-00001-of-00002.safetensors") + save_file(second, checkpoint / "model-00002-of-00002.safetensors") + weight_map = { + name: shard + for shard, tensors in ( + ("model-00001-of-00002.safetensors", first), + ("model-00002-of-00002.safetensors", second), + ) + for name in tensors + } + (checkpoint / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) + + +def test_validate_row_alignment_reports_incompatible_shapes(tmp_path: Path) -> None: + _write_indexed_checkpoint(tmp_path) + + report = _MODULE.validate_row_alignment( + tmp_path, + block_size=256, + includes=["*mixer.experts*.weight", "*mixer.shared_experts*.weight"], + excludes=["mtp.*"], + ) + + assert report["status"] == "failed" + assert report["selected_tensors"] == 3 + assert report["compatible_tensors"] == 2 + assert report["incompatible_tensors"] == 1 + assert report["incompatible"] == [ + { + "name": "model.layers.0.mixer.experts.0.down_proj.weight", + "shape": [8, 320], + "row_width": 320, + "remainder": 64, + "compatible": False, + } + ] + + +def test_validate_row_alignment_passes_aligned_selection(tmp_path: Path) -> None: + save_file( + {"model.layers.0.mixer.experts.0.up_proj.weight": torch.zeros((8, 512))}, + tmp_path / "model.safetensors", + ) + + report = _MODULE.validate_row_alignment( + tmp_path, + block_size=256, + includes=["*experts*.weight"], + excludes=[], + ) + + assert report["status"] == "passed" + assert report["compatible_tensors"] == 1 + assert report["incompatible"] == [] + + +def test_validate_row_alignment_requires_matches(tmp_path: Path) -> None: + save_file({"model.embed.weight": torch.zeros((8, 256))}, tmp_path / "model.safetensors") + + with pytest.raises(ValueError, match="No tensors matched"): + _MODULE.validate_row_alignment( + tmp_path, + block_size=256, + includes=["*experts*.weight"], + excludes=[], + ) diff --git a/tests/unit/examples/test_validate_iq2_xs_reconstruction.py b/tests/unit/examples/test_validate_iq2_xs_reconstruction.py new file mode 100644 index 00000000000..a20e91982cf --- /dev/null +++ b/tests/unit/examples/test_validate_iq2_xs_reconstruction.py @@ -0,0 +1,119 @@ +# 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. + +"""Tests for the IQ2_XS source-reconstruction validator.""" + +import importlib.util +import sys +from pathlib import Path + +import torch +from safetensors.torch import load_file, save_file + +from modelopt.torch.quantization.ggml import quantize_iq2_xs + +_EXAMPLES = Path(__file__).resolve().parents[3] / "examples" / "megatron_bridge" +_STRUCTURAL_SCRIPT = _EXAMPLES / "validate_mixed_quantized_hf.py" +_STRUCTURAL_SPEC = importlib.util.spec_from_file_location( + "validate_mixed_quantized_hf", _STRUCTURAL_SCRIPT +) +assert _STRUCTURAL_SPEC is not None and _STRUCTURAL_SPEC.loader is not None +structural_validator = importlib.util.module_from_spec(_STRUCTURAL_SPEC) +sys.modules[_STRUCTURAL_SPEC.name] = structural_validator +_STRUCTURAL_SPEC.loader.exec_module(structural_validator) + +_SCRIPT = _EXAMPLES / "validate_iq2_xs_reconstruction.py" +_SPEC = importlib.util.spec_from_file_location("validate_iq2_xs_reconstruction", _SCRIPT) +assert _SPEC is not None and _SPEC.loader is not None +validator = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = validator +_SPEC.loader.exec_module(validator) + + +def _write_fixture(tmp_path: Path) -> tuple[Path, Path, str]: + checkpoint = tmp_path / "checkpoint" + reference = tmp_path / "reference" + checkpoint.mkdir() + reference.mkdir() + name = "model.layers.0.mixer.experts.0.up_proj.weight" + generator = torch.Generator().manual_seed(1234) + source_weight = torch.randn((2, 257), generator=generator, dtype=torch.float32).bfloat16() + packed, logical_shape = quantize_iq2_xs(source_weight) + save_file({name: source_weight}, reference / "model.safetensors") + save_file( + { + name: packed, + name.removesuffix(".weight") + ".weight_logical_shape": logical_shape.cpu(), + name.removesuffix(".weight") + ".weight_padded_shape": torch.tensor([2, 512]), + }, + checkpoint / "model.safetensors", + ) + return checkpoint, reference, name + + +def test_reconstruction_matches_direct_packing_and_reports_partial_tail(tmp_path): + checkpoint, reference, name = _write_fixture(tmp_path) + + report = validator.validate_reconstruction( + checkpoint, + reference, + rows_per_tensor=2, + maximum_tensors=None, + require_repack_match=True, + ) + + assert report["status"] == "passed" + assert report["errors"] == [] + assert report["summary"]["available_iq2_xs_tensors"] == 1 + assert report["summary"]["sampled_tensors"] == 1 + assert report["summary"]["sampled_rows"] == 2 + assert report["summary"]["repack_matches"] == 2 + assert report["summary"]["repack_mismatches"] == 0 + assert report["summary"]["complete_blocks"]["values"] == 512 + assert report["summary"]["partial_tails"]["values"] == 2 + assert all(sample["tensor"] == name for sample in report["samples"]) + assert all(sample["partial_tail"]["values"] == 1 for sample in report["samples"]) + + +def test_reconstruction_detects_source_row_mismatch(tmp_path): + checkpoint, reference, name = _write_fixture(tmp_path) + tensors = load_file(reference / "model.safetensors") + tensors[name] = tensors[name].flip(0) + save_file(tensors, reference / "model.safetensors") + + report = validator.validate_reconstruction( + checkpoint, + reference, + rows_per_tensor=2, + maximum_tensors=None, + require_repack_match=True, + ) + + assert report["status"] == "failed" + assert report["summary"]["repack_matches"] == 0 + assert report["summary"]["repack_mismatches"] == 2 + assert all(sample["mismatched_bytes"] > 0 for sample in report["samples"]) + assert all("differs from direct packing" in error for error in report["errors"]) + + +def test_tensor_selection_is_even_and_honors_patterns(): + names = [f"layer.{index}.experts.weight" for index in range(10)] + + assert validator._select_tensors(names, (), 3) == [names[0], names[4], names[9]] + assert validator._select_tensors(names, ("layer.[135].*",), None) == [ + names[1], + names[3], + names[5], + ] diff --git a/tests/unit/examples/test_validate_mixed_quantized_hf.py b/tests/unit/examples/test_validate_mixed_quantized_hf.py new file mode 100644 index 00000000000..23a93f3f52c --- /dev/null +++ b/tests/unit/examples/test_validate_mixed_quantized_hf.py @@ -0,0 +1,269 @@ +# 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. + +"""Tests for the mixed unified-HF checkpoint validator.""" + +import importlib.util +import json +import sys +from pathlib import Path + +import numpy as np +import torch +from safetensors.torch import load_file, save_file + +_SCRIPT = ( + Path(__file__).resolve().parents[3] + / "examples" + / "megatron_bridge" + / "validate_mixed_quantized_hf.py" +) +_SPEC = importlib.util.spec_from_file_location("validate_mixed_quantized_hf", _SCRIPT) +assert _SPEC is not None and _SPEC.loader is not None +validator = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = validator +_SPEC.loader.exec_module(validator) + +_ORACLE_SCRIPT = ( + Path(__file__).resolve().parents[3] + / "examples" + / "megatron_bridge" + / "validate_iq2_xs_stock_ggml.py" +) +_ORACLE_SPEC = importlib.util.spec_from_file_location("validate_iq2_xs_stock_ggml", _ORACLE_SCRIPT) +assert _ORACLE_SPEC is not None and _ORACLE_SPEC.loader is not None +oracle = importlib.util.module_from_spec(_ORACLE_SPEC) +sys.modules[_ORACLE_SPEC.name] = oracle +_ORACLE_SPEC.loader.exec_module(oracle) + + +def _write_fixture(tmp_path: Path) -> tuple[Path, Path]: + source = tmp_path / "source" + checkpoint = tmp_path / "checkpoint" + source.mkdir() + checkpoint.mkdir() + + source_tensors = { + "model.embed_tokens.weight": torch.ones((8, 256), dtype=torch.bfloat16), + "model.layers.0.mixer.in_proj.weight": torch.ones((8, 256), dtype=torch.bfloat16), + "model.layers.0.mixer.out_proj.weight": torch.ones((8, 256), dtype=torch.bfloat16), + "model.layers.1.mixer.experts.0.up_proj.weight": torch.ones((8, 257), dtype=torch.bfloat16), + "model.layers.1.mixer.shared_experts.down_proj.weight": torch.ones( + (8, 257), dtype=torch.bfloat16 + ), + "model.layers.1.mixer.gate.weight": torch.ones((8, 256), dtype=torch.bfloat16), + "mtp.layers.0.mixer.experts.0.up_proj.weight": torch.ones((8, 256), dtype=torch.bfloat16), + } + save_file(source_tensors, source / "model.safetensors") + + checkpoint_tensors = { + "model.embed_tokens.weight": source_tensors["model.embed_tokens.weight"], + "model.layers.0.mixer.in_proj.weight": torch.arange(8 * 128, dtype=torch.int64) + .remainder(256) + .to(torch.uint8) + .reshape(8, 128), + "model.layers.0.mixer.in_proj.weight_scale": torch.ones((8, 16), dtype=torch.float8_e4m3fn), + "model.layers.0.mixer.in_proj.weight_scale_2": torch.tensor(1.0), + "model.layers.0.mixer.out_proj.weight": torch.zeros((8, 128), dtype=torch.uint8), + "model.layers.0.mixer.out_proj.weight_scale": torch.ones( + (8, 16), dtype=torch.float8_e4m3fn + ), + "model.layers.0.mixer.out_proj.weight_scale_2": torch.tensor(1.0), + "model.layers.1.mixer.experts.0.up_proj.weight": torch.arange(8 * 2 * 74, dtype=torch.int64) + .remainder(256) + .to(torch.uint8) + .reshape(8, 2, 74), + "model.layers.1.mixer.experts.0.up_proj.weight_logical_shape": torch.tensor([8, 257]), + "model.layers.1.mixer.experts.0.up_proj.weight_padded_shape": torch.tensor([8, 512]), + "model.layers.1.mixer.shared_experts.down_proj.weight": torch.zeros( + (8, 2, 74), dtype=torch.uint8 + ), + "model.layers.1.mixer.shared_experts.down_proj.weight_logical_shape": torch.tensor( + [8, 257] + ), + "model.layers.1.mixer.shared_experts.down_proj.weight_padded_shape": torch.tensor([8, 512]), + "model.layers.1.mixer.gate.weight": source_tensors["model.layers.1.mixer.gate.weight"], + } + save_file(checkpoint_tensors, checkpoint / "model.safetensors") + + quantized_layers = { + "model.layers.0.mixer.in_proj": { + "quant_algo": "W4A16_NVFP4", + "group_size": 16, + }, + "model.layers.0.mixer.out_proj": { + "quant_algo": "W4A16_NVFP4", + "group_size": 16, + }, + "model.layers.1.mixer.experts": { + "quant_algo": "IQ2_XS", + "group_size": 256, + "block_payload_bytes": 74, + "packing": "ggml", + "row_padding": "right", + "logical_shape_key": "weight_logical_shape", + "padded_shape_key": "weight_padded_shape", + }, + "model.layers.1.mixer.shared_experts.down_proj": { + "quant_algo": "IQ2_XS", + "group_size": 256, + "block_payload_bytes": 74, + "packing": "ggml", + "row_padding": "right", + "logical_shape_key": "weight_logical_shape", + "padded_shape_key": "weight_padded_shape", + }, + } + (checkpoint / "hf_quant_config.json").write_text( + json.dumps( + { + "producer": {"name": "modelopt", "version": "test"}, + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": None, + "quantized_layers": quantized_layers, + }, + } + ) + ) + return checkpoint, source + + +def test_validate_checkpoint_passes_and_hashes_iq2_payloads(tmp_path): + checkpoint, source = _write_fixture(tmp_path) + + report = validator.validate_checkpoint(checkpoint, source) + + assert report["status"] == "passed" + assert report["errors"] == [] + assert report["summary"] == { + "total_tensors": 14, + "total_weight_tensors": 6, + "iq2_xs_layers": 2, + "iq2_xs_tensors": 2, + "iq2_xs_blocks": 32, + "iq2_xs_logical_weights": 4112, + "iq2_xs_padded_weights": 8192, + "iq2_xs_payload_bytes": 2368, + "nvfp4_layers": 2, + "nvfp4_tensors": 2, + "unquantized_weight_tensors": 2, + } + assert len(report["iq2_xs"]["aggregate_sha256"]) == 64 + assert all(len(tensor["sha256"]) == 64 for tensor in report["iq2_xs"]["tensors"]) + + +def test_checkpoint_index_selects_authoritative_duplicate(tmp_path): + save_file( + { + "shared.weight": torch.ones((2, 2)), + "first.weight": torch.ones((1, 2)), + }, + tmp_path / "model-00001-of-00002.safetensors", + ) + save_file( + { + "shared.weight": torch.zeros((3, 2)), + "second.weight": torch.ones((1, 3)), + }, + tmp_path / "model-00002-of-00002.safetensors", + ) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "weight_map": { + "first.weight": "model-00001-of-00002.safetensors", + "second.weight": "model-00002-of-00002.safetensors", + "shared.weight": "model-00002-of-00002.safetensors", + } + } + ) + ) + + records = validator.read_checkpoint_index(tmp_path) + + assert set(records) == {"first.weight", "second.weight", "shared.weight"} + assert records["shared.weight"].shard == "model-00002-of-00002.safetensors" + assert records["shared.weight"].shape == (3, 2) + + +def test_validate_checkpoint_rejects_wrong_iq2_payload_shape(tmp_path): + checkpoint, source = _write_fixture(tmp_path) + state = load_file(checkpoint / "model.safetensors") + state["model.layers.1.mixer.experts.0.up_proj.weight"] = torch.zeros( + (8, 1, 73), dtype=torch.uint8 + ) + save_file(state, checkpoint / "model.safetensors") + + report = validator.validate_checkpoint(checkpoint, source, compute_digests=False) + + assert report["status"] == "failed" + assert any("expected [..., blocks, 74]" in error for error in report["errors"]) + + +def test_stock_ggml_oracle_samples_every_iq2_tensor(tmp_path, monkeypatch): + checkpoint, _ = _write_fixture(tmp_path) + + def modelopt_decode(_library, packed_bytes): + num_blocks = len(packed_bytes) // 74 + packed = torch.frombuffer(bytearray(packed_bytes), dtype=torch.uint8).reshape( + num_blocks, 1, 74 + ) + shape = torch.tensor([num_blocks, 256], dtype=torch.int64) + return oracle.dequantize_iq2_xs(packed, shape, dtype=torch.float32).numpy().reshape(-1) + + monkeypatch.setattr(oracle, "_stock_decode", modelopt_decode) + monkeypatch.setattr(oracle, "_llama_commit", lambda _source: "abc123") + + report = oracle.compare_checkpoint( + checkpoint, + tmp_path / "libggml-base.so", + tmp_path / "llama.cpp", + blocks_per_tensor=3, + ) + + assert report["status"] == "passed" + assert report["llama_cpp"]["commit"] == "abc123" + assert report["summary"] == { + "iq2_xs_tensors": 2, + "sampled_blocks": 6, + "decoded_values": 1536, + "bitwise_differences": 0, + "max_abs_difference": 0.0, + } + assert len(report["samples"]) == 6 + assert len(report["encoded_fields"]["sample_payload_sha256"]) == 64 + + +def test_stock_ggml_oracle_reports_bit_difference(tmp_path, monkeypatch): + checkpoint, _ = _write_fixture(tmp_path) + + def mismatching_decode(_library, packed_bytes): + output = np.zeros(len(packed_bytes) // 74 * 256, dtype=np.float32) + output[0] = 1.0 + return output + + monkeypatch.setattr(oracle, "_stock_decode", mismatching_decode) + monkeypatch.setattr(oracle, "_llama_commit", lambda _source: "abc123") + + report = oracle.compare_checkpoint( + checkpoint, + tmp_path / "libggml-base.so", + tmp_path / "llama.cpp", + blocks_per_tensor=1, + ) + + assert report["status"] == "failed" + assert report["summary"]["bitwise_differences"] > 0 diff --git a/tests/unit/recipe/test_nemotron_lightning_iq2_xs_recipe.py b/tests/unit/recipe/test_nemotron_lightning_iq2_xs_recipe.py new file mode 100644 index 00000000000..914bb197a32 --- /dev/null +++ b/tests/unit/recipe/test_nemotron_lightning_iq2_xs_recipe.py @@ -0,0 +1,180 @@ +# 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. + +"""Quantizer-placement tests for the Nemotron 3.5 Lightning mixed IQ2_XS recipe.""" + +import json + +import pytest +import torch +import transformers +from _test_utils.torch.transformers_models import get_tiny_nemotron_h +from safetensors.torch import load_file + +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.export import export_hf_checkpoint +from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format +from modelopt.torch.export.quant_utils import get_quant_config +from modelopt.torch.quantization.config import need_calibration +from modelopt.torch.quantization.conversion import set_quantizer_by_cfg +from modelopt.torch.quantization.nn import TensorQuantizer + +_RECIPE = "models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/iq2_xs_experts-nvfp4_mamba" + + +def _linear_with_weight_quantizer(*, per_expert: bool = False): + linear = torch.nn.Module() + quantizer = TensorQuantizer() + linear.weight_quantizer = torch.nn.ModuleList([quantizer]) if per_expert else quantizer + return linear + + +def test_nemotron_lightning_iq2_xs_recipe_matches_megatron_core_names(): + model = torch.nn.Module() + model.decoder = torch.nn.Module() + model.decoder.layers = torch.nn.ModuleList([torch.nn.Module()]) + layer = model.decoder.layers[0] + layer.mlp = torch.nn.Module() + layer.mlp.experts = torch.nn.Module() + layer.mlp.experts.linear_fc1 = _linear_with_weight_quantizer(per_expert=True) + layer.mlp.shared_experts = torch.nn.Module() + layer.mlp.shared_experts.linear_fc1 = _linear_with_weight_quantizer() + model.mtp = torch.nn.Module() + model.mtp.mlp = torch.nn.Module() + model.mtp.mlp.experts = torch.nn.Module() + model.mtp.mlp.experts.linear_fc1 = _linear_with_weight_quantizer(per_expert=True) + + config = load_recipe(_RECIPE).quantize.model_dump(exclude_unset=True) + set_quantizer_by_cfg(model, config["quant_cfg"]) + + routed = layer.mlp.experts.linear_fc1.weight_quantizer[0] + shared = layer.mlp.shared_experts.linear_fc1.weight_quantizer + mtp = model.mtp.mlp.experts.linear_fc1.weight_quantizer[0] + assert routed.is_enabled and routed.num_bits == "iq2_xs" + assert shared.is_enabled and shared.num_bits == "iq2_xs" + assert not mtp.is_enabled + + +@pytest.mark.skipif( + not hasattr(transformers, "NemotronHConfig"), + reason="NemotronH is not supported by this Transformers version", +) +def test_nemotron_lightning_iq2_xs_recipe_quantizer_placement(): + model = get_tiny_nemotron_h() + model.mtp = torch.nn.Module() + model.mtp.mixer = torch.nn.Module() + model.mtp.mixer.experts = torch.nn.Linear(256, 256, bias=False) + config = load_recipe(_RECIPE).quantize.model_dump(exclude_unset=True) + + assert not need_calibration(config) + mtq.quantize(model, config) + + quantizers = { + name: module + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) + } + iq2_xs = { + name: module + for name, module in quantizers.items() + if module.is_enabled and module.num_bits == "iq2_xs" + } + nvfp4 = { + name: module + for name, module in quantizers.items() + if module.is_enabled and module.num_bits == (2, 1) + } + + assert iq2_xs + assert all(".mixer.experts." in name or ".mixer.shared_experts." in name for name in iq2_xs) + assert any(".mixer.experts." in name for name in iq2_xs) + assert any(".mixer.shared_experts." in name for name in iq2_xs) + assert all(module.block_sizes == {-1: 256} for module in iq2_xs.values()) + + assert set(nvfp4) == { + "model.layers.0.mixer.in_proj.weight_quantizer", + "model.layers.0.mixer.out_proj.weight_quantizer", + } + assert all(module.block_sizes[-1] == 16 for module in nvfp4.values()) + assert all(module.block_sizes["type"] == "dynamic" for module in nvfp4.values()) + + enabled = {name for name, module in quantizers.items() if module.is_enabled} + assert enabled == set(iq2_xs) | set(nvfp4) + assert not quantizers["mtp.mixer.experts.weight_quantizer"].is_enabled + + hf_quant_config = get_quant_config(model) + quantization = hf_quant_config["quantization"] + assert quantization["quant_algo"] == "MIXED_PRECISION" + groups = convert_hf_quant_config_format(hf_quant_config)["config_groups"] + weights = [group["weights"] for group in groups.values()] + assert any(weight.get("packing") == "ggml" and weight["num_bits"] == 2 for weight in weights) + assert any(weight["num_bits"] == 4 and weight["group_size"] == 16 for weight in weights) + + +@pytest.mark.skipif( + not hasattr(transformers, "NemotronHConfig"), + reason="NemotronH is not supported by this Transformers version", +) +def test_nemotron_lightning_iq2_xs_recipe_direct_export(tmp_path): + model = get_tiny_nemotron_h( + hidden_size=256, + intermediate_size=256, + num_hidden_layers=2, + hybrid_override_pattern="ME", + num_attention_heads=8, + num_key_value_heads=4, + head_dim=32, + mamba_num_heads=8, + mamba_head_dim=32, + n_routed_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=256, + n_shared_experts=1, + moe_shared_expert_intermediate_size=256, + ) + model.config.architectures = [type(model).__name__] + config = load_recipe(_RECIPE).quantize.model_dump(exclude_unset=True) + mtq.quantize(model, config) + + export_dir = tmp_path / "checkpoint" + export_hf_checkpoint(model, export_dir=export_dir) + + state = load_file(export_dir / "model.safetensors") + iq2_xs_weights = { + name: tensor + for name, tensor in state.items() + if tensor.dtype == torch.uint8 and tensor.shape[-1] == 74 + } + assert iq2_xs_weights + assert all( + ".mixer.experts." in name or ".mixer.shared_experts." in name for name in iq2_xs_weights + ) + assert any(".mixer.experts." in name for name in iq2_xs_weights) + assert any(".mixer.shared_experts." in name for name in iq2_xs_weights) + for name, weight in iq2_xs_weights.items(): + base = name.removesuffix(".weight") + logical_shape = state[base + ".weight_logical_shape"] + padded_shape = state[base + ".weight_padded_shape"] + assert padded_shape[-1] == weight.shape[-2] * 256 + assert torch.all(padded_shape >= logical_shape) + + with open(export_dir / "config.json") as file: + quantization_config = json.load(file)["quantization_config"] + assert quantization_config["quant_algo"] == "MIXED_PRECISION" + weights = [group["weights"] for group in quantization_config["config_groups"].values()] + assert any(weight.get("packing") == "ggml" and weight["num_bits"] == 2 for weight in weights) + assert any(weight.get("row_padding") == "right" for weight in weights) + assert any(weight["num_bits"] == 4 and weight["group_size"] == 16 for weight in weights) diff --git a/tests/unit/torch/export/test_export_weight.py b/tests/unit/torch/export/test_export_weight.py index 6fc17d982e8..d932c0b8f66 100644 --- a/tests/unit/torch/export/test_export_weight.py +++ b/tests/unit/torch/export/test_export_weight.py @@ -20,10 +20,13 @@ from _test_utils.torch.export.utils import ToyModel, partial_fp8_config, partial_w4a8_config import modelopt.torch.quantization as mtq +from modelopt.torch.export.quant_utils import postprocess_state_dict from modelopt.torch.export.unified_export_hf import ( _export_quantized_weight, _process_quantized_modules, ) +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.quantization.utils import quantizer_attr_names @@ -102,6 +105,29 @@ def test_export_per_block_quantized_weight(): assert not hasattr(model.linears[2], quantizer_attrs.output_scale) +@pytest.mark.parametrize(("num_bits", "payload_bytes"), [("iq1_s", 50), ("iq2_xs", 74)]) +def test_export_iq_payload_as_weight(num_bits, payload_bytes): + linear = nn.Linear(257, 4, bias=False, dtype=torch.bfloat16) + linear.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=num_bits, + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + ) + + _export_quantized_weight(linear, torch.bfloat16) + state_dict = postprocess_state_dict(linear.state_dict(), maxbound=448, quantization=None) + + assert state_dict["weight"].shape == (4, 2, payload_bytes) + assert state_dict["weight"].dtype == torch.uint8 + assert "packed_weights" not in state_dict + assert "weight_shape" not in state_dict + assert state_dict["weight_logical_shape"].tolist() == [4, 257] + assert state_dict["weight_padded_shape"].tolist() == [4, 512] + + class QuantMoELinear(nn.Module): def __init__(self): super().__init__() diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index 90740587bf7..b33f50eb852 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -26,8 +26,11 @@ import modelopt.torch.export.unified_export_megatron as unified_export_megatron import modelopt.torch.quantization as mtq +from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format from modelopt.torch.export.quant_format import ( QUANTIZATION_FP8, + QUANTIZATION_IQ1_S, + QUANTIZATION_IQ2_XS, QUANTIZATION_NVFP4, QUANTIZATION_W4A8_AWQ, ) @@ -39,6 +42,48 @@ from modelopt.torch.quantization.nn import NVFP4StaticQuantizer +@pytest.mark.parametrize( + ("num_bits", "quantization_format", "payload_bytes", "effective_bits"), + [ + ("iq1_s", QUANTIZATION_IQ1_S, 50, 1.5625), + ("iq2_xs", QUANTIZATION_IQ2_XS, 74, 2.3125), + ], +) +def test_iq_quantization_config(num_bits, quantization_format, payload_bytes, effective_bits): + model = torch.nn.Sequential(torch.nn.Linear(256, 256, bias=False)) + mtq.quantize( + model, + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": { + "num_bits": num_bits, + "block_sizes": {-1: 256}, + "backend": "ggml", + "backend_extra_args": {"search_impl": "auto"}, + }, + }, + ], + "algorithm": None, + }, + ) + + assert get_quantization_format(model) == quantization_format + config = get_quant_config(model) + assert config["quantization"]["quant_algo"] == num_bits.upper() + assert config["quantization"]["block_payload_bytes"] == payload_bytes + hf_config = convert_hf_quant_config_format(config) + weights = hf_config["config_groups"]["group_0"]["weights"] + assert weights["group_size"] == 256 + assert weights["effective_bits"] == effective_bits + assert weights["packing"] == "ggml" + assert weights["row_padding"] == "right" + assert weights["logical_shape_key"] == "weight_logical_shape" + assert weights["padded_shape_key"] == "weight_padded_shape" + + class _FakeKVCacheQuantizer(torch.nn.Module): """Minimal FP8 KV cache quantizer for scaling-factor tests.""" diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py index 5ce964ca9a2..4b67fd214aa 100644 --- a/tests/unit/torch/export/test_offload_export.py +++ b/tests/unit/torch/export/test_offload_export.py @@ -35,18 +35,25 @@ wrap_in_parent_with_tied_keys, ) +import modelopt.torch.export.unified_export_hf as unified_export_hf +import modelopt.torch.export.unified_export_hf_streaming as unified_export_hf_streaming import modelopt.torch.quantization as mtq from modelopt.torch.export.model_utils import TiedWeightMap from modelopt.torch.export.quant_format import KV_CACHE_FP8 from modelopt.torch.export.quant_utils import _postprocess_single_tensor -from modelopt.torch.export.unified_export_hf import _export_quantized_weight +from modelopt.torch.export.unified_export_hf import ( + _export_quantized_weight, + _export_transformers_checkpoint, +) from modelopt.torch.export.unified_export_hf_streaming import ( + _export_transformers_checkpoint_streaming, _parse_shard_size, _StreamingShardWriter, name_shards_and_write_index, ) from modelopt.torch.quantization.nn.modules.quant_linear import RealQuantLinear from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector # --------------------------------------------------------------------------- # Helpers @@ -71,6 +78,47 @@ def _offload_module(module): set_module_tensor_to_device(module, "weight", "meta") +class _ExportConfig: + torch_dtype = torch.bfloat16 + tie_word_embeddings = False + _name_or_path = "" + + def save_pretrained(self, export_dir): + Path(export_dir, "config.json").write_text("{}") + + +class _SingleLayerModel(nn.Module): + def __init__(self, weight): + super().__init__() + self.layers = nn.ModuleList( + [nn.Linear(weight.shape[1], weight.shape[0], bias=False, dtype=weight.dtype)] + ) + self.layers[0].weight.data.copy_(weight) + self.config = _ExportConfig() + + +def _quantize_iq_weight(model, num_bits): + mtq.quantize( + model, + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": { + "num_bits": num_bits, + "block_sizes": {-1: 256}, + "backend": "ggml", + "backend_extra_args": {"search_impl": "auto"}, + }, + "enable": True, + }, + ], + "algorithm": "max", + }, + ) + + # --------------------------------------------------------------------------- # tied-weight alias map under offload # --------------------------------------------------------------------------- @@ -348,6 +396,57 @@ def test_postprocess_passthrough_normal_key(): assert val.shape == (4, 4) +@pytest.mark.parametrize("suffix", ["weight_logical_shape", "weight_padded_shape"]) +def test_postprocess_preserves_iq_shape_sidecar(suffix): + """Streaming export keeps the shape contract for row-padded IQ payloads.""" + shape = torch.tensor([4, 257 if suffix == "weight_logical_shape" else 512]) + key = f"model.layers.0.mlp.up_proj.{suffix}" + + exported_key, exported_shape = _postprocess_single_tensor(key, shape, 448.0, None) + + assert exported_key == key + assert torch.equal(exported_shape, shape) + + +@pytest.mark.parametrize("num_bits", ["iq1_s", "iq2_xs"]) +def test_streaming_offload_iq_payload_matches_resident_export(num_bits, tmp_path, monkeypatch): + """Offloaded streaming export writes the same row-padded IQ bytes as resident export.""" + generator = torch.Generator().manual_seed(5918) + weight = torch.randn((2, 257), generator=generator, dtype=torch.bfloat16) + resident_model = _SingleLayerModel(weight) + streaming_model = _SingleLayerModel(weight) + _quantize_iq_weight(resident_model, num_bits) + _quantize_iq_weight(streaming_model, num_bits) + + monkeypatch.setattr( + unified_export_hf, "requantize_resmooth_fused_llm_layers", lambda model: None + ) + monkeypatch.setattr( + unified_export_hf_streaming, + "requantize_resmooth_fused_llm_layers", + lambda model: None, + ) + + resident_state, _ = _export_transformers_checkpoint(resident_model, torch.bfloat16) + + _offload_module(streaming_model.layers[0]) + monkeypatch.setattr( + LayerActivationCollector, + "get_decoder_layers", + staticmethod(lambda model: model.layers), + ) + _export_transformers_checkpoint_streaming( + streaming_model, + torch.bfloat16, + export_dir=tmp_path, + ) + + with safe_open(str(tmp_path / "model.safetensors"), framework="pt") as exported: + for suffix in ("weight", "weight_logical_shape", "weight_padded_shape"): + key = f"layers.0.{suffix}" + assert torch.equal(exported.get_tensor(key), resident_state[key]) + + @pytest.mark.parametrize( "key", [ diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 49cd999a205..f6fe438a2d3 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -502,6 +502,37 @@ def forward_loop(m): self._cleanup_registry(expert_type) + def test_export_registers_packed_weight_buffers(self, monkeypatch): + """Packed expert weights must remain present in the exported state dict.""" + model = _TinyMoEModel() + expert_type = type(model.moe.experts) + self._cleanup_registry(expert_type) + register_fused_experts_on_the_fly(model) + + try: + converted = QuantModuleRegistry.convert(model.moe.experts) + + def _pack_weight_as_buffer(wrapper, dtype): + packed = torch.zeros((*wrapper.weight.shape, 1), dtype=torch.uint8) + del wrapper.weight + wrapper.register_buffer("weight", packed) + + monkeypatch.setattr( + "modelopt.torch.export.unified_export_hf._export_quantized_weight", + _pack_weight_as_buffer, + ) + + _export_fused_experts(converted, torch.float16) + + state_dict = converted.state_dict() + for idx in range(NUM_EXPERTS): + for projection in ("gate_proj", "up_proj", "down_proj"): + key = f"{idx}.{projection}.weight" + assert key in state_dict + assert state_dict[key].dtype == torch.uint8 + finally: + self._cleanup_registry(expert_type) + def test_uncalibrated_expert_gate_up_share_amax(self, monkeypatch): """gate_proj and up_proj must share weight_scale_2 even when an expert was never routed during calibration. diff --git a/tests/unit/torch/quantization/test_iq1_s.py b/tests/unit/torch/quantization/test_iq1_s.py new file mode 100644 index 00000000000..f3d96cbc845 --- /dev/null +++ b/tests/unit/torch/quantization/test_iq1_s.py @@ -0,0 +1,121 @@ +# 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 pytest +import torch + +from modelopt.torch.quantization.ggml.iq1_s import ( + IQ1_S_BLOCK_BYTES, + dequantize_iq1_s, + iq1_s_fake_quant, + iq1_s_grid, + quantize_iq1_s, +) + + +def test_iq1_s_canonical_grid(): + grid = iq1_s_grid() + + assert grid.shape == (2048, 8) + assert grid.dtype == torch.float32 + assert set(grid.unique().tolist()) == {-1.0, 0.0, 1.0} + assert grid[0].tolist() == [-1.0] * 8 + + +def test_iq1_s_zero_block_has_canonical_zero_encoding(): + weight = torch.zeros((2, 256), dtype=torch.bfloat16) + + packed, shape = quantize_iq1_s(weight) + + assert packed.shape == (2, 1, IQ1_S_BLOCK_BYTES) + assert packed.dtype == torch.uint8 + assert not packed.any() + assert shape.tolist() == [2, 256] + assert torch.equal(dequantize_iq1_s(packed, shape), weight) + + +def test_iq1_s_dequantizes_ggml_metadata_bit_fields(): + packed = torch.zeros((1, 1, 50), dtype=torch.uint8) + d = torch.tensor([2.0], dtype=torch.float16).view(torch.uint8) + packed[0, 0, :2] = d + entries = torch.tensor([0, 256, 511, 2047], dtype=torch.int64) + packed[0, 0, 2:6] = (entries & 0xFF).to(torch.uint8) + qh = ( + ((entries[0] >> 8) & 7) + | (((entries[1] >> 8) & 7) << 3) + | (((entries[2] >> 8) & 7) << 6) + | (((entries[3] >> 8) & 7) << 9) + | (3 << 12) + | (1 << 15) + ) + packed[0, 0, 34] = (qh & 0xFF).to(torch.uint8) + packed[0, 0, 35] = (qh >> 8).to(torch.uint8) + + decoded = dequantize_iq1_s(packed, torch.tensor([1, 256]), dtype=torch.float32) + expected = (iq1_s_grid()[entries] - 0.125) * 14.0 + + assert torch.equal(decoded[0, :32].reshape(4, 8), expected) + + +def test_iq1_s_round_trip_and_payload_fields(): + generator = torch.Generator().manual_seed(1234) + weight = torch.randn((2, 256), generator=generator, dtype=torch.bfloat16) + + packed, shape = quantize_iq1_s(weight, block_chunk_size=1) + reconstructed = dequantize_iq1_s(packed, shape) + + assert packed.shape == (2, 1, 50) + assert reconstructed.shape == weight.shape + assert reconstructed.dtype == torch.bfloat16 + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.25 + + blocks = packed.reshape(-1, 50) + qh = blocks[:, 34:50:2].to(torch.int64) | (blocks[:, 35:50:2].to(torch.int64) << 8) + assert torch.all(((qh >> 12) & 0x7) < 8) + assert torch.all((qh & 0xFFF) < 0x1000) + + +def test_iq1_s_right_pads_each_row_without_crossing_row_boundaries(): + generator = torch.Generator().manual_seed(4321) + weight = torch.randn((2, 257), generator=generator, dtype=torch.bfloat16) + explicitly_padded = torch.nn.functional.pad(weight, (0, 255)) + + packed, shape = quantize_iq1_s(weight, block_chunk_size=1) + expected, _ = quantize_iq1_s(explicitly_padded, block_chunk_size=1) + + assert packed.shape == (2, 2, IQ1_S_BLOCK_BYTES) + assert shape.tolist() == [2, 257] + assert torch.equal(packed, expected) + assert dequantize_iq1_s(packed, shape).shape == weight.shape + + +def test_iq1_s_rejects_scalar_weight(): + with pytest.raises(ValueError, match="at least one dimension"): + quantize_iq1_s(torch.tensor(1.0)) + + +def test_iq1_s_fake_quant_has_pass_through_gradient(): + class Quantizer: + num_bits = "iq1_s" + backend_extra_args = {"search_impl": "auto"} + + weight = torch.randn(1, 256, requires_grad=True) + output = iq1_s_fake_quant(weight, Quantizer()) + output.sum().backward() + + assert torch.equal(weight.grad, torch.ones_like(weight)) diff --git a/tests/unit/torch/quantization/test_iq2_xs.py b/tests/unit/torch/quantization/test_iq2_xs.py new file mode 100644 index 00000000000..1b495b30491 --- /dev/null +++ b/tests/unit/torch/quantization/test_iq2_xs.py @@ -0,0 +1,132 @@ +# 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 pytest +import torch + +import modelopt.torch.quantization.ggml.iq2_xs as iq2_xs_module +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.ggml.iq2_xs import ( + IQ2_XS_BLOCK_BYTES, + dequantize_iq2_xs, + iq2_xs_fake_quant, + iq2_xs_grid, + quantize_iq2_xs, +) +from modelopt.torch.quantization.nn import TensorQuantizer + + +def test_iq2_xs_canonical_grid(): + grid = iq2_xs_grid() + + assert grid.shape == (512, 8) + assert grid.dtype == torch.float32 + assert set(grid.unique().tolist()) == {8.0, 25.0, 43.0} + assert grid[0].tolist() == [8.0] * 8 + assert grid[-1].tolist() == [43.0] * 8 + + +def test_iq2_xs_zero_block_has_canonical_zero_encoding(): + weight = torch.zeros((2, 256), dtype=torch.bfloat16) + + packed, shape = quantize_iq2_xs(weight) + + assert packed.shape == (2, 1, IQ2_XS_BLOCK_BYTES) + assert packed.dtype == torch.uint8 + assert not packed.any() + assert shape.tolist() == [2, 256] + assert torch.equal(dequantize_iq2_xs(packed, shape), weight) + + +def test_iq2_xs_round_trip_and_payload_fields(): + generator = torch.Generator().manual_seed(1234) + weight = torch.randn((2, 512), generator=generator, dtype=torch.bfloat16) + + packed, shape = quantize_iq2_xs(weight, block_chunk_size=2) + reconstructed = dequantize_iq2_xs(packed, shape) + + assert packed.shape == (2, 2, 74) + assert reconstructed.shape == weight.shape + assert reconstructed.dtype == torch.bfloat16 + normalized_mse = ( + reconstructed.float() - weight.float() + ).square().mean() / weight.float().square().mean() + assert normalized_mse < 0.1 + + blocks = packed.reshape(-1, 74) + codes = blocks[:, 2:66:2].to(torch.int64) | (blocks[:, 3:66:2].to(torch.int64) << 8) + assert torch.all((codes & 0x1FF) < 512) + assert torch.all((codes >> 9) < 128) + + +def test_iq2_xs_right_pads_each_row_without_crossing_row_boundaries(): + generator = torch.Generator().manual_seed(4321) + weight = torch.randn((2, 257), generator=generator, dtype=torch.bfloat16) + explicitly_padded = torch.nn.functional.pad(weight, (0, 255)) + + packed, shape = quantize_iq2_xs(weight, block_chunk_size=2) + expected, _ = quantize_iq2_xs(explicitly_padded, block_chunk_size=2) + + assert packed.shape == (2, 2, IQ2_XS_BLOCK_BYTES) + assert shape.tolist() == [2, 257] + assert torch.equal(packed, expected) + assert dequantize_iq2_xs(packed, shape).shape == weight.shape + + +def test_iq2_xs_rejects_scalar_weight(): + with pytest.raises(ValueError, match="at least one dimension"): + quantize_iq2_xs(torch.tensor(1.0)) + + +def test_iq2_xs_fake_quant_has_pass_through_gradient(): + class Quantizer: + num_bits = "iq2_xs" + backend_extra_args = {"search_impl": "auto"} + + weight = torch.randn(1, 256, requires_grad=True) + output = iq2_xs_fake_quant(weight, Quantizer()) + output.sum().backward() + + assert torch.equal(weight.grad, torch.ones_like(weight)) + + +def test_iq2_xs_tensor_quantizer_matches_row_padded_backend(monkeypatch): + generator = torch.Generator().manual_seed(5918) + weight = torch.randn((2, 257), generator=generator, dtype=torch.bfloat16) + quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="iq2_xs", + block_sizes={-1: 256}, + backend="ggml", + backend_extra_args={"search_impl": "auto"}, + ) + ) + + packed_by_fake_quant = [] + original_quantize = quantize_iq2_xs + + def capture_quantize(inputs): + packed, shape = original_quantize(inputs) + packed_by_fake_quant.append(packed) + return packed, shape + + monkeypatch.setattr(iq2_xs_module, "quantize_iq2_xs", capture_quantize) + reconstructed = quantizer(weight) + packed, shape = original_quantize(weight) + expected = dequantize_iq2_xs(packed, shape, dtype=weight.dtype) + + assert len(packed_by_fake_quant) == 1 + assert torch.equal(packed_by_fake_quant[0].reshape_as(packed), packed) + torch.testing.assert_close(reconstructed, expected, rtol=0, atol=0.0078125) diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/iq2_xs_full_gguf_validate.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/iq2_xs_full_gguf_validate.yaml new file mode 100644 index 00000000000..f4a27127440 --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/iq2_xs_full_gguf_validate.yaml @@ -0,0 +1,57 @@ +# Materialize and validate the complete mixed checkpoint with a pinned stock GGUF stack. + +job_name: Nemotron-3.5-Lightning-IQ2XS-full-GGUF-validate +pipeline: + allow_to_fail: false + note: >- + Materialize the complete mixed checkpoint through an unmodified pinned llama.cpp + converter, require byte-identical IQ2_XS payloads in GGUF, and record the stock + loader result without treating full-model serving as a block-format requirement. + + global_vars: + hf_model: /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/cicd/cicd_1789316148520067000/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-IQ2XS-NVFP4/hf + output_dir: /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/users/hungyuehc/omniml-5918/evidence/full-gguf + + task_0: + inline: >- + unset HF_TOKEN GITHUB_TOKEN GH_TOKEN + && LLAMA_ROOT=/tmp/llama-cpp-$SLURM_JOB_ID + && BUILD_ROOT=/tmp/llama-cpp-build-$SLURM_JOB_ID + && LLAMA_COMMIT=37b3a9e0ccba261d1cc245a971deae0b18c201ab + && OUTPUT_DIR=<>/$SLURM_JOB_ID + && GGUF_PATH=$OUTPUT_DIR/nemotron-3.5-lightning-iq2-xs.gguf + && mkdir -p "$OUTPUT_DIR" + && git init "$LLAMA_ROOT" + && git -C "$LLAMA_ROOT" remote add origin https://github.com/ggml-org/llama.cpp.git + && git -C "$LLAMA_ROOT" fetch --depth=1 origin "$LLAMA_COMMIT" + && git -C "$LLAMA_ROOT" checkout --detach FETCH_HEAD + && test "$(git -C "$LLAMA_ROOT" rev-parse HEAD)" = "$LLAMA_COMMIT" + && python modules/Model-Optimizer/examples/megatron_bridge/materialize_mixed_iq_gguf.py + --checkpoint <> + --llama-source "$LLAMA_ROOT" + --outfile "$GGUF_PATH" + --report "$OUTPUT_DIR/payload-report.json" + && cmake -S "$LLAMA_ROOT" -B "$BUILD_ROOT" + -DGGML_CUDA=OFF -DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_EXAMPLES=OFF + -DLLAMA_BUILD_SERVER=ON + && cmake --build "$BUILD_ROOT" --target llama-cli --parallel 8 + && if "$BUILD_ROOT/bin/llama-cli" --model "$GGUF_PATH" --check-tensors + --n-gpu-layers 0 --no-display-prompt --prompt "2+2=" --n-predict 1 + >"$OUTPUT_DIR/stock-loader.log" 2>&1; then LOADER_STATUS=0; + else LOADER_STATUS=$?; fi + && export OUTPUT_DIR LLAMA_COMMIT LOADER_STATUS + && python -c 'import json, os, pathlib; output = pathlib.Path(os.environ["OUTPUT_DIR"]); payload = json.loads((output / "payload-report.json").read_text()); + report = {"schema_version": 1, "llama_cpp_commit": os.environ["LLAMA_COMMIT"], "payload_validation": payload["status"], "stock_loader_exit_code": + int(os.environ["LOADER_STATUS"]), "stock_loader_succeeded": os.environ["LOADER_STATUS"] == "0"}; (output / "loader-report.json").write_text(json.dumps(report, + indent=2, sort_keys=True) + "\n")' + slurm_config: + _factory_: slurm_factory + container: nvcr.io/nvidia/nemo:26.08 + container_mounts: + - /lustre:/lustre + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + nodes: 1 + ntasks_per_node: 1 + # This account's QoS requires four allocated GPUs; validation itself is CPU-only. + gpus_per_node: 4 + time: "02:00:00" diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/iq2_xs_stock_ggml_validate.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/iq2_xs_stock_ggml_validate.yaml new file mode 100644 index 00000000000..3151897ab00 --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/iq2_xs_stock_ggml_validate.yaml @@ -0,0 +1,50 @@ +# Validate exported IQ2_XS payloads with an unmodified pinned GGML decoder. + +job_name: Nemotron-3.5-Lightning-IQ2XS-stock-GGML-validate +pipeline: + allow_to_fail: false + note: >- + Sample canonical 74-byte blocks from every IQ2_XS tensor, require bitwise + agreement with the pinned stock GGML decoder, and round-trip the payloads + through a test-only GGUF container without changing their bytes. + + global_vars: + hf_model: /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/cicd/cicd_1789316148520067000/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-IQ2XS-NVFP4/hf + output_dir: /scratchspace/stock-ggml-iq2-xs-validation + + task_0: + inline: >- + unset HF_TOKEN GITHUB_TOKEN GH_TOKEN + && LLAMA_ROOT=/tmp/llama-cpp-$SLURM_JOB_ID + && BUILD_ROOT=/tmp/llama-cpp-build-$SLURM_JOB_ID + && LLAMA_COMMIT=37b3a9e0ccba261d1cc245a971deae0b18c201ab + && git init "$LLAMA_ROOT" + && git -C "$LLAMA_ROOT" remote add origin https://github.com/ggml-org/llama.cpp.git + && git -C "$LLAMA_ROOT" fetch --depth=1 origin "$LLAMA_COMMIT" + && git -C "$LLAMA_ROOT" checkout --detach FETCH_HEAD + && test "$(git -C "$LLAMA_ROOT" rev-parse HEAD)" = "$LLAMA_COMMIT" + && cmake -S "$LLAMA_ROOT" -B "$BUILD_ROOT" + -DGGML_CUDA=OFF -DBUILD_SHARED_LIBS=ON + -DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=OFF + && cmake --build "$BUILD_ROOT" --target ggml-cpu --parallel 8 + && GGML_LIBRARY=$(find "$BUILD_ROOT" -type f -name 'libggml-cpu.so*' | head -n 1) + && test -n "$GGML_LIBRARY" + && mkdir -p <> + && python modules/Model-Optimizer/examples/megatron_bridge/validate_iq2_xs_stock_ggml.py + --checkpoint <> + --ggml-library "$GGML_LIBRARY" + --llama-source "$LLAMA_ROOT" + --blocks-per-tensor 3 + --test-gguf <>/sampled-iq2-xs.gguf + --report <>/report.json + slurm_config: + _factory_: slurm_factory + container: nvcr.io/nvidia/nemo:26.08 + container_mounts: + - /lustre:/lustre + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + nodes: 1 + ntasks_per_node: 1 + # This account's QoS requires four allocated GPUs; validation itself is CPU-only. + gpus_per_node: 4 + time: "00:30:00" diff --git a/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_iq2_xs_export_validate.yaml b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_iq2_xs_export_validate.yaml new file mode 100644 index 00000000000..5cb8918d010 --- /dev/null +++ b/tools/launcher/examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_iq2_xs_export_validate.yaml @@ -0,0 +1,78 @@ +# Nemotron 3.5 Lightning mixed real-quant export and structural validation. +# +# The model path is expressed through /hf-local so the same YAML works with local Docker and +# Slurm. Point hf_local (local mode) or SLURM_HF_LOCAL (Slurm) at the parent model cache. +# +# Slurm: +# uv run launch.py --yaml \ +# examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_iq2_xs_export_validate.yaml \ +# --yes +# +# Local Docker: +# uv run launch.py --yaml \ +# examples/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/mbridge_iq2_xs_export_validate.yaml \ +# hf_local=/path/to/hf-local --yes + +job_name: Nemotron-3.5-Lightning-IQ2XS-export-validate +pipeline: + allow_to_fail: false + note: >- + Calibration-free mixed IQ2_XS expert and NVFP4 Mamba quantization, unified-HF export, + CPU/CUDA byte-parity check, and checkpoint layout/digest validation. + + global_vars: + hf_model: /hf-local/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 + output_dir: /scratchspace/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-IQ2XS-NVFP4 + + # Run as one sequential task so the launcher transfers one source archive. Later commands run + # only when every preceding validation or export command succeeds. + task_0: + environment: + - LAUNCH_SCRIPT: torchrun --nproc_per_node 4 + # Avoid stale shared-cache manifests whose compiled extension was evicted. + inline: >- + unset HF_TOKEN + && export TORCH_EXTENSIONS_DIR=/tmp/modelopt-torch-extensions-$SLURM_JOB_ID + MAX_JOBS=8 + && if [ -z "$SLURM_LOCALID" ] || [ "$SLURM_LOCALID" -eq 0 ]; then + python modules/Model-Optimizer/examples/megatron_bridge/validate_iq_cuda_parity.py + --rows 4 --row-width 257 + --report <>/iq_cpu_cuda_parity.json + && touch <>/.iq_parity_$SLURM_JOB_ID; + else for _ in $(seq 600); do + [ -f <>/.iq_parity_$SLURM_JOB_ID ] && break; sleep 1; + done; [ -f <>/.iq_parity_$SLURM_JOB_ID ]; fi + && + $LAUNCH_SCRIPT modules/Model-Optimizer/examples/megatron_bridge/quantize.py + --hf_model_name_or_path <> + --trust_remote_code + --tp_size 1 + --recipe models/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16/ptq/iq2_xs_experts-nvfp4_mamba + --skip_generate + --export_megatron_path <>/megatron + && + $LAUNCH_SCRIPT modules/Model-Optimizer/examples/megatron_bridge/export_quantized_megatron_to_hf.py + --hf_model_name_or_path <> + --megatron_path <>/megatron + --trust_remote_code + --pp_size 4 + --export_unified_hf_path <>/hf + && + if [ -z "$SLURM_LOCALID" ] || [ "$SLURM_LOCALID" -eq 0 ]; then + python modules/Model-Optimizer/examples/megatron_bridge/validate_mixed_quantized_hf.py + --checkpoint <>/hf + --reference-checkpoint <> + --report <>/checkpoint_validation.json; fi + slurm_config: + _factory_: "slurm_factory" + container: nvcr.io/nvidia/nemo:26.08 + container_mounts: + - /lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/hf-local:/hf-local + - /lustre:/lustre + modelopt_install_path: /opt/venv/lib/python3.12/site-packages/modelopt + nodes: 1 + # Slurm overrides LAUNCH_SCRIPT to python and supplies these four distributed ranks. + # Standalone validation commands above run only on rank zero. + ntasks_per_node: 4 + gpus_per_node: 4 + time: "04:00:00"