Replace fp16 linear with quantized linear
(model, weight_bit_width, backend="torch")
| 194 | |
| 195 | |
| 196 | def quantize(model, weight_bit_width, backend="torch"): |
| 197 | """Replace fp16 linear with quantized linear""" |
| 198 | |
| 199 | for i in range(len(model.language_model.transformer.layers) + 1): |
| 200 | if i == len(model.language_model.transformer.layers): |
| 201 | layer = model.language_model.transformer.topQueryLayer |
| 202 | else: |
| 203 | layer = model.language_model.transformer.layers[i] |
| 204 | |
| 205 | if backend == "torch": |
| 206 | layer.attention.query = QuantizedLinear( |
| 207 | in_features=layer.attention.query.in_features, |
| 208 | out_features=layer.attention.query.out_features, |
| 209 | weight_bit_width=weight_bit_width, |
| 210 | weight=layer.attention.query.weight.to(torch.cuda.current_device()), |
| 211 | bias=layer.attention.query.bias.to(torch.cuda.current_device()), |
| 212 | params_dtype=torch.half, |
| 213 | device=layer.attention.query.weight.device, |
| 214 | ) |
| 215 | layer.attention.value = QuantizedLinear( |
| 216 | in_features=layer.attention.value.in_features, |
| 217 | out_features=layer.attention.value.out_features, |
| 218 | weight_bit_width=weight_bit_width, |
| 219 | weight=layer.attention.value.weight.to(torch.cuda.current_device()), |
| 220 | bias=layer.attention.value.bias.to(torch.cuda.current_device()), |
| 221 | params_dtype=torch.half, |
| 222 | device=layer.attention.value.weight.device, |
| 223 | ) |
| 224 | layer.attention.key = QuantizedLinear( |
| 225 | in_features=layer.attention.key.in_features, |
| 226 | out_features=layer.attention.key.out_features, |
| 227 | weight_bit_width=weight_bit_width, |
| 228 | weight=layer.attention.key.weight.to(torch.cuda.current_device()), |
| 229 | bias=layer.attention.key.bias.to(torch.cuda.current_device()), |
| 230 | params_dtype=torch.half, |
| 231 | device=layer.attention.key.weight.device, |
| 232 | ) |
| 233 | layer.attention.dense = QuantizedLinear( |
| 234 | in_features=layer.attention.dense.in_features, |
| 235 | out_features=layer.attention.dense.out_features, |
| 236 | weight_bit_width=weight_bit_width, |
| 237 | weight=layer.attention.dense.weight.to(torch.cuda.current_device()), |
| 238 | bias=layer.attention.dense.bias.to(torch.cuda.current_device()), |
| 239 | params_dtype=torch.half, |
| 240 | device=layer.attention.dense.weight.device, |
| 241 | ) |
| 242 | layer.mlp.dense_h_to_4h = QuantizedLinear( |
| 243 | in_features=layer.mlp.dense_h_to_4h.in_features, |
| 244 | out_features=layer.mlp.dense_h_to_4h.out_features, |
| 245 | weight_bit_width=weight_bit_width, |
| 246 | weight=layer.mlp.dense_h_to_4h.weight.to(torch.cuda.current_device()), |
| 247 | bias=layer.mlp.dense_h_to_4h.bias.to(torch.cuda.current_device()), |
| 248 | params_dtype=torch.half, |
| 249 | device=layer.mlp.dense_h_to_4h.weight.device, |
| 250 | ) |
| 251 | layer.mlp.dense_4h_to_h = QuantizedLinear( |
| 252 | in_features=layer.mlp.dense_4h_to_h.in_features, |
| 253 | out_features=layer.mlp.dense_4h_to_h.out_features, |