# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai # SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai # SPDX-License-Identifier: Apache-2.0 # Contact: qubitium@modelcloud.ai, x.com/qubitium import copy import math import os import sys from concurrent.futures import ThreadPoolExecutor from functools import lru_cache from typing import Any, Dict, List, Optional, Tuple import numpy as np import torch as t # conflict with torch.py import torch.nn as nn import transformers from torch.nn.modules.conv import _ConvNd from ...adapter.adapter import LORA_MERGED_WEIGHT_PATHS, Adapter from ...models._const import DEVICE, PLATFORM from ...quantization import FORMAT, METHOD from ...utils.backend import BACKEND from ...utils.env import env_flag from ...utils.logger import setup_logger from ...utils.planar_packing import ( PLANAR_BITS, PLANAR_FORMAT_BITS, planar_pack_cols, planar_pack_rows, planar_unpack_cols, planar_unpack_rows, ) from ...utils.safe import THREADPOOLCTL log = setup_logger() def input_rows(x: t.Tensor) -> int: """Return the number of logical Linear rows in an arbitrary-rank input. Linear treats every dimension before the final feature dimension as a batch dimension, so a 4D input has ``B * S * T`` matrix rows. """ if x.ndim == 0: raise ValueError("Linear input must have at least one dimension") return math.prod(x.shape[:-1]) def empty_linear_output(x: t.Tensor, out_features: int) -> t.Tensor: """Allocate an empty Linear output while preserving rank, dtype, device.""" shape = x.shape[:-1] + (out_features,) if x.requires_grad: # Keep the empty path connected to the input so callers can backprop # through an empty batch just like through nn.Linear (with zero grad). return x.sum(dim=-1, keepdim=True).expand(shape) * 0 return t.empty( shape, dtype=x.dtype, device=x.device, ) # Packed quantized weights are unpacked through shift operations in several # kernels. Keep those shifts behind helpers so Ascend 910B/CANN 9.1 beta can use # arithmetic equivalents for operators that torch-npu does not expose as native # device kernels, while CPU/CUDA/XPU keep the standard bitwise ops. def _torch_shift_factor(shifts: int | t.Tensor, device: t.device) -> int | t.Tensor: if t.is_tensor(shifts): # Tensor shifts must be materialized on the target device; otherwise the # NPU arithmetic fallback would introduce cross-device operands. shifts_i64 = shifts.to(device=device, dtype=t.int64) return t.pow(t.full_like(shifts_i64, 2), shifts_i64) return 1 << int(shifts) def _torch_right_shift(values: t.Tensor, shifts: int | t.Tensor) -> t.Tensor: if values.device.type != "npu": return t.bitwise_right_shift(values, shifts) # CANN 9.1 beta on Ascend 910B may not provide bitwise_right_shift kernels # for these tensor paths. floor_divide by powers of two preserves arithmetic # right-shift behavior for signed packed int tensors and stays on-device. shifted = t.floor_divide(values.to(t.int64), _torch_shift_factor(shifts, values.device)) return shifted.to(values.dtype) def _torch_left_shift(values: t.Tensor, shifts: int | t.Tensor) -> t.Tensor: if values.device.type != "npu": return t.bitwise_left_shift(values, shifts) # Mirror left shift as multiplication by powers of two on NPU to avoid # missing-kernel or CPU-fallback paths in torch-npu. shifted = values.to(t.int64) * _torch_shift_factor(shifts, values.device) return shifted.to(values.dtype) class BaseQuantLinear(nn.Module): SUPPORTS_BACKENDS: List[BACKEND] = None # False for role-specific modules that are instantiated explicitly and # must never participate in general backend discovery. SUPPORTS_BACKEND_SELECTION: bool = True SUPPORTS_METHODS: List[METHOD] = None SUPPORTS_FORMATS: Dict[FORMAT, int] = None SUPPORTS_BITS: List[int] = None SUPPORTS_SHARDS: bool = None SUPPORTS_SHARDED_LOAD: bool = True SUPPORTS_TRAINING: bool = None # IPEX kernel will use Torch for training only and switches back to IPEX for eval/inference SUPPORTS_TRAINING_USE_TORCH_KERNEL: bool = False SUPPORTS_AUTO_PADDING: bool = None SUPPORTS_IN_FEATURES_DIVISIBLE_BY: List[int] = None SUPPORTS_OUT_FEATURES_DIVISIBLE_BY: List[int] = None SUPPORTS_PACK_DTYPES: List[t.dtype] = None SUPPORTS_ADAPTERS: List[Adapter] = None SUPPORTS_DEVICES: List[DEVICE] = None SUPPORTS_PLATFORM: List[PLATFORM] = None SUPPORTS_DTYPES: List[t.dtype] = None REQUIRES_FORMAT_V2: bool = False AUTOTUNE: bool = False def __init__(self, bits: int, in_features: int, out_features: int, bias: bool, backend: BACKEND, adapter: Adapter, name: str = None, register_buffers: bool = False, register_buffers_in_features: int = None, register_buffers_out_features: int = None, dtype: Optional[t.dtype] = None, validate_kwargs: Optional[Dict[str, Any]] = None, **kwargs): super().__init__() if name is None: name = f"{self.__class__.__module__}.{self.__class__.__qualname__}" self.name = name # full path module name in model weights self.in_features = in_features self.out_features = out_features self.bits = bits self.backend = backend # we need to clone the adapter since passed in adapter may be shared # adapter tensors are lodaed inside adapter so they must be unique per module self.adapter = copy.deepcopy(adapter) self.optimized = False self.autotune_enabled = self.AUTOTUNE self._autotune_complete = False self._autotune_result: Any = None # Rotation/online-Hadamard state for SpinQuant/QuaRot inference. self.online_full_had = False self.online_partial_had = False self.had_dim = -1 self.K = 1 # Not a registered buffer while None: accelerate's offload hooks cannot # handle None-valued buffers. `set_had_K` promotes it to a real buffer # when rotation assigns a tensor. self.had_K = None validate_args = { "bits": bits, "in_features": in_features, "out_features": out_features, "dtype": dtype, "adapter": adapter, } if validate_kwargs: validate_args.update(validate_kwargs) _, err = self.validate( **validate_args, ) if err: raise err # The root base only owns fields shared across all quantization methods. if register_buffers and bias: bias_shape = self.out_features if register_buffers_out_features is None else register_buffers_out_features self.register_buffer("bias", t.zeros(bias_shape, dtype=t.float16)) # load adapter if any if adapter is not None: if adapter.path in LORA_MERGED_WEIGHT_PATHS: print(f"Adapter (merged weights) lazy init: {self.adapter.name()}: {self.adapter}, module: {self.name}") # pre allocate buffers so accelerate can auto-bind merged weights in same tensor file as model self.register_buffer( "lora_A", t.zeros((in_features, adapter.rank), dtype=t.float16), ) self.register_buffer( "lora_B", t.zeros((adapter.rank, out_features), dtype=t.float16), ) else: pass # print(f"Adapter lazy init: {self.adapter.name()}: {self.adapter}, module: {self.name}") # TDOO: allow merged lora weights exist in gptq model safetensor file for direct loading # EoRA need to preallocate buffers for Lora_A and B weights so HF can load # self.register_buffer( # "lora_A", # torch.zeros((in_features, 128), dtype=torch.float16), # <-- EoRA lora_A shape needs to be calculated using pass in_features/out_features or other eora_test math # ) # # # EoRA need to preallocate buffers for Lora_A and B weights so HF can load # self.register_buffer( # "lora_B", # torch.zeros((128, out_features), dtype=torch.float16), # <-- EoRA lora_A shape needs to be calculated using pass in_features/out_features or other eora_test math # ) def list_buffers(self) -> List: tensors = [] seen = set() for state in (self._parameters, self._buffers): for tensor in state.values(): if tensor is None or not isinstance(tensor, t.Tensor): continue tensor_id = id(tensor) if tensor_id in seen: continue seen.add(tensor_id) tensors.append(tensor) return tensors def runtime_device(self) -> Optional[t.device]: # Prefer the real quantized storage tensor over adapter scratch buffers. for name in ("qweight", "weight", "B", "qzeros", "scales", "g_idx", "bias"): tensor = getattr(self, name, None) if isinstance(tensor, t.Tensor): return tensor.device buffers = self.list_buffers() return buffers[0].device if buffers else None def input_rows(self, x: t.Tensor) -> int: return input_rows(x) def empty_linear_output(self, x: t.Tensor) -> t.Tensor: """Return an empty output for any-rank Linear input. Applying an adapter is intentional: it preserves the same composition semantics as the non-empty path while still avoiding a native M=0 kernel launch. Bias addition has no observable effect for zero elements. """ out = empty_linear_output(x, self.out_features) if self.adapter is not None: out = self.adapter.apply(x=x, out=out) return out def smooth_block_size(self) -> int: return -1 # override me, to perform post-weight load to device init def post_init(self): self.clear_autotune() if self.adapter is not None: device = self.runtime_device() if device is None: raise RuntimeError(f"{self.__class__.__name__} cannot initialize adapters without any tensors.") self.adapter.post_init( weight_key=self.name, device=device, lora_A=getattr(self, "lora_A", None), lora_B=getattr(self, "lora_B", None)) def clear_autotune(self): self._autotune_complete = False self._autotune_result = None def get_autotune_result(self): return self._autotune_result def _autotune(self, *args, **kwargs): raise NotImplementedError(f"{self.__class__.__name__} does not implement `_autotune()`.") def maybe_autotune(self, *args, **kwargs): if not self.autotune_enabled or self.training: return self._autotune_result if self._autotune_complete: return self._autotune_result self._autotune_result = self._autotune(*args, **kwargs) self._autotune_complete = True return self._autotune_result @classmethod @lru_cache(maxsize=1024) def cached_validate_once(cls) -> Tuple[bool, Optional[Exception]]: ok, exp = cls.validate_once() if not ok and exp is not None: exp.with_traceback(None) return False, exp return True, None @classmethod def validate_once(cls) -> Tuple[bool, Optional[Exception]]: return True, None @classmethod # custom quant linear class can override this and add custom checks def validate( cls, bits: int, group_size: int = -1, desc_act: bool = False, sym: bool = True, in_features:int=None, out_features:int=None, pack_dtype:t.dtype=None, dtype: Optional[t.dtype]=None, dynamic:Optional[dict]=None, device:Optional[DEVICE | t.device]=None, trainable:Optional[bool]=None, adapter:Optional[Adapter]=None, ) -> Tuple[ bool, Optional[Exception]]: ok_once, exp_once = cls.cached_validate_once() if not ok_once: return False, exp_once return cls._validate(bits=bits, group_size=group_size, desc_act=desc_act, sym=sym, in_features=in_features, out_features=out_features, pack_dtype=pack_dtype, dtype=dtype, dynamic=dynamic, device=device, trainable=trainable, adapter=adapter) @classmethod # internal method and should not be overriden def verify_supports_params(cls): """ Validate that SUPPORTS parameters are not None or empty lists, raising an exception if the validation fails. """ base_supports_variables = [ (name, value) for name, value in BaseQuantLinear.__dict__.items() if name.startswith("SUPPORTS") and not callable(value) and value is None ] child_supports_variables = [ (name, value) for name, value in cls.__dict__.items() if name.startswith("SUPPORTS") and not callable(value) ] base_supports_variables.sort(key=lambda x: x[0]) child_supports_variables.sort(key=lambda x: x[0]) base_variable_names = {name for name, value in base_supports_variables} child_variable_names = {name for name, value in child_supports_variables} missing_variables = base_variable_names - child_variable_names if missing_variables: raise ValueError( f"{cls.__name__} these SUPPORTS variables are not overridden: {', '.join(sorted(missing_variables))}") for name, value in child_supports_variables: if not name.startswith("SUPPORTS") or callable(value): continue if value is None: raise ValueError(f"{cls.__name__}.{name} cannot be None.") # if isinstance(value, list) and not value: # raise ValueError(f"{cls.__name__}.{name} cannot be an empty list.") @classmethod def _validate_shared( cls, *, pack_dtype:t.dtype=None, dtype: Optional[t.dtype]=None, device:Optional[DEVICE | t.device]=None, trainable:Optional[bool]=None, adapter:Optional[Adapter]=None, ) -> Tuple[bool, Optional[Exception]]: cls.verify_supports_params() if adapter is not None and adapter.__class__ not in cls.SUPPORTS_ADAPTERS: err = f"{cls} does not support adapter: {adapter}" return False, NotImplementedError(err) if pack_dtype not in cls.SUPPORTS_PACK_DTYPES: err = f"{cls} does not support `pack_dtype`: {pack_dtype}" return False, NotImplementedError(err) if dtype is not None and dtype not in cls.SUPPORTS_DTYPES: err = f"{cls} only supports `{cls.SUPPORTS_DTYPES}` dtype: actual dtype = `{dtype}`" return False, NotImplementedError(err) if PLATFORM.ALL not in cls.SUPPORTS_PLATFORM and sys.platform not in cls.SUPPORTS_PLATFORM: err = f"{cls} does not support platform: {sys.platform}" return False, NotImplementedError(err) if DEVICE.ALL not in cls.SUPPORTS_DEVICES and device is not None: try: cls.validate_device(device) except NotImplementedError: e = f"{cls} does not support device: {device}" return False, NotImplementedError(e) if trainable and not cls.SUPPORTS_TRAINING: err = f"{cls} does not support training." return False, NotImplementedError(err) return True, None @classmethod def _validate_dynamic_bits( cls, *, bits: int, dynamic: Optional[dict], ) -> Tuple[bool, Optional[Exception]]: if dynamic is None: return True, None dynamic_bits = {} for pattern, pattern_dict in dynamic.items(): dynamic_bits[pattern] = pattern_dict.get("bits", bits) if len(cls.SUPPORTS_BITS) == 1: unsupported_dynamic_bits = { layer: dynamic_bits_value for layer, dynamic_bits_value in dynamic_bits.items() if dynamic_bits_value != bits } if unsupported_dynamic_bits: err = f"{cls} not supported dynamic_bits, only support `{cls.SUPPORTS_BITS}` bits" return False, NotImplementedError(err) else: for layer, dynamic_bits_value in dynamic_bits.items(): if dynamic_bits_value not in cls.SUPPORTS_BITS: err = f"{cls} only supports `{cls.SUPPORTS_BITS}` bits: actual dynamic_bits = `{dynamic_bits_value}` for layer `{layer}`" return False, NotImplementedError(err) return True, None @classmethod def _validate_shape_constraints( cls, *, in_features:int=None, out_features:int=None, ) -> Tuple[bool, Optional[Exception]]: if in_features is not None: validate = all(in_features % in_fea == 0 for in_fea in cls.SUPPORTS_IN_FEATURES_DIVISIBLE_BY) if not validate: err = f"{cls}: `in_features`: {in_features} must be divisible by {cls.SUPPORTS_IN_FEATURES_DIVISIBLE_BY}." return False, NotImplementedError(err) if out_features is not None: validate = all(out_features % out_fea == 0 for out_fea in cls.SUPPORTS_OUT_FEATURES_DIVISIBLE_BY) if not validate: err = f"{cls}: `out_features`: {out_features} must be divisible by {cls.SUPPORTS_OUT_FEATURES_DIVISIBLE_BY}." return False, NotImplementedError(err) return True, None @classmethod def _validate(cls, bits: int=4, group_size: int=128, desc_act: bool=False, sym: bool=False, pack_dtype:t.dtype=None, dtype: Optional[t.dtype]=None, dynamic:Optional[dict]=None, in_features:int=None, out_features:int=None, device:Optional[DEVICE | t.device]=None, trainable:Optional[bool]=None, adapter:Optional[Adapter]=None) -> Tuple[bool, Optional[Exception]]: ok, err = cls._validate_shared( pack_dtype=pack_dtype, dtype=dtype, device=device, trainable=trainable, adapter=adapter, ) if not ok: return ok, err if bits not in cls.SUPPORTS_BITS: err = f"{cls} only supports `{cls.SUPPORTS_BITS}` bits: actual bits = `{bits}`" return False, NotImplementedError(err) ok, err = cls._validate_dynamic_bits(bits=bits, dynamic=dynamic) if not ok: return ok, err return cls._validate_shape_constraints( in_features=in_features, out_features=out_features, ) @classmethod def validate_device(cls, device: DEVICE | t.device): # DEVICE declarations describe support families. Keep concrete # torch.device values intact for specialized kernels, while mapping # them to a family for this common declaration check. if isinstance(device, t.device): try: family = DEVICE(device.type) except ValueError: family = device.type else: family = device if family not in cls.SUPPORTS_DEVICES: raise NotImplementedError(f"{cls} only supports `{cls.SUPPORTS_DEVICES}`: actual device = `{device}`") # use optimize so we don't override native module.compile() # override me, to perform any torch.compile logic on the kernel pre forward def optimize(self, backend: str = "inductor", mode: str = None, fullgraph: bool = False): self.optimized = True log.info.once(f"Optimize: `{self.__class__.__name__}` compilation triggered.") pass # overrides nn.module.train() def train(self, mode=True): old_mode = self.training if old_mode == mode: return self # Custom behavior when switching to training mode if mode: if not self.SUPPORTS_TRAINING: err = f"{self.__class__.__name__}: `{self.name}` switching to training mode." log.error(err) raise NotImplementedError(err) else: pass # log.info(f"{self.__class__.__name__}: `{self.name}` switching to training mode.") else: pass # log.info(f"{self.__class__.__name__}: `{self.name}` switching to eval mode.") self.clear_autotune() return super().train(mode) def set_had_K(self, had_K: Optional[t.Tensor]) -> None: """Assign the rotation Hadamard tensor, registering it as a buffer. `had_K` starts life as a plain None attribute (None-valued buffers break accelerate's offload hooks), so a real tensor must be promoted to a registered buffer to follow module device/dtype moves. """ if "had_K" in self._buffers: if had_K is None: del self._buffers["had_K"] self.had_K = None else: self._buffers["had_K"] = had_K return if had_K is None: self.had_K = None return if hasattr(self, "had_K"): del self.had_K self.register_buffer("had_K", had_K, persistent=False) def _apply_rotation_to_input(self, x: t.Tensor) -> t.Tensor: """Apply online Hadamard transform to the input for SpinQuant/QuaRot inference.""" from ...quantization.rotation.hadamard_utils import apply_online_hadamard return apply_online_hadamard( x, online_full_had=self.online_full_had, online_partial_had=self.online_partial_had, had_K=self.had_K, K=self.K, had_dim=self.had_dim, ) class GroupedQuantLinear(BaseQuantLinear): SUPPORTS_GROUP_SIZE: List[int] = None SUPPORTS_DESC_ACT: List[bool] = None SUPPORTS_SYM: List[bool] = None @classmethod def verify_supports_params(cls): super().verify_supports_params() grouped_supports_variables = [ (name, value) for name, value in GroupedQuantLinear.__dict__.items() if name.startswith("SUPPORTS") and not callable(value) and value is None ] child_supports_variables = [ (name, value) for name, value in cls.__dict__.items() if name.startswith("SUPPORTS") and not callable(value) ] grouped_variable_names = {name for name, _ in grouped_supports_variables} child_variable_names = {name for name, _ in child_supports_variables} missing_variables = grouped_variable_names - child_variable_names if missing_variables: raise ValueError( f"{cls.__name__} these grouped SUPPORTS variables are not overridden: " f"{', '.join(sorted(missing_variables))}" ) for name, value in child_supports_variables: if name in grouped_variable_names and value is None: raise ValueError(f"{cls.__name__}.{name} cannot be None.") def __init__(self, bits: int, group_size: int, desc_act: bool, sym: bool, in_features: int, out_features: int, bias: bool, pack_dtype: t.dtype, backend: BACKEND, adapter: Adapter, name: str = None, register_buffers: bool = False, register_buffers_in_features: int = None, register_buffers_out_features: int = None, dtype: Optional[t.dtype] = None, **kwargs): super().__init__( bits=bits, in_features=in_features, out_features=out_features, bias=bias, backend=backend, adapter=adapter, name=name, register_buffers=False, register_buffers_in_features=register_buffers_in_features, register_buffers_out_features=register_buffers_out_features, dtype=dtype, validate_kwargs={ "group_size": group_size, "desc_act": desc_act, "sym": sym, "pack_dtype": pack_dtype, }, **kwargs, ) self.group_size = group_size if group_size != -1 else in_features self.requested_group_size = group_size self.desc_act = desc_act self.sym = sym def smooth_block_size(self) -> int: return -1 if self.requested_group_size == -1 else self.group_size @classmethod def _validate(cls, bits: int=4, group_size: int=128, desc_act: bool=False, sym: bool=False, pack_dtype:t.dtype=None, dtype: Optional[t.dtype]=None, dynamic:Optional[dict]=None, in_features:int=None, out_features:int=None, device:Optional[DEVICE | t.device]=None, trainable:Optional[bool]=None, adapter:Optional[Adapter]=None) -> Tuple[bool, Optional[Exception]]: ok, err = super()._validate( bits=bits, group_size=group_size, desc_act=desc_act, sym=sym, pack_dtype=pack_dtype, dtype=dtype, dynamic=dynamic, in_features=in_features, out_features=out_features, device=device, trainable=trainable, adapter=adapter, ) if not ok: return ok, err if group_size not in cls.SUPPORTS_GROUP_SIZE and group_size != in_features: err = f"{cls} only supports `{cls.SUPPORTS_GROUP_SIZE}` group_size: actual group_size = `{group_size}`" return False, NotImplementedError(err) if sym not in cls.SUPPORTS_SYM: err = f"{cls} only supports symmetric `{cls.SUPPORTS_SYM}` quantization: actual sym = `{sym}`" return False, NotImplementedError(err) if desc_act not in cls.SUPPORTS_DESC_ACT: err = f"{cls} only supports `{cls.SUPPORTS_DESC_ACT}` bits: actual desc_act = `{desc_act}`" return False, NotImplementedError(err) if dynamic is not None: dynamic_group_size = {} for pattern, pattern_dict in dynamic.items(): dynamic_group_size[pattern] = pattern_dict.get("group_size", group_size) for layer, layer_group_size in dynamic_group_size.items(): if layer_group_size not in cls.SUPPORTS_GROUP_SIZE: err = f"{cls} only supports `{cls.SUPPORTS_GROUP_SIZE}` group_size: actual group_size = `{layer_group_size}` for layer `{layer}`" return False, NotImplementedError(err) dynamic_sym = {} for pattern, pattern_dict in dynamic.items(): dynamic_sym[pattern] = pattern_dict.get("sym", sym) for layer, layer_sym in dynamic_sym.items(): if layer_sym not in cls.SUPPORTS_SYM: err = f"{cls} only supports `{cls.SUPPORTS_SYM}` bits: actual sym = `{layer_sym}` for layer `{layer}`" return False, NotImplementedError(err) dynamic_desc_act = {} for pattern, pattern_dict in dynamic.items(): dynamic_desc_act[pattern] = pattern_dict.get("desc_act", desc_act) for layer, layer_desc_act in dynamic_desc_act.items(): if layer_desc_act not in cls.SUPPORTS_DESC_ACT: err = f"{cls} only supports `{cls.SUPPORTS_DESC_ACT}` bits: actual desc_act = `{layer_desc_act}` for layer `{layer}`" return False, NotImplementedError(err) if in_features is not None: validate = in_features % group_size == 0 or cls.SUPPORTS_AUTO_PADDING if not validate: err = f"{cls}: `in_features`: {in_features} must be divisible by `group_size: {group_size}`." return False, NotImplementedError(err) return True, None class PackedGroupedQuantLinear(GroupedQuantLinear): def __init__(self, *args, pack_dtype: t.dtype, **kwargs): super().__init__(*args, pack_dtype=pack_dtype, **kwargs) # Packed GPTQ/AWQ layouts need explicit storage-word metadata. self.pack_dtype = pack_dtype self.maxq = 2 ** self.bits - 1 if self.pack_dtype == t.int8: self.pack_dtype_bits = 8 self.pack_np_dtype = np.int8 self.pack_np_math_dtype = np.uint8 elif self.pack_dtype == t.int16: self.pack_dtype_bits = 16 self.pack_np_dtype = np.int16 self.pack_np_math_dtype = np.uint16 elif self.pack_dtype == t.int32: self.pack_dtype_bits = 32 self.pack_np_dtype = np.int32 self.pack_np_math_dtype = np.uint32 elif self.pack_dtype == t.int64: self.pack_dtype_bits = 64 self.pack_np_dtype = np.int64 self.pack_np_math_dtype = np.uint64 else: raise ValueError(f"Unsupported pack_dtype: {self.pack_dtype}") # pack_factor is only meaningful for packed low-bit code storage. self.pack_factor = self.pack_dtype_bits // self.bits class WeightOnlyQuantLinear(BaseQuantLinear): def __init__(self, bits: int, in_features: int, out_features: int, bias: bool, backend: BACKEND, adapter: Adapter, name: str = None, register_buffers: bool = False, register_buffers_in_features: int = None, register_buffers_out_features: int = None, dtype: Optional[t.dtype] = None, pack_dtype: Optional[t.dtype] = None, **kwargs): super().__init__( bits=bits, in_features=in_features, out_features=out_features, bias=bias, backend=backend, adapter=adapter, name=name, register_buffers=register_buffers, register_buffers_in_features=register_buffers_in_features, register_buffers_out_features=register_buffers_out_features, dtype=dtype, validate_kwargs={"pack_dtype": pack_dtype}, **kwargs, ) class GPTQQuantLinear(PackedGroupedQuantLinear): def __init__(self, bits: int, group_size: int, desc_act: bool, sym: bool, in_features: int, out_features: int, bias: bool, pack_dtype: t.dtype, backend: BACKEND, adapter: Adapter, name: str = None, register_buffers: bool = False, register_buffers_in_features: int = None, register_buffers_out_features: int = None, dtype: Optional[t.dtype] = None, format: Optional[FORMAT] = None, **kwargs): super().__init__( bits=bits, group_size=group_size, desc_act=desc_act, sym=sym, in_features=in_features, out_features=out_features, bias=bias, pack_dtype=pack_dtype, backend=backend, adapter=adapter, name=name, register_buffers=False, register_buffers_in_features=register_buffers_in_features, register_buffers_out_features=register_buffers_out_features, dtype=dtype, **kwargs, ) self.format = format # Planar (split-plane, gptq_p) routing: 5/6/7-bit only exists planar; # 3-bit is planar only under gptq_p. 2/4/8-bit planar words are # bit-identical to the continuous layout, so they keep the fast paths. if format == FORMAT.GPTQ_P and self.bits not in PLANAR_FORMAT_BITS: raise ValueError( f"format={FORMAT.GPTQ_P} supports bits {PLANAR_FORMAT_BITS}, got bits={self.bits}" ) self.planar = self.bits in PLANAR_BITS or (format == FORMAT.GPTQ_P and self.bits == 3) # Planar packing stores whole 32-code blocks; `_validate` covers the # planar-only widths (5/6/7) but is format-blind, so format-dependent # planar 3-bit must be checked here where the format is known. # NotImplementedError keeps this recoverable for the kernel-selection # loop, which treats it as "unsupported by this kernel". if self.planar: for dim_name, dim in (("in_features", self.in_features), ("out_features", self.out_features)): if dim % 32 != 0: raise NotImplementedError( f"planar {self.bits}-bit requires `{dim_name}` divisible by 32, got {dim_name}={dim}." ) # GPTQ v1/v2 conversions only apply to GPTQ-style qzero storage. self._qzeros_format = 1 if register_buffers: self._register_gptq_buffers( bias=bias, register_buffers_in_features=register_buffers_in_features, register_buffers_out_features=register_buffers_out_features, ) @classmethod def _validate(cls, bits: int=4, group_size: int=128, desc_act: bool=False, sym: bool=False, pack_dtype:t.dtype=None, dtype: Optional[t.dtype]=None, dynamic:Optional[dict]=None, in_features:int=None, out_features:int=None, device:Optional[DEVICE | t.device]=None, trainable:Optional[bool]=None, adapter:Optional[Adapter]=None) -> Tuple[bool, Optional[Exception]]: ok, err = super()._validate( bits=bits, group_size=group_size, desc_act=desc_act, sym=sym, pack_dtype=pack_dtype, dtype=dtype, dynamic=dynamic, in_features=in_features, out_features=out_features, device=device, trainable=trainable, adapter=adapter, ) if not ok: return ok, err # Planar-only widths pack whole 32-code blocks, so the packed # dimensions must be 32-aligned regardless of the class-level # SUPPORTS_*_FEATURES_DIVISIBLE_BY declarations. if bits in PLANAR_BITS: for dim_name, dim in (("in_features", in_features), ("out_features", out_features)): if dim is not None and dim % 32 != 0: err = f"{cls}: planar {bits}-bit requires `{dim_name}` divisible by 32: actual {dim_name} = `{dim}`." return False, NotImplementedError(err) return True, None def _register_gptq_buffers( self, *, bias: bool, register_buffers_in_features: int = None, register_buffers_out_features: int = None, ) -> None: in_features = self.in_features if register_buffers_in_features is None else register_buffers_in_features out_features = self.out_features if register_buffers_out_features is None else register_buffers_out_features self.register_buffer( "qweight", t.zeros((math.ceil(in_features * self.bits / self.pack_dtype_bits), out_features), dtype=self.pack_dtype), ) self.register_buffer( "qzeros", t.zeros( ( math.ceil(in_features / self.group_size), math.ceil(out_features * self.bits / self.pack_dtype_bits), ), dtype=self.pack_dtype, ), ) self.register_buffer( "scales", t.zeros( (math.ceil(in_features / self.group_size), out_features), dtype=t.float16, ), ) self.register_buffer( "g_idx", t.tensor([i // self.group_size for i in range(in_features)], dtype=t.int32), ) if bias: self.register_buffer("bias", t.zeros(out_features, dtype=t.float16)) else: self.bias = None def qzero_format(self, format: int = None) -> int: if format is None: return self._qzeros_format if format not in [1, 2]: raise ValueError("Unsupported qzero format. Only 1 and 2 are supported.") self._qzeros_format = format return self._qzeros_format class PackableQuantLinear(GPTQQuantLinear): def __init__(self, *args, enable_wf_unsqueeze: bool = False, **kwargs): self.enable_wf_unsqueeze = enable_wf_unsqueeze super().__init__(*args, **kwargs) def _wf_device(self): if hasattr(self, "g_idx") and self.g_idx is not None: return self.g_idx.device if hasattr(self, "qweight") and self.qweight is not None: return self.qweight.device raise RuntimeError(f"{self.__class__.__name__} cannot initialize wf buffers without qweight or g_idx.") def _init_wf_unsqueeze_buffers(self): if not self.enable_wf_unsqueeze: return if ( getattr(self, "wf_unsqueeze_zero", None) is not None and getattr(self, "wf_unsqueeze_neg_one", None) is not None ): return if self.planar: # Planar dequantization unpacks bit planes directly and does not # use wf shift buffers. return device = self._wf_device() if self.bits in [2, 4, 8]: wf = t.tensor(list(range(0, self.pack_dtype_bits, self.bits)), dtype=t.int32).unsqueeze(0).to( device=device) elif self.bits == 3: wf = t.tensor( [ [0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 0], [0, 1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31], [0, 2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 0], ], dtype=t.int32, ).reshape(1, 3, 12).to(device=device) else: raise NotImplementedError(f"{self.__class__.__name__} does not support wf unpack buffers for {self.bits}-bit weights.") wf_zero = wf.unsqueeze(0).to(device=device) wf_neg_one = wf.unsqueeze(-1).to(device=device) if "wf_unsqueeze_zero" not in self._buffers: self.register_buffer("wf_unsqueeze_zero", wf_zero, persistent=False) else: self.wf_unsqueeze_zero = wf_zero if "wf_unsqueeze_neg_one" not in self._buffers: self.register_buffer("wf_unsqueeze_neg_one", wf_neg_one, persistent=False) else: self.wf_unsqueeze_neg_one = wf_neg_one def post_init(self, **kwargs): self._init_wf_unsqueeze_buffers() super().post_init(**kwargs) def list_buffers(self): buf = super().list_buffers() if hasattr(self, "wf_unsqueeze_zero") and self.wf_unsqueeze_zero is not None: buf.append(self.wf_unsqueeze_zero) if hasattr(self, "wf_unsqueeze_neg_one") and self.wf_unsqueeze_neg_one is not None: buf.append(self.wf_unsqueeze_neg_one) return buf def dequantize_weight(self, num_itr: int = 1): if self.planar: zeros = planar_unpack_cols(self.qzeros, self.bits).reshape(self.scales.shape) weight = planar_unpack_rows(self.qweight, self.bits) return self._dequantize_from_codes(weight, zeros, num_itr=num_itr) self._init_wf_unsqueeze_buffers() wf_zero = getattr(self, "wf_unsqueeze_zero", None) wf_neg_one = getattr(self, "wf_unsqueeze_neg_one", None) if wf_zero is None or wf_neg_one is None: raise RuntimeError( f"{self.__class__.__name__} requires wf unpack buffers for dequantization. " "Enable `enable_wf_unsqueeze` and call post_init()." ) if self.bits in [2, 4, 8]: zeros = _torch_right_shift( t.unsqueeze(self.qzeros, 2).expand(-1, -1, self.pack_factor), self.wf_unsqueeze_zero # self.wf.unsqueeze(0), ).to(self.dequant_dtype) zeros = t.bitwise_and(zeros, self.maxq).reshape(self.qzeros.shape[0], -1) zeros = zeros[:, :self.scales.shape[1]] weight = t.bitwise_and( _torch_right_shift( t.unsqueeze(self.qweight, 1).expand(-1, self.pack_factor, -1), self.wf_unsqueeze_neg_one # self.wf.unsqueeze(-1) ).to(self.dequant_dtype), self.maxq ) elif self.bits == 3: zeros = self.qzeros.reshape(self.qzeros.shape[0], self.qzeros.shape[1] // 3, 3, 1).expand( -1, -1, -1, 12 ) zeros = _torch_right_shift(zeros, self.wf_unsqueeze_zero) # self.wf.unsqueeze(0) zeros[:, :, 0, 10] = (zeros[:, :, 0, 10] & 0x3) | (_torch_left_shift(zeros[:, :, 1, 0], 2) & 0x4) zeros[:, :, 1, 11] = (zeros[:, :, 1, 11] & 0x1) | (_torch_left_shift(zeros[:, :, 2, 0], 1) & 0x6) zeros = zeros & 0x7 zeros = t.cat( [zeros[:, :, 0, :11], zeros[:, :, 1, 1:12], zeros[:, :, 2, 1:11]], dim=2, ).reshape(self.scales.shape) weight = self.qweight.reshape(self.qweight.shape[0] // 3, 3, 1, self.qweight.shape[1]).expand( -1, -1, 12, -1 ) weight = _torch_right_shift(weight, self.wf_unsqueeze_neg_one) & 0x7 # self.wf.unsqueeze(-1) weight[:, 0, 10] = (weight[:, 0, 10] & 0x3) | (_torch_left_shift(weight[:, 1, 0], 2) & 0x4) weight[:, 1, 11] = (weight[:, 1, 11] & 0x1) | (_torch_left_shift(weight[:, 2, 0], 1) & 0x6) weight = weight & 0x7 weight = t.cat([weight[:, 0, :11], weight[:, 1, 1:12], weight[:, 2, 1:11]], dim=1) weight = weight.reshape(weight.shape[0] * weight.shape[1], weight.shape[2]) return self._dequantize_from_codes(weight, zeros, num_itr=num_itr) def _dequantize_from_codes(self, weight: t.Tensor, zeros: t.Tensor, num_itr: int = 1): # Packed words may contain zero-filled tail codes when an embedding # vocabulary is not divisible by the packing width. The runtime # embedding contract remains the original vocabulary size. if weight.shape[0] > self.in_features: weight = weight[:self.in_features] if num_itr == 1: weights = self.scales[self.g_idx.long()] * (weight - zeros[self.g_idx.long()]) else: num_dim = self.g_idx.shape[0] // num_itr out_dim = weight.shape[1] // num_itr weights = [] for i in range(num_itr): row_start = i * num_dim row_end = (i + 1) * num_dim col_start = i * out_dim col_end = (i + 1) * out_dim if i < (num_itr - 1) else weight.shape[1] scale_i = self.scales[:, col_start:col_end] weight_i = weight[row_start:row_end, col_start:col_end] zeros_i = zeros[:, col_start:col_end] g_idx_i = self.g_idx[row_start:row_end].long() weights.append(scale_i[g_idx_i] * (weight_i - zeros_i[g_idx_i])) weights = t.cat(weights, dim=1) return weights # FIXME, optimum needs call pack(), we need to remove it def pack( self, linear: nn.Module, scales: t.Tensor, zeros: t.Tensor, g_idx: t.Tensor, block_in: int = 8192, workers: int = 8, ): self.pack_block(linear, scales, zeros, g_idx, block_in, workers) @t.inference_mode() def pack_block( self, linear: nn.Module, scales: t.Tensor, zeros: t.Tensor, g_idx: t.Tensor, block_in: int = 8192, workers: int = 1, ): """ Parallel qweight pack on CPU (threaded over input blocks). qzeros path = original logic. - qweight: streamed in 32-aligned blocks; each block packed by a worker into non-overlapping rows of the preallocated output tensor (no locks). - qzeros: EXACT original numpy-based packing to preserve layout/shape. - math done on CPU; no .numpy() round trips for qweight. Args: block_in: number of input channels per task (must be multiple of 32). workers: number of worker threads (None -> auto; 1 -> single-thread). """ MASK32 = (1 << 32) - 1 # for safe masking when packing via int64 # ---------- normalize weight to [out, in] (same as your code) ---------- W = linear.weight.detach() if isinstance(linear, _ConvNd): W = W.flatten(1) if isinstance(linear, transformers.pytorch_utils.Conv1D): W = W.T W = W.to("cpu", copy=False) out_features, in_features = W.shape # ---------- g_idx buffer ---------- if g_idx.numel() != in_features: raise ValueError(f"g_idx length {g_idx.numel()} != in_features {in_features}") g_idx = g_idx.to("cpu") self.register_buffer("g_idx", g_idx) # ---------- ORIGINAL scales/zeros logic (unchanged) ---------- scales = scales.T.contiguous() # [G, out] zeros = zeros.T.contiguous() # [G, out] scale_zeros = zeros * scales # [G, out] num_groups = scales.shape[0] # small buffers self.register_buffer("scales", scales.to(dtype=t.float16)) if linear.bias is not None: self.register_buffer("bias", linear.bias.detach().to("cpu", dtype=t.float16)) elif not hasattr(self, "bias"): self.bias = None # ---------- constants ---------- bits = int(self.bits) # 2,3,4,8 word_bits = int(self.pack_dtype_bits) # expect 32 assert word_bits == 32, "Only 32-bit packing words supported." if (in_features % word_bits) != 0: raise ValueError("in_features must be divisible by 32") if (out_features % word_bits) != 0: raise ValueError( "pack_block extension requires out_features to be divisible by 32" ) disable_ext = env_flag("GPTQMODEL_DISABLE_PACK_EXT") force_ext = env_flag("GPTQMODEL_FORCE_PACK_EXT") pack_block_threads = workers if workers and workers > 0 else 1 env_threads = os.getenv("GPTQMODEL_PACK_THREADS") if env_threads: try: pack_block_threads = max(int(env_threads), 1) except ValueError: log.warning( "pack_block: invalid GPTQMODEL_PACK_THREADS `%s`; defaulting to %d.", env_threads, pack_block_threads, ) planar = self.planar if not disable_ext and bits in (2, 4, 8): try: from .pack_block_ext import pack_block_cpu as pack_block_cpu_ext qweight_ext, qzeros_ext = pack_block_cpu_ext( W, scales, zeros, g_idx.to(dtype=t.int32), bits, word_bits, block_in, pack_block_threads, ) self.register_buffer("qweight", qweight_ext.to(dtype=self.pack_dtype)) self.register_buffer("qzeros", qzeros_ext.to(dtype=self.pack_dtype)) return except Exception as exc: if force_ext: raise log.debug("pack_block: native extension unavailable, falling back to Python path (%s)", exc) # NOTE: pack_factor is only meaningful for {2,4,8}. There is NO integer pack_factor for 3-bit. if bits in (2, 4, 8): pack_factor = word_bits // bits # 16, 8, 4 respectively # If the instance carries a different pack_factor, ignore it for safety. elif bits == 3 and not planar: pack_factor = None # sentinel: use the 10-1-10-1-10 scheme elif planar: pack_factor = None # sentinel: use split-plane packing else: raise NotImplementedError(f"Unsupported bits={bits}") # ---------- qweight allocation (storage dtype = int32) ---------- OUT_DTYPE = self.pack_dtype # should be torch.int32 per your dtype table if OUT_DTYPE != t.int32: raise ValueError("pack_block() expects self.pack_dtype == torch.int32 for 32-bit words.") # rows = (in // 32) * bits for every supported layout rows_per_group = bits qweight_rows = (in_features // word_bits) * rows_per_group qweight = t.empty((qweight_rows, out_features), dtype=OUT_DTYPE, device="cpu") # ---------- precompute shifts for 2/4/8-bit vectorized packing ---------- if bits in (2, 4, 8): shifts_pf64 = (t.arange(pack_factor, dtype=t.int64) * bits).view(1, pack_factor, 1) # [1,pf,1] # ---------- pack helpers (int64 shifts -> mask -> int32) ---------- @t.inference_mode() def _pack_rows_2_4_8(int32_blk_32xN: t.Tensor, dst: t.Tensor, dst_rows_base: int): # int32_blk_32xN: [32, N] r64 = int32_blk_32xN.to(t.int64) # [32,N] # interpret the 32 input rows as [bits, pack_factor, N] in row-major order R = r64.view(bits, pack_factor, -1) # [bits,pf,N] packed64 = (R << shifts_pf64).sum(dim=1, dtype=t.int64) # [bits,N] dst[dst_rows_base:dst_rows_base + bits] = ((packed64 & MASK32).to(t.int32)) @t.inference_mode() def _pack_rows_3(int32_blk_32xN: t.Tensor, dst: t.Tensor, dst_rows_base: int): # Layout: 32 inputs -> 3 output rows with pattern 10 | 1 | 10 | 1 | 10 (=> 32 scalars * 3 bits = 96 bits) x = int32_blk_32xN.to(t.int64) # [32,N] # A: first 10 values at bit offsets {0,3,6,...,27} plus low 2 bits of the 11th value at bit 30 A = ((x[0:10] << (t.arange(10, dtype=t.int64).view(-1, 1) * 3)).sum(dim=0, dtype=t.int64)) & MASK32 A = (A | ((x[10] << 30) & MASK32)) & MASK32 # B: carry bit 2 of the 11th into bit 0, then 10 values at offsets {1,4,7,...,28}, then bit 31 from 22nd B = (((x[10] >> 2) & 1) & MASK32) B = (B | ( (x[11:21] << (t.arange(10, dtype=t.int64).view(-1, 1) * 3 + 1)).sum(dim=0, dtype=t.int64))) & MASK32 B = (B | ((x[21] << 31) & MASK32)) & MASK32 # C: carry bits 1..0 of the 22nd into bits 0..1, then 10 values at offsets {2,5,8,...,29} C = (((x[21] >> 1) & 0x3) & MASK32) C = (C | ( (x[22:32] << (t.arange(10, dtype=t.int64).view(-1, 1) * 3 + 2)).sum(dim=0, dtype=t.int64))) & MASK32 dst[dst_rows_base + 0] = A.to(t.int32) dst[dst_rows_base + 1] = B.to(t.int32) dst[dst_rows_base + 2] = C.to(t.int32) # ---------- thread task: process a single [i0,i1) block ---------- block_in = max(word_bits, (block_in // word_bits) * word_bits) @t.inference_mode() def _process_block(i0: int, i1: int): blk = i1 - i0 # [out, blk] Wblk = W[:, i0:i1] # select group rows for these inputs from [G, out] -> [blk, out], then T -> [out, blk] gsel = g_idx[i0:i1].to(dtype=t.int64, copy=False) # [blk] if gsel.numel() == 0: return neg_mask = gsel < 0 if neg_mask.any(): gsel = gsel.clone() gsel[neg_mask] += num_groups gsel_max = int(gsel.max().item()) gsel_min = int(gsel.min().item()) if gsel_min < 0 or gsel_max >= num_groups: raise IndexError( f"pack_block: g_idx values out of range after normalization (min={gsel_min}, max={gsel_max}, groups={num_groups})." ) sz_blk_T = scale_zeros.index_select(0, gsel).T # [out, blk] s_blk_T = scales.index_select(0, gsel).T # [out, blk] # int_block = round((W + scale_zeros[g_idx]^T) / scales[g_idx]^T) int_block = t.round((Wblk + sz_blk_T) / s_blk_T) # [out, blk] # Values that drift outside the representable range must saturate. # Bit-packing an unchecked value wraps it modulo 2**bits and turns # a small reconstruction error into a full-range weight error. int_block.clamp_(0, self.maxq) int_block = int_block.to(t.int32).T.contiguous() # [blk, out] groups32 = blk // word_bits base_group = (i0 // word_bits) for g in range(groups32): sub = int_block[g * word_bits:(g + 1) * word_bits] # [32, out] int32 dst_rows = (base_group + g) * rows_per_group if bits in (2, 4, 8): _pack_rows_2_4_8(sub, qweight, dst_rows) elif bits == 3 and not planar: _pack_rows_3(sub, qweight, dst_rows) elif planar: qweight[dst_rows:dst_rows + bits] = planar_pack_rows(sub, bits) else: raise NotImplementedError(f"Unsupported bits={bits}") # free temps del Wblk, gsel, sz_blk_T, s_blk_T, int_block # ---------- schedule blocks across a thread pool ---------- starts = list(range(0, in_features, block_in)) ranges = [(i0, min(i0 + block_in, in_features)) for i0 in starts] len(ranges) total_blocks = len(ranges) # Each block writes a disjoint row range of the preallocated qweight and # only reads the shared W/scales/zeros tensors, so blocks can run on a # thread pool. Threads only help on free-threaded (GIL=0) Python builds; # keep the sequential default and enable the pool only through the # dedicated GPTQMODEL_PACK_PY_THREADS opt-in (GPTQMODEL_PACK_THREADS # keeps its existing meaning of sizing the native extension only). py_threads_env = os.getenv("GPTQMODEL_PACK_PY_THREADS") workers_eff = 1 if py_threads_env: try: workers_eff = max(1, min(int(py_threads_env), total_blocks)) except ValueError: log.warning( "pack_block: invalid GPTQMODEL_PACK_PY_THREADS `%s`; using sequential packing.", py_threads_env, ) if workers_eff == 1: for i0, i1 in ranges: _process_block(i0, i1) else: with ThreadPoolExecutor(max_workers=workers_eff, thread_name_prefix="pack qweight") as ex: futs = [ex.submit(_process_block, i0, i1) for (i0, i1) in ranges] for f in futs: f.result() # ---------- qzeros: EXACT original numpy path (unchanged) ---------- zeros_np = zeros.detach().cpu().numpy().astype(self.pack_np_math_dtype, copy=False) # [G, out] qzeros_np = np.zeros( (zeros_np.shape[0], (zeros_np.shape[1] // self.pack_dtype_bits) * bits), dtype=self.pack_np_math_dtype ) if bits in [2, 4, 8]: pf = (32 // bits) # compute locally to avoid relying on attribute for col in range(qzeros_np.shape[1]): base = col * pf for j in range(pf): qzeros_np[:, col] |= zeros_np[:, base + j] << (bits * j) elif bits == 3 and not planar: i = 0 col = 0 while col < qzeros_np.shape[1]: for j in range(i, i + 10): qzeros_np[:, col] |= zeros_np[:, j] << (3 * (j - i)) i += 10 qzeros_np[:, col] |= zeros_np[:, i] << 30 col += 1 qzeros_np[:, col] |= (zeros_np[:, i] >> 2) & 1 i += 1 for j in range(i, i + 10): qzeros_np[:, col] |= zeros_np[:, j] << (3 * (j - i) + 1) i += 10 qzeros_np[:, col] |= zeros_np[:, i] << 31 col += 1 qzeros_np[:, col] |= (zeros_np[:, i] >> 1) & 0x3 i += 1 for j in range(i, i + 10): qzeros_np[:, col] |= zeros_np[:, j] << (3 * (j - i) + 2) i += 10 col += 1 elif planar: qzeros_np = ( planar_pack_cols(zeros.to(t.int64), bits) .numpy() .astype(self.pack_np_math_dtype, copy=False) ) else: raise NotImplementedError(f"Unsupported bits={bits}") # ---------- register buffers ---------- self.register_buffer("qweight", qweight.to(dtype=self.pack_dtype)) self.register_buffer("qzeros", t.from_numpy(qzeros_np.astype(self.pack_np_dtype, copy=False))) @t.inference_mode() def pack_gpu( self, linear: nn.Module, scales: t.Tensor, zeros: t.Tensor, g_idx: t.Tensor, *, block_in: int = 8192, device: Optional[t.device] = None, ): """Pack quantized weights using CUDA kernels.""" if not t.cuda.is_available(): raise RuntimeError("pack_gpu requires CUDA availability but none was detected.") target_device = t.device(device) if device is not None else t.device("cuda") if target_device.type != "cuda": raise ValueError(f"pack_gpu expected a CUDA device, got `{target_device}`.") MASK32 = (1 << 32) - 1 mask_tensor = t.tensor(MASK32, dtype=t.int64, device=target_device) weight = linear.weight.detach().to(device=target_device, copy=True) # weight = linear.weight.detach() if isinstance(linear, _ConvNd): weight = weight.flatten(1) if isinstance(linear, transformers.pytorch_utils.Conv1D): weight = weight.T # weight = weight.to(device=target_device, copy=True) out_features, in_features = weight.shape if g_idx is None: raise ValueError("pack_gpu requires non-null g_idx") if g_idx.numel() != in_features: raise ValueError(f"g_idx length {g_idx.numel()} != in_features {in_features}") g_idx_cpu = g_idx.to(device="cpu", dtype=t.long, copy=False) g_idx_dev = g_idx_cpu.to(device=target_device) self.register_buffer("g_idx", g_idx_cpu.to(dtype=t.int32)) # Align layout with CPU paths: expect [groups, out_features] scales_g = scales.T.contiguous() zeros_g = zeros.T.contiguous() scale_zeros_g = zeros_g * scales_g scales_dev = scales_g.to(device=target_device) zeros_dev = zeros_g.to(device=target_device) scale_zeros_dev = scale_zeros_g.to(device=target_device) self.register_buffer("scales", scales_g.to(device="cpu", dtype=t.float16)) if linear.bias is not None: self.register_buffer("bias", linear.bias.detach().to(device="cpu", dtype=t.float16)) else: self.bias = None bits = int(self.bits) word_bits = int(self.pack_dtype_bits) if word_bits != 32: raise ValueError("pack_gpu currently supports 32-bit packing words only.") if in_features % word_bits != 0: raise ValueError("in_features must be divisible by 32 for pack_gpu") if bits in (2, 4, 8): pack_factor = word_bits // bits shifts_pf64 = ( t.arange(pack_factor, dtype=t.int64, device=target_device) .view(1, pack_factor, 1) * bits ) elif bits == 3 or self.planar: shifts_pf64 = None else: raise NotImplementedError(f"Unsupported bits={bits}") rows_per_group = bits qweight_dev = t.empty( (in_features // word_bits * rows_per_group, out_features), dtype=self.pack_dtype, device=target_device, ) range10 = t.arange(10, dtype=t.int64, device=target_device).view(-1, 1) @t.inference_mode() def _pack_rows_2_4_8(int32_blk_32xN: t.Tensor, dst: t.Tensor, dst_rows_base: int): r64 = int32_blk_32xN.to(t.int64) reshaped = r64.view(bits, pack_factor, -1) packed64 = (reshaped << shifts_pf64).sum(dim=1, dtype=t.int64) dst[dst_rows_base:dst_rows_base + bits] = ( (packed64 & mask_tensor).to(self.pack_dtype) ) @t.inference_mode() def _pack_rows_3_values(int32_blk_32xN: t.Tensor) -> t.Tensor: x = int32_blk_32xN.to(t.int64) A = ((x[0:10] << (range10 * 3)).sum(dim=0, dtype=t.int64)) & mask_tensor A = (A | ((x[10] << 30) & mask_tensor)) & mask_tensor B = (((x[10] >> 2) & 1) & mask_tensor) B = (B | ((x[11:21] << (range10 * 3 + 1)).sum(dim=0, dtype=t.int64))) & mask_tensor B = (B | ((x[21] << 31) & mask_tensor)) & mask_tensor C = (((x[21] >> 1) & 0x3) & mask_tensor) C = (C | ((x[22:32] << (range10 * 3 + 2)).sum(dim=0, dtype=t.int64))) & mask_tensor return t.stack([A, B, C], dim=0) def _pack_rows_3(int32_blk_32xN: t.Tensor, dst: t.Tensor, dst_rows_base: int): packed = _pack_rows_3_values(int32_blk_32xN) dst[dst_rows_base:dst_rows_base + 3] = packed.to(self.pack_dtype) block_in = max(word_bits, (block_in // word_bits) * word_bits) for i0 in range(0, in_features, block_in): i1 = min(i0 + block_in, in_features) blk = i1 - i0 Wblk = weight[:, i0:i1] gsel = g_idx_dev[i0:i1] sz_blk_T = scale_zeros_dev.index_select(0, gsel).T s_blk_T = scales_dev.index_select(0, gsel).T int_block = t.round((Wblk + sz_blk_T) / s_blk_T) int_block.clamp_(0, self.maxq) int_block = int_block.to(t.int32).T.contiguous() groups32 = blk // word_bits base_group = (i0 // word_bits) for g in range(groups32): sub = int_block[g * word_bits:(g + 1) * word_bits] dst_rows = (base_group + g) * rows_per_group if bits in (2, 4, 8): _pack_rows_2_4_8(sub, qweight_dev, dst_rows) elif bits == 3 and not self.planar: _pack_rows_3(sub, qweight_dev, dst_rows) elif self.planar: qweight_dev[dst_rows:dst_rows + bits] = planar_pack_rows(sub, bits).to(self.pack_dtype) else: raise NotImplementedError(f"Unsupported bits={bits}") zeros_int = zeros_dev.to(dtype=t.int64) if bits in (2, 4, 8): pack_factor = word_bits // bits if zeros_int.shape[1] % pack_factor != 0: raise ValueError( f"pack_gpu expected zeros second dimension divisible by pack_factor={pack_factor}, " f"got shape {zeros_int.shape}" ) zeros_view = zeros_int.view(zeros_int.shape[0], zeros_int.shape[1] // pack_factor, pack_factor) shifts = ( t.arange(pack_factor, dtype=t.int64, device=target_device) .view(1, 1, pack_factor) * bits ) packed = (zeros_view << shifts).sum(dim=-1, dtype=t.int64) & mask_tensor qzeros_dev = packed.to(dtype=self.pack_dtype) elif bits == 3 and not self.planar: groups = zeros_int.shape[1] // word_bits qzeros_chunks: List[t.Tensor] = [] for g in range(groups): block = zeros_int[:, g * word_bits:(g + 1) * word_bits] packed = _pack_rows_3_values(block.T).transpose(0, 1) qzeros_chunks.append(packed.to(dtype=self.pack_dtype)) qzeros_dev = t.cat(qzeros_chunks, dim=1) elif self.planar: qzeros_dev = planar_pack_cols(zeros_int, bits).to(dtype=self.pack_dtype) else: raise NotImplementedError(f"Unsupported bits={bits}") t.cuda.synchronize(device=target_device) self.register_buffer("qweight", qweight_dev.to(device="cpu", dtype=self.pack_dtype)) self.register_buffer("qzeros", qzeros_dev.to(device="cpu", dtype=self.pack_dtype)) del weight, scales_dev, zeros_dev, scale_zeros_dev, qweight_dev, qzeros_dev def pack_original(self, linear: nn.Module, scales: t.Tensor, zeros: t.Tensor, g_idx: t.Tensor=None): with THREADPOOLCTL.threadpool_limits(1): # TODO why did we need to clone? at packing, the original weight is no longer used by other processors? # W = linear.weight.data.clone() W = linear.weight.data is_embedding = isinstance(linear, nn.Embedding) if isinstance(linear, _ConvNd): W = W.flatten(1) if isinstance(linear, transformers.pytorch_utils.Conv1D): W = W.T # TODO why clone? # self.g_idx = g_idx.clone() if g_idx is not None else self.g_idx self.register_buffer("g_idx", g_idx if g_idx is not None else self.g_idx ) scales = scales.T.contiguous() zeros = zeros.T.contiguous() scale_zeros = zeros * scales # TODO why clone? # self.scales = scales.clone().to(dtype=t.float16) self.register_buffer("scales", scales.to(dtype=t.float16)) if hasattr(linear, "bias") and linear.bias is not None: # TODO why clone? # self.bias = linear.bias.clone().to(dtype=t.float16) self.register_buffer("bias", linear.bias.to(dtype=t.float16)) elif not hasattr(self, "bias"): self.bias = None if is_embedding: # Embedding weights already use [token, feature] orientation; # the GPTQ embedding result is transposed only for quantizing. int_weight = t.round((W + scale_zeros[self.g_idx]) / scales[self.g_idx]) else: int_weight = t.round((W + scale_zeros[self.g_idx].T) / scales[self.g_idx].T) int_weight.clamp_(0, self.maxq) int_weight = int_weight.to(t.int32) if not is_embedding: int_weight = int_weight.T.contiguous() else: int_weight = int_weight.contiguous() int_weight = int_weight.numpy().astype(self.pack_np_math_dtype) qweight = np.zeros((math.ceil(int_weight.shape[0] * self.bits / self.pack_dtype_bits), int_weight.shape[1]), dtype=self.pack_np_math_dtype) if self.bits in [2, 4, 8]: for row in range(qweight.shape[0]): for j in range(min(self.pack_factor, int_weight.shape[0] - row * self.pack_factor)): qweight[row] |= int_weight[row * self.pack_factor + j] << (self.bits * j) elif self.bits == 3 and not self.planar: i = 0 row = 0 while row < qweight.shape[0]: for j in range(i, i + 10): qweight[row] |= int_weight[j] << (3 * (j - i)) i += 10 qweight[row] |= int_weight[i] << 30 row += 1 qweight[row] |= (int_weight[i] >> 2) & 1 i += 1 for j in range(i, i + 10): qweight[row] |= int_weight[j] << (3 * (j - i) + 1) i += 10 qweight[row] |= int_weight[i] << 31 row += 1 qweight[row] |= (int_weight[i] >> 1) & 0x3 i += 1 for j in range(i, i + 10): qweight[row] |= int_weight[j] << (3 * (j - i) + 2) i += 10 row += 1 elif self.planar: qweight = ( planar_pack_rows( t.from_numpy(int_weight.astype(np.int64)), self.bits, pack_dtype=self.pack_dtype, ) .numpy() .astype(self.pack_np_math_dtype, copy=False) ) else: raise NotImplementedError(f"Unsupported bits={self.bits}") # self.qweight = t.from_numpy(qweight.astype(self.pack_np_dtype)) self.register_buffer("qweight", t.from_numpy(qweight.astype(self.pack_np_dtype))) zeros = zeros.numpy().astype(self.pack_np_math_dtype) qzeros = np.zeros((zeros.shape[0], math.ceil(zeros.shape[1] * self.bits / self.pack_dtype_bits)), dtype=self.pack_np_math_dtype) if self.bits in [2, 4, 8]: for col in range(qzeros.shape[1]): for j in range(min(self.pack_factor, zeros.shape[1] - col * self.pack_factor)): qzeros[:, col] |= zeros[:, col * self.pack_factor + j] << (self.bits * j) elif self.bits == 3 and not self.planar: i = 0 col = 0 while col < qzeros.shape[1]: for j in range(i, i + 10): qzeros[:, col] |= zeros[:, j] << (3 * (j - i)) i += 10 qzeros[:, col] |= zeros[:, i] << 30 col += 1 qzeros[:, col] |= (zeros[:, i] >> 2) & 1 i += 1 for j in range(i, i + 10): qzeros[:, col] |= zeros[:, j] << (3 * (j - i) + 1) i += 10 qzeros[:, col] |= zeros[:, i] << 31 col += 1 qzeros[:, col] |= (zeros[:, i] >> 1) & 0x3 i += 1 for j in range(i, i + 10): qzeros[:, col] |= zeros[:, j] << (3 * (j - i) + 2) i += 10 col += 1 elif self.planar: qzeros = ( planar_pack_cols( t.from_numpy(zeros.astype(np.int64)), self.bits, pack_dtype=self.pack_dtype, ) .numpy() .astype(self.pack_np_math_dtype, copy=False) ) else: raise NotImplementedError(f"Unsupported bits={self.bits}") # self.qzeros = t.from_numpy(qzeros.astype(self.pack_np_dtype)) self.register_buffer("qzeros", t.from_numpy(qzeros.astype(self.pack_np_dtype))) # assert # assert isinstance(self, TorchLinear), f"type: {self.__class_}" # wq = linear.weight.data # wq_dequantized = self.dequantize_weight().T # print(f"------ WQ -----") # print(wq) # print(f"------ WQ Dequantized -----") # print(wq_dequantized) # assert t.equal(wq, wq_dequantized) # print("self qw", self.qweight, self.scales, self.qzeros) class AWQuantLinear(PackedGroupedQuantLinear): def __init__(self, bias: bool = False, register_buffers: bool = False, **kwargs): super().__init__(bias=bias, register_buffers=False, **kwargs) in_features = self.in_features out_features = self.out_features if register_buffers: self.register_buffer( "qweight", t.zeros((in_features, out_features // (self.pack_dtype_bits // self.bits)), dtype=self.pack_dtype), ) self.register_buffer( "qzeros", t.zeros( (in_features // self.group_size, out_features // (self.pack_dtype_bits // self.bits)), dtype=self.pack_dtype, ), ) self.register_buffer( "scales", t.zeros( (in_features // self.group_size, out_features), dtype=t.float16, ), ) if bias: self.register_buffer("bias", t.zeros(out_features, dtype=t.float16)) else: self.bias = None # TODO FIX ME. this hack was needed because other part of code forgot to call nn.module register_buffer()! def list_buffers(self) -> List: buf = super().list_buffers() return buf