11import torch
22import torchao
3+ from torchao .quantization .quant_primitives import (
4+ MappingType ,
5+ ZeroPointDomain ,
6+ )
37from .subclass import ( # noqa
48 Int8DynamicallyQuantizedLinearWeight ,
59 Int8WeightOnlyQuantizedLinearWeight ,
610 QuantizedLinearWeightBase ,
711)
12+ from torchao .dtypes import AffineQuantizedTensor , PlainLayoutType
13+ from torchao .quantization .linear_activation_quantized_tensor import LinearActivationQuantizedTensor
814from torch .utils ._python_dispatch import return_and_correct_aliasing
915from .quant_primitives import (
1016 safe_int_mm ,
@@ -252,9 +258,9 @@ class AQMixin():
252258 def _autoquant_test (cls , act_mat , weight , bias , best_time , mode = ["relu" , None ]):
253259 w_qtensor = cls .from_float (weight )
254260 if _is_interpolate_mode (mode ):
255- q_c_op = torch .compile (cls ._quantized_op , mode = "max-autotune-no-cudagraphs" )
261+ q_c_op = torch .compile (cls ._quantized_linear_op , mode = "max-autotune-no-cudagraphs" )
256262 else :
257- func = lambda a ,b ,c : F .relu (cls ._quantized_op (F .relu (a ), b , c ))
263+ func = lambda a ,b ,c : F .relu (cls ._quantized_linear_op (F .relu (a ), b , c ))
258264 q_c_op = torch .compile (func , mode = "max-autotune-no-cudagraphs" )
259265 res = do_autoquant_bench (q_c_op , act_mat , w_qtensor , bias , warmup = 25 , rep = 100 )
260266 if res < best_time * 1.1 :
@@ -263,10 +269,53 @@ def _autoquant_test(cls, act_mat, weight, bias, best_time, mode=["relu", None]):
263269 print (f">>time: { res :0.3f} { cls } { best_time :0.3f} )
264270 return res
265271
266- class AQInt8DynamicallyQuantizedLinearWeight (AQMixin , Int8DynamicallyQuantizedLinearWeight ):
272+ ###### TODO !!!!!!!!!!!!!!!
273+ # 1) make class method from_float (just duplicate code)
274+ # 2) undo changes to quant_api?
275+ # 3) point to new quantized_op location
276+ # 4) rewrite the dynamic autoquant test
277+
278+ class AQInt8DynamicallyQuantizedLinearWeight (AQMixin , LinearActivationQuantizedTensor ):
267279 """
268280 AutoQuantizable version of Int8DynamicallyQuantizedLinearWeight
269281 """
282+ @classmethod
283+ def from_float (cls , weight ):
284+ in_features = weight .shape [1 ]
285+ # int8 dynamic quantization only has benefit when in_feature > 16
286+ # if in_features <= 16:
287+ # return weight
288+
289+ # avoid circular dep
290+ from torchao .dtypes import to_affine_quantized
291+ # weight settings
292+ mapping_type = MappingType .SYMMETRIC
293+ def get_weight_block_size (x ):
294+ return (1 , x .shape [1 ])
295+ target_dtype = torch .int8
296+ eps = torch .finfo (torch .float32 ).eps
297+ zero_point_dtype = torch .int64
298+
299+ # input settings
300+ def get_per_token_block_size (x ):
301+ block_size = list (x .shape )
302+ for i in range (len (block_size )- 1 ):
303+ block_size [i ] = 1
304+ return block_size
305+
306+ input_mapping_type = MappingType .SYMMETRIC
307+ input_target_dtype = torch .int8
308+ input_eps = 1e-5
309+ input_quant_min = - 127
310+ input_quant_max = 127
311+ layout_type = PlainLayoutType ()
312+ input_quant_func = lambda x : to_affine_quantized (x , input_mapping_type , get_per_token_block_size (x ), input_target_dtype , eps = input_eps , quant_min = input_quant_min , quant_max = input_quant_max , scale_dtype = torch .float32 if x .dtype == torch .float16 else None )
313+
314+ block_size = get_weight_block_size (weight )
315+ weight = to_affine_quantized (weight , mapping_type , block_size , target_dtype , eps = eps , zero_point_dtype = zero_point_dtype , layout_type = layout_type )
316+ weight = super (AQInt8DynamicallyQuantizedLinearWeight , cls ).from_float (weight , input_quant_func )
317+ return weight
318+
270319 @classmethod
271320 def _autoquant_test (cls , act_mat , weight , bias , best_time , mode = ["relu" , None ]):
272321 """
@@ -298,12 +347,13 @@ def _autoquant_test(cls, act_mat, weight, bias, best_time, mode=["relu", None]):
298347 )
299348 q_c_matmul = torch .compile (quantized_matmul , mode = "max-autotune-no-cudagraphs" )
300349 with torch .no_grad ():
301- res_matmul = do_autoquant_bench (q_c_matmul , x_vals_int8 , x_scales .reshape (- 1 ,1 ), w_qtensor .int_data )
350+ w_vals_int8 = w_qtensor .original_weight_tensor .layout_tensor .int_data .contiguous ().t ()
351+ res_matmul = do_autoquant_bench (q_c_matmul , x_vals_int8 , x_scales .reshape (- 1 ,1 ), w_vals_int8 )
302352 print (f">>time: { res_matmul :0.3f} { cls } { best_time :0.3f} )
303353
304354 # if the (much faster) matmul kernel is already beat, don't bother benchmarking full op
305- if res_matmul >= best_time :
306- return res_matmul
355+ # if res_matmul>=best_time:
356+ # return res_matmul
307357
308358 # calculate what time full op needs to beat for dynamic quant to be best given INTERPOLATION_CONSTANT
309359 to_beat = best_time + INTERPOLATION_CONSTANT / (1 - INTERPOLATION_CONSTANT )* (best_time - res_matmul )
@@ -313,18 +363,27 @@ def _autoquant_test(cls, act_mat, weight, bias, best_time, mode=["relu", None]):
313363 print (f">>time: { res_f :0.3f} { cls } { max_int_const_win :0.2f} )
314364 return res_f
315365
316- class AQWeightOnlyQuantizedLinearWeight (Int8WeightOnlyQuantizedLinearWeight , AQMixin ):
366+ class AQWeightOnlyQuantizedLinearWeight (AffineQuantizedTensor , AQMixin ):
317367 """
318368 AutoQuantizable version of Int8WeightOnlyQuantizedLinearWeight
319369 """
370+ @classmethod
371+ def from_float (cls , weight ):
372+ mapping_type = MappingType .SYMMETRIC
373+ target_dtype = torch .int8
374+ eps = torch .finfo (torch .float32 ).eps
375+ zero_point_dtype = torch .int64
376+ block_size = (1 , weight .shape [1 ])
377+ return super (AQWeightOnlyQuantizedLinearWeight , cls ).from_float (weight , mapping_type , block_size , target_dtype , eps = eps , zero_point_dtype = zero_point_dtype )
378+
320379
321- class AQWeightOnlyQuantizedLinearWeight2 (Int8WeightOnlyQuantizedLinearWeight , AQMixin ):
380+ class AQWeightOnlyQuantizedLinearWeight2 (AQWeightOnlyQuantizedLinearWeight , AQMixin ):
322381 """
323382 AutoQuantizable version of Int8WeightOnlyQuantizedLinearWeight that
324383 uses a different kernel
325384 """
326385 @staticmethod
327- def _quantized_op (act_mat , w_qtensor , bias ):
386+ def _quantized_linear_op (act_mat , w_qtensor , bias ):
328387 """
329388 Performs the quantized linear operations
330389
@@ -339,8 +398,8 @@ def _quantized_op(act_mat, w_qtensor, bias):
339398 orig_dtype = act_mat .dtype
340399 orig_shape = act_mat .shape
341400 act_mat = act_mat .reshape (- 1 , act_mat .shape [- 1 ], 1 )
342- y = (act_mat * w_qtensor .int_data .unsqueeze (0 )).sum (dim = - 2 )
343- y = y .reshape (* orig_shape [:- 1 ], y .shape [- 1 ]) * w_qtensor .q_scales
401+ y = (act_mat * w_qtensor .layout_tensor . int_data . t () .unsqueeze (0 )).sum (dim = - 2 )
402+ y = y .reshape (* orig_shape [:- 1 ], y .shape [- 1 ]) * w_qtensor .layout_tensor . scale
344403 if bias is not None :
345404 y += bias
346405 return y .to (orig_dtype )
@@ -352,14 +411,14 @@ def _autoquant_test(cls, act_mat, *args):
352411 return torch .inf
353412 return super ()._autoquant_test (act_mat , * args )
354413
355- class AQWeightOnlyQuantizedLinearWeight3 (Int8WeightOnlyQuantizedLinearWeight , AQMixin ):
414+ class AQWeightOnlyQuantizedLinearWeight3 (AQWeightOnlyQuantizedLinearWeight , AQMixin ):
356415 """
357416 AutoQuantizable version of Int8WeightOnlyQuantizedLinearWeight that
358417 uses a different kernel
359418 """
360- def _quantized_op (act_mat , w_qtensor , bias ):
419+ def _quantized_linear_op (act_mat , w_qtensor , bias ):
361420 orig_shape = act_mat .shape
362- y = torch .mm (act_mat .reshape (- 1 , orig_shape [- 1 ]), w_qtensor .int_data * w_qtensor .q_scales )
421+ y = torch .mm (act_mat .reshape (- 1 , orig_shape [- 1 ]), w_qtensor .layout_tensor . int_data . t () * w_qtensor .layout_tensor . scale )
363422 y = y .reshape (* orig_shape [:- 1 ], y .shape [- 1 ])
364423 if bias is not None :
365424 y += bias
@@ -377,7 +436,7 @@ def __init__(self):
377436 super ().__init__ ()
378437
379438 @staticmethod
380- def _quantized_op (act_mat , w_qtensor , bias ):
439+ def _quantized_linear_op (act_mat , w_qtensor , bias ):
381440 return torch .nn .functional .linear (act_mat , w_qtensor , bias )
382441
383442 @classmethod
@@ -389,7 +448,7 @@ def from_float(cls, weight):
389448 AQWeightOnlyQuantizedLinearWeight ,
390449 AQWeightOnlyQuantizedLinearWeight2 ,
391450 # AQWeightOnlyQuantizedLinearWeight3,
392- # TODO this gets picked in places where it makes perf worse, why?
451+ # # TODO this gets picked in places where it makes perf worse, why?
393452 AQInt8DynamicallyQuantizedLinearWeight ,
394453]
395454
0 commit comments