2121from .metrics import mse_loss
2222from paddle .distributed .fleet .meta_parallel import (
2323 ColumnParallelLinear ,
24- RowParallelLinear ,
25- )
24+ RowParallelLinear , )
2625__all__ = ['AutoClip' ]
2726
27+
2828class AutoClip (nn .Layer ):
2929 """
3030 AutoClip from AWQ[https://arxiv.org/abs/2306.00978]
3131 """
32+
3233 def __init__ (
3334 self ,
3435 model ,
@@ -39,8 +40,7 @@ def __init__(
3940 n_grid = 20 ,
4041 max_shrink = 0.5 ,
4142 n_sample_token = 512 ,
42- group_size = 128 ,
43- ):
43+ group_size = 128 , ):
4444 super (AutoClip , self ).__init__ ()
4545 self .model = model
4646 self .weight_bits = weight_bits
@@ -59,15 +59,17 @@ def __init__(
5959 def _apply_hook (self ):
6060 self ._forward_hook_list = []
6161 for _ , sub_layer in self .model .named_sublayers ():
62- if type (sub_layer ) in [ColumnParallelLinear , RowParallelLinear , paddle .nn .Linear ]:
62+ if type (sub_layer ) in [
63+ ColumnParallelLinear , RowParallelLinear , paddle .nn .Linear
64+ ]:
6365 forward_pre_hook_handle = sub_layer .register_forward_pre_hook (
6466 self ._forward_pre_hook )
6567 self ._forward_hook_list .append (forward_pre_hook_handle )
6668
6769 def _forward_pre_hook (self , layer , input ):
6870 self ._sample_scale (input , layer .full_name ())
6971 return input
70-
72+
7173 def _sample_scale (self , input , name ):
7274 input = input [0 ] if type (input ) == tuple else input
7375 input .stop_gradient = True
@@ -80,7 +82,6 @@ def _sample_scale(self, input, name):
8082 else :
8183 self .sampled_inputs [name ] = input
8284
83-
8485 def auto_clip (self , group_size = 128 , oc_batch_size = 256 ):
8586 """
8687 search clip scale for each layer and update the layer's weight
@@ -89,7 +90,7 @@ def auto_clip(self, group_size=128, oc_batch_size=256):
8990 name = sub_layer .full_name ()
9091 if name not in self .sampled_inputs or 'out_linear' in sub_name :
9192 continue
92-
93+
9394 weight = sub_layer .weight .cast ('float16' )
9495 weight_t = paddle .transpose (weight , perm = [1 , 0 ])
9596 x = self .sampled_inputs [name ].cast ('float16' )
@@ -98,33 +99,41 @@ def auto_clip(self, group_size=128, oc_batch_size=256):
9899 x = x .reshape ([1 , x .shape [0 ], - 1 , group_size ])
99100 x = x [:, 0 ::x .shape [1 ] // self .n_sample_token ]
100101 weight_t = weight_t .reshape ([weight_t .shape [0 ], 1 , - 1 , group_size ])
101- oc_batch_size = oc_batch_size if weight_t .shape [0 ] % oc_batch_size == 0 else 128 # prevent OOM
102+ oc_batch_size = oc_batch_size if weight_t .shape [
103+ 0 ] % oc_batch_size == 0 else 128 # prevent OOM
102104 assert weight_t .shape [0 ] % oc_batch_size == 0
103105
104106 w_all = weight_t
105107 best_max_val_all = []
106108
107109 for i_b in range (weight_t .shape [0 ] // oc_batch_size ):
108- w = w_all [i_b * oc_batch_size : (i_b + 1 ) * oc_batch_size ]
110+ w = w_all [i_b * oc_batch_size :(i_b + 1 ) * oc_batch_size ]
109111
110- org_max_val = w .abs ().max (axis = - 1 , keepdim = True ) # co, 1, n_group, 1
112+ org_max_val = w .abs ().max (
113+ axis = - 1 , keepdim = True ) # co, 1, n_group, 1
111114 best_max_val = org_max_val .clone ()
112115 min_errs = paddle .ones_like (org_max_val , dtype = 'float16' ) * 1e9
113116 org_out = (x * w ).sum (axis = - 1 ) # co, n_token, n_group
114117 for i_s in range (int (self .max_shrink * self .n_grid )):
115118 max_val = org_max_val * (1 - i_s / self .n_grid )
116119 max_val_tmp = max_val
117120 cur_w = paddle .where (w > max_val_tmp , max_val_tmp , w )
118- cur_w = paddle .where (cur_w < - max_val_tmp , - max_val_tmp , cur_w )
121+ cur_w = paddle .where (cur_w < - max_val_tmp , - max_val_tmp ,
122+ cur_w )
119123 org_w_shape = cur_w .shape
120- cur_w_r = cur_w .reshape ([- 1 , self .group_size ]).transpose ([1 , 0 ])
121- quant_dequant_weight = fake_quant (cur_w_r , method = 'abs_max_channel_wise' , weight_bits = 4 )
122- quant_dequant_weight = quant_dequant_weight .transpose ([1 , 0 ]).reshape (org_w_shape )
124+ cur_w_r = cur_w .reshape ([- 1 ,
125+ self .group_size ]).transpose ([1 , 0 ])
126+ quant_dequant_weight = fake_quant (
127+ cur_w_r , method = 'abs_max_channel_wise' , weight_bits = 4 )
128+ quant_dequant_weight = quant_dequant_weight .transpose (
129+ [1 , 0 ]).reshape (org_w_shape )
123130 cur_out = (x * quant_dequant_weight ).sum (axis = - 1 )
124131 # co, 1, n_group, 1
125132 tmp = (cur_out - org_out ).detach ().clone ()
126- err = paddle .pow (tmp , 2 ).mean (axis = 1 ).reshape (min_errs .shape )
127- print ('block {} search s {} err {}' .format (i_b , i_s , err .mean ().item ()))
133+ err = paddle .pow (tmp ,
134+ 2 ).mean (axis = 1 ).reshape (min_errs .shape )
135+ print ('block {} search s {} err {}' .format (
136+ i_b , i_s , err .mean ().item ()))
128137 del cur_w , cur_out , quant_dequant_weight , tmp , cur_w_r
129138 paddle .device .cuda .empty_cache ()
130139
@@ -143,16 +152,21 @@ def auto_clip(self, group_size=128, oc_batch_size=256):
143152 if 'w_0' in param .name :
144153 param_tmp = param .transpose (perm = [1 , 0 ]).cast ('float16' )
145154 tmp_shape = param_tmp .shape
146- param_tmp = param_tmp .reshape ([best_max_val .shape [0 ], best_max_val .shape [1 ], - 1 ])
147- best_max_val = paddle .tile (best_max_val , repeat_times = (1 , 1 , param_tmp .shape [- 1 ]))
148- param_tmp = paddle .where (param_tmp > best_max_val , best_max_val , param_tmp )
149- param_tmp = paddle .where (param_tmp < - best_max_val , - best_max_val , param_tmp )
155+ param_tmp = param_tmp .reshape (
156+ [best_max_val .shape [0 ], best_max_val .shape [1 ], - 1 ])
157+ best_max_val = paddle .tile (
158+ best_max_val , repeat_times = (1 , 1 , param_tmp .shape [- 1 ]))
159+ param_tmp = paddle .where (param_tmp > best_max_val ,
160+ best_max_val , param_tmp )
161+ param_tmp = paddle .where (param_tmp < - best_max_val ,
162+ - best_max_val , param_tmp )
150163 param_tmp = param_tmp .reshape (tmp_shape ).cast (param .dtype )
151164 param_tmp = param_tmp .transpose (perm = [1 , 0 ])
152165 paddle .assign (param_tmp , output = param )
153166 del param_tmp
154167 paddle .device .cuda .empty_cache ()
155168 break
156169
157- del best_max_val , weight_t , x , weight , self .sampled_inputs [name ], w_all , best_max_val_all
170+ del best_max_val , weight_t , x , weight , self .sampled_inputs [
171+ name ], w_all , best_max_val_all
158172 paddle .device .cuda .empty_cache ()
0 commit comments