From 5f32cfb66335a74420510d012a41a7e1c843823b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 06:10:40 +0000 Subject: [PATCH 1/4] Add a tool to compare the quantization scales of two PTQ runs Diffs the input, weight and KV scales of two exported safetensors checkpoints (or any two name -> tensor mappings) per scale kind, reports how many tensors differ and the median and worst relative difference, and exits non-zero when they differ beyond a tolerance. It is the measuring stick for how reproducible PTQ is across repeated runs, batch sizes and parallel layouts. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + .../torch/quantization/utils/scale_compare.py | 196 ++++++++++++++++++ .../torch/quantization/test_scale_compare.py | 115 ++++++++++ 3 files changed, 312 insertions(+) create mode 100644 modelopt/torch/quantization/utils/scale_compare.py create mode 100644 tests/unit/torch/quantization/test_scale_compare.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d6397b993a0..27a57326667 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -14,6 +14,7 @@ Changelog - GPTQ no longer runs a max-calibration forward when no activation quantizer needs one, and ``hf_ptq.py`` now sizes the default layerwise calibration batch at 8,192 tokens (16 samples at ``--calib_seq 512``) instead of 1. Layerwise GPTQ on HF fused-MoE models runs over 10x faster; pass ``--batch_size 1`` for the old batching, which pads no sample. - Add a checkpoint-specific NVIDIA Nemotron 3.5 Super VL AutoQuantize recipe with a 4.9-effective-bit NVFP4-MSE/FP8/BF16 search and fixed NVFP4 MTP experts. Select it under ``models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/auto_quantize/``. - Add ``examples/hf_ptq/convert_ggml_to_int8.py`` to expand a checkpoint exported with a GGML recipe (IQ formats, ``q8_0``) into an INT8 weight-only compressed-tensors checkpoint. Stock vLLM serves it with its integer kernels at about half the BF16 weight memory, reproducing the GGML weights up to rounding their scales to the model dtype. +- Add ``python -m modelopt.torch.quantization.utils.scale_compare`` to diff the quantization scales of two exported checkpoints within a tolerance, e.g. to check how reproducible a PTQ run is across repeats or parallel layouts. *Speculative Decoding* diff --git a/modelopt/torch/quantization/utils/scale_compare.py b/modelopt/torch/quantization/utils/scale_compare.py new file mode 100644 index 00000000000..5a305396d6c --- /dev/null +++ b/modelopt/torch/quantization/utils/scale_compare.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compare the quantization scales of two PTQ runs. + +Export both runs as safetensors checkpoints and diff every scale tensor:: + + python -m modelopt.torch.quantization.utils.scale_compare ckpt_a ckpt_b --rtol 1e-3 + +The exit status is 0 when all scales agree within the tolerance, so the command doubles as a +reproducibility check for repeated runs or different parallel layouts. :func:`compare_scales` +takes any two ``{name: tensor}`` mappings, e.g. flattened ``get_quantizer_state_dict`` outputs. +""" + +import argparse +import re +import statistics +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path + +import torch +from safetensors import safe_open + +__all__ = ["ScaleDiff", "TensorDiff", "compare_scales", "load_checkpoint_scales", "main"] + +# Exported scales (input_scale, weight_scale, weight_scale_2, pre_quant_scale, k_scale, v_scale) +# and the quantizer buffers (_amax) they come from. +_SCALE_NAME = re.compile(r"(_scale|_scale_2|_amax)$") + + +@dataclass +class TensorDiff: + """Difference between one scale tensor of the two runs.""" + + name: str + elements: int + outside_tolerance: int + max_rel: float # largest |a - b| / max(|a|, |b|) over the elements + + @property + def kind(self) -> str: + """Last component of the name, e.g. ``input_scale``.""" + return self.name.rsplit(".", 1)[-1] + + +@dataclass +class ScaleDiff: + """Result of :func:`compare_scales`.""" + + tensors: list[TensorDiff] = field(default_factory=list) + only_in_a: list[str] = field(default_factory=list) + only_in_b: list[str] = field(default_factory=list) + shape_mismatch: list[str] = field(default_factory=list) + + @property + def within_tolerance(self) -> bool: + """Whether tensors were compared and none is missing, reshaped or outside the tolerance.""" + return ( + bool(self.tensors) + and not (self.only_in_a or self.only_in_b or self.shape_mismatch) + and all(t.outside_tolerance == 0 for t in self.tensors) + ) + + def summary(self, top: int = 5) -> str: + """Per-kind table (kind = last component of the tensor name) and the worst tensors.""" + by_kind: dict[str, list[TensorDiff]] = {} + for t in self.tensors: + by_kind.setdefault(t.kind, []).append(t) + lines = [ + f"{'kind':<16}{'tensors':>9}{'differing':>11}{'diff elems':>12}" + f"{'median rel':>12}{'max rel':>10}" + ] + for kind, ts in sorted(by_kind.items()): + differing = [t.max_rel for t in ts if t.outside_tolerance] + elements = sum(t.elements for t in ts) + outside = sum(t.outside_tolerance for t in ts) + median = statistics.median(differing) if differing else 0.0 + lines.append( + f"{kind:<16}{len(ts):>9}{len(differing):>11}{outside / max(elements, 1):>12.2%}" + f"{median:>12.2e}{max(t.max_rel for t in ts):>10.2e}" + ) + worst = sorted((t for t in self.tensors if t.outside_tolerance), key=lambda t: -t.max_rel) + if worst: + lines.append("worst:") + lines += [f" {t.max_rel:.2e} {t.name}" for t in worst[:top]] + for label, names in ( + ("only in A", self.only_in_a), + ("only in B", self.only_in_b), + ("shape mismatch", self.shape_mismatch), + ): + if names: + lines.append(f"{label}: {len(names)} tensors, e.g. {names[0]}") + lines.append("result: " + ("within tolerance" if self.within_tolerance else "DIFFERENT")) + return "\n".join(lines) + + +def _diff(a: torch.Tensor, b: torch.Tensor, rtol: float, atol: float) -> tuple[int, float]: + """Return the number of elements outside the tolerance and the largest relative difference.""" + dtype = torch.float64 if torch.float64 in (a.dtype, b.dtype) else torch.float32 + a, b = a.to(dtype), b.to(dtype) # fp8 and bf16 scales have no arithmetic of their own + same = (a == b) | (a.isnan() & b.isnan()) + delta = (a - b).abs() + magnitude = torch.maximum(a.abs(), b.abs()) + rel = torch.nan_to_num(delta / magnitude, nan=torch.inf).masked_fill(same, 0.0) + close = same | (a.isfinite() & b.isfinite() & (delta <= atol + rtol * magnitude)) + return int((~close).sum()), float(rel.max()) if rel.numel() else 0.0 + + +def compare_scales( + a: Mapping[str, torch.Tensor], + b: Mapping[str, torch.Tensor], + rtol: float = 0.0, + atol: float = 0.0, +) -> ScaleDiff: + """Diff the tensors of ``a`` and ``b`` by name. + + An element is outside the tolerance when ``|a - b| > atol + rtol * max(|a|, |b|)``. The default + tolerance of zero asks for identical scales. + """ + # Plain name sets: a lazy mapping's `in` would load the tensor. + names_a, names_b = set(a), set(b) + result = ScaleDiff(only_in_a=sorted(names_a - names_b), only_in_b=sorted(names_b - names_a)) + for name in sorted(names_a & names_b): + ta, tb = a[name], b[name] + if ta.shape != tb.shape: + result.shape_mismatch.append(name) + continue + outside, max_rel = _diff(ta, tb, rtol, atol) + result.tensors.append(TensorDiff(name, ta.numel(), outside, max_rel)) + return result + + +class _CheckpointScales(Mapping[str, torch.Tensor]): + """Lazy ``name -> tensor`` view of the scale tensors of a safetensors checkpoint.""" + + def __init__(self, path: str | Path): + path = Path(path) + files = sorted(path.glob("*.safetensors")) if path.is_dir() else [path] + if not files: + raise FileNotFoundError(f"No .safetensors files in {path}") + self._handles: dict[str, safe_open] = {} + for file in files: + handle = safe_open(file, framework="pt") + self._handles.update(dict.fromkeys(filter(_SCALE_NAME.search, handle.keys()), handle)) + if not self._handles: + raise ValueError(f"No quantization scale tensors in {path}") + + def __getitem__(self, name: str) -> torch.Tensor: + return self._handles[name].get_tensor(name) + + def __iter__(self) -> Iterator[str]: + return iter(self._handles) + + def __len__(self) -> int: + return len(self._handles) + + +def load_checkpoint_scales(path: str | Path) -> Mapping[str, torch.Tensor]: + """Scale tensors of a safetensors file or of every ``*.safetensors`` file in a directory.""" + return _CheckpointScales(path) + + +def main(argv: Sequence[str] | None = None) -> int: + """Compare two checkpoints; the exit status is 0 when their scales agree within tolerance.""" + parser = argparse.ArgumentParser(description="Compare the quantization scales of two PTQ runs.") + parser.add_argument("checkpoint_a", help="Safetensors file or checkpoint directory.") + parser.add_argument("checkpoint_b", help="Safetensors file or checkpoint directory.") + parser.add_argument("--rtol", type=float, default=0.0, help="Relative tolerance (default 0).") + parser.add_argument("--atol", type=float, default=0.0, help="Absolute tolerance (default 0).") + parser.add_argument("--top", type=int, default=5, help="Number of worst tensors to list.") + args = parser.parse_args(argv) + result = compare_scales( + load_checkpoint_scales(args.checkpoint_a), + load_checkpoint_scales(args.checkpoint_b), + args.rtol, + args.atol, + ) + print(result.summary(args.top)) + return 0 if result.within_tolerance else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/unit/torch/quantization/test_scale_compare.py b/tests/unit/torch/quantization/test_scale_compare.py new file mode 100644 index 00000000000..8df65bdf449 --- /dev/null +++ b/tests/unit/torch/quantization/test_scale_compare.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections import Counter +from collections.abc import Mapping + +import pytest +import torch +from safetensors.torch import save_file + +from modelopt.torch.quantization.utils.scale_compare import compare_scales, main + + +class _LoadCounting(Mapping): + """Mapping that counts tensor loads, like the lazy checkpoint view does on disk reads.""" + + def __init__(self, tensors): + self.tensors = tensors + self.loads = Counter() + + def __getitem__(self, name): + self.loads[name] += 1 + return self.tensors[name] + + def __iter__(self): + return iter(self.tensors) + + def __len__(self): + return len(self.tensors) + + +@pytest.mark.parametrize(("rtol", "input_scale_outside"), [(0.0, 1), (0.02, 0)]) +def test_compare_scales(rtol, input_scale_outside): + fp8 = lambda *values: torch.tensor(values).to(torch.float8_e4m3fn) # noqa: E731 + a = { + "l0.q.input_scale": torch.tensor(0.5), + "l1.q.input_scale": torch.tensor(1.0), + "l0.q.weight_scale": fp8(1.0, 2.0, 4.0, 8.0), + "gone.input_scale": torch.tensor(1.0), + "reshaped.input_scale": torch.zeros(2), + } + b = { + "l0.q.input_scale": torch.tensor(0.505), # 1% off + "l1.q.input_scale": torch.tensor(1.0), + "l0.q.weight_scale": fp8(1.0, 2.0, 4.0, 9.0), # one fp8 step off + "new.input_scale": torch.tensor(1.0), + "reshaped.input_scale": torch.zeros(3), + } + + diff = compare_scales(a, b, rtol=rtol) + + assert diff.only_in_a == ["gone.input_scale"] + assert diff.only_in_b == ["new.input_scale"] + assert diff.shape_mismatch == ["reshaped.input_scale"] + by_name = {t.name: t for t in diff.tensors} + assert by_name["l1.q.input_scale"].outside_tolerance == 0 + assert by_name["l0.q.input_scale"].outside_tolerance == input_scale_outside + assert by_name["l0.q.input_scale"].max_rel == pytest.approx(0.005 / 0.505, rel=1e-4) + assert by_name["l0.q.weight_scale"].outside_tolerance == 1 + assert by_name["l0.q.weight_scale"].max_rel == pytest.approx(1 / 9) + assert not diff.within_tolerance + + +@pytest.mark.parametrize(("args", "exit_code"), [([], 1), (["--rtol", "0.02"], 0)]) +def test_main_compares_checkpoint_directories(tmp_path, capsys, args, exit_code): + run_a, run_b = tmp_path / "a", tmp_path / "b" + run_a.mkdir() + run_b.mkdir() + # Run A is sharded by layer, run B is one file; only the scales are compared, not the weights. + save_file( + {"l0.q.input_scale": torch.tensor(0.5), "l0.q.weight": torch.ones(2, 2)}, + run_a / "model-1.safetensors", + ) + save_file({"l1.q.input_scale": torch.tensor(1.0)}, run_a / "model-2.safetensors") + save_file( + { + "l0.q.input_scale": torch.tensor(0.5), + "l1.q.input_scale": torch.tensor(1.01), + "l0.q.weight": torch.zeros(2, 2), + }, + run_b / "model.safetensors", + ) + + assert main([str(run_a), str(run_b), *args]) == exit_code + + assert "input_scale" in capsys.readouterr().out + + +def test_compare_scales_loads_each_tensor_once(): + a = _LoadCounting({f"l{i}.q.input_scale": torch.tensor(float(i + 1)) for i in range(3)}) + b = _LoadCounting(dict(a.tensors)) + + assert compare_scales(a, b).within_tolerance + + assert set(a.loads.values()) == set(b.loads.values()) == {1} + + +def test_nothing_to_compare_is_not_a_pass(tmp_path): + assert not compare_scales({}, {}).within_tolerance + + save_file({"l0.q.weight": torch.ones(2, 2)}, tmp_path / "model.safetensors") + with pytest.raises(ValueError, match="No quantization scale tensors"): + main([str(tmp_path), str(tmp_path)]) From 19dd30297f3b219454e68d6cb1c6cf6972d8cadf Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 06:11:39 +0000 Subject: [PATCH 2/4] Add --deterministic to hf_ptq for best-effort reproducible PTQ hf_ptq seeded only python and numpy, and its default --batch_size 0 probes the calibration batch size from free GPU memory, so another GPU or memory state calibrates with another batch size. The new switch also seeds torch, selects deterministic torch kernels (an op without one only warns), sets CUBLAS_WORKSPACE_CONFIG, and rejects --batch_size 0, through a small enable_determinism helper. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + examples/hf_ptq/README.md | 12 ++++++ examples/hf_ptq/hf_ptq.py | 26 ++++++++++-- modelopt/torch/utils/__init__.py | 1 + modelopt/torch/utils/determinism.py | 42 +++++++++++++++++++ tests/examples/hf_ptq/test_hf_ptq_args.py | 36 +++++++++++++++++ tests/unit/torch/utils/test_determinism.py | 47 ++++++++++++++++++++++ 7 files changed, 162 insertions(+), 3 deletions(-) create mode 100644 modelopt/torch/utils/determinism.py create mode 100644 tests/unit/torch/utils/test_determinism.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 27a57326667..5edc498e426 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -15,6 +15,7 @@ Changelog - Add a checkpoint-specific NVIDIA Nemotron 3.5 Super VL AutoQuantize recipe with a 4.9-effective-bit NVFP4-MSE/FP8/BF16 search and fixed NVFP4 MTP experts. Select it under ``models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/auto_quantize/``. - Add ``examples/hf_ptq/convert_ggml_to_int8.py`` to expand a checkpoint exported with a GGML recipe (IQ formats, ``q8_0``) into an INT8 weight-only compressed-tensors checkpoint. Stock vLLM serves it with its integer kernels at about half the BF16 weight memory, reproducing the GGML weights up to rounding their scales to the model dtype. - Add ``python -m modelopt.torch.quantization.utils.scale_compare`` to diff the quantization scales of two exported checkpoints within a tolerance, e.g. to check how reproducible a PTQ run is across repeats or parallel layouts. +- Add ``--deterministic`` to ``examples/hf_ptq/hf_ptq.py`` for best-effort reproducible PTQ: it seeds every random number generator and selects deterministic torch kernels, and requires an explicit ``--batch_size`` because the default probes it from free GPU memory. *Speculative Decoding* diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 422729277e1..2ee3419eb32 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -788,6 +788,18 @@ After the TensorRT-LLM checkpoint export, you can use the `trtllm-build` build c - Deployable on [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM), [vLLM](https://github.com/vllm-project/vllm) and [SGLang](https://github.com/sgl-project/sglang) - More models coming soon! +## Reproducible PTQ + +Repeating a run with the same model, data, recipe, batch size, hardware and software gives the same scales. What moves them is the calibration batch size: the default `--batch_size 0` probes it from free GPU memory, so another GPU or memory state silently calibrates with another batch size. `--deterministic` pins what the run still depends on: it requires an explicit `--batch_size`, seeds every random number generator and selects deterministic torch kernels (an op without one only warns). It costs some run time. + +```bash +python hf_ptq.py ... --batch_size 8 --deterministic --export_path run_a +python hf_ptq.py ... --batch_size 8 --deterministic --export_path run_b +python -m modelopt.torch.quantization.utils.scale_compare run_a run_b +``` + +`scale_compare` lists, per scale kind (`input_scale`, `weight_scale`, ...), how many tensors differ and by how much, and exits non-zero when a scale differs by more than `--rtol` / `--atol` (default 0: identical). Run it on checkpoints that differ in batch size or GPU layout to see how far the scales move. + ## Tracking runs with MLflow Set MLflow's own `MLFLOW_TRACKING_URI`, or pass `--mlflow `, to record a PTQ diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 27e8fb3e8d5..9b0309fc7b9 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -93,7 +93,7 @@ EagleOfflineDataCollator, OfflineSupervisedDataset, ) -from modelopt.torch.utils import print_rank_0 +from modelopt.torch.utils import enable_determinism, print_rank_0 from modelopt.torch.utils.dataset_utils import ( create_forward_loop, get_dataset_dataloader, @@ -1208,6 +1208,17 @@ def parse_args() -> argparse.Namespace: type=int, default=512, ) + parser.add_argument( + "--deterministic", + help=( + "Best-effort reproducible PTQ: seed every random number generator and use " + "deterministic torch kernels (an op without one only warns). Needs an explicit " + "--batch_size, since the default probes it from free GPU memory and so differs " + "between GPUs and runs." + ), + default=False, + action="store_true", + ) parser.add_argument("--export_path", default="exported_model") parser.add_argument( "--dataset", @@ -1436,6 +1447,12 @@ def parse_args() -> argparse.Namespace: if args.calib_seq <= 0: parser.error("--calib_seq must be positive.") + if args.deterministic and args.batch_size == 0: + parser.error( + "--deterministic needs an explicit --batch_size: the default probes it from free GPU " + "memory, so it differs between GPUs and runs." + ) + if args.specdec_offline_dataset is not None and args.sparsity_fmt != "dense": parser.error("--specdec_offline_dataset is only supported with --sparsity_fmt dense (PTQ).") @@ -1499,8 +1516,11 @@ def main(args: argparse.Namespace): if not torch.cuda.is_available(): raise OSError("GPU is required for inference.") - random.seed(RAND_SEED) - np.random.seed(RAND_SEED) + if args.deterministic: + enable_determinism(RAND_SEED) + else: + random.seed(RAND_SEED) + np.random.seed(RAND_SEED) setup_distributed_args(args) diff --git a/modelopt/torch/utils/__init__.py b/modelopt/torch/utils/__init__.py index a38c80cac01..e33865b9c53 100644 --- a/modelopt/torch/utils/__init__.py +++ b/modelopt/torch/utils/__init__.py @@ -18,6 +18,7 @@ from ._pytree import * from .cpp_extension import * from .dataset_utils import * +from .determinism import * from .graph import * from .import_utils import * from .list import * diff --git a/modelopt/torch/utils/determinism.py b/modelopt/torch/utils/determinism.py new file mode 100644 index 00000000000..c6e8defb478 --- /dev/null +++ b/modelopt/torch/utils/determinism.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Best-effort determinism for reproducible runs.""" + +import os +import random + +import numpy as np +import torch + +from .logging import print_rank_0 + +__all__ = ["enable_determinism"] + + +def enable_determinism(seed: int) -> None: + """Seed every random number generator and ask PyTorch for deterministic kernels. + + An operation without a deterministic kernel only warns, so repeated runs on the same hardware, + software and parallel layout agree as closely as the remaining kernels allow. Call this at + startup, before the first CUDA matmul: cuBLAS reads ``CUBLAS_WORKSPACE_CONFIG`` when it creates + its handle. + """ + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + torch.use_deterministic_algorithms(True, warn_only=True) + print_rank_0(f"Deterministic mode: seeded every RNG with {seed}, deterministic torch kernels.") diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index b441dcc3ce7..c14623a1282 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -964,6 +964,42 @@ def load_model(args): ] +def test_deterministic_needs_an_explicit_batch_size(monkeypatch): + flags = ["--pyt_ckpt_path", "/models/Qwen3-0.6B", "--deterministic"] + with pytest.raises(SystemExit): + _parse_hf_ptq_args(monkeypatch, *flags) + + _, args = _parse_hf_ptq_args(monkeypatch, *flags, "--batch_size", "8") + + assert args.deterministic + + +@pytest.mark.usefixtures("clean_env") +def test_main_enables_determinism_before_loading(monkeypatch, tmp_path): + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, "--pyt_ckpt_path", "org/model", "--deterministic", "--batch_size", "8" + ) + seeds = [] + + class _LoadReachedError(Exception): + pass + + def load_model(args): + raise _LoadReachedError + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.compiler, "set_stance", lambda *args, **kwargs: None) + monkeypatch.setattr(hf_ptq, "launch_memory_monitor", lambda: None) + monkeypatch.setattr(hf_ptq, "ensure_local_checkpoint", lambda path: (path, str(tmp_path))) + monkeypatch.setattr(hf_ptq, "enable_determinism", seeds.append) + monkeypatch.setattr(hf_ptq, "load_model", load_model) + + with pytest.raises(_LoadReachedError): + hf_ptq.main(args) + + assert seeds == [hf_ptq.RAND_SEED] + + def test_tensorrt_llm_export_copies_non_model_files_from_the_local_copy(monkeypatch, tmp_path): """The deprecated TensorRT-LLM exporter does not copy them itself; hf_ptq does, from disk.""" hf_ptq, args = _parse_hf_ptq_args( diff --git a/tests/unit/torch/utils/test_determinism.py b/tests/unit/torch/utils/test_determinism.py new file mode 100644 index 00000000000..1f0bed42778 --- /dev/null +++ b/tests/unit/torch/utils/test_determinism.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import random +from unittest import mock + +import numpy as np +import pytest +import torch + +from modelopt.torch.utils import enable_determinism + + +@pytest.fixture +def restore_global_state(): + deterministic = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + with mock.patch.dict(os.environ): + yield + torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) + + +@pytest.mark.usefixtures("restore_global_state") +def test_enable_determinism(): + os.environ.pop("CUBLAS_WORKSPACE_CONFIG", None) + + enable_determinism(7) + draws = (random.random(), np.random.rand(), torch.rand(1).item()) + enable_determinism(7) + + assert (random.random(), np.random.rand(), torch.rand(1).item()) == draws + assert torch.are_deterministic_algorithms_enabled() + assert torch.is_deterministic_algorithms_warn_only_enabled() + assert os.environ["CUBLAS_WORKSPACE_CONFIG"] == ":4096:8" From d73aa56381ec15d1641ccbc3006648a857f10ae2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 18:53:28 +0000 Subject: [PATCH 3/4] Reject a scale in two shards and qualify the reproducibility claim A scale name that appeared in two files of a checkpoint directory was silently taken from the later file, so a stale or overlapping export could pass a comparison without every tensor being checked; it is an error now. The README presented identical scales as a guarantee, but --deterministic only warns for an op without a deterministic kernel, so it now says the result was observed on the measured configurations. The determinism test also restores the RNG states it reseeds. Signed-off-by: Shengliang Xu --- examples/hf_ptq/README.md | 2 +- modelopt/torch/quantization/utils/scale_compare.py | 5 ++++- tests/unit/torch/quantization/test_scale_compare.py | 8 ++++++++ tests/unit/torch/utils/test_determinism.py | 4 ++++ 4 files changed, 17 insertions(+), 2 deletions(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 2ee3419eb32..8d85c4ca85e 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -790,7 +790,7 @@ After the TensorRT-LLM checkpoint export, you can use the `trtllm-build` build c ## Reproducible PTQ -Repeating a run with the same model, data, recipe, batch size, hardware and software gives the same scales. What moves them is the calibration batch size: the default `--batch_size 0` probes it from free GPU memory, so another GPU or memory state silently calibrates with another batch size. `--deterministic` pins what the run still depends on: it requires an explicit `--batch_size`, seeds every random number generator and selects deterministic torch kernels (an op without one only warns). It costs some run time. +On the configurations we measured (Qwen3-30B-A3B NVFP4 on one and two GPUs, Qwen3-8B AWQ), repeating a run with the same model, data, recipe, batch size, hardware and software gave identical scales. What moved them was the calibration batch size: the default `--batch_size 0` probes it from free GPU memory, so another GPU or memory state silently calibrates with another batch size. `--deterministic` pins what the run still depends on: it requires an explicit `--batch_size`, seeds every random number generator and selects deterministic torch kernels. This is best effort, not a guarantee: an op without a deterministic kernel still runs and only warns. It costs some run time. ```bash python hf_ptq.py ... --batch_size 8 --deterministic --export_path run_a diff --git a/modelopt/torch/quantization/utils/scale_compare.py b/modelopt/torch/quantization/utils/scale_compare.py index 5a305396d6c..ec04c3b7598 100644 --- a/modelopt/torch/quantization/utils/scale_compare.py +++ b/modelopt/torch/quantization/utils/scale_compare.py @@ -154,7 +154,10 @@ def __init__(self, path: str | Path): self._handles: dict[str, safe_open] = {} for file in files: handle = safe_open(file, framework="pt") - self._handles.update(dict.fromkeys(filter(_SCALE_NAME.search, handle.keys()), handle)) + names = list(filter(_SCALE_NAME.search, handle.keys())) + if duplicates := self._handles.keys() & set(names): + raise ValueError(f"{sorted(duplicates)[0]} is in more than one file of {path}") + self._handles.update(dict.fromkeys(names, handle)) if not self._handles: raise ValueError(f"No quantization scale tensors in {path}") diff --git a/tests/unit/torch/quantization/test_scale_compare.py b/tests/unit/torch/quantization/test_scale_compare.py index 8df65bdf449..03e327cc796 100644 --- a/tests/unit/torch/quantization/test_scale_compare.py +++ b/tests/unit/torch/quantization/test_scale_compare.py @@ -113,3 +113,11 @@ def test_nothing_to_compare_is_not_a_pass(tmp_path): save_file({"l0.q.weight": torch.ones(2, 2)}, tmp_path / "model.safetensors") with pytest.raises(ValueError, match="No quantization scale tensors"): main([str(tmp_path), str(tmp_path)]) + + +def test_a_scale_in_two_shards_is_an_error(tmp_path): + save_file({"l0.q.input_scale": torch.ones(1)}, tmp_path / "model-1.safetensors") + save_file({"l0.q.input_scale": torch.ones(1)}, tmp_path / "model-2.safetensors") + + with pytest.raises(ValueError, match="more than one file"): + main([str(tmp_path), str(tmp_path)]) diff --git a/tests/unit/torch/utils/test_determinism.py b/tests/unit/torch/utils/test_determinism.py index 1f0bed42778..70e8512a970 100644 --- a/tests/unit/torch/utils/test_determinism.py +++ b/tests/unit/torch/utils/test_determinism.py @@ -28,9 +28,13 @@ def restore_global_state(): deterministic = torch.are_deterministic_algorithms_enabled() warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + rng_states = (random.getstate(), np.random.get_state(), torch.get_rng_state()) with mock.patch.dict(os.environ): yield torch.use_deterministic_algorithms(deterministic, warn_only=warn_only) + random.setstate(rng_states[0]) + np.random.set_state(rng_states[1]) + torch.set_rng_state(rng_states[2]) @pytest.mark.usefixtures("restore_global_state") From 133fb1d2a01d9a2010feaffbdb75e9da7694bf25 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 19:40:58 +0000 Subject: [PATCH 4/4] Compare weight_scale_inv and warn on a non-deterministic cuBLAS config The scale pattern missed weight_scale_inv, which the unified HF export emits for block-scaled FP8, so a comparison of such checkpoints skipped those scales and could pass without having looked at them. A missing checkpoint path now says so, and the summary documents what its columns count. enable_determinism keeps a CUBLAS_WORKSPACE_CONFIG the user already set but warns when cuBLAS does not treat it as deterministic. hf_ptq seeds after setup_distributed_args, so the "Deterministic mode" line is printed once and not on every rank. Signed-off-by: Shengliang Xu --- examples/hf_ptq/hf_ptq.py | 4 ++-- .../torch/quantization/utils/scale_compare.py | 15 +++++++++++---- modelopt/torch/utils/determinism.py | 16 +++++++++++++--- .../torch/quantization/test_scale_compare.py | 9 +++++++++ tests/unit/torch/utils/test_determinism.py | 4 ++++ 5 files changed, 39 insertions(+), 9 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 9b0309fc7b9..cca885d59cd 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -1516,14 +1516,14 @@ def main(args: argparse.Namespace): if not torch.cuda.is_available(): raise OSError("GPU is required for inference.") + setup_distributed_args(args) + if args.deterministic: enable_determinism(RAND_SEED) else: random.seed(RAND_SEED) np.random.seed(RAND_SEED) - setup_distributed_args(args) - try: # Entered inside the try: opening the run is fatal by design, and skipping # cleanup_distributed would leave the other ranks blocked on the first collective diff --git a/modelopt/torch/quantization/utils/scale_compare.py b/modelopt/torch/quantization/utils/scale_compare.py index ec04c3b7598..b2962440e97 100644 --- a/modelopt/torch/quantization/utils/scale_compare.py +++ b/modelopt/torch/quantization/utils/scale_compare.py @@ -36,9 +36,9 @@ __all__ = ["ScaleDiff", "TensorDiff", "compare_scales", "load_checkpoint_scales", "main"] -# Exported scales (input_scale, weight_scale, weight_scale_2, pre_quant_scale, k_scale, v_scale) -# and the quantizer buffers (_amax) they come from. -_SCALE_NAME = re.compile(r"(_scale|_scale_2|_amax)$") +# Exported scales (input_scale, weight_scale, weight_scale_2, weight_scale_inv, pre_quant_scale, +# k_scale, v_scale) and the quantizer buffers (_amax) they come from. +_SCALE_NAME = re.compile(r"(_scale(_2|_inv)?|_amax)$") @dataclass @@ -75,7 +75,12 @@ def within_tolerance(self) -> bool: ) def summary(self, top: int = 5) -> str: - """Per-kind table (kind = last component of the tensor name) and the worst tensors.""" + """Per-kind table (kind = last component of the tensor name) and the worst tensors. + + ``differing`` counts the tensors outside the tolerance and ``median rel`` is the median of + their largest relative differences; ``max rel`` is the largest over all tensors of the + kind, including those inside the tolerance. + """ by_kind: dict[str, list[TensorDiff]] = {} for t in self.tensors: by_kind.setdefault(t.kind, []).append(t) @@ -148,6 +153,8 @@ class _CheckpointScales(Mapping[str, torch.Tensor]): def __init__(self, path: str | Path): path = Path(path) + if not path.exists(): + raise FileNotFoundError(f"{path} does not exist") files = sorted(path.glob("*.safetensors")) if path.is_dir() else [path] if not files: raise FileNotFoundError(f"No .safetensors files in {path}") diff --git a/modelopt/torch/utils/determinism.py b/modelopt/torch/utils/determinism.py index c6e8defb478..81d7aa9e566 100644 --- a/modelopt/torch/utils/determinism.py +++ b/modelopt/torch/utils/determinism.py @@ -21,10 +21,12 @@ import numpy as np import torch -from .logging import print_rank_0 +from .logging import print_rank_0, warn_rank_0 __all__ = ["enable_determinism"] +_DETERMINISTIC_CUBLAS_CONFIGS = (":4096:8", ":16:8") + def enable_determinism(seed: int) -> None: """Seed every random number generator and ask PyTorch for deterministic kernels. @@ -32,11 +34,19 @@ def enable_determinism(seed: int) -> None: An operation without a deterministic kernel only warns, so repeated runs on the same hardware, software and parallel layout agree as closely as the remaining kernels allow. Call this at startup, before the first CUDA matmul: cuBLAS reads ``CUBLAS_WORKSPACE_CONFIG`` when it creates - its handle. + its handle. A value of that variable that is already set is kept, with a warning if cuBLAS does + not treat it as deterministic. """ random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) - os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + cublas_config = os.environ.setdefault( + "CUBLAS_WORKSPACE_CONFIG", _DETERMINISTIC_CUBLAS_CONFIGS[0] + ) + if cublas_config not in _DETERMINISTIC_CUBLAS_CONFIGS: + warn_rank_0( + f"CUBLAS_WORKSPACE_CONFIG={cublas_config} is not one of {_DETERMINISTIC_CUBLAS_CONFIGS}: " + "cuBLAS reductions are not made deterministic." + ) torch.use_deterministic_algorithms(True, warn_only=True) print_rank_0(f"Deterministic mode: seeded every RNG with {seed}, deterministic torch kernels.") diff --git a/tests/unit/torch/quantization/test_scale_compare.py b/tests/unit/torch/quantization/test_scale_compare.py index 03e327cc796..7b23df80136 100644 --- a/tests/unit/torch/quantization/test_scale_compare.py +++ b/tests/unit/torch/quantization/test_scale_compare.py @@ -113,6 +113,8 @@ def test_nothing_to_compare_is_not_a_pass(tmp_path): save_file({"l0.q.weight": torch.ones(2, 2)}, tmp_path / "model.safetensors") with pytest.raises(ValueError, match="No quantization scale tensors"): main([str(tmp_path), str(tmp_path)]) + with pytest.raises(FileNotFoundError, match="does not exist"): + main([str(tmp_path / "missing"), str(tmp_path)]) def test_a_scale_in_two_shards_is_an_error(tmp_path): @@ -121,3 +123,10 @@ def test_a_scale_in_two_shards_is_an_error(tmp_path): with pytest.raises(ValueError, match="more than one file"): main([str(tmp_path), str(tmp_path)]) + + +def test_weight_scale_inv_is_compared(tmp_path): + save_file({"l0.q.weight_scale_inv": torch.tensor(1.0)}, tmp_path / "a.safetensors") + save_file({"l0.q.weight_scale_inv": torch.tensor(1.1)}, tmp_path / "b.safetensors") + + assert main([str(tmp_path / "a.safetensors"), str(tmp_path / "b.safetensors")]) == 1 diff --git a/tests/unit/torch/utils/test_determinism.py b/tests/unit/torch/utils/test_determinism.py index 70e8512a970..119f9050dc7 100644 --- a/tests/unit/torch/utils/test_determinism.py +++ b/tests/unit/torch/utils/test_determinism.py @@ -49,3 +49,7 @@ def test_enable_determinism(): assert torch.are_deterministic_algorithms_enabled() assert torch.is_deterministic_algorithms_warn_only_enabled() assert os.environ["CUBLAS_WORKSPACE_CONFIG"] == ":4096:8" + + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":0:0" + with pytest.warns(UserWarning, match="CUBLAS_WORKSPACE_CONFIG=:0:0"): + enable_determinism(7)