compressed-tensors 0.13.1a20260123__tar.gz → 0.13.1a20260130__tar.gz

{compressed_tensors-0.13.1a20260123 → compressed_tensors-0.13.1a20260130}/.github/actions/test/action.yml

@@ -23,7 +23,7 @@ runs:
       with:
           venv: ${{ inputs.venv }}
           name: compressed
-          extra: "[dev,accelerate]"
+          extra: "[dev]"
 
     - name: clean up
       run: |

{compressed_tensors-0.13.1a20260123 → compressed_tensors-0.13.1a20260130}/.github/workflows/test-check.yaml

@@ -12,7 +12,7 @@ on:
 
 jobs:
   python-tests:
-    runs-on: ibm-wdc-k8s-vllm-h100-solo
+    runs-on: gcp-k8s-vllm-l4-duo
     env:
         HF_TOKEN: ${{ secrets.HF_RED_HAT_READ_ONLY }}
     steps:
@@ -30,7 +30,7 @@ jobs:
         - name: Set Env
           run: pip3 install --upgrade pip setuptools
         - name: "⚙️ Install dependencies"
-          run: pip3 install .[dev,accelerate]
+          run: pip3 install .[dev]
         - name: clean up
           run: |
             echo "cleaning up disk space as GHA runner has limited disk size."

{compressed_tensors-0.13.1a20260123 → compressed_tensors-0.13.1a20260130}/Makefile

@@ -23,7 +23,7 @@ style:
 # run tests for the repo
 test:
 	@echo "Running python tests";
-	pytest tests;
+	pytest -ra tests;
 
 # creates wheel file
 build:

{compressed_tensors-0.13.1a20260123/src/compressed_tensors.egg-info → compressed_tensors-0.13.1a20260130}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.13.1a20260123
+Version: 0.13.1a20260130
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/vllm-project/compressed-tensors
 Author: Neuralmagic, Inc.
@@ -9,7 +9,7 @@ License: Apache 2.0
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch<=2.9.1,>=1.7.0
-Requires-Dist: transformers
+Requires-Dist: transformers<5.0.0
 Requires-Dist: pydantic>=2.0
 Requires-Dist: loguru
 Provides-Extra: dev
@@ -19,6 +19,8 @@ Requires-Dist: wheel>=0.36.2; extra == "dev"
 Requires-Dist: flake8>=3.8.3; extra == "dev"
 Requires-Dist: pytest>=6.0.0; extra == "dev"
 Requires-Dist: nbconvert>=7.16.3; extra == "dev"
+Requires-Dist: transformers<5.0; extra == "dev"
+Requires-Dist: accelerate; extra == "dev"
 Provides-Extra: accelerate
 Requires-Dist: accelerate; extra == "accelerate"
 Dynamic: author

{compressed_tensors-0.13.1a20260123 → compressed_tensors-0.13.1a20260130}/setup.py

@@ -88,11 +88,11 @@ def _setup_packages() -> List:
     )
 
 def _setup_install_requires() -> List:
-    return ["torch>=1.7.0,<=2.9.1", "transformers", "pydantic>=2.0", "loguru"]
+    return ["torch>=1.7.0,<=2.9.1", "transformers<5.0.0", "pydantic>=2.0", "loguru"]
 
 def _setup_extras() -> Dict:
     return {
-        "dev": ["black==22.12.0", "isort==5.8.0", "wheel>=0.36.2", "flake8>=3.8.3", "pytest>=6.0.0", "nbconvert>=7.16.3"],
+        "dev": ["black==22.12.0", "isort==5.8.0", "wheel>=0.36.2", "flake8>=3.8.3", "pytest>=6.0.0", "nbconvert>=7.16.3", "transformers<5.0", "accelerate"],
         "accelerate": ["accelerate"]
     }
 

{compressed_tensors-0.13.1a20260123 → compressed_tensors-0.13.1a20260130}/src/compressed_tensors/linear/compressed_linear.py

@@ -87,12 +87,6 @@ class CompressedLinear(Linear):
         # mark module as compressed
         module.quantization_status = QuantizationStatus.COMPRESSED
 
-        # handles case where forward is wrapped in new_forward by accelerate hooks
-        if hasattr(module, "_old_forward"):
-            module._old_forward = CompressedLinear.forward.__get__(
-                module, CompressedLinear
-            )
-
         return module
 
     def forward(self, input: Tensor) -> Tensor:

{compressed_tensors-0.13.1a20260123 → compressed_tensors-0.13.1a20260130}/src/compressed_tensors/offload/__init__.py

