File size: 7,663 Bytes
2528275 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 |
import torch
class TruncateFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, input, threshold):
ctx.save_for_backward(input)
ctx.threshold = threshold
truncated_tensor = input.clone()
truncated_tensor = torch.where(truncated_tensor.abs() < threshold, truncated_tensor.sign() * threshold, truncated_tensor)
return truncated_tensor
@staticmethod
def backward(ctx, grad_output):
input, = ctx.saved_tensors
threshold = ctx.threshold
grad_input = grad_output.clone()
grad_input[input.abs() < threshold] = 0
return grad_input, None
def truncate_number(number, threshold=1e-2):
# avoid overflow with AMP training
return TruncateFunction.apply(number, threshold)
def smooth_ln_fcs_temporary(ln, fcs, scales,shifts,weight_in_cpu=False):
ln.use_temporary_parameter = True
if weight_in_cpu is True:
ln.weight = ln.weight.to(scales.device)
for fc in fcs:
fc.weight = fc.weight.to(scales.device)
if not isinstance(fcs, list):
fcs = [fcs]
if hasattr(ln, 'bias') and ln.bias is not None:
ln.temp_bias = (ln.bias - shifts) / scales
else:
ln.temp_bias = (-1*shifts) / scales
ln.temp_weight = ln.weight / scales
for fc in fcs:
fc.use_temporary_parameter = True
if hasattr(fc, 'bias') and fc.bias is not None:
fc.temp_bias = fc.bias + torch.matmul(fc.weight,shifts)
else:
fc.temp_bias = torch.matmul(fc.weight,shifts)
fc.temp_weight = fc.weight * scales.view(1,-1)
if weight_in_cpu is True:
ln.weight = ln.weight.cpu()
for fc in fcs:
fc.weight = fc.weight.cpu()
def smooth_ln_fcs_inplace(ln, fcs, scales,shifts):
ln.use_temporary_parameter = False
if not isinstance(fcs, list):
fcs = [fcs]
if hasattr(ln, 'bias') and ln.bias is not None:
ln.bias.sub_(shifts)
ln.bias.div_(scales)
else:
del ln.bias
ln.register_buffer('bias',(-1*shifts)/scales)
ln.weight.div_(scales)
for fc in fcs:
fc.use_temporary_parameter = False
if hasattr(fc, 'bias') and fc.bias is not None:
fc.bias.add_(fc.weight@shifts)
else:
del fc.bias
# import ipdb;ipdb.set_trace()
fc.register_buffer('bias',fc.weight@shifts)
fc.weight.mul_(scales.view(1,-1))
def smooth_fc_fc_temporary(fc1, fc2, scales,shifts=None,num_key_value_groups=1,head_dim=128,weight_in_cpu=False,args=None):
# only support for v_proj and out_proh now.
fc1.use_temporary_parameter = True
fc2.use_temporary_parameter = True
if weight_in_cpu is True:
fc1.weight = fc1.weight.to(scales.device)
fc2.weight = fc2.weight.to(scales.device)
# import ipdb;ipdb.set_trace()
if num_key_value_groups > 1:
# import ipdb;ipdb.set_trace()
if args.gqa_scales == "copy":
kv_scales = scales
kv_shift = shifts
scales = scales.view(-1,head_dim).repeat_interleave(num_key_value_groups,dim=0).view(-1)
shifts = shifts.view(-1,head_dim).repeat_interleave(num_key_value_groups,dim=0).view(-1)
elif args.gqa_scales == "mean":
kv_scales = scales.view(-1,num_key_value_groups,head_dim).mean(dim=1).view(-1)
kv_shift = shifts.view(-1,num_key_value_groups,head_dim).mean(dim=1).view(-1)
else:
raise NotImplementedError("Only implemented copy and mean for gqa")
else:
kv_scales = scales
kv_shift = shifts
if hasattr(fc1, 'temp_weight'):
fc1.temp_bias = fc1.temp_bias - kv_shift
fc1.temp_bias = fc1.temp_bias/kv_scales.view(-1)
fc1.temp_weight = fc1.temp_weight/kv_scales.view(-1,1)
else:
fc1.temp_bias = fc1.bias/kv_scales.view(-1)
fc1.temp_weight = fc1.weight/kv_scales.view(-1,1)
if hasattr(fc2, 'bias') and fc2.bias is not None:
fc2.temp_bias = fc2.bias + fc2.weight@shifts
else:
fc2.temp_bias = fc2.weight@shifts
fc2.temp_weight = fc2.weight * scales.view(1,-1)
if weight_in_cpu is True:
fc1.weight = fc1.weight.cpu()
fc2.weight = fc2.weight.cpu()
def smooth_fc_fc_inplace(fc1, fc2, scales,shifts=None,num_key_value_groups=1,head_dim=128,args=None):
# only support for v_proj and out_proh now.
fc1.use_temporary_parameter = False
fc2.use_temporary_parameter = False
if num_key_value_groups > 1:
# import ipdb;ipdb.set_trace()
if args.gqa_scales == "copy":
kv_scales = scales
kv_shift = shifts
scales = scales.view(-1,head_dim).repeat_interleave(num_key_value_groups,dim=0).view(-1)
shifts = shifts.view(-1,head_dim).repeat_interleave(num_key_value_groups,dim=0).view(-1)
elif args.gqa_scales == "mean":
kv_scales = scales.view(-1,num_key_value_groups,head_dim).mean(dim=1).view(-1)
kv_shift = shifts.view(-1,num_key_value_groups,head_dim).mean(dim=1).view(-1)
else:
raise NotImplementedError("Only implemented copy and mean for gqa")
else:
kv_scales = scales
kv_shift = shifts
fc1.bias.sub_(kv_shift)
fc1.bias.div_(kv_scales.view(-1))
fc1.weight.div_(kv_scales.view(-1,1))
if hasattr(fc2, 'bias') and fc2.bias is not None:
fc2.bias.add_(fc2.weight@shifts)
else:
del fc2.bias
fc2.register_buffer('bias',fc2.weight@shifts)
fc2.weight.mul_(scales.view(1,-1))
def smooth_q_k_temporary(q_proj, k_proj,scales,num_key_value_groups=1,head_dim=128,weight_in_cpu=False,args=None):
q_proj.use_temporary_parameter = True
k_proj.use_temporary_parameter = True
if weight_in_cpu is True:
q_proj.weight = q_proj.weight.to(scales.device)
k_proj.weight = k_proj.weight.to(scales.device)
if num_key_value_groups > 1:
# import ipdb;ipdb.set_trace()
if args.gqa_scales == "copy":
kv_scales = scales
scales = scales.view(-1,head_dim).repeat_interleave(num_key_value_groups,dim=0).view(-1)
elif args.gqa_scales == "mean":
kv_scales = scales.view(-1,num_key_value_groups,head_dim).mean(dim=1).view(-1)
else:
raise NotImplementedError("Only implemented copy and mean for gqa")
else:
kv_scales = scales
q_proj.temp_weight = q_proj.temp_weight/scales.view(-1,1)
q_proj.temp_bias = q_proj.temp_bias/scales.view(-1)
k_proj.temp_weight = k_proj.temp_weight*kv_scales.view(-1,1)
k_proj.temp_bias = k_proj.temp_bias*kv_scales.view(-1)
if weight_in_cpu is True:
q_proj.weight = q_proj.weight.cpu()
k_proj.weight = k_proj.weight.cpu()
def smooth_q_k_inplace(q_proj, k_proj, scales,num_key_value_groups=1,head_dim=128,args=None):
q_proj.use_temporary_parameter = False
k_proj.use_temporary_parameter = False
if num_key_value_groups > 1:
# import ipdb;ipdb.set_trace()
if args.gqa_scales == "copy":
kv_scales = scales
scales = scales.view(-1,head_dim).repeat_interleave(num_key_value_groups,dim=0).view(-1)
elif args.gqa_scales == "mean":
kv_scales = scales.view(-1,num_key_value_groups,head_dim).mean(dim=1).view(-1)
else:
raise NotImplementedError("Only implemented copy and mean for gqa")
else:
kv_scales = scales
q_proj.weight.div_(scales.view(-1,1))
q_proj.bias.div_(scales.view(-1))
k_proj.weight.mul_(kv_scales.view(-1,1))
k_proj.bias.mul_(kv_scales.view(-1)) |