(self, input_)
| 175 | self.weight_scale = Parameter(self.weight_scale.to(kwargs["device"]), requires_grad=False) |
| 176 | |
| 177 | def forward(self, input_): |
| 178 | # Set up backprop all-reduce. |
| 179 | if self.input_is_parallel: |
| 180 | input_parallel = input_ |
| 181 | else: |
| 182 | input_parallel = scatter_to_tensor_model_parallel_region(input_) |
| 183 | # Matrix multiply. |
| 184 | output_parallel = W8A16Linear.apply(input_parallel, self.weight, self.weight_scale, self.weight_bit_width) |
| 185 | # All-reduce across all the partitions. |
| 186 | output_ = reduce_from_tensor_model_parallel_region(output_parallel) |
| 187 | if self.bias is not None and not self.skip_bias_add: |
| 188 | output = output_ + self.bias |
| 189 | else: |
| 190 | output = output_ |
| 191 | output_bias = self.bias if self.skip_bias_add else None |
| 192 | |
| 193 | return output, output_bias |
| 194 | |
| 195 | |
| 196 | def quantize(model, weight_bit_width, backend="torch"): |
nothing calls this directly
no test coverage detected