(self, input_)
| 116 | self.weight_scale = Parameter(self.weight_scale.to(kwargs["device"]), requires_grad=False) |
| 117 | |
| 118 | def forward(self, input_): |
| 119 | # Set up backprop all-reduce. |
| 120 | input_parallel = copy_to_tensor_model_parallel_region(input_) |
| 121 | # Matrix multiply. |
| 122 | output_parallel = W8A16Linear.apply(input_parallel, self.weight, self.weight_scale, self.weight_bit_width) |
| 123 | if self.bias is not None and not self.skip_bias_add: |
| 124 | output_parallel = output_parallel + self.bias |
| 125 | if self.gather_output: |
| 126 | # All-gather across the partitions. |
| 127 | output = gather_from_tensor_model_parallel_region(output_parallel) |
| 128 | else: |
| 129 | output = output_parallel |
| 130 | |
| 131 | output_bias = self.bias if self.skip_bias_add else None |
| 132 | |
| 133 | return output, output_bias |
| 134 | |
| 135 | |
| 136 | class QuantizedRowParallelLinear(RowParallelLinear): |
nothing calls this directly
no test coverage detected