|
| 1 | +import torch |
| 2 | + |
| 3 | +from .subclass import ( # noqa |
| 4 | + Int8DynamicallyQuantizedLinearWeight, |
| 5 | + Int8WeightOnlyQuantizedLinearWeight, |
| 6 | + QuantizedLinearWeightBase, |
| 7 | +) |
| 8 | +from torch.utils._python_dispatch import return_and_correct_aliasing |
| 9 | +from .utils import benchmark |
| 10 | + |
| 11 | +aten = torch.ops.aten |
| 12 | + |
| 13 | +AUTOQUANT_CACHE = {} |
| 14 | + |
| 15 | +def check_cache(cls, shape, dtype): |
| 16 | + return AUTOQUANT_CACHE.get((cls, shape, dtype), None) |
| 17 | + |
| 18 | +def update_cache(cls, shape, dtype, res): |
| 19 | + AUTOQUANT_CACHE[(cls, shape, dtype)] = res |
| 20 | + |
| 21 | +class AutoQuantizableLinearWeight(torch.Tensor): |
| 22 | + """ |
| 23 | + when run, finds best type of quantization for this tensor and swaps itself with that |
| 24 | + """ |
| 25 | + @staticmethod |
| 26 | + def __new__(cls, weight, qtensor_class_list, *args, **kwargs): |
| 27 | + kwargs["device"] = weight.device |
| 28 | + kwargs["layout"] = ( |
| 29 | + kwargs.get("layout") if kwargs.get("layout", False) else weight.layout |
| 30 | + ) |
| 31 | + kwargs["dtype"] = ( |
| 32 | + kwargs.get("dtype") if kwargs.get("dtype", False) else weight.dtype |
| 33 | + ) |
| 34 | + kwargs["requires_grad"] = False |
| 35 | + shape = kwargs.pop("shape", weight.shape) |
| 36 | + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) # type: ignore[attr-defined] |
| 37 | + |
| 38 | + def __init__(self, weight, qtensor_class_list, *args, **kwargs): |
| 39 | + self.weight = weight |
| 40 | + self.qtensor_class_list = qtensor_class_list |
| 41 | + self.logged_shape = None |
| 42 | + self.logged_dtype = None |
| 43 | + |
| 44 | + def __repr__(self): |
| 45 | + return ( |
| 46 | + f"{self.__class__.__name__}(data={self.weight}, shape={self.shape}, " |
| 47 | + f"device={self.device}, dtype={self.dtype}, qtensor_class_list={self.qtensor_class_list})" |
| 48 | + ) |
| 49 | + |
| 50 | + @staticmethod |
| 51 | + def log_shape(act_mat, w_autoquant, bias): |
| 52 | + orig_shape = act_mat.shape |
| 53 | + act_mat = act_mat.reshape(-1, act_mat.shape[-1]) |
| 54 | + logged_shape = (act_mat.shape, w_autoquant.shape, None if bias is None else bias.shape) |
| 55 | + logged_dtype = act_mat.dtype |
| 56 | + w_autoquant.logged_shape = logged_shape |
| 57 | + w_autoquant.logged_dtype = logged_dtype |
| 58 | + for q_cls in w_autoquant.qtensor_class_list: |
| 59 | + if check_cache(q_cls, logged_shape, logged_dtype) is None: |
| 60 | + update_cache(q_cls, logged_shape, logged_dtype, None) |
| 61 | + y = torch.mm(act_mat, w_autoquant.weight.t()) |
| 62 | + y = y.reshape(*orig_shape[:-1], y.shape[-1]) |
| 63 | + if bias is not None: |
| 64 | + y += bias |
| 65 | + return y |
| 66 | + |
| 67 | + def tune_autoquant(self, q_cls): |
| 68 | + act_shape, w_shape, bias_shape = self.logged_shape |
| 69 | + if check_cache(q_cls, self.logged_shape, self.logged_dtype) is None: |
| 70 | + with torch.no_grad(): |
| 71 | + act_mat = torch.randn(act_shape, dtype=self.logged_dtype, device=self.device) |
| 72 | + bias = None if bias_shape is None else torch.randn(bias_shape, dtype=self.logged_dtype, device=self.device) |
| 73 | + print(q_cls, self.logged_shape, self.logged_dtype) |
| 74 | + print("mem", torch.cuda.max_memory_allocated()/1e6, torch.cuda.memory_usage()) |
| 75 | + res = q_cls._autoquant_test(act_mat, self.weight, bias) |
| 76 | + update_cache(q_cls, self.logged_shape, self.logged_dtype, res) |
| 77 | + |
| 78 | + def to_quantized(self): |
| 79 | + if self.logged_shape is None or self.logged_dtype is None: |
| 80 | + raise RuntimeError("must run module normally to get shape, dtype info for autoquant") |
| 81 | + best_time = torch.inf |
| 82 | + best_cls = None |
| 83 | + for q_cls in self.qtensor_class_list: |
| 84 | + if check_cache(q_cls, self.logged_shape, self.logged_dtype) is None: |
| 85 | + self.tune_autoquant(q_cls) |
| 86 | + cls_res = AUTOQUANT_CACHE.get((q_cls, self.logged_shape, self.logged_dtype), torch.inf) |
| 87 | + if best_time >= cls_res: |
| 88 | + best_time = cls_res |
| 89 | + best_cls = q_cls |
| 90 | + # TODO handle random cls args/kwargs? or should they be curried |
| 91 | + self = best_cls.from_float(self.weight) |
| 92 | + return self |
| 93 | + |
| 94 | + def _apply_fn_to_data(self, fn): |
| 95 | + return self.__class__( |
| 96 | + fn(self.weight), self.qtensor_class_list, dtype=self.dtype |
| 97 | + ) |
| 98 | + |
| 99 | + def __tensor_flatten__(self): |
| 100 | + return ["weight"], [self.qtensor_class_list, self.dtype, self.shape] |
| 101 | + |
| 102 | + @classmethod |
| 103 | + def __tensor_unflatten__(cls, tensor_data_dict, tensor_attributes, outer_size=None, outer_stride=None): |
| 104 | + weight = tensor_data_dict["weight"] |
| 105 | + qtensor_class_list, dtype, shape = tensor_attributes[0] |
| 106 | + return cls(weight, qtensor_class_list, shape=shape if outer_size is None else outer_size, dtype=dtype, strides=outer_stride) |
| 107 | + |
| 108 | + @classmethod |
| 109 | + def from_float(cls, weight, qtensor_class_list): |
| 110 | + return cls(weight, qtensor_class_list) |
| 111 | + |
| 112 | + @classmethod |
| 113 | + def __torch_function__(cls, func, types, args=(), kwargs=None): |
| 114 | + kwargs = {} if kwargs is None else kwargs |
| 115 | + |
| 116 | + if func is torch.nn.functional.linear: |
| 117 | + mat1, w_autoquant, bias = ( |
| 118 | + args[0], |
| 119 | + args[1], |
| 120 | + args[2] if len(args)>2 else None |
| 121 | + ) |
| 122 | + return cls.log_shape(mat1, w_autoquant, bias) |
| 123 | + |
| 124 | + try: |
| 125 | + with torch._C.DisableTorchFunctionSubclass(): |
| 126 | + return func(*args, **kwargs) |
| 127 | + except: |
| 128 | + print(f"ERR: subclass doesn't implement {func}") |
| 129 | + |
| 130 | + @classmethod |
| 131 | + def __torch_dispatch__(cls, func, types, args, kwargs): |
| 132 | + if func is aten.detach.default: |
| 133 | + return return_and_correct_aliasing(func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)) |
| 134 | + |
| 135 | + |
| 136 | +class DefaultLinear(torch.Tensor): |
| 137 | + """ |
| 138 | + An class to be used in concert with AutoQuantizableLinearWeight to provide a |
| 139 | + default/non-quantized option. Only implements the bare minimum needed to work with the |
| 140 | + AutoQuantizableLinearWeight class using the same interfaces that would normally be |
| 141 | + used by QTensor subclasses but for a default linear op instead. |
| 142 | + """ |
| 143 | + def __init__(self): |
| 144 | + super().__init__() |
| 145 | + |
| 146 | + @classmethod |
| 147 | + def _autoquant_test(cls, act_mat, weight, bias): |
| 148 | + w_qtensor = cls.from_float(weight) |
| 149 | + q_c_op = torch.compile(cls._quantized_op, mode="max-autotune") |
| 150 | + with torch.no_grad(): |
| 151 | + res=benchmark(q_c_op, act_mat, w_qtensor, bias) |
| 152 | + print(cls, res) |
| 153 | + return res |
| 154 | + |
| 155 | + @staticmethod |
| 156 | + def _quantized_op(act_mat, w_qtensor, bias): |
| 157 | + return torch.nn.functional.linear(act_mat, w_qtensor, bias) |
| 158 | + |
| 159 | + @classmethod |
| 160 | + def from_float(cls, weight): |
| 161 | + return weight |
| 162 | + |
| 163 | +DEFAULT_CLASS_LIST = [ |
| 164 | + Int8DynamicallyQuantizedLinearWeight, |
| 165 | + DefaultLinear, |
| 166 | + Int8WeightOnlyQuantizedLinearWeight, |
| 167 | + |
| 168 | +] |
| 169 | + |
| 170 | +if False: |
| 171 | + # def _get_to_kwargs(self, *args, **kwargs): |
| 172 | + # device, dtype, _, memory_format = torch._C._nn._parse_to(*args, **kwargs) |
| 173 | + # device = self.device if device is None else device |
| 174 | + # dtype = self.dtype if dtype is None else dtype |
| 175 | + # memory_format = ( |
| 176 | + # memory_format if memory_format is not None else torch.preserve_format |
| 177 | + # ) |
| 178 | + # kwargs = { |
| 179 | + # "device": device, |
| 180 | + # "dtype": dtype, |
| 181 | + # "memory_format": memory_format, |
| 182 | + # } |
| 183 | + # return kwargs |
| 184 | + |
| 185 | + # def to(self, *args, **kwargs): |
| 186 | + # kwargs = self._get_to_kwargs(*args, **kwargs) |
| 187 | + # return self.__class__( |
| 188 | + # self.int_data.to(kwargs["device"]), |
| 189 | + # self.q_scales.to(kwargs["device"]), |
| 190 | + # self.transposed, |
| 191 | + # self.shape, |
| 192 | + # **kwargs, |
| 193 | + # ) |
| 194 | + |
| 195 | + # def _apply_fn_to_data(self, fn): |
| 196 | + # return self.__class__( |
| 197 | + # fn(self.int_data), fn(self.q_scales), self.transposed, self.shape, dtype=self.dtype |
| 198 | + # ) |
| 199 | + |
| 200 | + # def _change_shape(self, shape): |
| 201 | + # return self.__class__( |
| 202 | + # self.int_data, self.q_scales, self.transposed, shape, dtype=self.dtype |
| 203 | + # ) |
| 204 | + |
| 205 | + # def half(self): |
| 206 | + # return self.to(torch.float16) |
| 207 | + pass |
0 commit comments