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

{compressed_tensors-0.13.1a20260127 → 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:

{compressed_tensors-0.13.1a20260127 → 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.1a20260127/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.1a20260127
+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.

{compressed_tensors-0.13.1a20260127 → 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.1a20260127 → 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.1a20260127 → 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.1a20260127 → 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 (
@@ -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,
 ):
     """
@@ -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.1a20260127 → 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.1a20260127 → 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)

{compressed_tensors-0.13.1a20260127 → compressed_tensors-0.13.1a20260130}/src/compressed_tensors/quantization/quant_scheme.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 import warnings
 from copy import deepcopy
-from typing import List, Optional
 
 import torch
 from compressed_tensors.config import CompressionFormat
@@ -47,11 +46,11 @@ class QuantizationScheme(BaseModel):
     :param format: CompressionFormat for the layer
     """
 
-    targets: List[str]
-    weights: Optional[QuantizationArgs] = None
-    input_activations: Optional[QuantizationArgs] = None
-    output_activations: Optional[QuantizationArgs] = None
-    format: Optional[str] = None
+    targets: list[str]
+    weights: QuantizationArgs | None = None
+    input_activations: QuantizationArgs | None = None
+    output_activations: QuantizationArgs | None = None
+    format: str | None = None
 
     @model_validator(mode="after")
     def validate_model_after(model: "QuantizationScheme") -> "QuantizationScheme":
@@ -121,7 +120,7 @@ Pre-Set Quantization Scheme Args
 """
 
 
-def preset_name_to_scheme(name: str, targets: List[str]) -> QuantizationScheme:
+def preset_name_to_scheme(name: str, targets: list[str]) -> QuantizationScheme:
     """
     :param name: preset quantization settings name. must exist in upper case in
         PRESET_SCHEMES
@@ -175,7 +174,6 @@ NVFP4 = dict(
         symmetric=True,
         dynamic=False,
         group_size=16,
-        observer="static_minmax",
         scale_dtype=FP8_E4M3_DATA.dtype,
         zp_dtype=FP8_E4M3_DATA.dtype,
     ),
@@ -244,7 +242,6 @@ INT8_W8A8 = dict(
         strategy=QuantizationStrategy.TOKEN,
         symmetric=True,
         dynamic=True,
-        observer=None,
     ),
 )
 
@@ -299,7 +296,6 @@ INT8_W4A8 = dict(
         strategy=QuantizationStrategy.TOKEN,
         symmetric=True,
         dynamic=True,
-        observer=None,
     ),
 )
 
@@ -356,7 +352,6 @@ FP8_DYNAMIC = dict(
         strategy=QuantizationStrategy.TOKEN,
         symmetric=True,
         dynamic=True,
-        observer=None,
     ),
 )
 
@@ -378,7 +373,6 @@ FP8_BLOCK = dict(
         strategy=QuantizationStrategy.GROUP,
         symmetric=True,
         dynamic=True,
-        observer=None,
         group_size=128,
     ),
 )

{compressed_tensors-0.13.1a20260127 → compressed_tensors-0.13.1a20260130}/src/compressed_tensors/quantization/utils/helpers.py

@@ -14,7 +14,7 @@
 
 import logging
 import math
-from typing import Generator, Optional, Tuple
+from collections.abc import Generator
 
 import torch
 from compressed_tensors.quantization.quant_args import (
@@ -66,8 +66,8 @@ def calculate_qparams(
     min_vals: Tensor,
     max_vals: Tensor,
     quantization_args: QuantizationArgs,
-    global_scale: Optional[Tensor] = None,
-) -> Tuple[FloatTensor, IntTensor]:
+    global_scale: Tensor | None = None,
+) -> tuple[FloatTensor, IntTensor]:
     """
     :param min_vals: tensor of min value(s) to calculate scale(s) and zero point(s)
         from
@@ -152,7 +152,7 @@ def compute_dynamic_scales_and_zp(
     value: Tensor,
     args: QuantizationArgs,
     module: torch.nn.Module,
-    global_scale: Optional[Tensor] = None,
+    global_scale: Tensor | None = None,
 ):
     """
     Returns the computed scales and zero points for dynamic activation
@@ -207,7 +207,9 @@ def compute_dynamic_scales_and_zp(
     return calculate_qparams(min_val, max_val, args, global_scale=global_scale)
 
 
-def calculate_range(quantization_args: QuantizationArgs, device: str) -> Tuple:
+def calculate_range(
+    quantization_args: QuantizationArgs, device: str
+) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Calculated the effective quantization range for the given Quantization Args
 
@@ -285,7 +287,7 @@ def module_type(module: Module) -> str:
     "Please use `model.named_modules()` and filter by "
     "compressed_tensors.InternalModule if neceessary"
 )
