Replace fp16 linear with quantized linear
(model, weight_bit_width)
| 101 | return output |
| 102 | |
| 103 | def quantize_oneflow(model, weight_bit_width): |
| 104 | """Replace fp16 linear with quantized linear""" |
| 105 | |
| 106 | for i in range(len(model.language_model.transformer.layers) + 1): |
| 107 | if i == len(model.language_model.transformer.layers): |
| 108 | layer = model.language_model.transformer.topQueryLayer |
| 109 | else: |
| 110 | layer = model.language_model.transformer.layers[i] |
| 111 | |
| 112 | layer.attention.query = QuantizedLinear( |
| 113 | in_features=layer.attention.query.in_features, |
| 114 | out_features=layer.attention.query.out_features, |
| 115 | weight_bit_width=weight_bit_width, |
| 116 | weight=layer.attention.query.weight.to(torch.cuda.current_device()), |
| 117 | bias=layer.attention.query.bias.to(torch.cuda.current_device()), |
| 118 | params_dtype=torch.half, |
| 119 | device=layer.attention.query.weight.device, |
| 120 | ) |
| 121 | layer.attention.value = QuantizedLinear( |
| 122 | in_features=layer.attention.value.in_features, |
| 123 | out_features=layer.attention.value.out_features, |
| 124 | weight_bit_width=weight_bit_width, |
| 125 | weight=layer.attention.value.weight.to(torch.cuda.current_device()), |
| 126 | bias=layer.attention.value.bias.to(torch.cuda.current_device()), |
| 127 | params_dtype=torch.half, |
| 128 | device=layer.attention.value.weight.device, |
| 129 | ) |
| 130 | layer.attention.key = QuantizedLinear( |
| 131 | in_features=layer.attention.key.in_features, |
| 132 | out_features=layer.attention.key.out_features, |
| 133 | weight_bit_width=weight_bit_width, |
| 134 | weight=layer.attention.key.weight.to(torch.cuda.current_device()), |
| 135 | bias=layer.attention.key.bias.to(torch.cuda.current_device()), |
| 136 | params_dtype=torch.half, |
| 137 | device=layer.attention.key.weight.device, |
| 138 | ) |
| 139 | layer.attention.dense = QuantizedLinear( |
| 140 | in_features=layer.attention.dense.in_features, |
| 141 | out_features=layer.attention.dense.out_features, |
| 142 | weight_bit_width=weight_bit_width, |
| 143 | weight=layer.attention.dense.weight.to(torch.cuda.current_device()), |
| 144 | bias=layer.attention.dense.bias.to(torch.cuda.current_device()), |
| 145 | params_dtype=torch.half, |
| 146 | device=layer.attention.dense.weight.device, |
| 147 | ) |
| 148 | layer.mlp.dense_h_to_4h = QuantizedLinear( |
| 149 | in_features=layer.mlp.dense_h_to_4h.in_features, |
| 150 | out_features=layer.mlp.dense_h_to_4h.out_features, |
| 151 | weight_bit_width=weight_bit_width, |
| 152 | weight=layer.mlp.dense_h_to_4h.weight.to(torch.cuda.current_device()), |
| 153 | bias=layer.mlp.dense_h_to_4h.bias.to(torch.cuda.current_device()), |
| 154 | params_dtype=torch.half, |
| 155 | device=layer.mlp.dense_h_to_4h.weight.device, |
| 156 | ) |
| 157 | layer.mlp.dense_4h_to_h = QuantizedLinear( |
| 158 | in_features=layer.mlp.dense_4h_to_h.in_features, |
| 159 | out_features=layer.mlp.dense_4h_to_h.out_features, |
| 160 | weight_bit_width=weight_bit_width, |