|
| 1 | +import logging |
| 2 | +from typing import Any, Callable, Dict, Optional, Type |
| 3 | + |
| 4 | +import torch |
| 5 | +import torch.nn.functional as F |
| 6 | + |
| 7 | +from torchao.quantization.GPTQ import _check_linear_int4_k |
| 8 | +from torchao.quantization.unified import Quantizer |
| 9 | +from torchao.quantization.utils import groupwise_affine_quantize_tensor |
| 10 | + |
| 11 | + |
| 12 | +# This module is copied from torchao.quantization.GPTQ.WeightOnlyInt4Linear with |
| 13 | +# changes at the annotated lines. |
| 14 | +class VkWeightOnlyInt4Linear(torch.nn.Module): |
| 15 | + __constants__ = ["in_features", "out_features"] |
| 16 | + in_features: int |
| 17 | + out_features: int |
| 18 | + weight: torch.Tensor |
| 19 | + |
| 20 | + def __init__( |
| 21 | + self, |
| 22 | + in_features: int, |
| 23 | + out_features: int, |
| 24 | + # TODO: remove dtype field, not used |
| 25 | + bias=False, |
| 26 | + device=None, |
| 27 | + dtype=None, |
| 28 | + groupsize: int = 128, |
| 29 | + inner_k_tiles: int = 8, |
| 30 | + precision: torch.dtype = torch.bfloat16, |
| 31 | + scales_precision: torch.dtype = torch.bfloat16, |
| 32 | + ) -> None: |
| 33 | + super().__init__() |
| 34 | + self.padding = not _check_linear_int4_k(in_features, groupsize, inner_k_tiles) |
| 35 | + if self.padding: |
| 36 | + from torchao.quantization.utils import find_multiple |
| 37 | + |
| 38 | + self.origin_in_features = in_features |
| 39 | + in_features = find_multiple(in_features, (1024,)) |
| 40 | + |
| 41 | + self.in_features = in_features |
| 42 | + self.out_features = out_features |
| 43 | + assert not bias, "require bias=False" |
| 44 | + self.device = device |
| 45 | + self.groupsize = groupsize |
| 46 | + self.inner_k_tiles = inner_k_tiles |
| 47 | + self.precision = precision |
| 48 | + self.scales_precision = scales_precision |
| 49 | + |
| 50 | + if dtype is not None: |
| 51 | + raise ValueError("Please specify 'precision' instead of 'dtype'") |
| 52 | + |
| 53 | + assert out_features % 8 == 0, "require out_features % 8 == 0" |
| 54 | + assert ( |
| 55 | + in_features % (inner_k_tiles * 16) == 0 |
| 56 | + ), "require in_features % (innerKTiles * 16) == 0" |
| 57 | + # In the original implementation, the weight buffer is registered with the packed |
| 58 | + # sizes, i.e. the result of calling the _convert_weight_to_int4pack operator. |
| 59 | + # However, the Vulkan implementation does not expect the weights to be packed |
| 60 | + # therefore the weight tensor is registered with the unpacked sizes instead. |
| 61 | + # Note that in_features is divided by 2 because each `uint8` tensor element |
| 62 | + # contains 2 4-bit packed values. |
| 63 | + self.register_buffer( |
| 64 | + "weight", |
| 65 | + torch.empty( |
| 66 | + (out_features, in_features // 2), |
| 67 | + dtype=torch.uint8, |
| 68 | + device=device, |
| 69 | + ), |
| 70 | + ) |
| 71 | + self.dtype = dtype |
| 72 | + self.register_buffer( |
| 73 | + "scales_and_zeros", |
| 74 | + torch.empty( |
| 75 | + (in_features // groupsize, out_features, 2), |
| 76 | + dtype=self.scales_precision, |
| 77 | + device=device, |
| 78 | + ), |
| 79 | + ) |
| 80 | + |
| 81 | + def forward(self, input: torch.Tensor) -> torch.Tensor: |
| 82 | + if self.padding: |
| 83 | + input = F.pad(input, pad=(0, self.in_features - self.origin_in_features)) |
| 84 | + # The forward method is replaced. In the original implementation, the forward |
| 85 | + # method is torchao.quantization.GPTQ.linear_forward_int4; here a Vulkan custom |
| 86 | + # operator is called instead. |
| 87 | + return torch.ops.et_vk.linear_weight_int4( |
| 88 | + input, |
| 89 | + self.weight, |
| 90 | + self.groupsize, |
| 91 | + self.scales_and_zeros, |
| 92 | + self.inner_k_tiles, |
| 93 | + ) |
| 94 | + |
| 95 | + |
| 96 | +# This function is coped from torchao.quantization.GPTQ._replace_linear_int4 |
| 97 | +# with small changes at the annotated locations. |
| 98 | +def _vk_replace_linear_int4( |
| 99 | + module: torch.nn.Module, |
| 100 | + groupsize: int, |
| 101 | + inner_k_tiles: Optional[int], |
| 102 | + padding_allowed: bool, |
| 103 | + skip_layer_func: Optional[Callable] = None, |
| 104 | + precision: torch.dtype = torch.bfloat16, |
| 105 | + scales_precision: torch.dtype = torch.bfloat16, |
| 106 | + # Use custom vulkan linear layer as default |
| 107 | + linear_class: Type[torch.nn.Module] = VkWeightOnlyInt4Linear, |
| 108 | + copy_weights: bool = False, |
| 109 | + # Serves the same purpose as `tensor_dim_limit` in |
| 110 | + # executorch.backends.vulkan.partitioner.VulkanSupportedOperators |
| 111 | + feature_limit: int = 16384, |
| 112 | +): |
| 113 | + for name, child in module.named_children(): |
| 114 | + if isinstance(child, torch.nn.Linear) and ( |
| 115 | + skip_layer_func is None or not skip_layer_func(child.weight) |
| 116 | + ): |
| 117 | + # Add an additional condition that the out/in features must not exceed the |
| 118 | + # `feature_limit` argument. |
| 119 | + if ( |
| 120 | + _check_linear_int4_k(child.in_features, groupsize, inner_k_tiles) |
| 121 | + or padding_allowed |
| 122 | + ) and ( |
| 123 | + child.out_features < feature_limit and child.in_features < feature_limit |
| 124 | + ): |
| 125 | + new_linear = linear_class( |
| 126 | + child.in_features, |
| 127 | + child.out_features, |
| 128 | + bias=False, |
| 129 | + device=child.weight.device, |
| 130 | + groupsize=groupsize, |
| 131 | + inner_k_tiles=inner_k_tiles, |
| 132 | + precision=precision, |
| 133 | + scales_precision=scales_precision, |
| 134 | + ) |
| 135 | + if copy_weights and child.weight.device != torch.device("meta"): |
| 136 | + new_linear.weight = child.weight |
| 137 | + setattr(module, name, new_linear) |
| 138 | + else: |
| 139 | + _vk_replace_linear_int4( |
| 140 | + child, |
| 141 | + groupsize, |
| 142 | + inner_k_tiles, |
| 143 | + padding_allowed, |
| 144 | + skip_layer_func, |
| 145 | + precision, |
| 146 | + scales_precision, |
| 147 | + linear_class, |
| 148 | + copy_weights, |
| 149 | + ) |
| 150 | + |
| 151 | + |
| 152 | +# This module is copied from torchao.quantization.GPTQ.Int4WeightOnlyQuantizer |
| 153 | +# with some changes at the annotated lines. |
| 154 | +class VkInt4WeightOnlyQuantizer(Quantizer): |
| 155 | + def __init__( |
| 156 | + self, |
| 157 | + groupsize: int = 256, |
| 158 | + padding_allowed: bool = True, |
| 159 | + inner_k_tiles: Optional[int] = 8, |
| 160 | + device: torch.device = torch.device("cpu"), # noqa |
| 161 | + precision: torch.dtype = torch.float32, |
| 162 | + feature_limit: int = 16384, |
| 163 | + ) -> None: |
| 164 | + super().__init__() |
| 165 | + assert inner_k_tiles in [2, 4, 8] |
| 166 | + assert groupsize in [32, 64, 128, 256] |
| 167 | + |
| 168 | + self.inner_k_tiles = inner_k_tiles |
| 169 | + self.groupsize: int = groupsize |
| 170 | + self.padding_allowed: bool = padding_allowed |
| 171 | + self.device: torch.device = device |
| 172 | + self.precision: torch.dtype = precision |
| 173 | + # Serves the same purpose as `tensor_dim_limit` in |
| 174 | + # executorch.backends.vulkan.partitioner.VulkanSupportedOperators |
| 175 | + self.feature_limit = feature_limit |
| 176 | + |
| 177 | + @torch.no_grad() |
| 178 | + def _create_quantized_state_dict( |
| 179 | + self, model: torch.nn.Module |
| 180 | + ) -> Dict[str, torch.Tensor]: |
| 181 | + cur_state_dict = model.state_dict() |
| 182 | + for fqn, mod in model.named_modules(): |
| 183 | + # Add additional check to make sure features do not exceed feature limit |
| 184 | + if isinstance(mod, torch.nn.Linear) and ( |
| 185 | + mod.out_features < self.feature_limit |
| 186 | + and mod.in_features < self.feature_limit |
| 187 | + ): |
| 188 | + assert not mod.bias |
| 189 | + out_features = mod.out_features |
| 190 | + in_features = mod.in_features |
| 191 | + logging.info(f"linear: {fqn}, in={in_features}, out={out_features}") |
| 192 | + |
| 193 | + assert ( |
| 194 | + in_features % self.groupsize == 0 |
| 195 | + ), f"require in_features:{in_features} % self.groupsize:{self.groupsize} == 0" |
| 196 | + |
| 197 | + weight = mod.weight.data |
| 198 | + if not _check_linear_int4_k( |
| 199 | + in_features, self.groupsize, self.inner_k_tiles |
| 200 | + ): |
| 201 | + if self.padding_allowed: |
| 202 | + import torch.nn.functional as F |
| 203 | + |
| 204 | + from torchao.quantization.utils import find_multiple |
| 205 | + |
| 206 | + logging.warn( |
| 207 | + f"warning: {fqn} is padded to satisfy in_features % 1024 == 0" |
| 208 | + ) |
| 209 | + padded_in_features = find_multiple(in_features, (1024,)) |
| 210 | + weight = F.pad( |
| 211 | + weight, pad=(0, padded_in_features - in_features) |
| 212 | + ) |
| 213 | + else: |
| 214 | + logging.warn( |
| 215 | + f"warning: {fqn} is skipped, int4 requires that in_features is 32, 64, or is divisible by 1024, " |
| 216 | + + "and that groupsize and inner_k_tiles*16 evenly divide into it" |
| 217 | + ) |
| 218 | + continue |
| 219 | + (w_int4x8, scales_and_zeros) = groupwise_affine_quantize_tensor( |
| 220 | + weight, |
| 221 | + 4, # n_bit |
| 222 | + self.groupsize, |
| 223 | + self.precision, # dtype for scales_and_zeros |
| 224 | + ) |
| 225 | + # In the original implementation, w_int4x8 is packed via calling the |
| 226 | + # _convert_weight_to_int4pack operator before storing the weight. However |
| 227 | + # the Vulkan implementation does not expect the weights to be packed, so |
| 228 | + # the w_int4x8 tensor is stored as the weight instead. |
| 229 | + cur_state_dict[f"{fqn}.weight"] = w_int4x8.to(self.device) |
| 230 | + cur_state_dict[f"{fqn}.scales_and_zeros"] = scales_and_zeros.to( |
| 231 | + self.device |
| 232 | + ) |
| 233 | + return cur_state_dict |
| 234 | + |
| 235 | + def _convert_for_runtime(self, model: torch.nn.Module) -> torch.nn.Module: |
| 236 | + _vk_replace_linear_int4( |
| 237 | + model, |
| 238 | + self.groupsize, |
| 239 | + self.inner_k_tiles, |
| 240 | + self.padding_allowed, |
| 241 | + skip_layer_func=None, |
| 242 | + precision=self.precision, |
| 243 | + scales_precision=self.precision, |
| 244 | + ) |
| 245 | + return model |
| 246 | + |
| 247 | + def quantize( |
| 248 | + self, model: torch.nn.Module, *args: Any, **kwargs: Any |
| 249 | + ) -> torch.nn.Module: |
| 250 | + state_dict = self._create_quantized_state_dict(model) |
| 251 | + model = self._convert_for_runtime(model) |
| 252 | + model.load_state_dict(state_dict, strict=False) |
| 253 | + return model |
0 commit comments