From a3eec6dfbf180b258501304f1a3008bc301ddffa Mon Sep 17 00:00:00 2001 From: guptaishaan Date: Mon, 27 Jul 2026 16:09:33 -0700 Subject: [PATCH] Fix uintx safe-globals import for torchao 0.17.0 `_update_torch_safe_globals` imported `UintxAQTTensorImpl`/`UintxTensor` from the hardcoded path `torchao.dtypes.uintx.uintx_layout`, which torchao removed in 0.17.0 -- the classes moved to `torchao.prototype.dtypes.uintx.uintx_utils` (0.16.0 still shipped the old path as a deprecation shim). The import shares a `try` block with `from torchao.dtypes import NF4Tensor`, so the whole block raised, was swallowed into a warning, and the uintx tensor subclasses were never registered with `torch.serialization.add_safe_globals`. Loading a uintx-quantized checkpoint with `weights_only=True` then failed with `UnpicklingError`. Version-gate the import with the same `is_torchao_version(">", ...)` idiom already used below it, and add a test asserting the uintx subclasses end up in `torch.serialization.get_safe_globals()`. Verified against the released torchao 0.15.0, 0.16.0 and 0.17.0 wheels. torchao main is not covered: it deletes `UintxAQTTensorImpl` outright and moves `NF4Tensor` to `torchao.quantization`, and the layout is still in flux ahead of 0.18.0. --- .../quantizers/torchao/torchao_quantizer.py | 8 +++++++- tests/quantization/torchao/test_torchao.py | 13 +++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/src/diffusers/quantizers/torchao/torchao_quantizer.py b/src/diffusers/quantizers/torchao/torchao_quantizer.py index b33a18cd142c..238313ee1b32 100644 --- a/src/diffusers/quantizers/torchao/torchao_quantizer.py +++ b/src/diffusers/quantizers/torchao/torchao_quantizer.py @@ -97,7 +97,13 @@ def _update_torch_safe_globals(): ] try: from torchao.dtypes import NF4Tensor - from torchao.dtypes.uintx.uintx_layout import UintxAQTTensorImpl, UintxTensor + + # note: is_torchao_version(">=", "0.17.0") does not work correctly + # with torchao nightly, so using a ">" check which does work correctly + if is_torchao_version(">", "0.16.0"): + from torchao.prototype.dtypes.uintx.uintx_utils import UintxAQTTensorImpl, UintxTensor + else: + from torchao.dtypes.uintx.uintx_layout import UintxAQTTensorImpl, UintxTensor safe_globals.extend([UintxTensor, UintxAQTTensorImpl, NF4Tensor]) diff --git a/tests/quantization/torchao/test_torchao.py b/tests/quantization/torchao/test_torchao.py index ea2dad559b49..1ff2841e3009 100644 --- a/tests/quantization/torchao/test_torchao.py +++ b/tests/quantization/torchao/test_torchao.py @@ -86,6 +86,19 @@ def _is_xpu_or_cuda_capability_atleast_8_9() -> bool: from torchao.utils import TorchAOBaseTensor, get_model_size_in_bytes +@require_torch +@require_torchao_version_greater_or_equal("0.15.0") +class TestTorchAoSafeGlobals: + def test_uintx_tensor_subclasses_are_registered(self): + """ + `UintxTensor`/`UintxAQTTensorImpl` moved between torchao releases. If the import in + `_update_torch_safe_globals` goes stale, they are silently skipped and loading a uintx + checkpoint with `weights_only=True` fails. + """ + registered = {getattr(safe_global, "__name__", None) for safe_global in torch.serialization.get_safe_globals()} + assert {"UintxTensor", "UintxAQTTensorImpl"}.issubset(registered) + + @require_torch @require_torch_accelerator @require_torchao_version_greater_or_equal("0.15.0")