-def iter_named_leaf_modules(model: Module) -> Generator[Tuple[str, Module], None, None]:
+def iter_named_leaf_modules(model: Module) -> Generator[tuple[str, Module], None, None]:
     """
     Yields modules that do not have any submodules except observers. The observers
     themselves are not yielded
@@ -321,7 +323,7 @@ def iter_named_quantizable_modules(
     include_children: bool = True,
     include_attn: bool = False,
     include_mlp: bool = False,
-) -> Generator[Tuple[str, Module], None, None]:
+) -> Generator[tuple[str, Module], None, None]:
     """
     Yield name and submodule of
     - leaf modules, set by include_children
@@ -416,9 +418,9 @@ def is_kv_cache_quant_scheme(scheme: QuantizationScheme) -> bool:
 def generate_gparam(
     updated_min_val: torch.Tensor,
     updated_max_val: torch.Tensor,
-    scale_data: Optional[FloatArgs] = FP8_E4M3_DATA,
-    quant_data: Optional[FloatArgs] = FP4_E2M1_DATA,
-    dtype: Optional[torch.dtype] = torch.float32,
+    scale_data: FloatArgs | None = FP8_E4M3_DATA,
+    quant_data: FloatArgs | None = FP4_E2M1_DATA,
+    dtype: torch.dtype | None = torch.float32,
 ):
     """
     Generate a global scale for an entire tensor (input_tensor).
@@ -439,7 +441,7 @@ def generate_gparam(
 def strategy_cdiv(
     value: int,
     divisor: int,
-    strategy: Optional[QuantizationStrategy],
+    strategy: QuantizationStrategy | None,
     strict: bool = False,
 ) -> int:
     dividend = math.ceil(value / divisor)

{compressed_tensors-0.13.1a20260127 → compressed_tensors-0.13.1a20260130}/src/compressed_tensors/utils/helpers.py

@@ -14,19 +14,10 @@
 
 import contextlib
 import warnings
+from collections.abc import Callable, Iterable, Mapping
 from functools import wraps
 from types import MappingProxyType
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    Iterable,
-    List,
-    Mapping,
-    Optional,
-    TypeVar,
-)
+from typing import TYPE_CHECKING, Any, TypeVar
 
 import numpy
 import torch
@@ -66,7 +57,7 @@ FSDP_WRAPPER_NAME = "_fsdp_wrapped_module"
 
 def infer_compressor_from_model_config(
     pretrained_model_name_or_path: str,
-) -> Optional["ModelCompressor"]:  # noqa: F821
+) -> "ModelCompressor | None":  # noqa: F821
     """
     Given a path to a model config, extract a sparsity config if it exists and return
     the associated ModelCompressor
@@ -185,7 +176,7 @@ def getattr_chain(obj: Any, chain_str: str, *args, **kwargs) -> Any:
 
 
 def deprecated(
-    future_name: Optional[str] = None, message: Optional[str] = None
+    future_name: str | None = None, message: str | None = None
 ) -> Callable[[T], T]:
     """
     Decorator to mark functions as deprecated
@@ -224,7 +215,7 @@ class Aliasable:
     """
 
     @staticmethod
-    def get_aliases() -> Dict[str, str]:
+    def get_aliases() -> dict[str, str]:
         raise NotImplementedError()
 
     def __eq__(self, other):
@@ -246,8 +237,8 @@ class Aliasable:
 
 
 def shard_tensor(
-    tensor: torch.Tensor, shard_sizes: List[int], dim: int = 0
-) -> List[torch.Tensor]:
+    tensor: torch.Tensor, shard_sizes: list[int], dim: int = 0
+) -> list[torch.Tensor]:
     """
     Shards a tensor into a list of tensors along a given dimension.
 
@@ -277,7 +268,7 @@ def shard_tensor(
     return shards
 
 
-def combine_shards(shards, dim=0):
+def combine_shards(shards: list[torch.Tensor], dim: int = 0) -> torch.Tensor:
     """
     Combine decompressed shards along a given dimension using `narrow`.
 
@@ -325,7 +316,7 @@ def pack_bitmasks(bytemasks: torch.Tensor) -> torch.Tensor:
 
 
 def unpack_bitmasks(
-    packed_bitmasks: torch.Tensor, original_shape: List[int]
+    packed_bitmasks: torch.Tensor, original_shape: list[int]
 ) -> torch.Tensor:
     """
     Converts a bitmask tensor back to a bytemask tensor for use during decompression