Repository navigation
[OMNIML-6284] Add a PTQ scale-comparison tool and hf_ptq --deterministic #2995
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,206 @@ | ||
| # 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, 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 | ||
| 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. | ||
|
|
||
| ``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) | ||
| 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) | ||
|
Comment on lines
+87
to
+100
Contributor
There was a problem hiding this comment. [SUGGESTION] Two columns in the table mean something narrower than their headers suggest, and the docstring on
The PR description spells both out; the tool's own output is what a user will have in front of them months later. One extra clause on the def summary(self, top: int = 5) -> str:
"""Per-kind table (kind = last component of the tensor name) and the worst tensors.
``median rel`` is the median over the differing tensors of each tensor's largest
relative difference; ``max rel`` is the largest over all tensors of the kind.
"""
Collaborator
Author
There was a problem hiding this comment. Done in 133fb1d: the |
||
| 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) | ||
| 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}") | ||
| self._handles: dict[str, safe_open] = {} | ||
| for file in files: | ||
| handle = safe_open(file, framework="pt") | ||
| 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}") | ||
|
|
||
| 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) | ||
|
Comment on lines
+153
to
+183
Contributor
There was a problem hiding this comment. [SUGGESTION] For the CLI that's fine — the process exits. For the documented library use ( Minor, separate: when
Collaborator
Author
There was a problem hiding this comment. Partly done in 133fb1d: a path that does not exist now raises |
||
|
|
||
|
|
||
| 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()) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,52 @@ | ||
| # 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, 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. | ||
|
|
||
| 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. 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) | ||
| 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.") |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[SUGGESTION]
enable_determinismruns beforesetup_distributed_args(args), so itsprint_rank_0is not yet rank-aware:dist.is_master()returnsTrueon every rank untiltorch.distributedis initialized, and undertorchrunall N ranks print the "Deterministic mode" line.Harmless, but it's the one line of this feature's output and duplicating it per rank is the opposite of what
print_rank_0is there for. Moving the call just aftersetup_distributed_args(args)keeps it comfortably before the first CUDA matmul (the constraint the docstring actually cares about) and makes the message appear once.The ordering in
tests/examples/hf_ptq/test_hf_ptq_args.py::test_main_enables_determinism_before_loading("before loading") still holds either way, sinceload_modelcomes well after.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Moved in 133fb1d: the seeding block now runs right after
setup_distributed_args(args), so the "Deterministic mode" line is printed once instead of on every rank under torchrun. It is still beforeensure_local_checkpoint, the model load and the first matmul.