diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a073a683af9..08b1293ea12 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -15,6 +15,7 @@ Changelog - GPTQ now also quantizes HuggingFace fused MoE experts, each with its own Hessian from the tokens routed to it; they were previously skipped. The built-in IQ1 and IQ2 recipes therefore run GPTQ on the experts too, and large MoEs need layerwise calibration to hold one layer's expert Hessians at a time. - 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 to what one decoder layer's forward fits in GPU memory, up to 32,768 tokens (64 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. - 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 ``hessian_input_dtype`` to the GPTQ and ``local_hessian`` algorithm configs. ``bfloat16`` accumulates the float32 Hessian with a tensor-core GEMM on bfloat16 activations, several times faster and without a float32 copy of each batch; the default ``float32`` keeps the existing float32 GEMM. - 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. *Speculative Decoding* diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index a43e031f135..88cab998218 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -999,7 +999,26 @@ class MseCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): ) -class LocalHessianCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): +class _HessianInputDtypeConfig(ModeloptBaseConfig): + """The ``hessian_input_dtype`` knob, shared by GPTQ and local-Hessian calibration.""" + + hessian_input_dtype: Literal["float32", "bfloat16"] = ModeloptField( + default="float32", + title="Dtype the Hessian GEMM reads the activations in.", + description="""The Hessian is accumulated and stored in float32 either way. ``float32`` + copies each batch of activations to float32 for the GEMM. ``bfloat16`` feeds them to a + tensor-core GEMM in bfloat16 that accumulates in float32: for bfloat16 activations it + forms the same products and only rounds the sums differently. At 32K tokens per batch + that is up to about 3e-5 of the Hessian's diagonal for GPTQ's full Hessian and 1.5e-4 + for local Hessian's 16-column blocks, and the GEMM is 6-16x faster. It needs CUDA GEMMs + that take ``out_dtype``; without them the Hessian falls back to float32 with a + warning.""", + ) + + +class LocalHessianCalibConfig( + _SharedStatesConfig, _HessianInputDtypeConfig, QuantizeAlgorithmConfig +): """Configuration for local Hessian-weighted MSE calibration. This algorithm uses activation information to optimize per-block scales for weight @@ -1210,7 +1229,7 @@ class SVDQuantConfig(QuantizeAlgorithmConfig): ) -class GPTQCalibConfig(QuantizeAlgorithmConfig): +class GPTQCalibConfig(_HessianInputDtypeConfig, QuantizeAlgorithmConfig): """The config for GPTQ quantization. GPTQ minimizes the layer-wise quantization error by using second-order (Hessian) information diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index ac20729da19..047e4443189 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -74,6 +74,7 @@ _GPTQ_HELPER_REGISTRY, _IMATRIX_ENCODER_REGISTRY, GPTQHelper, + bfloat16_hessian_gemm, track_calib_forward_memory, ) @@ -957,10 +958,13 @@ class _LocalHessianAccumulator: scale; the buffer is allocated lazily so never-routed experts cost nothing. """ - def __init__(self, cout: int, cin: int, block_size: int): + def __init__( + self, cout: int, cin: int, block_size: int, input_dtype: torch.dtype = torch.float32 + ): self.cout = cout self.cin = cin self.block_size = block_size + self.input_dtype = input_dtype self.num_blocks_per_cin = cin // block_size # Not block-divisible -> no Hessian (falls back to plain MSE). self.is_enabled = cin % block_size == 0 @@ -973,10 +977,15 @@ def accumulate(self, input_tensor: torch.Tensor) -> None: """Accumulate ``XᵀX`` per block from an activation of shape ``(..., cin)``.""" if not self.is_enabled: return - # fp32 GEMM avoids bf16/fp16 precision loss; (cin, tokens) -> (n_blocks, bs, tokens). - x = input_tensor.reshape(-1, self.cin).to(torch.float32).T + # (cin, tokens) -> (n_blocks, bs, tokens). A float32 GEMM avoids bf16/fp16 precision loss; + # a bfloat16 one, accumulating in float32, is faster and needs no float32 copy. + bf16 = bfloat16_hessian_gemm(self.input_dtype, input_tensor) + x = input_tensor.reshape(-1, self.cin).to(torch.bfloat16 if bf16 else torch.float32).T x = x.reshape(self.num_blocks_per_cin, self.block_size, -1) - hessian_batch = x @ x.transpose(-1, -2) + if bf16: + hessian_batch = torch.bmm(x, x.transpose(-1, -2), out_dtype=torch.float32) + else: + hessian_batch = x @ x.transpose(-1, -2) if self.hessian_per_block is None: self.hessian_per_block = hessian_batch else: @@ -1178,6 +1187,7 @@ def local_hessian_calibrate( block_size: int = 16, debug: bool = False, shared_states: Mapping[str, Mapping[str, Sequence[str]]] | None = None, + hessian_input_dtype: str = "float32", ): """Calibrate weight quantizers by minimizing the Hessian-weighted error. @@ -1206,6 +1216,8 @@ def local_hessian_calibrate( block_size: Block size for local Hessian computation (default: 16). debug: If True, retain the per-quantizer Hessian accumulators on the model (``model._local_hessian_accumulators``) for inspection. + hessian_input_dtype: ``"float32"`` or ``"bfloat16"``, the dtype the Hessian GEMM reads + the activations in. See :class:`LocalHessianCalibConfig ` for details on the configuration options. @@ -1227,7 +1239,9 @@ def capture(weight_quantizer, weight, input_tensor): input_local = input_tensor.to_local() if hasattr(input_tensor, "to_local") else input_tensor acc = accumulators.get(id(weight_quantizer)) if acc is None: - acc = _LocalHessianAccumulator(weight.shape[0], weight.shape[1], block_size) + acc = _LocalHessianAccumulator( + weight.shape[0], weight.shape[1], block_size, getattr(torch, hessian_input_dtype) + ) accumulators[id(weight_quantizer)] = acc acc.accumulate(input_local) @@ -2425,6 +2439,7 @@ def gptq( block_size: int = 128, fused: bool = False, importance_weighted: bool = False, + hessian_input_dtype: str = "float32", ): """GPTQ quantization. @@ -2463,6 +2478,8 @@ def gptq( fused: If True, use fused Triton kernel for NVFP4 static quantization. importance_weighted: If True, formats whose codebook search takes a per-input-column importance weight it by the square root of the Hessian diagonal. + hessian_input_dtype: ``"float32"`` or ``"bfloat16"``, the dtype the Hessian GEMM reads + the activations in; see :class:`GPTQCalibConfig `. """ total_start = time.time() @@ -2513,6 +2530,7 @@ def _make_gptq_handle(name, m, weight_quantizer=None, weight=None, log_mse=True) offload_to_cpu=True, fused=fused, importance_weighted=importance_weighted, + hessian_input_dtype=hessian_input_dtype, weight_quantizer=weight_quantizer, weight=weight, log_mse=log_mse, diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index 4895a31103b..8b045de3f15 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -38,12 +38,18 @@ """GPTQ helper and Hessian utilities for calibration.""" +import functools from contextlib import contextmanager import torch import torch.distributed as dist -from modelopt.torch.utils import get_available_gpu_memory, get_process_gpus, print_rank_0 +from modelopt.torch.utils import ( + get_available_gpu_memory, + get_process_gpus, + print_rank_0, + warn_rank_0, +) from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method # A GPU takes a new Hessian only if enough memory stays available after it for the calibration @@ -107,20 +113,23 @@ def _hessian_offload_device(nbytes: int) -> torch.device: return torch.device("cpu") -def update_hessian(input, hessian, n_samples): +def update_hessian(input, hessian, n_samples, input_dtype=torch.float32): """Update hessian matrix with new input samples using incremental formula. Args: input: Input tensor (batch_size, ..., features) hessian: Current Hessian matrix to update in-place n_samples: Number of samples already processed + input_dtype: The dtype the GEMM reads ``input`` in: ``torch.float32``, or + ``torch.bfloat16`` for a tensor-core GEMM that accumulates in float32. + Returns: Tuple of (updated_hessian, new_sample_count) Note: input must be non-empty (batch_size > 0); a zero-sized input causes division by zero. """ # Flatten to 2D (total_tokens, features) first, so batch_size counts tokens - input_flat = input.reshape(-1, input.shape[-1]).float() + input_flat = input.reshape(-1, input.shape[-1]) batch_size = input_flat.shape[0] if batch_size == 0: # in MOEs some experts receive no tokens @@ -130,16 +139,71 @@ def update_hessian(input, hessian, n_samples): # once per expert per forward, and separate scale and add passes each rewrite all of H. scale = n_samples / (n_samples + batch_size) n_samples += batch_size + alpha = 2 / n_samples + if bfloat16_hessian_gemm(input_dtype, input_flat): + # BF16 products are exact in float32, so a tensor-core GEMM writing float32 differs from a + # float32 GEMM only in how it rounds the sums: up to ~3e-5 of the Hessian's diagonal at 32K + # tokens per batch. It is 9-16x faster on GB300 and needs no float32 copy of the inputs. + input_flat = input_flat.to(torch.bfloat16) + if hessian.device == input_flat.device: + torch.addmm( + hessian, + input_flat.t(), + input_flat, + beta=scale, + alpha=alpha, + out_dtype=torch.float32, + out=hessian, + ) + else: + gram = torch.mm(input_flat.t(), input_flat, out_dtype=torch.float32) + hessian.mul_(scale).add_(gram.to(hessian.device), alpha=alpha) + return hessian, n_samples + + input_flat = input_flat.float() if hessian.device == input_flat.device: - hessian.addmm_(input_flat.t(), input_flat, beta=scale, alpha=2 / n_samples) + hessian.addmm_(input_flat.t(), input_flat, beta=scale, alpha=alpha) else: - hessian.mul_(scale).add_( - (input_flat.t() @ input_flat).to(hessian.device), alpha=2 / n_samples - ) + hessian.mul_(scale).add_((input_flat.t() @ input_flat).to(hessian.device), alpha=alpha) return hessian, n_samples +def bfloat16_hessian_gemm(input_dtype: torch.dtype, inputs: torch.Tensor) -> bool: + """Whether a Hessian GEMM over ``inputs`` runs on bfloat16 with float32 output. + + True when ``input_dtype`` asks for bfloat16 and the inputs' GPU has such GEMMs; when they are + missing, warns once per device and returns False, so the caller computes in float32. + """ + if input_dtype != torch.bfloat16: + return False + if inputs.is_cuda and _has_float32_out_gemm(inputs.device): + return True + _warn_bfloat16_hessian_unavailable(inputs.device) + return False + + +@functools.cache +def _has_float32_out_gemm(device: torch.device) -> bool: + """Whether GEMMs on CUDA ``device`` take bfloat16 inputs and write float32 (``out_dtype``).""" + try: + x = torch.zeros(1, 1, device=device, dtype=torch.bfloat16) + h = torch.mm(x, x, out_dtype=torch.float32) + torch.addmm(h, x, x, out_dtype=torch.float32, out=h) + torch.bmm(x[None], x[None], out_dtype=torch.float32) + except (TypeError, RuntimeError): + return False + return True + + +@functools.cache +def _warn_bfloat16_hessian_unavailable(device: torch.device) -> None: + warn_rank_0( + f"hessian_input_dtype='bfloat16' needs CUDA GEMMs that take out_dtype, which {device} lacks " + "with this torch; computing its Hessians in float32." + ) + + def compute_hessian_inverse(hessian, weight, perc_damp): """Compute damped upper-Cholesky inverse Hessian. @@ -202,6 +266,7 @@ def __init__( weight_quantizer=None, weight=None, log_mse=True, + hessian_input_dtype="float32", ): """Initialize GPTQHelper with module state and Hessian storage. @@ -209,11 +274,13 @@ def __init__( one expert's slice of an HF fused-MoE parameter; by default the helper owns ``module.weight`` and ``module.weight_quantizer``. ``log_mse=False`` skips the logged Hessian-weighted error, which costs two weight-by-Hessian matmuls. + ``hessian_input_dtype`` is the dtype the Hessian GEMM reads the inputs in, by name. """ self.module = module self.name = name self.fused = fused self.importance_weighted = importance_weighted + self.hessian_input_dtype = getattr(torch, hessian_input_dtype) self._weight_quantizer = weight_quantizer self.log_mse = log_mse weight = module.weight if weight is None else weight @@ -259,7 +326,9 @@ def hessian_forward(self, input, *args, **kwargs): def accumulate(self, inputs): """Add a batch of this weight's inputs, after input quantization, to the Hessian.""" inputs = inputs.to_local() if hasattr(inputs, "to_local") else inputs - self.hessian, self.n_samples = update_hessian(inputs, self.hessian, self.n_samples) + self.hessian, self.n_samples = update_hessian( + inputs, self.hessian, self.n_samples, self.hessian_input_dtype + ) def cleanup(self): """Unpatch the module's forward method.""" diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 286883969f9..d2733293234 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -81,6 +81,52 @@ def test_hessian_of_a_full_gpu_falls_back_to_the_cpu(need_2_gpus, monkeypatch, b assert helper.hessian.device == torch.device("cpu") +def test_bf16_hessian_input_gives_the_float32_hessian(): + """``input_dtype=torch.bfloat16`` gives the float32 Hessian, on the inputs' GPU or the CPU, + and with float32-output GEMMs makes no float32 copy of the inputs.""" + torch.manual_seed(0) + batches = [torch.randn(4096, 1024, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + inputs = torch.cat(batches).double() + expected = 2 * inputs.t() @ inputs / inputs.shape[0] + fast = calib_utils._has_float32_out_gemm(batches[0].device) + + for device in ("cuda", "cpu"): + hessian, n_samples = torch.zeros(1024, 1024, device=device), 0 + for batch in batches: + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + hessian, n_samples = update_hessian(batch, hessian, n_samples, torch.bfloat16) + growth = torch.cuda.max_memory_allocated() - base + # Tensor cores round float32 sums their own way: ~1e-5 of the diagonal (2), not BF16's 1e-2. + torch.testing.assert_close(hessian.cuda().double(), expected, rtol=0, atol=1e-4) + if device == "cuda" and fast: + assert growth < batch.numel() * 4 + + +def test_gptq_hessian_input_dtype_reaches_the_hessian(monkeypatch): + """A recipe's ``hessian_input_dtype: bfloat16`` drops the float32 copy of each batch that the + default float32 Hessian GEMM makes.""" + if not calib_utils._has_float32_out_gemm(torch.device("cuda", torch.cuda.current_device())): + pytest.skip("needs CUDA GEMMs that take out_dtype") + inputs = torch.randn(4, 1024, 2048, device="cuda", dtype=torch.bfloat16) + peaks = {} + for hessian_input_dtype in ("float32", "bfloat16"): + monkeypatch.setattr(calib_utils, "_calib_forward_peak", {}) + model = torch.nn.Linear(2048, 16, bias=False, device="cuda", dtype=torch.bfloat16) + quant_cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + ], + "algorithm": {"method": "gptq", "hessian_input_dtype": hessian_input_dtype}, + } + mtq.quantize(model, quant_cfg, forward_loop=lambda m: m(inputs)) + peaks[hessian_input_dtype] = calib_utils._calib_forward_peak[torch.cuda.current_device()] + # The float32 GEMM reads a float32 copy of the batch: 4 bytes per element. + assert peaks["float32"] - peaks["bfloat16"] >= inputs.numel() * 2 + + @pytest.mark.parametrize( ("measured_fraction", "stays"), [(None, False), (0.05, True), (0.25, False)], diff --git a/tests/gpu/torch/quantization/test_local_hessian_gpu.py b/tests/gpu/torch/quantization/test_local_hessian_gpu.py new file mode 100644 index 00000000000..25318cff2cb --- /dev/null +++ b/tests/gpu/torch/quantization/test_local_hessian_gpu.py @@ -0,0 +1,77 @@ +# 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. + +"""GPU tests for local-Hessian calibration's bfloat16 Hessian GEMM.""" + +import pytest +import torch + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.model_calib import _LocalHessianAccumulator +from modelopt.torch.quantization.utils import calib_utils + + +def _fast(): + return calib_utils._has_float32_out_gemm(torch.device("cuda", torch.cuda.current_device())) + + +def test_bf16_local_hessian_input_gives_the_float32_hessian(): + """``input_dtype=torch.bfloat16`` gives the float32 per-block Hessian, and with float32-output + GEMMs makes no float32 copy of the inputs.""" + torch.manual_seed(0) + batches = [torch.randn(4096, 1024, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + inputs = torch.cat(batches).double().T.reshape(64, 16, -1) + expected = inputs @ inputs.transpose(-1, -2) + accumulator = _LocalHessianAccumulator(8, 1024, 16, torch.bfloat16) + for batch in batches: + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + accumulator.accumulate(batch) + growth = torch.cuda.max_memory_allocated() - base + error = (accumulator.hessian_per_block.double() - expected).abs().max() + # Tensor cores round float32 sums their own way: ~1e-4 of the diagonal, not BF16's 1e-2. + assert error < 1e-3 * expected.diagonal(dim1=-2, dim2=-1).mean() + if _fast(): + assert growth < batch.numel() * 4 + + +def test_local_hessian_input_dtype_reaches_the_hessian(): + """A recipe's ``hessian_input_dtype: bfloat16`` reaches local Hessian's accumulator: its Hessian + matches the float32 one to bfloat16-GEMM rounding, without being the same sums.""" + if not _fast(): + pytest.skip("needs CUDA GEMMs that take out_dtype") + torch.manual_seed(0) + inputs = torch.randn(4, 2048, 1024, device="cuda", dtype=torch.bfloat16) + hessians = {} + for hessian_input_dtype in ("float32", "bfloat16"): + model = torch.nn.Linear(1024, 32, bias=False, device="cuda", dtype=torch.bfloat16) + quant_cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + ], + "algorithm": { + "method": "local_hessian", + "hessian_input_dtype": hessian_input_dtype, + "debug": True, + }, + } + mtq.quantize(model, quant_cfg, forward_loop=lambda m: m(inputs)) + (accumulator,) = model._local_hessian_accumulators.values() + hessians[hessian_input_dtype] = accumulator.hessian_per_block + diagonal = hessians["float32"].diagonal(dim1=-2, dim2=-1).mean() + assert not torch.equal(hessians["float32"], hessians["bfloat16"]) + assert (hessians["float32"] - hessians["bfloat16"]).abs().max() < 1e-3 * diagonal diff --git a/tests/unit/torch/quantization/test_gptq.py b/tests/unit/torch/quantization/test_gptq.py index 5744348d67e..6e5b01190ec 100644 --- a/tests/unit/torch/quantization/test_gptq.py +++ b/tests/unit/torch/quantization/test_gptq.py @@ -15,6 +15,8 @@ """CPU unit tests for GPTQ utilities.""" +import copy +import warnings from functools import partial import pytest @@ -153,6 +155,36 @@ def forward_loop(model): assert len(calls) == forwards +def test_bf16_hessian_input_on_cpu_gives_the_float32_result_and_warns_once(): + """CPU GEMMs take no ``out_dtype``, so ``hessian_input_dtype: bfloat16`` falls back to the + float32 Hessian, warning once however many batches it reads.""" + calib_utils._warn_bfloat16_hessian_unavailable.cache_clear() + torch.manual_seed(0) + model = torch.nn.Linear(32, 8, bias=False) + batches = [torch.randn(4, 32) for _ in range(3)] + weights = {} + for hessian_input_dtype in ("float32", "bfloat16"): + quantized = copy.deepcopy(model) + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + ], + "algorithm": { + "method": "gptq", + "block_size": 16, + "hessian_input_dtype": hessian_input_dtype, + }, + } + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + mtq.quantize(quantized, config, lambda m: [m(batch) for batch in batches]) + fallbacks = [w for w in caught if "hessian_input_dtype" in str(w.message)] + assert len(fallbacks) == (hessian_input_dtype == "bfloat16") + weights[hessian_input_dtype] = quantized.weight + assert torch.equal(weights["float32"], weights["bfloat16"]) + + def _gptq_hessian_dp_sync_helper(rank, size, token_counts): torch.manual_seed(0) inputs = [torch.randn(n, 8) for n in token_counts] diff --git a/tests/unit/torch/quantization/test_local_hessian.py b/tests/unit/torch/quantization/test_local_hessian.py index d8b4c4eecec..8cf89707126 100644 --- a/tests/unit/torch/quantization/test_local_hessian.py +++ b/tests/unit/torch/quantization/test_local_hessian.py @@ -39,6 +39,7 @@ _QUANT_FUNCTIONAL_BACKENDS, register_quant_backend, ) +from modelopt.torch.quantization.utils import calib_utils # Weight-only INT8 per-channel; calibration is re-run explicitly per test. INT8_WEIGHT_CFG = { @@ -132,6 +133,29 @@ def test_refines_amax_beyond_max_and_plain_mse(self): assert any(not torch.allclose(lh[n], max_amax[n]) for n in lh) # refined past max-cal assert any(not torch.allclose(lh[n], mse[n]) for n in lh) # Hessian changed the choice + def test_bf16_hessian_input_on_cpu_gives_the_float32_result_and_warns_once(self): + """CPU GEMMs take no ``out_dtype``, so ``hessian_input_dtype="bfloat16"`` falls back to + the float32 Hessians, warning once however many batches they read.""" + calib_utils._warn_bfloat16_hessian_unavailable.cache_clear() + amaxes = {} + for hessian_input_dtype in ("float32", "bfloat16"): + torch.manual_seed(0) + model = SimpleLinear() + mtq.quantize(model, INT8_WEIGHT_CFG, forward_loop=_make_forward_loop()) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + local_hessian_calibrate( + model, + _make_forward_loop(), + fp8_scale_sweep=False, + hessian_input_dtype=hessian_input_dtype, + ) + fallbacks = [w for w in caught if "hessian_input_dtype" in str(w.message)] + assert len(fallbacks) == (hessian_input_dtype == "bfloat16") + amaxes[hessian_input_dtype] = _weight_amaxes(model) + for name, amax in amaxes["float32"].items(): + assert torch.equal(amax, amaxes["bfloat16"][name]) + def test_warns_with_module_name_when_cin_not_divisible(self): class _OddModel(nn.Module): def __init__(self):