1212
1313AUTOQUANT_CACHE = {}
1414
15- def check_cache (shape , cls ):
16- if shape in AUTOQUANT_CACHE :
17- return AUTOQUANT_CACHE [shape ].get (cls , None )
18- else :
19- return None
15+ def check_cache (cls , shape , dtype ):
16+ return AUTOQUANT_CACHE .get ((cls , shape , dtype ), None )
2017
21- def update_cache (shape , cls , res ):
22- if not shape in AUTOQUANT_CACHE :
23- AUTOQUANT_CACHE [shape ] = {}
24- AUTOQUANT_CACHE [shape ][cls ] = res
18+ def update_cache (cls , shape , dtype , res ):
19+ AUTOQUANT_CACHE [(cls , shape , dtype )] = res
2520
2621class AutoQuantizableLinearWeight (torch .Tensor ):
2722 """
@@ -43,7 +38,8 @@ def __new__(cls, weight, qtensor_class_list, *args, **kwargs):
4338 def __init__ (self , weight , qtensor_class_list , * args , ** kwargs ):
4439 self .weight = weight
4540 self .qtensor_class_list = qtensor_class_list
46- self .cache_shape = None
41+ self .logged_shape = None
42+ self .logged_dtype = None
4743
4844 def __repr__ (self ):
4945 return (
@@ -52,36 +48,46 @@ def __repr__(self):
5248 )
5349
5450 @staticmethod
55- def tune_autoquant (act_mat , w_autoquant , bias ):
51+ def log_shape (act_mat , w_autoquant , bias ):
5652 orig_shape = act_mat .shape
5753 act_mat = act_mat .reshape (- 1 , act_mat .shape [- 1 ])
58- cache_shape = (act_mat .shape , w_autoquant .shape , None if bias is None else bias .shape )
59- w_autoquant .cache_shape = cache_shape
60- for cur_cls in w_autoquant .qtensor_class_list :
61- if check_cache (cache_shape , cur_cls ) is None :
62- with torch .no_grad ():
63- print (cur_cls , cache_shape )
64- print (torch .cuda .max_memory_allocated ()/ 1e6 , torch .cuda .memory_usage ())
65- res = cur_cls ._autoquant_test (act_mat .clone (), w_autoquant .weight .clone (), None if bias is None else bias .clone ())
66- update_cache (cache_shape , cur_cls , res )
67- print (torch .cuda .max_memory_allocated ()/ 1e6 , torch .cuda .memory_usage ())
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 )
6861 y = torch .mm (act_mat , w_autoquant .weight .t ())
6962 y = y .reshape (* orig_shape [:- 1 ], y .shape [- 1 ])
7063 if bias is not None :
7164 y += bias
7265 return y
7366
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+
7478 def to_quantized (self ):
75- if self .cache_shape is None or self .cache_shape not in AUTOQUANT_CACHE :
76- raise RuntimeError ("must run module normally to find best quantization option " )
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 " )
7781 best_time = torch .inf
7882 best_cls = None
79- for cur_cls in self .qtensor_class_list :
80- cls_res = AUTOQUANT_CACHE [self .cache_shape ].get (cur_cls , torch .inf )
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 )
8187 if best_time >= cls_res :
8288 best_time = cls_res
83- best_cls = cur_cls
84- # need to handle random cls args/kwargs?
89+ best_cls = q_cls
90+ # TODO handle random cls args/kwargs? or should they be curried
8591 self = best_cls .from_float (self .weight )
8692 return self
8793
@@ -113,7 +119,7 @@ def __torch_function__(cls, func, types, args=(), kwargs=None):
113119 args [1 ],
114120 args [2 ] if len (args )> 2 else None
115121 )
116- return cls .tune_autoquant (mat1 , w_autoquant , bias )
122+ return cls .log_shape (mat1 , w_autoquant , bias )
117123
118124 try :
119125 with torch ._C .DisableTorchFunctionSubclass ():
@@ -155,9 +161,10 @@ def from_float(cls, weight):
155161 return weight
156162
157163DEFAULT_CLASS_LIST = [
164+ Int8DynamicallyQuantizedLinearWeight ,
158165 DefaultLinear ,
159166 Int8WeightOnlyQuantizedLinearWeight ,
160- Int8DynamicallyQuantizedLinearWeight ,
167+
161168]
162169
163170if False :
0 commit comments