@@ -135,9 +135,7 @@ def register_offload_module(base: torch.nn.Module, name: str, module: torch.nn.M
     """
     cache = base._parameters
     if isinstance(cache, OffloadCache):
-        offload_module(
-            module, cache.onload_device, cache.offload_device, no_split=False
-        )
+        offload_module(module, cache.onload_device, cache.offload_device)
 
     base.register_module(name, module)
 
@@ -178,9 +176,12 @@ def align_module_device(
     if isinstance(module._parameters, OffloadCache):
         assert isinstance(module._buffers, OffloadCache)
         with module._parameters.disable_offloading():
-            with patch_attr(
-                module._parameters, "onload_device", execution_device
-            ), patch_attr(module._buffers, "onload_device", execution_device):
+            if execution_device is not None:
+                with patch_attr(
+                    module._parameters, "onload_device", execution_device
+                ), patch_attr(module._buffers, "onload_device", execution_device):
+                    yield
+            else:
                 yield
 
     else:

{compressed_tensors-0.13.1a20260123 → compressed_tensors-0.13.1a20260130}/src/compressed_tensors/offload/cache/base.py

@@ -67,6 +67,7 @@ class OffloadCache(MutableMapping, ABC):
         """
         from compressed_tensors.offload.cache.cpu import CPUCache
         from compressed_tensors.offload.cache.device import DeviceCache
+        from compressed_tensors.offload.cache.dist_cpu import DistributedCPUCache
 
         device_type = torch.device(device).type if device != "disk" else "disk"
         distributed = dist.is_available() and dist.is_initialized()
@@ -74,6 +75,8 @@ class OffloadCache(MutableMapping, ABC):
         match (device_type, distributed):
             case ("cpu", False):
                 return CPUCache
+            case ("cpu", True):
+                return DistributedCPUCache
             case ("cuda", False):
                 return DeviceCache
             case _:

{compressed_tensors-0.13.1a20260123 → compressed_tensors-0.13.1a20260130}/src/compressed_tensors/offload/cache/device.py

@@ -35,8 +35,8 @@ class DeviceCache(OffloadCache):
         :param key: cpu tensor to onload
         :return: device tensor
         """
-        assert offloaded.device == self.onload_device
-        return offloaded
+        # move because onload_device might be modified after init
+        return send_tensors(offloaded, device=self.onload_device, copy=False)
 
     def offload(self, tensor: torch.Tensor | None) -> torch.Tensor:
         """

compressed_tensors-0.13.1a20260130/src/compressed_tensors/offload/cache/dist_cpu.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# 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 torch
+import torch.distributed as dist
+from compressed_tensors.offload.cache.cpu import CPUCache
+
+
+class DistributedCPUCache(CPUCache):
+    """
+    Handles offloading and onloading tensors from/to cpu memory shared across processes
+    """
+
+    offload_device = torch.device("cpu")
+
+    def offload(self, tensor: torch.Tensor | None) -> torch.Tensor:
+        if tensor is None:
+            return None
+
+        # slight runtime cost for views
+        tensor = tensor.contiguous()
+
+        if dist.get_rank() == 0:
+            # create shared memory cpu tensor
+            tensor = super().offload(tensor).share_memory_()
+            (handle, filename, nbytes) = tensor.untyped_storage()._share_filename_cpu_()
+            broadcast_obj = [handle, filename, nbytes]
+        else:
+            broadcast_obj = [None, None, None]
+
+        # receive shared memory file handle
+        dist.broadcast_object_list(broadcast_obj, src=0)
+
+        if dist.get_rank() != 0:
+            # reconstruct tensor from shared memory file handle
+            tensor = torch.empty_like(tensor, device=self.offload_device)
+            tensor.set_(torch.UntypedStorage._new_shared_filename_cpu(*broadcast_obj))
+
+        # ensure that rank 0 does not garbage collect before other ranks reconstruct
+        dist.barrier()
+
+        return tensor

{compressed_tensors-0.13.1a20260123 → compressed_tensors-0.13.1a20260130}/src/compressed_tensors/offload/dispatch.py

@@ -39,7 +39,7 @@ ModelType = TypeVar("ModelType", bound=torch.nn.Module)
 def offload_model(
     model: ModelType,
     onload_device: torch.device | str,
-    offload_device: Optional[torch.device | str | Literal["disk"]] = None,
+    offload_device: torch.device | str | Literal["disk"] = torch.device("cpu"),
 ) -> ModelType:
     """
     Offload a model to the `offload_device`. During forward passes, model weights will

{compressed_tensors-0.13.1a20260123 → compressed_tensors-0.13.1a20260130}/src/compressed_tensors/quantization/lifecycle/apply.py

@@ -14,9 +14,6 @@
 
 from collections import OrderedDict
 from copy import deepcopy
-from typing import Dict, List, Optional
-from typing import OrderedDict as OrderedDictType
-from typing import Union
 
 import torch
 from compressed_tensors.config import CompressionFormat
@@ -60,8 +57,8 @@ from compressed_tensors.utils.safetensors_load import (
 
 def load_pretrained_quantization_parameters(
     model: Module,
-    model_name_or_path: Optional[str] = None,
-    load_weight_qparams: Optional[bool] = False,
+    model_name_or_path: str | None = None,
+    load_weight_qparams: bool = False,
 ):
     """
     Loads the quantization parameters (scale and zero point) from model_name_or_path to
@@ -110,7 +107,7 @@ def load_pretrained_quantization_parameters(
 
 
 def apply_quantization_config(
-    model: Module, config: Union[QuantizationConfig, None], run_compressed: bool = False
+    model: Module, config: QuantizationConfig | None, run_compressed: bool = False
 ):
     """
     Initializes the model for quantization in-place based on the given config.
@@ -207,7 +204,7 @@ def _apply_kv_cache_scheme(
 
 
 def _load_quant_args_from_mapping(
-    base_name: str, module_name: str, module: Module, mapping: Dict
+    base_name: str, module_name: str, module: Module, mapping: dict
 ):
     # TODO: skip update and just register here, don't do it in initialize
     """
@@ -251,8 +248,8 @@ def _load_quant_args_from_mapping(
 
 
 def _scheme_from_targets(
-    target_to_scheme: OrderedDictType[str, QuantizationScheme],
-    targets: List[str],
+    target_to_scheme: OrderedDict[str, QuantizationScheme],
+    targets: list[str],
     name: str,
 ) -> QuantizationScheme:
     # return the first scheme (the prioritized one,

{compressed_tensors-0.13.1a20260123 → compressed_tensors-0.13.1a20260130}/src/compressed_tensors/quantization/lifecycle/forward.py

@@ -14,7 +14,6 @@
 
 from functools import wraps
 from math import ceil
-from typing import Optional
 
 import torch
 from compressed_tensors.quantization.quant_args import (
@@ -47,9 +46,9 @@ def quantize(
     scale: torch.Tensor,
     zero_point: torch.Tensor,
     args: QuantizationArgs,
-    dtype: Optional[torch.dtype] = None,
-    g_idx: Optional[torch.Tensor] = None,
-    global_scale: Optional[torch.Tensor] = None,
+    dtype: torch.dtype | None = None,
+    g_idx: torch.Tensor | None = None,
+    global_scale: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """
     Quantize the input tensor x using the QuantizationStrategy specified in args.
@@ -85,11 +84,11 @@ def quantize(
 def dequantize(
     x_q: torch.Tensor,
     scale: torch.Tensor,
-    zero_point: Optional[torch.Tensor] = None,
-    args: Optional[QuantizationArgs] = None,
-    dtype: Optional[torch.dtype] = None,
-    g_idx: Optional[torch.Tensor] = None,
-    global_scale: Optional[torch.Tensor] = None,
+    zero_point: torch.Tensor | None = None,
+    args: QuantizationArgs | None = None,
+    dtype: torch.dtype | None = None,
+    g_idx: torch.Tensor | None = None,
+    global_scale: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """
     Dequantize a quantized input tensor x_q based on the strategy specified in args. If
@@ -159,8 +158,8 @@ def fake_quantize(
     scale: torch.Tensor,
     zero_point: torch.Tensor,
     args: QuantizationArgs,
-    g_idx: Optional[torch.Tensor] = None,
-    global_scale: Optional[torch.Tensor] = None,
+    g_idx: torch.Tensor | None = None,
+    global_scale: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """
     Fake quantize the input tensor x by quantizing then dequantizing with
@@ -195,11 +194,11 @@ def _process_quantization(
     scale: torch.Tensor,
     zero_point: torch.Tensor,
     args: QuantizationArgs,
-    g_idx: Optional[torch.Tensor] = None,
-    dtype: Optional[torch.dtype] = None,
+    g_idx: torch.Tensor | None = None,
+    dtype: torch.dtype | None = None,
     do_quantize: bool = True,
     do_dequantize: bool = True,
-    global_scale: Optional[torch.Tensor] = None,
+    global_scale: torch.Tensor | None = None,
 ) -> torch.Tensor:
     q_min, q_max = calculate_range(args, x.device)
     group_size = args.group_size
@@ -457,8 +456,8 @@ def _quantize(
     q_min: torch.Tensor,
     q_max: torch.Tensor,
     args: QuantizationArgs,
-    dtype: Optional[torch.dtype] = None,
-    global_scale: Optional[torch.Tensor] = None,
+    dtype: torch.dtype | None = None,
+    global_scale: torch.Tensor | None = None,
 ) -> torch.Tensor:
 
     # if a global scale is optionally provided, use it
@@ -486,9 +485,9 @@ def _quantize(
 def _dequantize(
     x_q: torch.Tensor,
     scale: torch.Tensor,
-    zero_point: torch.Tensor = None,
-    dtype: Optional[torch.dtype] = None,
-    global_scale: Optional[torch.Tensor] = None,
+    zero_point: torch.Tensor | None = None,
+    dtype: torch.dtype | None = None,
+    global_scale: torch.Tensor | None = None,
 ) -> torch.Tensor:
 
     # if a global scale is optionally provided, use it

{compressed_tensors-0.13.1a20260123 → compressed_tensors-0.13.1a20260130}/src/compressed_tensors/quantization/lifecycle/initialize.py

@@ -14,7 +14,6 @@
 
 
 import logging
-from typing import Optional, Tuple, Union
 
 import torch
 from compressed_tensors.modeling import (
@@ -23,6 +22,7 @@ from compressed_tensors.modeling import (
     QuantizedAttentionImpl,
     QuantizedKVCache,
 )
+from compressed_tensors.offload import unwrap_offload_forward
 from compressed_tensors.quantization import (
     ActivationOrdering,
     DynamicType,
@@ -37,7 +37,6 @@ from compressed_tensors.quantization.lifecycle.forward import (
 )
 from compressed_tensors.quantization.utils import strategy_cdiv
 from compressed_tensors.utils import (
-    disable_hf_hook,
     get_execution_device,
     get_head_dim,
     get_num_attn_heads,
@@ -60,7 +59,7 @@ _LOGGER = logging.getLogger(__name__)
 
 def initialize_module_for_quantization(
     module: Module,
-    scheme: Optional[QuantizationScheme] = None,
+    scheme: QuantizationScheme | None = None,
     force_zero_point: bool = True,
 ):
     """
@@ -134,7 +133,7 @@ def initialize_module_for_quantization(
                 force_zero_point=force_zero_point,
             )
 
-        with disable_hf_hook(module):
+        with unwrap_offload_forward(module):
             # wrap forward call of module to perform
             # quantized actions based on calltime status
             wrap_module_forward_quantized(module, scheme)
@@ -148,6 +147,7 @@ def is_attention_module(module: Module):
         hasattr(module, "k_proj")
         or hasattr(module, "v_proj")
         or hasattr(module, "qkv_proj")
+        or hasattr(module, "kv_b_proj")
     )
 
 
@@ -155,7 +155,7 @@ def initialize_qparams(
     module: Module,
     base_name: str,
     quantization_args: QuantizationArgs,
-    observed_shape: Tuple[Union[int, None]],
+    observed_shape: tuple[int | None, ...],
     observed_dtype: torch.dtype,
     force_zero_point: bool = True,
 ):
@@ -279,8 +279,8 @@ def initialize_attn_qparams(
 ):
     """Initlaize k_scale, v_scale for self_attn"""
 
-    impl: Optional[QuantizedAttentionImpl] = getattr(module, IMPL_ATTR, None)
-    kv_cache: Optional[QuantizedKVCache] = getattr(module, KV_CACHE_ATTR, None)
+    impl: QuantizedAttentionImpl | None = getattr(module, IMPL_ATTR, None)
+    kv_cache: QuantizedKVCache | None = getattr(module, KV_CACHE_ATTR, None)
 
     if impl is None and kv_cache is None:
         raise ValueError(

{compressed_tensors-0.13.1a20260123 → compressed_tensors-0.13.1a20260130}/src/compressed_tensors/quantization/quant_args.py

@@ -14,7 +14,7 @@
 
 import warnings
 from enum import Enum
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 import torch
 from compressed_tensors.utils import Aliasable
@@ -48,10 +48,10 @@ __all__ = [
 class FloatArgs:
     exponent: int
     mantissa: int
-    bits: Optional[int] = None
-    max: Optional[float] = None
-    min: Optional[float] = None
-    dtype: Optional[torch.dtype] = None
+    bits: int | None = None
+    max: float | None = None
+    min: float | None = None
+    dtype: torch.dtype | None = None
 
 
 class FP4_E2M1_DATA(FloatArgs):
@@ -147,7 +147,7 @@ class ActivationOrdering(Aliasable, str, Enum):
     STATIC = "static"
 
     @staticmethod
-    def get_aliases() -> Dict[str, str]:
+    def get_aliases() -> dict[str, str]:
         return {
             "dynamic": "group",
             "static": "weight",
@@ -178,21 +178,21 @@ class QuantizationArgs(BaseModel, use_enum_values=True):
     num_bits: int = 8
     type: QuantizationType = QuantizationType.INT
     symmetric: bool = True
-    group_size: Optional[int] = None
-    strategy: Optional[QuantizationStrategy] = None
-    block_structure: Optional[List[int]] = None
-    dynamic: Union[DynamicType, bool] = False
-    actorder: Union[ActivationOrdering, bool, None] = None
-    scale_dtype: Optional[TorchDtype] = None
-    zp_dtype: Optional[TorchDtype] = None
-    observer: Optional[str] = Field(
+    group_size: int | None = None
+    strategy: QuantizationStrategy | None = None
+    block_structure: list[int] | None = None
+    dynamic: DynamicType | bool = False
+    actorder: ActivationOrdering | bool | None = None
+    scale_dtype: TorchDtype | None = None
+    zp_dtype: TorchDtype | None = None
+    observer: str | None = Field(
         default=None,
         description=(
             "Determines the method of computing quantization parameters (scales and "
             "zero-points). Defaults to min-max when not using dynamic quantization"
         ),
     )
-    observer_kwargs: Dict[str, Any] = Field(
+    observer_kwargs: dict[str, Any] = Field(
         default_factory=dict,
         description=(
             "optional dict of kwargs to be passed directly to torch quantization "
@@ -214,7 +214,7 @@ class QuantizationArgs(BaseModel, use_enum_values=True):
         return value
 
     @field_validator("group_size", mode="before")
-    def validate_group(cls, value) -> Union[int, None]:
+    def validate_group(cls, value) -> int | None:
         if value is None:
             return value
 
@@ -227,7 +227,7 @@ class QuantizationArgs(BaseModel, use_enum_values=True):
         return value
 
     @field_validator("block_structure", mode="before")
-    def validate_block_structure(cls, value) -> Optional[List[int]]:
+    def validate_block_structure(cls, value) -> list[int] | None:
         if value is None:
             return value
         # For backward compatibility, allow string format "2x4", "8x16", etc.
@@ -251,14 +251,14 @@ class QuantizationArgs(BaseModel, use_enum_values=True):
         )
 
     @field_validator("strategy", mode="before")
-    def validate_strategy(cls, value) -> Union[QuantizationStrategy, None]:
+    def validate_strategy(cls, value) -> QuantizationStrategy | None:
         if isinstance(value, str):
             return QuantizationStrategy(value.lower())
 
         return value
 
     @field_validator("actorder", mode="before")
-    def validate_actorder(cls, value) -> Optional[ActivationOrdering]:
+    def validate_actorder(cls, value) -> ActivationOrdering | None:
         if isinstance(value, bool):
             return ActivationOrdering.GROUP if value else None
 
@@ -268,7 +268,7 @@ class QuantizationArgs(BaseModel, use_enum_values=True):
         return value
 
     @field_validator("dynamic", mode="before")
-    def validate_dynamic(cls, value) -> Union[DynamicType, bool]:
+    def validate_dynamic(cls, value) -> DynamicType | bool:
         if isinstance(value, str):
             return DynamicType(value.lower())
         return value
@@ -329,10 +329,13 @@ class QuantizationArgs(BaseModel, use_enum_values=True):
             raise ValueError(f"Block structure requires block strategy\n{model}")
 
         # validate activation ordering and strategy
-        if actorder is not None and strategy != QuantizationStrategy.GROUP:
+        if actorder is not None and strategy not in (
+            QuantizationStrategy.GROUP,
+            QuantizationStrategy.TENSOR_GROUP,
+        ):
             raise ValueError(
-                "Must use group quantization strategy in order to apply "
-                "activation ordering"
+                "Must use group or tensor_group quantization strategy in "
+                "order to apply activation ordering"
             )
 
         # infer observer w.r.t. dynamic
@@ -369,7 +372,7 @@ class QuantizationArgs(BaseModel, use_enum_values=True):
 
         elif observer is None:
             # default to minmax for non-dynamic cases
-            observer = "minmax"
+            observer = "memoryless_minmax"
 
         if zp_dtype is None:
             if model.num_bits == 4 and model.type == QuantizationType.FLOAT:
@@ -409,7 +412,7 @@ class QuantizationArgs(BaseModel, use_enum_values=True):
 def round_to_quantized_type_dtype(
     tensor: torch.Tensor,
     dtype: torch.dtype,
-    cast_to_original_dtype: Optional[bool] = True,
+    cast_to_original_dtype: bool = True,
 ) -> torch.Tensor:
     """
     Rounds an input tensor to the nearest quantized representation given a dtype.
@@ -439,7 +442,7 @@ def round_to_quantized_type_args(
     args: QuantizationArgs,
     min: torch.Tensor,
     max: torch.Tensor,
-    cast_to_original_dtype: Optional[bool] = True,
+    cast_to_original_dtype: bool = True,
 ) -> torch.Tensor:
     """
     Rounds an input tensor to the nearest quantized representation given

{compressed_tensors-0.13.1a20260123 → compressed_tensors-0.13.1a20260130}/src/compressed_tensors/quantization/quant_config.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 from collections import defaultdict
 from enum import Enum
-from typing import Annotated, Any, Dict, List, Optional, Set, Union
+from typing import Annotated, Any
 
 from compressed_tensors.config import CompressionFormat
 from compressed_tensors.quantization.quant_args import DynamicType, QuantizationArgs
@@ -55,7 +55,7 @@ class QuantizationStatus(str, Enum):
     COMPRESSED = "compressed"
 
     @classmethod
-    def lifecycle_order(cls) -> List["QuantizationStatus"]:
+    def lifecycle_order(cls) -> list["QuantizationStatus"]:
         """
         :return: list of correct quantization lifecycle order
         """
@@ -131,13 +131,13 @@ class QuantizationConfig(BaseModel):
         are not quantized even if they match up with a target in config_groups
     """
 
-    config_groups: Dict[str, Union[QuantizationScheme, List[str]]]
+    config_groups: dict[str, QuantizationScheme | list[str]]
     quant_method: str = DEFAULT_QUANTIZATION_METHOD
-    kv_cache_scheme: Optional[QuantizationArgs] = None
+    kv_cache_scheme: QuantizationArgs | None = None
     format: str = DEFAULT_QUANTIZATION_FORMAT
     quantization_status: QuantizationStatus = QuantizationStatus.INITIALIZED
-    global_compression_ratio: Optional[float] = None
-    ignore: Optional[List[str]] = Field(default_factory=list)
+    global_compression_ratio: float | None = None
+    ignore: list[str] | None = Field(default_factory=list)
     # `run_compressed` is a dummy, unused arg for backwards compatibility
     # see: https://github.com/huggingface/transformers/pull/39324
     run_compressed: Annotated[Any, Field(exclude=True)] = None
@@ -161,8 +161,8 @@ class QuantizationConfig(BaseModel):
 
     @staticmethod
     def from_pretrained(
-        model: Module, format: Optional[Union[str, list]] = None
-    ) -> Optional["QuantizationConfig"]:
+        model: Module, format: str | list | None = None
+    ) -> "QuantizationConfig | None":
         """
         Converts a model into its associated QuantizationConfig based on the
         QuantizationScheme attached to each quantized module
@@ -177,21 +177,21 @@ class QuantizationConfig(BaseModel):
 
         # set of all quantization schemes
         # TODO: make quant config/scheme/args frozen/hashable and use a set
-        quantization_schemes: List[QuantizationScheme] = list()
+        quantization_schemes: list[QuantizationScheme] = list()
 
         # use any status from modules (in practice, use the last module)
         model_status = None
 
         # set of all quantized types
         # this is later used to create the ignore list
-        quantization_type_names: Set[str] = set()
+        quantization_type_names: set[str] = set()
 
         # maps types to names which are not quantized
         # this is later used to create the ignore list
-        ignore: Dict[str, List[str]] = defaultdict(list)
+        ignore: dict[str, list[str]] = defaultdict(list)
 
         # this keeps track of any kvcache schemes
-        kv_cache_scheme: Optional[QuantizationArgs] = None
+        kv_cache_scheme: QuantizationArgs | None = None
 
         for name, submodule in model.named_modules():
             layer_type: str = module_type(submodule)