diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d6397b993a0..5edc498e426 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -14,6 +14,8 @@ 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. +- 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..8d85c4ca85e 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 + +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 +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..cca885d59cd 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,11 +1516,14 @@ 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) - setup_distributed_args(args) + if args.deterministic: + enable_determinism(RAND_SEED) + else: + random.seed(RAND_SEED) + np.random.seed(RAND_SEED) + 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 new file mode 100644 index 00000000000..b2962440e97 --- /dev/null +++ b/modelopt/torch/quantization/utils/scale_compare.py @@ -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) + 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) + + +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/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..81d7aa9e566 --- /dev/null +++ b/modelopt/torch/utils/determinism.py @@ -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.") 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/quantization/test_scale_compare.py b/tests/unit/torch/quantization/test_scale_compare.py new file mode 100644 index 00000000000..7b23df80136 --- /dev/null +++ b/tests/unit/torch/quantization/test_scale_compare.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +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)]) + 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): + 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)]) + + +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 new file mode 100644 index 00000000000..119f9050dc7 --- /dev/null +++ b/tests/unit/torch/utils/test_determinism.py @@ -0,0 +1,55 @@ +# 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() + 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") +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" + + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":0:0" + with pytest.warns(UserWarning, match="CUBLAS_WORKSPACE_CONFIG=:0:0"): + enable_determinism(7)