From c98933dc1866ce4a8fc4bb08740dc1d81ca2b033 Mon Sep 17 00:00:00 2001 From: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Date: Thu, 9 Apr 2026 15:44:27 +0000 Subject: [PATCH 1/9] GPTQ vector and unit test Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> --- .../torch/quantization/utils/calib_utils.py | 91 ++- tests/gpu/torch/quantization/test_gptq_vq.py | 544 ++++++++++++++++++ 2 files changed, 634 insertions(+), 1 deletion(-) create mode 100644 tests/gpu/torch/quantization/test_gptq_vq.py diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index e52a8438d55..6d05f2f51ac 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -47,6 +47,30 @@ from modelopt.torch.utils.perf import get_used_gpu_mem_fraction +def load_vector_lut_codebook(quantizer): + """Load vector LUT codebook and quantizer params from a weight_quantizer. + + Returns: + Tuple of (codebook, quant_block_size, scale_type). + """ + from luts import encode as luts_encode + + extra_args = quantizer.backend_extra_args + encode_format = quantizer.num_bits + encode_path = extra_args.get("encode_path", "") + if encode_path and not encode_path.endswith("/"): + encode_path += "/" + + if "sorted" in encode_format: + cb = torch.load(encode_path + encode_format + ".pt", map_location="cpu") + codebook = cb["sorted_values"].cuda().float() + else: + codebook, _ = luts_encode(encode_format, path=encode_path, norm=False, cuda=True) + codebook = codebook.float() + + return codebook, extra_args.get("block_sizes"), extra_args.get("scale_type") + + def update_hessian(input, hessian, n_samples): """Update hessian matrix with new input samples using incremental formula. @@ -140,11 +164,18 @@ def update_weights(self, block_size, perc_damp): Populates ``self.weight`` and ``self.h_inv``, runs the blockwise update, logs MSE, and writes the result back to the module. """ + backend_extra_args = getattr(self.module.weight_quantizer, "backend_extra_args", None) + is_vector_lut = bool( + backend_extra_args and backend_extra_args.get("lut_type") == "vector_lut" + ) hessian = self.hessian.to(self.module.weight.device) self.weight = self.module.weight.data.float().clone() self._prepare_hessian_inverse(hessian, perc_damp) - self._blockwise_update(block_size) + if is_vector_lut: + self._blockwise_vector_update(block_size) + else: + self._blockwise_update(block_size) self._print_mse_error(hessian) self.module.weight.data = self.weight.reshape(self.module.weight.shape).to( @@ -224,6 +255,64 @@ def _blockwise_update(self, block_size): errs, self.h_inv[block_start:block_end, block_end:], alpha=-1 ) + def _blockwise_vector_update(self, block_size): + """GPTQ blockwise update for vector LUT quantizers. + + Pre-computes scales once, then runs the standard GPTQ 3-loop + with per-vector-group static quantization via clip_vector_prescaled. + """ + import torch.nn.functional as F + from luts import clip_vector_prescaled, clip_vector_scalesign_fast + + codebook, quant_block_size, scale_type = load_vector_lut_codebook( + self.module.weight_quantizer + ) + + # Get vector size from codebook + vector_size = codebook.shape[1] + + assert self.weight is not None and self.h_inv is not None + num_cols = self.weight.shape[1] + assert block_size % quant_block_size == 0 + + # Pre-compute scales once outside the GPTQ loop + _, scales = clip_vector_scalesign_fast( + self.weight, + codebook, + quant_block_size, + scale_type, + scale_algo="max", + sign_scale=True, + return_scales=True, + ) + scales_2d = scales.reshape(self.weight.shape[0], -1) + + w = self.weight.clone() + h_inv = self.h_inv + + for blk_start in range(0, num_cols, block_size): + blk_end = min(blk_start + block_size, num_cols) + errs = torch.zeros_like(w[:, blk_start:blk_end]) + + for j in range(blk_start, blk_end, vector_size): + d = min(vector_size, blk_end - j) + s = scales_2d[:, j // quant_block_size].contiguous() + + sub = w[:, j : j + d].contiguous() + if d < vector_size: + sub = F.pad(sub, (0, vector_size - d)) + q_sub = clip_vector_prescaled(sub, codebook, s) + + for k in range(d): + col = j + k + self.weight[:, col] = q_sub[:, k] + err = (w[:, col] - q_sub[:, k]) / h_inv[col, col] + errs[:, col - blk_start] = err + w[:, col:blk_end].addr_(err, h_inv[col, col:blk_end], alpha=-1) + + if blk_end < num_cols: + w[:, blk_end:] -= errs @ h_inv[blk_start:blk_end, blk_end:] + def _print_mse_error(self, hessian): """Log Hessian-weighted relative MSE between ``self.weight`` and original weights.""" w_orig = self.module.weight.float() diff --git a/tests/gpu/torch/quantization/test_gptq_vq.py b/tests/gpu/torch/quantization/test_gptq_vq.py new file mode 100644 index 00000000000..0fd3c8da624 --- /dev/null +++ b/tests/gpu/torch/quantization/test_gptq_vq.py @@ -0,0 +1,544 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Test _blockwise_vector_update against gptq_quantize_scaled_vq from adaptive_rounding.py.""" +# ruff: noqa: N803, N806 — uppercase names (W, A, Q, H, etc.) match math notation in the reference. + +import tempfile +from types import SimpleNamespace + +import torch +import torch.nn.functional as F + +from modelopt.torch.quantization.utils.calib_utils import GPTQHelper + +# --------------------------------------------------------------------------- +# Exact copies from modelopt-internal/docs/adaptive_rounding.py +# --------------------------------------------------------------------------- + + +def compute_output_nmse( + W: torch.Tensor, + W_q: torch.Tensor, + A: torch.Tensor, +) -> float: + """ + Compute the output Normalized MSE: + + NMSE = || (W - W_q) @ A ||_F^2 / || W @ A ||_F^2 + + where @ denotes matrix multiplication. This measures how much + weight quantization error propagates through the activations. + + Args: + W: Original weight matrix, shape (out_features, in_features) + W_q: Quantized weight matrix, same shape as W + A: Activation matrix, shape (in_features, n_samples) + + Returns: + Scalar NMSE value + """ + error_output = (W - W_q) @ A + ref_output = W @ A + nmse = error_output.norm() ** 2 / ref_output.norm() ** 2 + return nmse.item() + + +def _round_to_e4m3(s: torch.Tensor) -> torch.Tensor: + """Round a float32 scale tensor to the nearest E4M3 representable value.""" + return s.to(dtype=torch.float8_e4m3fn).to(dtype=torch.float32) + + +def _vq_nearest(vecs: torch.Tensor, codebook: torch.Tensor, chunk_size: int = 1) -> torch.Tensor: + """ + Find the nearest codeword for each row-vector in vecs. + + Args: + vecs: (N, D) float tensor + codebook: (K, D) float tensor + chunk_size: Process this many vectors at a time to bound memory. + Default 1 to avoid torch.cdist batch precision issues. + + Returns: + (N, D) tensor of nearest codewords + """ + N = vecs.shape[0] + result = torch.empty_like(vecs) + for start in range(0, N, chunk_size): + end = min(start + chunk_size, N) + dists = torch.cdist(vecs[start:end], codebook) + result[start:end] = codebook[dists.argmin(dim=1)] + return result + + +def scaled_vector_quantize( + W: torch.Tensor, + codebook: torch.Tensor, + quant_block_size: int, + n_scale_iters: int = 3, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Block-scaled vector quantization with iterative scale optimisation. + + For each row-block of quant_block_size elements we find a scalar + scale s and codeword assignments that (approximately) minimise + + || block - s * reconstruct(block / s) ||^2 + + via alternating minimisation: + 1. Fix s -> assign each sub-vector to nearest codeword in codebook + 2. Fix assignments -> solve for optimal s in closed form: + s* = sum_i / sum_i ||c_i||^2 + + Initialised with s = max(|block|) / max(|codebook|) (amax), then + refined for n_scale_iters iterations. + + Effective bitrate: log2(n_codewords) / codebook_dim + 16 / quant_block_size + + Args: + W: (out_features, in_features) + codebook: (n_codewords, codebook_dim) — must evenly divide quant_block_size + quant_block_size: Elements per scaling block + n_scale_iters: Number of assign-then-rescale iterations (default 3) + + Returns: + W_q: Dequantised weight matrix + scales: (out_features, n_blocks) + """ + codebook_dim = codebook.shape[1] + out_features, in_features = W.shape + assert quant_block_size % codebook_dim == 0, ( + f"quant_block_size ({quant_block_size}) must be a multiple of codebook_dim ({codebook_dim})" + ) + assert in_features % codebook_dim == 0, ( + f"in_features ({in_features}) must be a multiple of codebook_dim ({codebook_dim})" + ) + + cb_absmax = codebook.abs().max().item() + n_blocks = (in_features + quant_block_size - 1) // quant_block_size + + # Initial scales: amax heuristic + pad = n_blocks * quant_block_size - in_features + if pad > 0: + W_padded = torch.nn.functional.pad(W, (0, pad)) + else: + W_padded = W + W_blocks = W_padded.reshape(out_features, n_blocks, quant_block_size) + block_max = W_blocks.abs().amax(dim=2) + scales = (block_max / cb_absmax).clamp(min=1e-10) + + W_q = torch.empty_like(W) + + for b in range(n_blocks): + start = b * quant_block_size + end = min(start + quant_block_size, in_features) + block = W[:, start:end] # (out_features, block_len) + s = scales[:, b] # (out_features,) + + for _ in range(n_scale_iters): + # Step 1: assign sub-vectors to nearest codeword at current scale + normalized = block / s.unsqueeze(1) + vecs = normalized.reshape(-1, codebook_dim) + q_vecs = _vq_nearest(vecs, codebook) + q_block = q_vecs.reshape(out_features, end - start) + + # Step 2: optimal scale given assignments + # s* = sum_i / sum_i ||c_i||^2 + # where x_i are original sub-vectors, c_i are assigned codewords + numerator = (block * q_block).sum(dim=1) # (out_features,) + denominator = (q_block * q_block).sum(dim=1) # (out_features,) + s = (numerator / denominator.clamp(min=1e-20)).clamp(min=1e-10) + + # Round converged scale to E4M3 and do final assignment + s = _round_to_e4m3(s) + normalized = block / s.unsqueeze(1) + vecs = normalized.reshape(-1, codebook_dim) + q_vecs = _vq_nearest(vecs, codebook) + W_q[:, start:end] = q_vecs.reshape(out_features, end - start) * s.unsqueeze(1) + scales[:, b] = s + + return W_q, scales + + +def gptq_quantize_scaled_vq( + W: torch.Tensor, + A: torch.Tensor, + codebook: torch.Tensor, + quant_block_size: int, + gptq_block_size: int = 128, + damp_pct: float = 0.01, + n_scale_iters: int = 3, + h_inv: torch.Tensor | None = None, +) -> tuple[torch.Tensor, float]: + """ + GPTQ with block-scaled vector quantization. + + Scales are pre-computed via the same iterative optimisation as + scaled_vector_quantize. During the GPTQ pass, columns are processed + in groups of codebook_dim: each group is VQ-quantized to the nearest + codeword, and per-column errors are compensated into remaining columns + using the inverse Hessian. + + Args: + W: (out_features, in_features) + A: (in_features, n_samples) + codebook: (n_codewords, codebook_dim) + quant_block_size: Elements per scaling block + gptq_block_size: Columns per GPTQ lazy batch update (must be + a multiple of codebook_dim) + damp_pct: Dampening fraction + n_scale_iters: Iterations for scale optimisation + h_inv: Optional pre-computed upper-triangular Cholesky factor of the + damped inverse Hessian. If None, computed from A. + + Returns: + W_q: Quantized weight matrix (dequantised) + nmse: Output NMSE achieved + """ + codebook_dim = codebook.shape[1] + assert gptq_block_size % codebook_dim == 0, ( + f"gptq_block_size ({gptq_block_size}) must be a multiple of codebook_dim ({codebook_dim})" + ) + + W_orig = W.clone() + W = W.clone() + out_features, in_features = W.shape + + _, scales = scaled_vector_quantize( + W_orig, codebook, quant_block_size, n_scale_iters=n_scale_iters + ) + + if h_inv is None: + H = 2.0 * (A @ A.T) + damp = damp_pct * torch.diag(H).mean() + H.diagonal().add_(damp) + + H_inv = torch.linalg.inv(H) + try: + L = torch.linalg.cholesky(H_inv) + except torch.linalg.LinAlgError: + H.diagonal().add_(damp * 10) + H_inv = torch.linalg.inv(H) + L = torch.linalg.cholesky(H_inv) + Hinv = L.T + else: + Hinv = h_inv + + Q = torch.zeros_like(W) + + for i in range(0, in_features, gptq_block_size): + j_end = min(i + gptq_block_size, in_features) + E = torch.zeros(out_features, j_end - i, dtype=W.dtype, device=W.device) + + for j in range(i, j_end, codebook_dim): + d = min(codebook_dim, j_end - j) + sb = j // quant_block_size + s = scales[:, sb] # (out_features,) + + sub_vec = W[:, j : j + d] / s.unsqueeze(1) + q_vecs = _vq_nearest( + sub_vec.reshape(-1, codebook_dim) if d == codebook_dim else sub_vec.reshape(-1, d), + codebook[:, :d], + ) + q_block = q_vecs.reshape(out_features, d) * s.unsqueeze(1) + Q[:, j : j + d] = q_block + + for k in range(d): + col = j + k + err = (W[:, col] - Q[:, col]) / Hinv[col, col] + E[:, col - i] = err + W[:, col:j_end] -= err.unsqueeze(1) * Hinv[col, col:j_end].unsqueeze(0) + + if j_end < in_features: + W[:, j_end:] -= E @ Hinv[i:j_end, j_end:] + + nmse = compute_output_nmse(W_orig, Q, A) + return Q, nmse + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_codebook_file(tmp_dir, name, n_codewords, vector_size): + """Save a codebook .pt file in the format load_vector_lut_codebook expects.""" + codebook = torch.randn(n_codewords, vector_size) + torch.save({"sorted_values": codebook}, f"{tmp_dir}/{name}.pt") + return codebook.cuda().float() + + +def _make_hessian(activations): + """Build Hessian H = 2 * X @ X^T from activations (batch, seq, features).""" + X = activations.reshape(-1, activations.shape[-1]).t().float() + return 2.0 * (X @ X.t()) + + +def _attach_mock_quantizer(module, encode_path, encode_format, quant_block_size, scale_type): + """Attach a mock weight_quantizer with the right attributes for _blockwise_vector_update.""" + module.weight_quantizer = SimpleNamespace( + num_bits=encode_format, + backend="psx_luts", + backend_extra_args={ + "encode_path": encode_path, + "lut_type": "vector_lut", + "block_sizes": quant_block_size, + "scale_type": scale_type, + }, + ) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +SEED = 42 +OUT_FEATURES = 64 +IN_FEATURES = 128 +GPTQ_BLOCK_SIZE = 128 +QUANT_BLOCK_SIZE = 16 +VECTOR_SIZE = 8 +N_CODEWORDS = 256 +PERC_DAMP = 0.01 +SCALE_TYPE = "e4m3" + + +def test_clip_vector_prescaled_vs_vq_nearest(): + """clip_vector_prescaled and _vq_nearest must produce identical quantized vectors.""" + from luts import clip_vector_prescaled, clip_vector_scalesign_fast + + torch.manual_seed(SEED) + + with tempfile.TemporaryDirectory() as tmp_dir: + codebook = _make_codebook_file(tmp_dir, "test_sorted", N_CODEWORDS, VECTOR_SIZE) + + W = torch.randn(OUT_FEATURES, IN_FEATURES, device="cuda", dtype=torch.float32) + + # Compute scales with sign_scale=True (scales can be negative) + _, scales = clip_vector_scalesign_fast( + W, + codebook, + QUANT_BLOCK_SIZE, + SCALE_TYPE, + scale_algo="max", + sign_scale=True, + return_scales=True, + ) + scales_2d = scales.reshape(OUT_FEATURES, -1) + + for j in range(0, IN_FEATURES, VECTOR_SIZE): + d = min(VECTOR_SIZE, IN_FEATURES - j) + s = scales_2d[:, j // QUANT_BLOCK_SIZE].contiguous() + sub = W[:, j : j + d].contiguous() + + # clip_vector_prescaled path + if d < VECTOR_SIZE: + sub_padded = F.pad(sub, (0, VECTOR_SIZE - d)) + else: + sub_padded = sub + q_luts = clip_vector_prescaled(sub_padded, codebook, s)[:, :d] + + # _vq_nearest path (manual normalize -> lookup -> denormalize) + normalized = sub / s.unsqueeze(1) + vecs = ( + normalized.reshape(-1, VECTOR_SIZE) + if d == VECTOR_SIZE + else normalized.reshape(-1, d) + ) + cb_slice = codebook if d == VECTOR_SIZE else codebook[:, :d] + q_vecs = _vq_nearest(vecs, cb_slice) + q_ref = q_vecs.reshape(OUT_FEATURES, d) * s.unsqueeze(1) + + assert torch.allclose(q_luts, q_ref, atol=1e-5), ( + f"VQ mismatch at col {j}: max diff={(q_luts - q_ref).abs().max().item():.2e}" + ) + + +def test_blockwise_vector_update_vs_gptq_quantize_scaled_vq(): + """_blockwise_vector_update must produce identical weights to gptq_quantize_scaled_vq. + + Both paths share the same pre-computed scales (from clip_vector_scalesign_fast) + and the same h_inv, so the only variable is the GPTQ loop itself. + The reference's scaled_vector_quantize is patched to return the shared scales. + """ + from unittest.mock import patch + + from luts import clip_vector_scalesign_fast + + torch.manual_seed(SEED) + + with tempfile.TemporaryDirectory() as tmp_dir: + encode_format = "test_sorted" + codebook = _make_codebook_file(tmp_dir, encode_format, N_CODEWORDS, VECTOR_SIZE) + + W = torch.randn(OUT_FEATURES, IN_FEATURES, device="cuda", dtype=torch.float32) + A = torch.randn(4, 16, IN_FEATURES, device="cuda", dtype=torch.float32) + A_flat = A.reshape(-1, IN_FEATURES).t() # (in_features, n_samples) for reference + hessian = _make_hessian(A) + + # Shared scales — computed once, used by both paths + Q_rtn, scales = clip_vector_scalesign_fast( + W, + codebook, + QUANT_BLOCK_SIZE, + SCALE_TYPE, + scale_algo="max", + sign_scale=True, + return_scales=True, + ) + scales_2d = scales.reshape(OUT_FEATURES, -1) + + # --- Reference: exact gptq_quantize_scaled_vq from adaptive_rounding.py --- + # Patch scaled_vector_quantize to return the shared scales (Q_rtn unused by caller) + with patch( + f"{__name__}.scaled_vector_quantize", + return_value=(Q_rtn, scales_2d), + ): + Q_ref, _ = gptq_quantize_scaled_vq( + W, + A_flat, + codebook, + QUANT_BLOCK_SIZE, + GPTQ_BLOCK_SIZE, + damp_pct=PERC_DAMP, + ) + + # --- Modelopt: _blockwise_vector_update --- + module = torch.nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False).cuda() + module.weight.data = W.clone() + _attach_mock_quantizer(module, tmp_dir, encode_format, QUANT_BLOCK_SIZE, SCALE_TYPE) + + helper = GPTQHelper(module, "test_layer") + helper.hessian = hessian.clone() + helper.n_samples = 1 + helper.update_weights(GPTQ_BLOCK_SIZE, PERC_DAMP) + Q_modelopt = module.weight.data.float() + + # Same scales + same GPTQ loop structure => weights must match + assert torch.allclose(Q_modelopt, Q_ref, atol=1e-5), ( + f"Weights differ: max diff={(Q_modelopt - Q_ref).abs().max().item():.2e}" + ) + + +def test_mtq_quantize_gptq_vs_gptq_quantize_scaled_vq(): + """End-to-end: mtq.quantize with GPTQ vs gptq_quantize_scaled_vq. + + Both paths compute their own hessian and h_inv independently. + Shared scales ensure the only variable is the GPTQ loop. The hessian is + injected into GPTQHelper so both paths compute identical h_inv. + """ + from unittest.mock import patch + + from luts import clip_vector_scalesign_fast + + # Register psx_luts backend (workaround for _default_disabled_quantizer_cfg list/dict change) + import modelopt.torch.quantization.config as _mtq_cfg + + _orig = _mtq_cfg._default_disabled_quantizer_cfg + if isinstance(_orig, list): + _mtq_cfg._default_disabled_quantizer_cfg = {} + import modelopt_internal.torch.quantization.plugins.psx_luts # noqa: F401 + + _mtq_cfg._default_disabled_quantizer_cfg = _orig + + import modelopt.torch.quantization as mtq + + torch.manual_seed(SEED) + + with tempfile.TemporaryDirectory() as tmp_dir: + encode_format = "test_sorted" + codebook = _make_codebook_file(tmp_dir, encode_format, N_CODEWORDS, VECTOR_SIZE) + + W = torch.randn(OUT_FEATURES, IN_FEATURES, device="cuda", dtype=torch.float32) + A = torch.randn(4, 16, IN_FEATURES, device="cuda", dtype=torch.float32) + A_flat = A.reshape(-1, IN_FEATURES).t() + + # --- Modelopt: mtq.quantize with GPTQ (list-based config) --- + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": { + "backend": "psx_luts", + "num_bits": encode_format, + "pass_through_bwd": True, + "backend_extra_args": { + "encode_path": tmp_dir, + "lut_type": "vector_lut", + "block_sizes": QUANT_BLOCK_SIZE, + "scale_type": SCALE_TYPE, + "scale_algo": "max", + "round_mode": "rne", + "sign_scale": True, + }, + }, + }, + ], + "algorithm": { + "method": "gptq", + "use_sequential": False, + "perc_damp": PERC_DAMP, + "block_size": GPTQ_BLOCK_SIZE, + }, + } + + model = torch.nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False).cuda() + model.weight.data = W.clone() + + def forward_loop(m): + m(A) + + # Capture h_inv from the modelopt GPTQ run + captured = {} + _orig_update = GPTQHelper.update_weights + + def _capturing_update(self, *args, **kwargs): + _orig_update(self, *args, **kwargs) + captured["h_inv"] = self.h_inv.clone() + + with patch.object(GPTQHelper, "update_weights", _capturing_update): + mtq.quantize(model, config, forward_loop=forward_loop) + Q_modelopt = model.weight.data.float() + + # --- Reference: gptq_quantize_scaled_vq with shared scales + captured h_inv --- + Q_rtn, scales = clip_vector_scalesign_fast( + W, + codebook, + QUANT_BLOCK_SIZE, + SCALE_TYPE, + scale_algo="max", + sign_scale=True, + return_scales=True, + ) + scales_2d = scales.reshape(OUT_FEATURES, -1) + + with patch( + f"{__name__}.scaled_vector_quantize", + return_value=(Q_rtn, scales_2d), + ): + Q_ref, _ = gptq_quantize_scaled_vq( + W, + A_flat, + codebook, + QUANT_BLOCK_SIZE, + GPTQ_BLOCK_SIZE, + damp_pct=PERC_DAMP, + h_inv=captured["h_inv"], + ) + + assert torch.allclose(Q_modelopt, Q_ref, atol=1e-5), ( + f"Weights differ: max diff={(Q_modelopt - Q_ref).abs().max().item():.2e}" + ) From 50d700904016b062c65368db2031c172600218ea Mon Sep 17 00:00:00 2001 From: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Date: Mon, 13 Apr 2026 16:38:09 +0000 Subject: [PATCH 2/9] latest tested on Qwen3-8B Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> --- examples/llm_ptq/example_utils.py | 4 +- examples/llm_ptq/hf_ptq.py | 85 +++- modelopt/torch/quantization/model_calib.py | 19 + modelopt/torch/utils/dataset_utils.py | 383 +++++++++++++++++-- tests/gpu/torch/quantization/test_gptq_vq.py | 10 - 5 files changed, 438 insertions(+), 63 deletions(-) diff --git a/examples/llm_ptq/example_utils.py b/examples/llm_ptq/example_utils.py index c2d4d4bfca8..fb07e1016ac 100755 --- a/examples/llm_ptq/example_utils.py +++ b/examples/llm_ptq/example_utils.py @@ -583,7 +583,7 @@ def get_model( model_kwargs = config_kwargs.copy() # Don't set torch_dtype for VILA models as they handle it explicitly in their builder if "vila" not in ckpt_path.lower(): - model_kwargs.setdefault("dtype", "auto") + model_kwargs.setdefault("torch_dtype", "auto") if "vila" in ckpt_path.lower(): hf_vila = AutoModel.from_pretrained( @@ -666,7 +666,7 @@ def has_pack_quantized_config(config): model_kwargs2 = model_kwargs.copy() if auto_model_module not in [AutoModelForCausalLM, AutoModel]: model_kwargs2.pop("trust_remote_code", None) - model_kwargs2["dtype"] = torch_dtype + model_kwargs2["torch_dtype"] = torch_dtype model_kwargs2.pop("max_memory", None) model = from_config(hf_config, **model_kwargs2) diff --git a/examples/llm_ptq/hf_ptq.py b/examples/llm_ptq/hf_ptq.py index 327605406c4..fcd735e0575 100755 --- a/examples/llm_ptq/hf_ptq.py +++ b/examples/llm_ptq/hf_ptq.py @@ -71,6 +71,7 @@ ) from modelopt.torch.utils.dataset_utils import ( create_forward_loop, + get_calib_and_holdout_dataloaders, get_dataset_dataloader, get_max_batch_size, get_supported_datasets, @@ -203,9 +204,10 @@ def make_calib_dataloader( tokenizer: PreTrainedTokenizerBase | None, device: torch.device, model_type: str | None, -) -> tuple[DataLoader | _DeviceDataLoader, str | None]: +) -> tuple[DataLoader | _DeviceDataLoader, str | None, Path | None]: calib_dataloader = None first_text_speech_dataset = None + holdout_path = None if args.specdec_offline_dataset is not None: offline_data_path = Path(args.specdec_offline_dataset) dumped_files = sorted(str(p) for p in offline_data_path.glob("*.pt")) @@ -283,15 +285,29 @@ def make_calib_dataloader( include_labels = ( args.auto_quantize_bits is not None and args.auto_quantize_method == "gradient" ) - calib_dataloader = get_dataset_dataloader( - dataset_name=args.dataset, - tokenizer=tokenizer, - batch_size=args.batch_size, - num_samples=args.calib_size, - device=device, - include_labels=include_labels, - ) - return calib_dataloader, first_text_speech_dataset + + if args.holdout_size > 0: + calib_dataloader, holdout_path = get_calib_and_holdout_dataloaders( + dataset_name=args.dataset, + tokenizer=tokenizer, + batch_size=args.batch_size, + calib_size=args.calib_size, + holdout_size=args.holdout_size, + max_sample_length=args.calib_seq, + device=device, + include_labels=include_labels, + save_dir=args.calib_data_dir, + ) + else: + calib_dataloader = get_dataset_dataloader( + dataset_name=args.dataset, + tokenizer=tokenizer, + batch_size=args.batch_size, + num_samples=args.calib_size, + device=device, + include_labels=include_labels, + ) + return calib_dataloader, first_text_speech_dataset, holdout_path def auto_quantize( @@ -419,10 +435,15 @@ def load_model(args: argparse.Namespace): attn_implementation=args.attn_implementation, ) else: - assert args.qformat in QUANT_CFG_CHOICES, ( - f"Quantization format is not supported for low memory mode. Supported formats: {QUANT_CFG_CHOICES.keys()}" - ) - quant_cfg = QUANT_CFG_CHOICES[args.qformat] + if args.qformat in QUANT_CFG_CHOICES: + quant_cfg = QUANT_CFG_CHOICES[args.qformat] + elif hasattr(mtq, args.qformat): + quant_cfg = getattr(mtq, args.qformat) + else: + raise AssertionError( + f"Quantization format is not supported for low memory mode. " + f"Supported formats: {QUANT_CFG_CHOICES.keys()}" + ) if args.kv_cache_qformat != "none": quant_cfg = mtq.utils.update_quant_cfg_with_kv_cache_quant( quant_cfg, @@ -1028,7 +1049,7 @@ def quantize_main( print(f"Use calib batch_size {args.batch_size}") - calib_dataloader, first_text_speech_dataset = make_calib_dataloader( + calib_dataloader, first_text_speech_dataset, holdout_path = make_calib_dataloader( args, language_model, processor, tokenizer, device, model_type ) @@ -1066,10 +1087,14 @@ def quantize_main( "Plain quantization supports only one quantization format." ) - assert args.qformat in QUANT_CFG_CHOICES, ( - f"Unsupported quantization format: {args.qformat}, choices are: {list(QUANT_CFG_CHOICES.keys())}" - ) - quant_cfg = QUANT_CFG_CHOICES[args.qformat] + if args.qformat in QUANT_CFG_CHOICES: + quant_cfg = QUANT_CFG_CHOICES[args.qformat] + elif hasattr(mtq, args.qformat): + quant_cfg = getattr(mtq, args.qformat) + else: + raise AssertionError( + f"Unsupported quantization format: {args.qformat}, choices are: {list(QUANT_CFG_CHOICES.keys())}" + ) quant_cfg = build_quant_cfg( args.qformat, @@ -1104,7 +1129,7 @@ def quantize_main( quant_cfg = copy.deepcopy(quant_cfg) _set_kv_cache_constant_amax(quant_cfg["quant_cfg"]) - if args.qformat in QUANT_CFG_CHOICES: + if args.qformat in QUANT_CFG_CHOICES or hasattr(mtq, args.qformat): mono_quantize( args, quant_cfg, @@ -1180,6 +1205,26 @@ def parse_args() -> argparse.Namespace: type=str, default="512", ) + parser.add_argument( + "--holdout_size", + help=( + "Number of holdout samples to save as a .pt file for evaluation. " + "Holdout samples are drawn from the same dataset immediately after " + "the calibration samples so there is no overlap. 0 disables holdout." + ), + type=int, + default=0, + ) + parser.add_argument( + "--calib_data_dir", + help=( + "Directory to save/load calib.pt and holdout.pt. " + "If both files exist, data is reloaded from disk instead of re-downloading. " + "Defaults to --export_path if not specified." + ), + type=str, + default=None, + ) parser.add_argument( "--calib_seq", help="Maximum sequence length for calibration.", diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 35a0e931c9d..faa19862a0c 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1589,6 +1589,21 @@ def sequential_calibrate( def _layer_forward_loop(m, _inputs=layer_inputs): for args, kwargs_input in _inputs: + # Reset past_key_values to prevent the KV cache from + # accumulating across multiple forward replays (e.g. + # max_calibrate then Hessian collection in GPTQ). + # The layer doesn't need stale KV data — each replay + # should start with a fresh cache. + if ( + "past_key_values" in kwargs_input + and kwargs_input["past_key_values"] is not None + ): + kwargs_input = dict(kwargs_input) + cache = kwargs_input["past_key_values"] + if hasattr(cache, "reset"): + cache.reset() + else: + kwargs_input["past_key_values"] = None m(*args, **kwargs_input) calib_func(layer, _layer_forward_loop, **calib_kwargs) @@ -1665,6 +1680,10 @@ def gptq( print_rank_0("Updating weights using GPTQ algorithm...") for handle in gptq_handles.values(): handle.update_weights(block_size, perc_damp) + + # Disable weight quantizer after running GPTQ update since weights are already QDQ'ed + if hasattr(handle.module, "weight_quantizer"): + handle.module.weight_quantizer.disable() handle.free() del gptq_handles diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 1e9a7fbbbd8..d15dedb5d12 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -108,8 +108,10 @@ __all__ = [ "create_forward_loop", "download_hf_dataset_as_jsonl", + "get_calib_and_holdout_dataloaders", "get_dataset_dataloader", "get_dataset_samples", + "get_dataset_samples_with_holdout", "get_jsonl_text_samples", "get_max_batch_size", "get_supported_datasets", @@ -211,43 +213,23 @@ def _auto_preprocess_sample( ) -def get_dataset_samples( +def _load_streaming_dataset( dataset_name: str, - num_samples: int, *, apply_chat_template: bool = False, tokenizer: "PreTrainedTokenizerBase | None" = None, split: str | list[str] | None = None, -) -> list[str]: - """Load a portion of a dataset with the dataset name and a given size. +) -> tuple[list, Callable[[dict], str]]: + """Resolve dataset config and return streaming splits with a preprocessing function. - Supports both registered datasets (in ``SUPPORTED_DATASET_CONFIG``) and arbitrary - HuggingFace datasets. Unregistered datasets are auto-detected by column names: - ``messages``/``conversations`` (chat), ``prompt``, ``text``, or ``input``. - - Args: - dataset_name: Name or HuggingFace path of the dataset to load, a local directory path, - or a path to a ``.jsonl`` file. For local directory paths, the - predefined config from ``SUPPORTED_DATASET_CONFIG`` is matched if the base folder name - matches a registered key (e.g. ``/hf-local/abisee/cnn_dailymail`` matches ``cnn_dailymail`` key). - num_samples: Number of samples to load from the dataset. - apply_chat_template: Whether to apply the chat template to the samples - (if supported by the dataset). For unregistered datasets with a - ``messages`` column, chat template is always applied regardless of - this flag. - tokenizer: Tokenizer to use for applying the chat template to the samples. - No tokenization is done and plain text is still returned. - split: Override the split(s) to load. Accepts a single split name or a list. - If ``None``, uses the splits defined in ``SUPPORTED_DATASET_CONFIG`` for - registered datasets, or ``["train"]`` for unregistered datasets. + This is a shared helper for :func:`get_dataset_samples` and + :func:`get_dataset_samples_with_holdout`. Returns: - Samples: The list of samples. + A tuple of ``(dataset_splits, preprocess_fn)`` where *dataset_splits* is a list of + HuggingFace ``IterableDataset`` objects and *preprocess_fn* maps a raw sample dict + to a plain-text string (empty string signals a sample to skip). """ - # Local JSONL file path support (each line is a JSON object with a `text` field). - if dataset_name.endswith(".jsonl"): - return get_jsonl_text_samples(dataset_name, num_samples, key="text") - from datasets import load_dataset local_dataset_path = None @@ -301,21 +283,158 @@ def _preprocess(sample: dict) -> str: print(f"Loading dataset with {config=} and {splits=}") dataset_splits = [load_dataset(streaming=True, **config, split=s) for s in splits] + return dataset_splits, _preprocess + + +def get_dataset_samples( + dataset_name: str, + num_samples: int, + *, + apply_chat_template: bool = False, + tokenizer: "PreTrainedTokenizerBase | None" = None, + split: str | list[str] | None = None, +) -> list[str]: + """Load a portion of a dataset with the dataset name and a given size. + + Supports both registered datasets (in ``SUPPORTED_DATASET_CONFIG``) and arbitrary + HuggingFace datasets. Unregistered datasets are auto-detected by column names: + ``messages``/``conversations`` (chat), ``prompt``, ``text``, or ``input``. + + Args: + dataset_name: Name or HuggingFace path of the dataset to load, a local directory path, + or a path to a ``.jsonl`` file. For local directory paths, the + predefined config from ``SUPPORTED_DATASET_CONFIG`` is matched if the base folder name + matches a registered key (e.g. ``/hf-local/abisee/cnn_dailymail`` matches ``cnn_dailymail`` key). + num_samples: Number of samples to load from the dataset. + apply_chat_template: Whether to apply the chat template to the samples + (if supported by the dataset). For unregistered datasets with a + ``messages`` column, chat template is always applied regardless of + this flag. + tokenizer: Tokenizer to use for applying the chat template to the samples. + No tokenization is done and plain text is still returned. + split: Override the split(s) to load. Accepts a single split name or a list. + If ``None``, uses the splits defined in ``SUPPORTED_DATASET_CONFIG`` for + registered datasets, or ``["train"]`` for unregistered datasets. + + Returns: + Samples: The list of samples. + """ + # Local JSONL file path support (each line is a JSON object with a `text` field). + if dataset_name.endswith(".jsonl"): + return get_jsonl_text_samples(dataset_name, num_samples, key="text") + + dataset_splits, _preprocess = _load_streaming_dataset( + dataset_name, + apply_chat_template=apply_chat_template, + tokenizer=tokenizer, + split=split, + ) + num_per_split = [num_samples // len(dataset_splits)] * len(dataset_splits) num_per_split[-1] += num_samples - sum(num_per_split) samples: list[str] = [] for dataset, n in zip(dataset_splits, num_per_split): - for i, sample in enumerate(dataset): - if i >= n: + split_samples: list[str] = [] + for sample in dataset: + if len(split_samples) >= n: break text = _preprocess(sample) if text: - samples.append(text) + split_samples.append(text) + samples.extend(split_samples) return samples +def get_dataset_samples_with_holdout( + dataset_name: str, + num_samples: int, + num_holdout: int, + *, + apply_chat_template: bool = False, + tokenizer: "PreTrainedTokenizerBase | None" = None, + split: str | list[str] | None = None, +) -> tuple[list[str], list[str]]: + """Load calibration and holdout samples from a dataset in a single streaming pass. + + Samples are loaded sequentially: the first ``num_samples`` non-empty samples go to the + calibration set, and the next ``num_holdout`` non-empty samples go to the holdout set. + Empty samples are skipped without counting toward either quota. This guarantees zero + overlap between the two sets by construction. + + Args: + dataset_name: Name or HuggingFace path of the dataset to load, a local directory path, + or a path to a ``.jsonl`` file. + num_samples: Number of non-empty samples for the calibration set. + num_holdout: Number of non-empty samples for the holdout set. + apply_chat_template: Whether to apply the chat template to the samples. + tokenizer: Tokenizer to use for applying the chat template. + split: Override the split(s) to load. + + Returns: + A tuple ``(calib_samples, holdout_samples)``. + """ + if dataset_name.endswith(".jsonl"): + all_samples = get_jsonl_text_samples(dataset_name, num_samples + num_holdout, key="text") + calib = all_samples[:num_samples] + holdout = all_samples[num_samples:] + if len(calib) < num_samples: + warn( + f"JSONL file only had {len(calib)} samples, " + f"fewer than the requested {num_samples} calibration samples." + ) + if len(holdout) < num_holdout: + warn( + f"JSONL file only had {len(holdout)} holdout samples after calibration, " + f"fewer than the requested {num_holdout}." + ) + return calib, holdout + + dataset_splits, _preprocess = _load_streaming_dataset( + dataset_name, + apply_chat_template=apply_chat_template, + tokenizer=tokenizer, + split=split, + ) + + calib_per_split = [num_samples // len(dataset_splits)] * len(dataset_splits) + calib_per_split[-1] += num_samples - sum(calib_per_split) + + holdout_per_split = [num_holdout // len(dataset_splits)] * len(dataset_splits) + holdout_per_split[-1] += num_holdout - sum(holdout_per_split) + + calib_samples: list[str] = [] + holdout_samples: list[str] = [] + for dataset, calib_n, holdout_n in zip(dataset_splits, calib_per_split, holdout_per_split): + split_calib: list[str] = [] + split_holdout: list[str] = [] + for sample in dataset: + text = _preprocess(sample) + if not text: + continue + if len(split_calib) < calib_n: + split_calib.append(text) + elif len(split_holdout) < holdout_n: + split_holdout.append(text) + else: + break + if len(split_calib) < calib_n: + warn( + f"Dataset split exhausted: only {len(split_calib)} calibration samples " + f"collected, fewer than the requested {calib_n}." + ) + if len(split_holdout) < holdout_n: + warn( + f"Dataset split exhausted: only {len(split_holdout)} holdout samples " + f"collected, fewer than the requested {holdout_n}." + ) + calib_samples.extend(split_calib) + holdout_samples.extend(split_holdout) + + return calib_samples, holdout_samples + + class _CustomDataset(torch.utils.data.Dataset): def __init__(self, encodings): self.encodings = encodings @@ -331,6 +450,208 @@ def __len__(self): return len(next(iter(self.encodings.values()))) +def _validate_no_overlap(calib_ids: torch.Tensor, holdout_ids: torch.Tensor) -> None: + """Assert that calibration and holdout ``input_ids`` have no shared rows.""" + calib_set = {tuple(row.tolist()) for row in calib_ids} + holdout_set = {tuple(row.tolist()) for row in holdout_ids} + overlap = calib_set & holdout_set + assert not overlap, f"Calibration and holdout data overlap: {len(overlap)} shared sequences." + print(f"Validated no overlap: {len(calib_set)} calib, {len(holdout_set)} holdout, 0 shared.") + + +def _dedup_tokenized( + calib_encoded: dict[str, torch.Tensor], + holdout_encoded: dict[str, torch.Tensor], +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + """Remove duplicate tokenized sequences within calib and between calib/holdout. + + After truncation, different texts can produce identical ``input_ids`` (e.g. + same problem statement with different solutions). Feeding duplicates to + calibration wastes samples and skews statistics; keeping them in holdout + contaminates evaluation. + + Dedup order: + 1. Remove duplicates within calibration (keep first occurrence). + 2. Remove holdout rows that match any remaining calibration row. + + Returns filtered copies of both encoded dicts. + """ + # Dedup within calib. + seen: set[tuple] = set() + calib_keep: list[int] = [] + for i, row in enumerate(calib_encoded["input_ids"]): + key = tuple(row.tolist()) + if key not in seen: + seen.add(key) + calib_keep.append(i) + n_calib_dupes = len(calib_encoded["input_ids"]) - len(calib_keep) + + # Remove holdout rows that collide with calib. + holdout_keep: list[int] = [] + for i, row in enumerate(holdout_encoded["input_ids"]): + if tuple(row.tolist()) not in seen: + holdout_keep.append(i) + n_holdout_overlap = len(holdout_encoded["input_ids"]) - len(holdout_keep) + + if n_calib_dupes or n_holdout_overlap: + print( + f"Dedup: removed {n_calib_dupes} duplicate calib and " + f"{n_holdout_overlap} overlapping holdout sequences " + f"({len(calib_keep)} calib, {len(holdout_keep)} holdout remaining)." + ) + else: + print(f"No duplicates: {len(calib_keep)} calib, {len(holdout_keep)} holdout, 0 shared.") + + def _index(encoded: dict[str, torch.Tensor], indices: list[int]) -> dict[str, torch.Tensor]: + idx = torch.tensor(indices, dtype=torch.long) + return {k: v[idx] for k, v in encoded.items()} + + return _index(calib_encoded, calib_keep), _index(holdout_encoded, holdout_keep) + + +def get_calib_and_holdout_dataloaders( + dataset_name: str | list[str], + tokenizer: "PreTrainedTokenizerBase", + batch_size: int, + calib_size: int | list[int], + holdout_size: int, + max_sample_length: int = 512, + device: torch.device | None = None, + include_labels: bool = False, + apply_chat_template: bool = False, + save_dir: str | None = None, +) -> tuple[DataLoader, Path | None]: + """Create a calibration DataLoader and save holdout data to disk. + + On the first call the function streams ``calib_size + holdout_size`` samples + from ``dataset_name``, tokenizes both splits, saves them as ``.pt`` files + under ``save_dir``, and returns the calibration DataLoader. + + On subsequent calls, if both ``calib.pt`` and ``holdout.pt`` already exist + under ``save_dir``, the function reloads directly from disk (skipping the + dataset download) and validates that calibration and holdout have zero + overlapping ``input_ids`` rows. + + Args: + dataset_name: Name(s) of the dataset(s) to load. + tokenizer: HuggingFace tokenizer instance. + batch_size: Batch size for the calibration DataLoader. + calib_size: Number of calibration samples (per dataset if list). + holdout_size: Total number of holdout samples. + max_sample_length: Maximum token length per sample. + device: Target device for the calibration DataLoader tensors. + include_labels: Whether to include labels in the calibration DataLoader. + apply_chat_template: Whether to apply the chat template. + save_dir: Directory to save/load ``calib.pt`` and ``holdout.pt``. + + Returns: + A ``(calib_dataloader, holdout_path)`` tuple. + """ + tokenizer = copy.deepcopy(tokenizer) + + if isinstance(dataset_name, str): + dataset_name = [dataset_name] + if isinstance(calib_size, int): + calib_size = [calib_size] + + save_path = Path(save_dir) if save_dir else None + calib_pt = save_path / "calib.pt" if save_path else None + holdout_pt = save_path / "holdout.pt" if save_path else None + + if calib_pt and holdout_pt and calib_pt.is_file() and holdout_pt.is_file(): + # ---- Reload from previously saved .pt files ---- + print(f"Loading calibration data from {calib_pt}") + calib_data = torch.load(calib_pt, map_location="cpu", weights_only=True) + print(f"Loading holdout data from {holdout_pt}") + holdout_data = torch.load(holdout_pt, map_location="cpu", weights_only=True) + + _validate_no_overlap(calib_data["input_ids"], holdout_data["input_ids"]) + + if device: + calib_data = {k: v.to(device) for k, v in calib_data.items()} + if include_labels: + calib_data["labels"] = torch.where( + calib_data["attention_mask"] > 0.5, + calib_data["input_ids"], + -100, + ) + else: + calib_data = {"input_ids": calib_data["input_ids"]} + + calib_dl = DataLoader(_CustomDataset(calib_data), batch_size=batch_size, shuffle=False) + return calib_dl, holdout_pt + + # ---- Fresh load: stream calib + holdout, tokenize, save ---- + all_calib_text: list[str] = [] + all_holdout_text: list[str] = [] + holdout_per_ds = [holdout_size // len(dataset_name)] * len(dataset_name) + holdout_per_ds[-1] += holdout_size - sum(holdout_per_ds) + + for ds_name, n_calib, n_holdout in zip(dataset_name, calib_size, holdout_per_ds): + calib_text, holdout_text = get_dataset_samples_with_holdout( + ds_name, + num_samples=n_calib, + num_holdout=n_holdout, + apply_chat_template=apply_chat_template, + tokenizer=tokenizer, + ) + all_calib_text.extend(calib_text) + all_holdout_text.extend(holdout_text) + + calib_encoded = tokenizer( + all_calib_text, + return_tensors="pt", + padding=True, + truncation=True, + max_length=max_sample_length, + ) + holdout_encoded = tokenizer( + all_holdout_text, + return_tensors="pt", + padding=True, + truncation=True, + max_length=max_sample_length, + ) + + calib_encoded, holdout_encoded = _dedup_tokenized(calib_encoded, holdout_encoded) + _validate_no_overlap(calib_encoded["input_ids"], holdout_encoded["input_ids"]) + + # Save to disk. + if save_path: + save_path.mkdir(parents=True, exist_ok=True) + torch.save( + { + "input_ids": calib_encoded["input_ids"], + "attention_mask": calib_encoded["attention_mask"], + }, + calib_pt, + ) + print(f"Saved {len(calib_encoded['input_ids'])} calib samples to {calib_pt}") + torch.save( + { + "input_ids": holdout_encoded["input_ids"], + "attention_mask": holdout_encoded["attention_mask"], + }, + holdout_pt, + ) + print(f"Saved {len(holdout_encoded['input_ids'])} holdout samples to {holdout_pt}") + + # Build calib dataloader. + if device: + calib_encoded = {k: v.to(device) for k, v in calib_encoded.items()} + if include_labels: + calib_encoded["labels"] = torch.where( + calib_encoded["attention_mask"] > 0.5, + calib_encoded["input_ids"], + -100, + ) + else: + calib_encoded = {"input_ids": calib_encoded["input_ids"]} + + calib_dl = DataLoader(_CustomDataset(calib_encoded), batch_size=batch_size, shuffle=False) + return calib_dl, holdout_pt + + def get_dataset_dataloader( dataset_name: str | list[str] = "cnn_dailymail", tokenizer: "PreTrainedTokenizerBase | None" = None, diff --git a/tests/gpu/torch/quantization/test_gptq_vq.py b/tests/gpu/torch/quantization/test_gptq_vq.py index 0fd3c8da624..9741e5f4853 100644 --- a/tests/gpu/torch/quantization/test_gptq_vq.py +++ b/tests/gpu/torch/quantization/test_gptq_vq.py @@ -443,16 +443,6 @@ def test_mtq_quantize_gptq_vs_gptq_quantize_scaled_vq(): from luts import clip_vector_scalesign_fast - # Register psx_luts backend (workaround for _default_disabled_quantizer_cfg list/dict change) - import modelopt.torch.quantization.config as _mtq_cfg - - _orig = _mtq_cfg._default_disabled_quantizer_cfg - if isinstance(_orig, list): - _mtq_cfg._default_disabled_quantizer_cfg = {} - import modelopt_internal.torch.quantization.plugins.psx_luts # noqa: F401 - - _mtq_cfg._default_disabled_quantizer_cfg = _orig - import modelopt.torch.quantization as mtq torch.manual_seed(SEED) From adafec5b059f7cd09fda9d9854e5c6985c309bcc Mon Sep 17 00:00:00 2001 From: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Date: Tue, 14 Apr 2026 22:48:10 +0000 Subject: [PATCH 3/9] removed dataset holdout logic Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> --- examples/llm_ptq/hf_ptq.py | 57 +---- modelopt/torch/utils/dataset_utils.py | 295 +------------------------- 2 files changed, 12 insertions(+), 340 deletions(-) diff --git a/examples/llm_ptq/hf_ptq.py b/examples/llm_ptq/hf_ptq.py index fcd735e0575..2ffd3ef3950 100755 --- a/examples/llm_ptq/hf_ptq.py +++ b/examples/llm_ptq/hf_ptq.py @@ -71,7 +71,6 @@ ) from modelopt.torch.utils.dataset_utils import ( create_forward_loop, - get_calib_and_holdout_dataloaders, get_dataset_dataloader, get_max_batch_size, get_supported_datasets, @@ -204,10 +203,9 @@ def make_calib_dataloader( tokenizer: PreTrainedTokenizerBase | None, device: torch.device, model_type: str | None, -) -> tuple[DataLoader | _DeviceDataLoader, str | None, Path | None]: +) -> tuple[DataLoader | _DeviceDataLoader, str | None]: calib_dataloader = None first_text_speech_dataset = None - holdout_path = None if args.specdec_offline_dataset is not None: offline_data_path = Path(args.specdec_offline_dataset) dumped_files = sorted(str(p) for p in offline_data_path.glob("*.pt")) @@ -286,28 +284,15 @@ def make_calib_dataloader( args.auto_quantize_bits is not None and args.auto_quantize_method == "gradient" ) - if args.holdout_size > 0: - calib_dataloader, holdout_path = get_calib_and_holdout_dataloaders( - dataset_name=args.dataset, - tokenizer=tokenizer, - batch_size=args.batch_size, - calib_size=args.calib_size, - holdout_size=args.holdout_size, - max_sample_length=args.calib_seq, - device=device, - include_labels=include_labels, - save_dir=args.calib_data_dir, - ) - else: - calib_dataloader = get_dataset_dataloader( - dataset_name=args.dataset, - tokenizer=tokenizer, - batch_size=args.batch_size, - num_samples=args.calib_size, - device=device, - include_labels=include_labels, - ) - return calib_dataloader, first_text_speech_dataset, holdout_path + calib_dataloader = get_dataset_dataloader( + dataset_name=args.dataset, + tokenizer=tokenizer, + batch_size=args.batch_size, + num_samples=args.calib_size, + device=device, + include_labels=include_labels, + ) + return calib_dataloader, first_text_speech_dataset def auto_quantize( @@ -1049,7 +1034,7 @@ def quantize_main( print(f"Use calib batch_size {args.batch_size}") - calib_dataloader, first_text_speech_dataset, holdout_path = make_calib_dataloader( + calib_dataloader, first_text_speech_dataset = make_calib_dataloader( args, language_model, processor, tokenizer, device, model_type ) @@ -1205,26 +1190,6 @@ def parse_args() -> argparse.Namespace: type=str, default="512", ) - parser.add_argument( - "--holdout_size", - help=( - "Number of holdout samples to save as a .pt file for evaluation. " - "Holdout samples are drawn from the same dataset immediately after " - "the calibration samples so there is no overlap. 0 disables holdout." - ), - type=int, - default=0, - ) - parser.add_argument( - "--calib_data_dir", - help=( - "Directory to save/load calib.pt and holdout.pt. " - "If both files exist, data is reloaded from disk instead of re-downloading. " - "Defaults to --export_path if not specified." - ), - type=str, - default=None, - ) parser.add_argument( "--calib_seq", help="Maximum sequence length for calibration.", diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index d15dedb5d12..b0bab281476 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -108,10 +108,8 @@ __all__ = [ "create_forward_loop", "download_hf_dataset_as_jsonl", - "get_calib_and_holdout_dataloaders", "get_dataset_dataloader", "get_dataset_samples", - "get_dataset_samples_with_holdout", "get_jsonl_text_samples", "get_max_batch_size", "get_supported_datasets", @@ -222,8 +220,7 @@ def _load_streaming_dataset( ) -> tuple[list, Callable[[dict], str]]: """Resolve dataset config and return streaming splits with a preprocessing function. - This is a shared helper for :func:`get_dataset_samples` and - :func:`get_dataset_samples_with_holdout`. + This is a shared helper for :func:`get_dataset_samples`. Returns: A tuple of ``(dataset_splits, preprocess_fn)`` where *dataset_splits* is a list of @@ -347,94 +344,6 @@ def get_dataset_samples( return samples -def get_dataset_samples_with_holdout( - dataset_name: str, - num_samples: int, - num_holdout: int, - *, - apply_chat_template: bool = False, - tokenizer: "PreTrainedTokenizerBase | None" = None, - split: str | list[str] | None = None, -) -> tuple[list[str], list[str]]: - """Load calibration and holdout samples from a dataset in a single streaming pass. - - Samples are loaded sequentially: the first ``num_samples`` non-empty samples go to the - calibration set, and the next ``num_holdout`` non-empty samples go to the holdout set. - Empty samples are skipped without counting toward either quota. This guarantees zero - overlap between the two sets by construction. - - Args: - dataset_name: Name or HuggingFace path of the dataset to load, a local directory path, - or a path to a ``.jsonl`` file. - num_samples: Number of non-empty samples for the calibration set. - num_holdout: Number of non-empty samples for the holdout set. - apply_chat_template: Whether to apply the chat template to the samples. - tokenizer: Tokenizer to use for applying the chat template. - split: Override the split(s) to load. - - Returns: - A tuple ``(calib_samples, holdout_samples)``. - """ - if dataset_name.endswith(".jsonl"): - all_samples = get_jsonl_text_samples(dataset_name, num_samples + num_holdout, key="text") - calib = all_samples[:num_samples] - holdout = all_samples[num_samples:] - if len(calib) < num_samples: - warn( - f"JSONL file only had {len(calib)} samples, " - f"fewer than the requested {num_samples} calibration samples." - ) - if len(holdout) < num_holdout: - warn( - f"JSONL file only had {len(holdout)} holdout samples after calibration, " - f"fewer than the requested {num_holdout}." - ) - return calib, holdout - - dataset_splits, _preprocess = _load_streaming_dataset( - dataset_name, - apply_chat_template=apply_chat_template, - tokenizer=tokenizer, - split=split, - ) - - calib_per_split = [num_samples // len(dataset_splits)] * len(dataset_splits) - calib_per_split[-1] += num_samples - sum(calib_per_split) - - holdout_per_split = [num_holdout // len(dataset_splits)] * len(dataset_splits) - holdout_per_split[-1] += num_holdout - sum(holdout_per_split) - - calib_samples: list[str] = [] - holdout_samples: list[str] = [] - for dataset, calib_n, holdout_n in zip(dataset_splits, calib_per_split, holdout_per_split): - split_calib: list[str] = [] - split_holdout: list[str] = [] - for sample in dataset: - text = _preprocess(sample) - if not text: - continue - if len(split_calib) < calib_n: - split_calib.append(text) - elif len(split_holdout) < holdout_n: - split_holdout.append(text) - else: - break - if len(split_calib) < calib_n: - warn( - f"Dataset split exhausted: only {len(split_calib)} calibration samples " - f"collected, fewer than the requested {calib_n}." - ) - if len(split_holdout) < holdout_n: - warn( - f"Dataset split exhausted: only {len(split_holdout)} holdout samples " - f"collected, fewer than the requested {holdout_n}." - ) - calib_samples.extend(split_calib) - holdout_samples.extend(split_holdout) - - return calib_samples, holdout_samples - - class _CustomDataset(torch.utils.data.Dataset): def __init__(self, encodings): self.encodings = encodings @@ -450,208 +359,6 @@ def __len__(self): return len(next(iter(self.encodings.values()))) -def _validate_no_overlap(calib_ids: torch.Tensor, holdout_ids: torch.Tensor) -> None: - """Assert that calibration and holdout ``input_ids`` have no shared rows.""" - calib_set = {tuple(row.tolist()) for row in calib_ids} - holdout_set = {tuple(row.tolist()) for row in holdout_ids} - overlap = calib_set & holdout_set - assert not overlap, f"Calibration and holdout data overlap: {len(overlap)} shared sequences." - print(f"Validated no overlap: {len(calib_set)} calib, {len(holdout_set)} holdout, 0 shared.") - - -def _dedup_tokenized( - calib_encoded: dict[str, torch.Tensor], - holdout_encoded: dict[str, torch.Tensor], -) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: - """Remove duplicate tokenized sequences within calib and between calib/holdout. - - After truncation, different texts can produce identical ``input_ids`` (e.g. - same problem statement with different solutions). Feeding duplicates to - calibration wastes samples and skews statistics; keeping them in holdout - contaminates evaluation. - - Dedup order: - 1. Remove duplicates within calibration (keep first occurrence). - 2. Remove holdout rows that match any remaining calibration row. - - Returns filtered copies of both encoded dicts. - """ - # Dedup within calib. - seen: set[tuple] = set() - calib_keep: list[int] = [] - for i, row in enumerate(calib_encoded["input_ids"]): - key = tuple(row.tolist()) - if key not in seen: - seen.add(key) - calib_keep.append(i) - n_calib_dupes = len(calib_encoded["input_ids"]) - len(calib_keep) - - # Remove holdout rows that collide with calib. - holdout_keep: list[int] = [] - for i, row in enumerate(holdout_encoded["input_ids"]): - if tuple(row.tolist()) not in seen: - holdout_keep.append(i) - n_holdout_overlap = len(holdout_encoded["input_ids"]) - len(holdout_keep) - - if n_calib_dupes or n_holdout_overlap: - print( - f"Dedup: removed {n_calib_dupes} duplicate calib and " - f"{n_holdout_overlap} overlapping holdout sequences " - f"({len(calib_keep)} calib, {len(holdout_keep)} holdout remaining)." - ) - else: - print(f"No duplicates: {len(calib_keep)} calib, {len(holdout_keep)} holdout, 0 shared.") - - def _index(encoded: dict[str, torch.Tensor], indices: list[int]) -> dict[str, torch.Tensor]: - idx = torch.tensor(indices, dtype=torch.long) - return {k: v[idx] for k, v in encoded.items()} - - return _index(calib_encoded, calib_keep), _index(holdout_encoded, holdout_keep) - - -def get_calib_and_holdout_dataloaders( - dataset_name: str | list[str], - tokenizer: "PreTrainedTokenizerBase", - batch_size: int, - calib_size: int | list[int], - holdout_size: int, - max_sample_length: int = 512, - device: torch.device | None = None, - include_labels: bool = False, - apply_chat_template: bool = False, - save_dir: str | None = None, -) -> tuple[DataLoader, Path | None]: - """Create a calibration DataLoader and save holdout data to disk. - - On the first call the function streams ``calib_size + holdout_size`` samples - from ``dataset_name``, tokenizes both splits, saves them as ``.pt`` files - under ``save_dir``, and returns the calibration DataLoader. - - On subsequent calls, if both ``calib.pt`` and ``holdout.pt`` already exist - under ``save_dir``, the function reloads directly from disk (skipping the - dataset download) and validates that calibration and holdout have zero - overlapping ``input_ids`` rows. - - Args: - dataset_name: Name(s) of the dataset(s) to load. - tokenizer: HuggingFace tokenizer instance. - batch_size: Batch size for the calibration DataLoader. - calib_size: Number of calibration samples (per dataset if list). - holdout_size: Total number of holdout samples. - max_sample_length: Maximum token length per sample. - device: Target device for the calibration DataLoader tensors. - include_labels: Whether to include labels in the calibration DataLoader. - apply_chat_template: Whether to apply the chat template. - save_dir: Directory to save/load ``calib.pt`` and ``holdout.pt``. - - Returns: - A ``(calib_dataloader, holdout_path)`` tuple. - """ - tokenizer = copy.deepcopy(tokenizer) - - if isinstance(dataset_name, str): - dataset_name = [dataset_name] - if isinstance(calib_size, int): - calib_size = [calib_size] - - save_path = Path(save_dir) if save_dir else None - calib_pt = save_path / "calib.pt" if save_path else None - holdout_pt = save_path / "holdout.pt" if save_path else None - - if calib_pt and holdout_pt and calib_pt.is_file() and holdout_pt.is_file(): - # ---- Reload from previously saved .pt files ---- - print(f"Loading calibration data from {calib_pt}") - calib_data = torch.load(calib_pt, map_location="cpu", weights_only=True) - print(f"Loading holdout data from {holdout_pt}") - holdout_data = torch.load(holdout_pt, map_location="cpu", weights_only=True) - - _validate_no_overlap(calib_data["input_ids"], holdout_data["input_ids"]) - - if device: - calib_data = {k: v.to(device) for k, v in calib_data.items()} - if include_labels: - calib_data["labels"] = torch.where( - calib_data["attention_mask"] > 0.5, - calib_data["input_ids"], - -100, - ) - else: - calib_data = {"input_ids": calib_data["input_ids"]} - - calib_dl = DataLoader(_CustomDataset(calib_data), batch_size=batch_size, shuffle=False) - return calib_dl, holdout_pt - - # ---- Fresh load: stream calib + holdout, tokenize, save ---- - all_calib_text: list[str] = [] - all_holdout_text: list[str] = [] - holdout_per_ds = [holdout_size // len(dataset_name)] * len(dataset_name) - holdout_per_ds[-1] += holdout_size - sum(holdout_per_ds) - - for ds_name, n_calib, n_holdout in zip(dataset_name, calib_size, holdout_per_ds): - calib_text, holdout_text = get_dataset_samples_with_holdout( - ds_name, - num_samples=n_calib, - num_holdout=n_holdout, - apply_chat_template=apply_chat_template, - tokenizer=tokenizer, - ) - all_calib_text.extend(calib_text) - all_holdout_text.extend(holdout_text) - - calib_encoded = tokenizer( - all_calib_text, - return_tensors="pt", - padding=True, - truncation=True, - max_length=max_sample_length, - ) - holdout_encoded = tokenizer( - all_holdout_text, - return_tensors="pt", - padding=True, - truncation=True, - max_length=max_sample_length, - ) - - calib_encoded, holdout_encoded = _dedup_tokenized(calib_encoded, holdout_encoded) - _validate_no_overlap(calib_encoded["input_ids"], holdout_encoded["input_ids"]) - - # Save to disk. - if save_path: - save_path.mkdir(parents=True, exist_ok=True) - torch.save( - { - "input_ids": calib_encoded["input_ids"], - "attention_mask": calib_encoded["attention_mask"], - }, - calib_pt, - ) - print(f"Saved {len(calib_encoded['input_ids'])} calib samples to {calib_pt}") - torch.save( - { - "input_ids": holdout_encoded["input_ids"], - "attention_mask": holdout_encoded["attention_mask"], - }, - holdout_pt, - ) - print(f"Saved {len(holdout_encoded['input_ids'])} holdout samples to {holdout_pt}") - - # Build calib dataloader. - if device: - calib_encoded = {k: v.to(device) for k, v in calib_encoded.items()} - if include_labels: - calib_encoded["labels"] = torch.where( - calib_encoded["attention_mask"] > 0.5, - calib_encoded["input_ids"], - -100, - ) - else: - calib_encoded = {"input_ids": calib_encoded["input_ids"]} - - calib_dl = DataLoader(_CustomDataset(calib_encoded), batch_size=batch_size, shuffle=False) - return calib_dl, holdout_pt - - def get_dataset_dataloader( dataset_name: str | list[str] = "cnn_dailymail", tokenizer: "PreTrainedTokenizerBase | None" = None, From b2f5a6ec1d90615219ba90635b9ea13716c9901a Mon Sep 17 00:00:00 2001 From: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Date: Tue, 14 Apr 2026 22:53:34 +0000 Subject: [PATCH 4/9] reverted dataset_utils Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> --- modelopt/torch/utils/dataset_utils.py | 90 +++++++++------------------ 1 file changed, 31 insertions(+), 59 deletions(-) diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index b0bab281476..1e9a7fbbbd8 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -211,22 +211,43 @@ def _auto_preprocess_sample( ) -def _load_streaming_dataset( +def get_dataset_samples( dataset_name: str, + num_samples: int, *, apply_chat_template: bool = False, tokenizer: "PreTrainedTokenizerBase | None" = None, split: str | list[str] | None = None, -) -> tuple[list, Callable[[dict], str]]: - """Resolve dataset config and return streaming splits with a preprocessing function. +) -> list[str]: + """Load a portion of a dataset with the dataset name and a given size. - This is a shared helper for :func:`get_dataset_samples`. + Supports both registered datasets (in ``SUPPORTED_DATASET_CONFIG``) and arbitrary + HuggingFace datasets. Unregistered datasets are auto-detected by column names: + ``messages``/``conversations`` (chat), ``prompt``, ``text``, or ``input``. + + Args: + dataset_name: Name or HuggingFace path of the dataset to load, a local directory path, + or a path to a ``.jsonl`` file. For local directory paths, the + predefined config from ``SUPPORTED_DATASET_CONFIG`` is matched if the base folder name + matches a registered key (e.g. ``/hf-local/abisee/cnn_dailymail`` matches ``cnn_dailymail`` key). + num_samples: Number of samples to load from the dataset. + apply_chat_template: Whether to apply the chat template to the samples + (if supported by the dataset). For unregistered datasets with a + ``messages`` column, chat template is always applied regardless of + this flag. + tokenizer: Tokenizer to use for applying the chat template to the samples. + No tokenization is done and plain text is still returned. + split: Override the split(s) to load. Accepts a single split name or a list. + If ``None``, uses the splits defined in ``SUPPORTED_DATASET_CONFIG`` for + registered datasets, or ``["train"]`` for unregistered datasets. Returns: - A tuple of ``(dataset_splits, preprocess_fn)`` where *dataset_splits* is a list of - HuggingFace ``IterableDataset`` objects and *preprocess_fn* maps a raw sample dict - to a plain-text string (empty string signals a sample to skip). + Samples: The list of samples. """ + # Local JSONL file path support (each line is a JSON object with a `text` field). + if dataset_name.endswith(".jsonl"): + return get_jsonl_text_samples(dataset_name, num_samples, key="text") + from datasets import load_dataset local_dataset_path = None @@ -280,66 +301,17 @@ def _preprocess(sample: dict) -> str: print(f"Loading dataset with {config=} and {splits=}") dataset_splits = [load_dataset(streaming=True, **config, split=s) for s in splits] - return dataset_splits, _preprocess - - -def get_dataset_samples( - dataset_name: str, - num_samples: int, - *, - apply_chat_template: bool = False, - tokenizer: "PreTrainedTokenizerBase | None" = None, - split: str | list[str] | None = None, -) -> list[str]: - """Load a portion of a dataset with the dataset name and a given size. - - Supports both registered datasets (in ``SUPPORTED_DATASET_CONFIG``) and arbitrary - HuggingFace datasets. Unregistered datasets are auto-detected by column names: - ``messages``/``conversations`` (chat), ``prompt``, ``text``, or ``input``. - - Args: - dataset_name: Name or HuggingFace path of the dataset to load, a local directory path, - or a path to a ``.jsonl`` file. For local directory paths, the - predefined config from ``SUPPORTED_DATASET_CONFIG`` is matched if the base folder name - matches a registered key (e.g. ``/hf-local/abisee/cnn_dailymail`` matches ``cnn_dailymail`` key). - num_samples: Number of samples to load from the dataset. - apply_chat_template: Whether to apply the chat template to the samples - (if supported by the dataset). For unregistered datasets with a - ``messages`` column, chat template is always applied regardless of - this flag. - tokenizer: Tokenizer to use for applying the chat template to the samples. - No tokenization is done and plain text is still returned. - split: Override the split(s) to load. Accepts a single split name or a list. - If ``None``, uses the splits defined in ``SUPPORTED_DATASET_CONFIG`` for - registered datasets, or ``["train"]`` for unregistered datasets. - - Returns: - Samples: The list of samples. - """ - # Local JSONL file path support (each line is a JSON object with a `text` field). - if dataset_name.endswith(".jsonl"): - return get_jsonl_text_samples(dataset_name, num_samples, key="text") - - dataset_splits, _preprocess = _load_streaming_dataset( - dataset_name, - apply_chat_template=apply_chat_template, - tokenizer=tokenizer, - split=split, - ) - num_per_split = [num_samples // len(dataset_splits)] * len(dataset_splits) num_per_split[-1] += num_samples - sum(num_per_split) samples: list[str] = [] for dataset, n in zip(dataset_splits, num_per_split): - split_samples: list[str] = [] - for sample in dataset: - if len(split_samples) >= n: + for i, sample in enumerate(dataset): + if i >= n: break text = _preprocess(sample) if text: - split_samples.append(text) - samples.extend(split_samples) + samples.append(text) return samples From 23aeaf9ead576a5b04a7f7e419ee8a3695dc5ce1 Mon Sep 17 00:00:00 2001 From: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Date: Thu, 16 Apr 2026 01:11:08 +0000 Subject: [PATCH 5/9] update Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> --- .../torch/quantization/utils/calib_utils.py | 85 ++++++++----------- 1 file changed, 34 insertions(+), 51 deletions(-) diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index 6d05f2f51ac..5da7649a1cb 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -39,6 +39,7 @@ """GPTQ helper and Hessian utilities for calibration.""" import math +from contextlib import contextmanager import torch @@ -47,30 +48,6 @@ from modelopt.torch.utils.perf import get_used_gpu_mem_fraction -def load_vector_lut_codebook(quantizer): - """Load vector LUT codebook and quantizer params from a weight_quantizer. - - Returns: - Tuple of (codebook, quant_block_size, scale_type). - """ - from luts import encode as luts_encode - - extra_args = quantizer.backend_extra_args - encode_format = quantizer.num_bits - encode_path = extra_args.get("encode_path", "") - if encode_path and not encode_path.endswith("/"): - encode_path += "/" - - if "sorted" in encode_format: - cb = torch.load(encode_path + encode_format + ".pt", map_location="cpu") - codebook = cb["sorted_values"].cuda().float() - else: - codebook, _ = luts_encode(encode_format, path=encode_path, norm=False, cuda=True) - codebook = codebook.float() - - return codebook, extra_args.get("block_sizes"), extra_args.get("scale_type") - - def update_hessian(input, hessian, n_samples): """Update hessian matrix with new input samples using incremental formula. @@ -158,21 +135,33 @@ def free(self): self.weight = None self.h_inv = None + def is_vector_lut(self) -> bool: + """Check if this module's weight quantizer is configured for vector LUT quantization.""" + extra_args = getattr(self.module.weight_quantizer, "backend_extra_args", None) + return bool(extra_args and extra_args.get("lut_type") == "vector_lut") + + @contextmanager + def weight_slice_scales(self, scales): + """Temporarily replace _psx_scales with per-slice scales, then restore.""" + quantizer = self.module.weight_quantizer + original = quantizer._psx_scales + quantizer._psx_scales = scales + try: + yield + finally: + quantizer._psx_scales = original + def update_weights(self, block_size, perc_damp): """Run GPTQ blockwise weight update on this module. Populates ``self.weight`` and ``self.h_inv``, runs the blockwise update, logs MSE, and writes the result back to the module. """ - backend_extra_args = getattr(self.module.weight_quantizer, "backend_extra_args", None) - is_vector_lut = bool( - backend_extra_args and backend_extra_args.get("lut_type") == "vector_lut" - ) hessian = self.hessian.to(self.module.weight.device) self.weight = self.module.weight.data.float().clone() self._prepare_hessian_inverse(hessian, perc_damp) - if is_vector_lut: + if self.is_vector_lut(): self._blockwise_vector_update(block_size) else: self._blockwise_update(block_size) @@ -256,36 +245,28 @@ def _blockwise_update(self, block_size): ) def _blockwise_vector_update(self, block_size): - """GPTQ blockwise update for vector LUT quantizers. + """GPTQ blockwise update for vector quantizers. - Pre-computes scales once, then runs the standard GPTQ 3-loop - with per-vector-group static quantization via clip_vector_prescaled. + A single ``quantizer(weight)`` call computes and caches per-block + scales (``_psx_scales``) on the quantizer via the backend's + ``static_scales`` path. The GPTQ loop then slices per-vector-group + scales from ``_psx_scales`` for each sub-vector quantization call. """ import torch.nn.functional as F - from luts import clip_vector_prescaled, clip_vector_scalesign_fast - codebook, quant_block_size, scale_type = load_vector_lut_codebook( - self.module.weight_quantizer + quantizer = self.module.weight_quantizer + assert quantizer.backend_extra_args.get("static_scales", False), ( + "GPTQ vector update requires static_scales=True in backend_extra_args." ) - - # Get vector size from codebook - vector_size = codebook.shape[1] + vector_size = quantizer.backend_extra_args["vector_size"] + quant_block_size = quantizer.backend_extra_args["block_sizes"] assert self.weight is not None and self.h_inv is not None num_cols = self.weight.shape[1] assert block_size % quant_block_size == 0 - # Pre-compute scales once outside the GPTQ loop - _, scales = clip_vector_scalesign_fast( - self.weight, - codebook, - quant_block_size, - scale_type, - scale_algo="max", - sign_scale=True, - return_scales=True, - ) - scales_2d = scales.reshape(self.weight.shape[0], -1) + # Compute and cache _psx_scales on the quantizer via the backend. + quantizer(self.weight) w = self.weight.clone() h_inv = self.h_inv @@ -296,12 +277,14 @@ def _blockwise_vector_update(self, block_size): for j in range(blk_start, blk_end, vector_size): d = min(vector_size, blk_end - j) - s = scales_2d[:, j // quant_block_size].contiguous() + s = quantizer._psx_scales[:, j // quant_block_size].contiguous() sub = w[:, j : j + d].contiguous() if d < vector_size: sub = F.pad(sub, (0, vector_size - d)) - q_sub = clip_vector_prescaled(sub, codebook, s) + + with self.weight_slice_scales(s): + q_sub = quantizer(sub) for k in range(d): col = j + k From ad03449fd935f69d2092d08cb254a83d807a1329 Mon Sep 17 00:00:00 2001 From: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Date: Fri, 17 Apr 2026 16:10:54 +0000 Subject: [PATCH 6/9] update Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> --- modelopt/torch/quantization/model_calib.py | 9 +- .../torch/quantization/utils/calib_utils.py | 89 +++---------------- 2 files changed, 21 insertions(+), 77 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index faa19862a0c..5110b0c9c39 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -49,7 +49,7 @@ reduce_amax, weight_attr_names, ) -from .utils.calib_utils import GPTQHelper +from .utils.calib_utils import _GPTQ_HELPER_REGISTRY, GPTQHelper __all__ = [ "awq", @@ -1663,7 +1663,12 @@ def gptq( print_rank_0("No quantized linear layers found, skipping GPTQ") return - gptq_handles = {name: GPTQHelper(m, name, offload_to_cpu=True) for name, m in quantized_layers} + def _make_gptq_handle(name, m): + backend = getattr(m.weight_quantizer, "backend", None) + cls = _GPTQ_HELPER_REGISTRY.get(backend, GPTQHelper) + return cls(m, name, offload_to_cpu=True) + + gptq_handles = {name: _make_gptq_handle(name, m) for name, m in quantized_layers} for handle in gptq_handles.values(): handle.setup() diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index 5da7649a1cb..8c0ae1ee6ee 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -39,7 +39,6 @@ """GPTQ helper and Hessian utilities for calibration.""" import math -from contextlib import contextmanager import torch @@ -135,22 +134,6 @@ def free(self): self.weight = None self.h_inv = None - def is_vector_lut(self) -> bool: - """Check if this module's weight quantizer is configured for vector LUT quantization.""" - extra_args = getattr(self.module.weight_quantizer, "backend_extra_args", None) - return bool(extra_args and extra_args.get("lut_type") == "vector_lut") - - @contextmanager - def weight_slice_scales(self, scales): - """Temporarily replace _psx_scales with per-slice scales, then restore.""" - quantizer = self.module.weight_quantizer - original = quantizer._psx_scales - quantizer._psx_scales = scales - try: - yield - finally: - quantizer._psx_scales = original - def update_weights(self, block_size, perc_damp): """Run GPTQ blockwise weight update on this module. @@ -160,12 +143,7 @@ def update_weights(self, block_size, perc_damp): hessian = self.hessian.to(self.module.weight.device) self.weight = self.module.weight.data.float().clone() self._prepare_hessian_inverse(hessian, perc_damp) - - if self.is_vector_lut(): - self._blockwise_vector_update(block_size) - else: - self._blockwise_update(block_size) - + self._blockwise_update(block_size) self._print_mse_error(hessian) self.module.weight.data = self.weight.reshape(self.module.weight.shape).to( self.module.weight.data.dtype @@ -244,58 +222,6 @@ def _blockwise_update(self, block_size): errs, self.h_inv[block_start:block_end, block_end:], alpha=-1 ) - def _blockwise_vector_update(self, block_size): - """GPTQ blockwise update for vector quantizers. - - A single ``quantizer(weight)`` call computes and caches per-block - scales (``_psx_scales``) on the quantizer via the backend's - ``static_scales`` path. The GPTQ loop then slices per-vector-group - scales from ``_psx_scales`` for each sub-vector quantization call. - """ - import torch.nn.functional as F - - quantizer = self.module.weight_quantizer - assert quantizer.backend_extra_args.get("static_scales", False), ( - "GPTQ vector update requires static_scales=True in backend_extra_args." - ) - vector_size = quantizer.backend_extra_args["vector_size"] - quant_block_size = quantizer.backend_extra_args["block_sizes"] - - assert self.weight is not None and self.h_inv is not None - num_cols = self.weight.shape[1] - assert block_size % quant_block_size == 0 - - # Compute and cache _psx_scales on the quantizer via the backend. - quantizer(self.weight) - - w = self.weight.clone() - h_inv = self.h_inv - - for blk_start in range(0, num_cols, block_size): - blk_end = min(blk_start + block_size, num_cols) - errs = torch.zeros_like(w[:, blk_start:blk_end]) - - for j in range(blk_start, blk_end, vector_size): - d = min(vector_size, blk_end - j) - s = quantizer._psx_scales[:, j // quant_block_size].contiguous() - - sub = w[:, j : j + d].contiguous() - if d < vector_size: - sub = F.pad(sub, (0, vector_size - d)) - - with self.weight_slice_scales(s): - q_sub = quantizer(sub) - - for k in range(d): - col = j + k - self.weight[:, col] = q_sub[:, k] - err = (w[:, col] - q_sub[:, k]) / h_inv[col, col] - errs[:, col - blk_start] = err - w[:, col:blk_end].addr_(err, h_inv[col, col:blk_end], alpha=-1) - - if blk_end < num_cols: - w[:, blk_end:] -= errs @ h_inv[blk_start:blk_end, blk_end:] - def _print_mse_error(self, hessian): """Log Hessian-weighted relative MSE between ``self.weight`` and original weights.""" w_orig = self.module.weight.float() @@ -303,3 +229,16 @@ def _print_mse_error(self, hessian): mse = (delta).mm(hessian).mul(delta).mean() / (w_orig.mm(hessian).mul(w_orig).mean() + 1e-6) suffix = f", n_hessian_samples: {self.n_samples}" if self.n_samples else "" print_rank_0(f"[{self.name}] Relative MSE error: {mse.item():.2e}{suffix}") + + +_GPTQ_HELPER_REGISTRY: dict[str, type[GPTQHelper]] = {} + + +def register_gptq_helper(backend: str, factory: type[GPTQHelper]) -> None: + """Register a :class:`GPTQHelper` subclass for a quantizer backend. + + When :func:`modelopt.torch.quantization.model_calib.gptq` encounters a + module whose ``weight_quantizer.backend`` matches ``backend``, it will + construct ``factory`` instead of the default ``GPTQHelper``. + """ + _GPTQ_HELPER_REGISTRY[backend] = factory From 5870cc128532fe07be54b0620619f8db50743b3e Mon Sep 17 00:00:00 2001 From: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Date: Fri, 17 Apr 2026 16:34:22 +0000 Subject: [PATCH 7/9] update Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> --- examples/llm_ptq/example_utils.py | 4 ++-- examples/llm_ptq/hf_ptq.py | 13 ++++--------- modelopt/torch/quantization/model_calib.py | 9 ++++----- 3 files changed, 10 insertions(+), 16 deletions(-) diff --git a/examples/llm_ptq/example_utils.py b/examples/llm_ptq/example_utils.py index fb07e1016ac..c2d4d4bfca8 100755 --- a/examples/llm_ptq/example_utils.py +++ b/examples/llm_ptq/example_utils.py @@ -583,7 +583,7 @@ def get_model( model_kwargs = config_kwargs.copy() # Don't set torch_dtype for VILA models as they handle it explicitly in their builder if "vila" not in ckpt_path.lower(): - model_kwargs.setdefault("torch_dtype", "auto") + model_kwargs.setdefault("dtype", "auto") if "vila" in ckpt_path.lower(): hf_vila = AutoModel.from_pretrained( @@ -666,7 +666,7 @@ def has_pack_quantized_config(config): model_kwargs2 = model_kwargs.copy() if auto_model_module not in [AutoModelForCausalLM, AutoModel]: model_kwargs2.pop("trust_remote_code", None) - model_kwargs2["torch_dtype"] = torch_dtype + model_kwargs2["dtype"] = torch_dtype model_kwargs2.pop("max_memory", None) model = from_config(hf_config, **model_kwargs2) diff --git a/examples/llm_ptq/hf_ptq.py b/examples/llm_ptq/hf_ptq.py index 2ffd3ef3950..82f5f814698 100755 --- a/examples/llm_ptq/hf_ptq.py +++ b/examples/llm_ptq/hf_ptq.py @@ -420,15 +420,10 @@ def load_model(args: argparse.Namespace): attn_implementation=args.attn_implementation, ) else: - if args.qformat in QUANT_CFG_CHOICES: - quant_cfg = QUANT_CFG_CHOICES[args.qformat] - elif hasattr(mtq, args.qformat): - quant_cfg = getattr(mtq, args.qformat) - else: - raise AssertionError( - f"Quantization format is not supported for low memory mode. " - f"Supported formats: {QUANT_CFG_CHOICES.keys()}" - ) + assert args.qformat in QUANT_CFG_CHOICES, ( + f"Quantization format is not supported for low memory mode. Supported formats: {QUANT_CFG_CHOICES.keys()}" + ) + quant_cfg = QUANT_CFG_CHOICES[args.qformat] if args.kv_cache_qformat != "none": quant_cfg = mtq.utils.update_quant_cfg_with_kv_cache_quant( quant_cfg, diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 5110b0c9c39..2336cb6a01b 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1665,7 +1665,10 @@ def gptq( def _make_gptq_handle(name, m): backend = getattr(m.weight_quantizer, "backend", None) - cls = _GPTQ_HELPER_REGISTRY.get(backend, GPTQHelper) + if backend is None: + cls = GPTQHelper + else: + cls = _GPTQ_HELPER_REGISTRY.get(backend, GPTQHelper) return cls(m, name, offload_to_cpu=True) gptq_handles = {name: _make_gptq_handle(name, m) for name, m in quantized_layers} @@ -1685,10 +1688,6 @@ def _make_gptq_handle(name, m): print_rank_0("Updating weights using GPTQ algorithm...") for handle in gptq_handles.values(): handle.update_weights(block_size, perc_damp) - - # Disable weight quantizer after running GPTQ update since weights are already QDQ'ed - if hasattr(handle.module, "weight_quantizer"): - handle.module.weight_quantizer.disable() handle.free() del gptq_handles From 956d44b07b1ace30a9b5ef569a8da0330ef19e98 Mon Sep 17 00:00:00 2001 From: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Date: Fri, 17 Apr 2026 16:37:21 +0000 Subject: [PATCH 8/9] minor cleanup Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> --- examples/llm_ptq/hf_ptq.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/examples/llm_ptq/hf_ptq.py b/examples/llm_ptq/hf_ptq.py index 82f5f814698..37a88f97c74 100755 --- a/examples/llm_ptq/hf_ptq.py +++ b/examples/llm_ptq/hf_ptq.py @@ -1067,14 +1067,10 @@ def quantize_main( "Plain quantization supports only one quantization format." ) - if args.qformat in QUANT_CFG_CHOICES: - quant_cfg = QUANT_CFG_CHOICES[args.qformat] - elif hasattr(mtq, args.qformat): - quant_cfg = getattr(mtq, args.qformat) - else: - raise AssertionError( - f"Unsupported quantization format: {args.qformat}, choices are: {list(QUANT_CFG_CHOICES.keys())}" - ) + assert args.qformat in QUANT_CFG_CHOICES, ( + f"Unsupported quantization format: {args.qformat}, choices are: {list(QUANT_CFG_CHOICES.keys())}" + ) + quant_cfg = QUANT_CFG_CHOICES[args.qformat] quant_cfg = build_quant_cfg( args.qformat, @@ -1109,7 +1105,7 @@ def quantize_main( quant_cfg = copy.deepcopy(quant_cfg) _set_kv_cache_constant_amax(quant_cfg["quant_cfg"]) - if args.qformat in QUANT_CFG_CHOICES or hasattr(mtq, args.qformat): + if args.qformat in QUANT_CFG_CHOICES: mono_quantize( args, quant_cfg, From f2f28258ab4fda55eddeb1a3699b7b1bee7f7f71 Mon Sep 17 00:00:00 2001 From: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Date: Fri, 17 Apr 2026 18:33:13 +0000 Subject: [PATCH 9/9] removed stray unit tests Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> --- tests/gpu/torch/quantization/test_gptq_vq.py | 534 ------------------- 1 file changed, 534 deletions(-) delete mode 100644 tests/gpu/torch/quantization/test_gptq_vq.py diff --git a/tests/gpu/torch/quantization/test_gptq_vq.py b/tests/gpu/torch/quantization/test_gptq_vq.py deleted file mode 100644 index 9741e5f4853..00000000000 --- a/tests/gpu/torch/quantization/test_gptq_vq.py +++ /dev/null @@ -1,534 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 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. - -"""Test _blockwise_vector_update against gptq_quantize_scaled_vq from adaptive_rounding.py.""" -# ruff: noqa: N803, N806 — uppercase names (W, A, Q, H, etc.) match math notation in the reference. - -import tempfile -from types import SimpleNamespace - -import torch -import torch.nn.functional as F - -from modelopt.torch.quantization.utils.calib_utils import GPTQHelper - -# --------------------------------------------------------------------------- -# Exact copies from modelopt-internal/docs/adaptive_rounding.py -# --------------------------------------------------------------------------- - - -def compute_output_nmse( - W: torch.Tensor, - W_q: torch.Tensor, - A: torch.Tensor, -) -> float: - """ - Compute the output Normalized MSE: - - NMSE = || (W - W_q) @ A ||_F^2 / || W @ A ||_F^2 - - where @ denotes matrix multiplication. This measures how much - weight quantization error propagates through the activations. - - Args: - W: Original weight matrix, shape (out_features, in_features) - W_q: Quantized weight matrix, same shape as W - A: Activation matrix, shape (in_features, n_samples) - - Returns: - Scalar NMSE value - """ - error_output = (W - W_q) @ A - ref_output = W @ A - nmse = error_output.norm() ** 2 / ref_output.norm() ** 2 - return nmse.item() - - -def _round_to_e4m3(s: torch.Tensor) -> torch.Tensor: - """Round a float32 scale tensor to the nearest E4M3 representable value.""" - return s.to(dtype=torch.float8_e4m3fn).to(dtype=torch.float32) - - -def _vq_nearest(vecs: torch.Tensor, codebook: torch.Tensor, chunk_size: int = 1) -> torch.Tensor: - """ - Find the nearest codeword for each row-vector in vecs. - - Args: - vecs: (N, D) float tensor - codebook: (K, D) float tensor - chunk_size: Process this many vectors at a time to bound memory. - Default 1 to avoid torch.cdist batch precision issues. - - Returns: - (N, D) tensor of nearest codewords - """ - N = vecs.shape[0] - result = torch.empty_like(vecs) - for start in range(0, N, chunk_size): - end = min(start + chunk_size, N) - dists = torch.cdist(vecs[start:end], codebook) - result[start:end] = codebook[dists.argmin(dim=1)] - return result - - -def scaled_vector_quantize( - W: torch.Tensor, - codebook: torch.Tensor, - quant_block_size: int, - n_scale_iters: int = 3, -) -> tuple[torch.Tensor, torch.Tensor]: - """ - Block-scaled vector quantization with iterative scale optimisation. - - For each row-block of quant_block_size elements we find a scalar - scale s and codeword assignments that (approximately) minimise - - || block - s * reconstruct(block / s) ||^2 - - via alternating minimisation: - 1. Fix s -> assign each sub-vector to nearest codeword in codebook - 2. Fix assignments -> solve for optimal s in closed form: - s* = sum_i / sum_i ||c_i||^2 - - Initialised with s = max(|block|) / max(|codebook|) (amax), then - refined for n_scale_iters iterations. - - Effective bitrate: log2(n_codewords) / codebook_dim + 16 / quant_block_size - - Args: - W: (out_features, in_features) - codebook: (n_codewords, codebook_dim) — must evenly divide quant_block_size - quant_block_size: Elements per scaling block - n_scale_iters: Number of assign-then-rescale iterations (default 3) - - Returns: - W_q: Dequantised weight matrix - scales: (out_features, n_blocks) - """ - codebook_dim = codebook.shape[1] - out_features, in_features = W.shape - assert quant_block_size % codebook_dim == 0, ( - f"quant_block_size ({quant_block_size}) must be a multiple of codebook_dim ({codebook_dim})" - ) - assert in_features % codebook_dim == 0, ( - f"in_features ({in_features}) must be a multiple of codebook_dim ({codebook_dim})" - ) - - cb_absmax = codebook.abs().max().item() - n_blocks = (in_features + quant_block_size - 1) // quant_block_size - - # Initial scales: amax heuristic - pad = n_blocks * quant_block_size - in_features - if pad > 0: - W_padded = torch.nn.functional.pad(W, (0, pad)) - else: - W_padded = W - W_blocks = W_padded.reshape(out_features, n_blocks, quant_block_size) - block_max = W_blocks.abs().amax(dim=2) - scales = (block_max / cb_absmax).clamp(min=1e-10) - - W_q = torch.empty_like(W) - - for b in range(n_blocks): - start = b * quant_block_size - end = min(start + quant_block_size, in_features) - block = W[:, start:end] # (out_features, block_len) - s = scales[:, b] # (out_features,) - - for _ in range(n_scale_iters): - # Step 1: assign sub-vectors to nearest codeword at current scale - normalized = block / s.unsqueeze(1) - vecs = normalized.reshape(-1, codebook_dim) - q_vecs = _vq_nearest(vecs, codebook) - q_block = q_vecs.reshape(out_features, end - start) - - # Step 2: optimal scale given assignments - # s* = sum_i / sum_i ||c_i||^2 - # where x_i are original sub-vectors, c_i are assigned codewords - numerator = (block * q_block).sum(dim=1) # (out_features,) - denominator = (q_block * q_block).sum(dim=1) # (out_features,) - s = (numerator / denominator.clamp(min=1e-20)).clamp(min=1e-10) - - # Round converged scale to E4M3 and do final assignment - s = _round_to_e4m3(s) - normalized = block / s.unsqueeze(1) - vecs = normalized.reshape(-1, codebook_dim) - q_vecs = _vq_nearest(vecs, codebook) - W_q[:, start:end] = q_vecs.reshape(out_features, end - start) * s.unsqueeze(1) - scales[:, b] = s - - return W_q, scales - - -def gptq_quantize_scaled_vq( - W: torch.Tensor, - A: torch.Tensor, - codebook: torch.Tensor, - quant_block_size: int, - gptq_block_size: int = 128, - damp_pct: float = 0.01, - n_scale_iters: int = 3, - h_inv: torch.Tensor | None = None, -) -> tuple[torch.Tensor, float]: - """ - GPTQ with block-scaled vector quantization. - - Scales are pre-computed via the same iterative optimisation as - scaled_vector_quantize. During the GPTQ pass, columns are processed - in groups of codebook_dim: each group is VQ-quantized to the nearest - codeword, and per-column errors are compensated into remaining columns - using the inverse Hessian. - - Args: - W: (out_features, in_features) - A: (in_features, n_samples) - codebook: (n_codewords, codebook_dim) - quant_block_size: Elements per scaling block - gptq_block_size: Columns per GPTQ lazy batch update (must be - a multiple of codebook_dim) - damp_pct: Dampening fraction - n_scale_iters: Iterations for scale optimisation - h_inv: Optional pre-computed upper-triangular Cholesky factor of the - damped inverse Hessian. If None, computed from A. - - Returns: - W_q: Quantized weight matrix (dequantised) - nmse: Output NMSE achieved - """ - codebook_dim = codebook.shape[1] - assert gptq_block_size % codebook_dim == 0, ( - f"gptq_block_size ({gptq_block_size}) must be a multiple of codebook_dim ({codebook_dim})" - ) - - W_orig = W.clone() - W = W.clone() - out_features, in_features = W.shape - - _, scales = scaled_vector_quantize( - W_orig, codebook, quant_block_size, n_scale_iters=n_scale_iters - ) - - if h_inv is None: - H = 2.0 * (A @ A.T) - damp = damp_pct * torch.diag(H).mean() - H.diagonal().add_(damp) - - H_inv = torch.linalg.inv(H) - try: - L = torch.linalg.cholesky(H_inv) - except torch.linalg.LinAlgError: - H.diagonal().add_(damp * 10) - H_inv = torch.linalg.inv(H) - L = torch.linalg.cholesky(H_inv) - Hinv = L.T - else: - Hinv = h_inv - - Q = torch.zeros_like(W) - - for i in range(0, in_features, gptq_block_size): - j_end = min(i + gptq_block_size, in_features) - E = torch.zeros(out_features, j_end - i, dtype=W.dtype, device=W.device) - - for j in range(i, j_end, codebook_dim): - d = min(codebook_dim, j_end - j) - sb = j // quant_block_size - s = scales[:, sb] # (out_features,) - - sub_vec = W[:, j : j + d] / s.unsqueeze(1) - q_vecs = _vq_nearest( - sub_vec.reshape(-1, codebook_dim) if d == codebook_dim else sub_vec.reshape(-1, d), - codebook[:, :d], - ) - q_block = q_vecs.reshape(out_features, d) * s.unsqueeze(1) - Q[:, j : j + d] = q_block - - for k in range(d): - col = j + k - err = (W[:, col] - Q[:, col]) / Hinv[col, col] - E[:, col - i] = err - W[:, col:j_end] -= err.unsqueeze(1) * Hinv[col, col:j_end].unsqueeze(0) - - if j_end < in_features: - W[:, j_end:] -= E @ Hinv[i:j_end, j_end:] - - nmse = compute_output_nmse(W_orig, Q, A) - return Q, nmse - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _make_codebook_file(tmp_dir, name, n_codewords, vector_size): - """Save a codebook .pt file in the format load_vector_lut_codebook expects.""" - codebook = torch.randn(n_codewords, vector_size) - torch.save({"sorted_values": codebook}, f"{tmp_dir}/{name}.pt") - return codebook.cuda().float() - - -def _make_hessian(activations): - """Build Hessian H = 2 * X @ X^T from activations (batch, seq, features).""" - X = activations.reshape(-1, activations.shape[-1]).t().float() - return 2.0 * (X @ X.t()) - - -def _attach_mock_quantizer(module, encode_path, encode_format, quant_block_size, scale_type): - """Attach a mock weight_quantizer with the right attributes for _blockwise_vector_update.""" - module.weight_quantizer = SimpleNamespace( - num_bits=encode_format, - backend="psx_luts", - backend_extra_args={ - "encode_path": encode_path, - "lut_type": "vector_lut", - "block_sizes": quant_block_size, - "scale_type": scale_type, - }, - ) - - -# --------------------------------------------------------------------------- -# Tests -# --------------------------------------------------------------------------- - -SEED = 42 -OUT_FEATURES = 64 -IN_FEATURES = 128 -GPTQ_BLOCK_SIZE = 128 -QUANT_BLOCK_SIZE = 16 -VECTOR_SIZE = 8 -N_CODEWORDS = 256 -PERC_DAMP = 0.01 -SCALE_TYPE = "e4m3" - - -def test_clip_vector_prescaled_vs_vq_nearest(): - """clip_vector_prescaled and _vq_nearest must produce identical quantized vectors.""" - from luts import clip_vector_prescaled, clip_vector_scalesign_fast - - torch.manual_seed(SEED) - - with tempfile.TemporaryDirectory() as tmp_dir: - codebook = _make_codebook_file(tmp_dir, "test_sorted", N_CODEWORDS, VECTOR_SIZE) - - W = torch.randn(OUT_FEATURES, IN_FEATURES, device="cuda", dtype=torch.float32) - - # Compute scales with sign_scale=True (scales can be negative) - _, scales = clip_vector_scalesign_fast( - W, - codebook, - QUANT_BLOCK_SIZE, - SCALE_TYPE, - scale_algo="max", - sign_scale=True, - return_scales=True, - ) - scales_2d = scales.reshape(OUT_FEATURES, -1) - - for j in range(0, IN_FEATURES, VECTOR_SIZE): - d = min(VECTOR_SIZE, IN_FEATURES - j) - s = scales_2d[:, j // QUANT_BLOCK_SIZE].contiguous() - sub = W[:, j : j + d].contiguous() - - # clip_vector_prescaled path - if d < VECTOR_SIZE: - sub_padded = F.pad(sub, (0, VECTOR_SIZE - d)) - else: - sub_padded = sub - q_luts = clip_vector_prescaled(sub_padded, codebook, s)[:, :d] - - # _vq_nearest path (manual normalize -> lookup -> denormalize) - normalized = sub / s.unsqueeze(1) - vecs = ( - normalized.reshape(-1, VECTOR_SIZE) - if d == VECTOR_SIZE - else normalized.reshape(-1, d) - ) - cb_slice = codebook if d == VECTOR_SIZE else codebook[:, :d] - q_vecs = _vq_nearest(vecs, cb_slice) - q_ref = q_vecs.reshape(OUT_FEATURES, d) * s.unsqueeze(1) - - assert torch.allclose(q_luts, q_ref, atol=1e-5), ( - f"VQ mismatch at col {j}: max diff={(q_luts - q_ref).abs().max().item():.2e}" - ) - - -def test_blockwise_vector_update_vs_gptq_quantize_scaled_vq(): - """_blockwise_vector_update must produce identical weights to gptq_quantize_scaled_vq. - - Both paths share the same pre-computed scales (from clip_vector_scalesign_fast) - and the same h_inv, so the only variable is the GPTQ loop itself. - The reference's scaled_vector_quantize is patched to return the shared scales. - """ - from unittest.mock import patch - - from luts import clip_vector_scalesign_fast - - torch.manual_seed(SEED) - - with tempfile.TemporaryDirectory() as tmp_dir: - encode_format = "test_sorted" - codebook = _make_codebook_file(tmp_dir, encode_format, N_CODEWORDS, VECTOR_SIZE) - - W = torch.randn(OUT_FEATURES, IN_FEATURES, device="cuda", dtype=torch.float32) - A = torch.randn(4, 16, IN_FEATURES, device="cuda", dtype=torch.float32) - A_flat = A.reshape(-1, IN_FEATURES).t() # (in_features, n_samples) for reference - hessian = _make_hessian(A) - - # Shared scales — computed once, used by both paths - Q_rtn, scales = clip_vector_scalesign_fast( - W, - codebook, - QUANT_BLOCK_SIZE, - SCALE_TYPE, - scale_algo="max", - sign_scale=True, - return_scales=True, - ) - scales_2d = scales.reshape(OUT_FEATURES, -1) - - # --- Reference: exact gptq_quantize_scaled_vq from adaptive_rounding.py --- - # Patch scaled_vector_quantize to return the shared scales (Q_rtn unused by caller) - with patch( - f"{__name__}.scaled_vector_quantize", - return_value=(Q_rtn, scales_2d), - ): - Q_ref, _ = gptq_quantize_scaled_vq( - W, - A_flat, - codebook, - QUANT_BLOCK_SIZE, - GPTQ_BLOCK_SIZE, - damp_pct=PERC_DAMP, - ) - - # --- Modelopt: _blockwise_vector_update --- - module = torch.nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False).cuda() - module.weight.data = W.clone() - _attach_mock_quantizer(module, tmp_dir, encode_format, QUANT_BLOCK_SIZE, SCALE_TYPE) - - helper = GPTQHelper(module, "test_layer") - helper.hessian = hessian.clone() - helper.n_samples = 1 - helper.update_weights(GPTQ_BLOCK_SIZE, PERC_DAMP) - Q_modelopt = module.weight.data.float() - - # Same scales + same GPTQ loop structure => weights must match - assert torch.allclose(Q_modelopt, Q_ref, atol=1e-5), ( - f"Weights differ: max diff={(Q_modelopt - Q_ref).abs().max().item():.2e}" - ) - - -def test_mtq_quantize_gptq_vs_gptq_quantize_scaled_vq(): - """End-to-end: mtq.quantize with GPTQ vs gptq_quantize_scaled_vq. - - Both paths compute their own hessian and h_inv independently. - Shared scales ensure the only variable is the GPTQ loop. The hessian is - injected into GPTQHelper so both paths compute identical h_inv. - """ - from unittest.mock import patch - - from luts import clip_vector_scalesign_fast - - import modelopt.torch.quantization as mtq - - torch.manual_seed(SEED) - - with tempfile.TemporaryDirectory() as tmp_dir: - encode_format = "test_sorted" - codebook = _make_codebook_file(tmp_dir, encode_format, N_CODEWORDS, VECTOR_SIZE) - - W = torch.randn(OUT_FEATURES, IN_FEATURES, device="cuda", dtype=torch.float32) - A = torch.randn(4, 16, IN_FEATURES, device="cuda", dtype=torch.float32) - A_flat = A.reshape(-1, IN_FEATURES).t() - - # --- Modelopt: mtq.quantize with GPTQ (list-based config) --- - config = { - "quant_cfg": [ - {"quantizer_name": "*", "enable": False}, - { - "quantizer_name": "*weight_quantizer", - "cfg": { - "backend": "psx_luts", - "num_bits": encode_format, - "pass_through_bwd": True, - "backend_extra_args": { - "encode_path": tmp_dir, - "lut_type": "vector_lut", - "block_sizes": QUANT_BLOCK_SIZE, - "scale_type": SCALE_TYPE, - "scale_algo": "max", - "round_mode": "rne", - "sign_scale": True, - }, - }, - }, - ], - "algorithm": { - "method": "gptq", - "use_sequential": False, - "perc_damp": PERC_DAMP, - "block_size": GPTQ_BLOCK_SIZE, - }, - } - - model = torch.nn.Linear(IN_FEATURES, OUT_FEATURES, bias=False).cuda() - model.weight.data = W.clone() - - def forward_loop(m): - m(A) - - # Capture h_inv from the modelopt GPTQ run - captured = {} - _orig_update = GPTQHelper.update_weights - - def _capturing_update(self, *args, **kwargs): - _orig_update(self, *args, **kwargs) - captured["h_inv"] = self.h_inv.clone() - - with patch.object(GPTQHelper, "update_weights", _capturing_update): - mtq.quantize(model, config, forward_loop=forward_loop) - Q_modelopt = model.weight.data.float() - - # --- Reference: gptq_quantize_scaled_vq with shared scales + captured h_inv --- - Q_rtn, scales = clip_vector_scalesign_fast( - W, - codebook, - QUANT_BLOCK_SIZE, - SCALE_TYPE, - scale_algo="max", - sign_scale=True, - return_scales=True, - ) - scales_2d = scales.reshape(OUT_FEATURES, -1) - - with patch( - f"{__name__}.scaled_vector_quantize", - return_value=(Q_rtn, scales_2d), - ): - Q_ref, _ = gptq_quantize_scaled_vq( - W, - A_flat, - codebook, - QUANT_BLOCK_SIZE, - GPTQ_BLOCK_SIZE, - damp_pct=PERC_DAMP, - h_inv=captured["h_inv"], - ) - - assert torch.allclose(Q_modelopt, Q_ref, atol=1e-5), ( - f"Weights differ: max diff={(Q_modelopt - Q_ref).abs().max().item():.2e}" - )