| 134 | |
| 135 | |
| 136 | class QuantizedRowParallelLinear(RowParallelLinear): |
| 137 | def __init__( |
| 138 | self, |
| 139 | input_size: int, |
| 140 | output_size: int, |
| 141 | weight_bit_width: int, |
| 142 | weight: torch.Tensor = None, |
| 143 | bias: torch.Tensor = None, |
| 144 | *args, |
| 145 | **kwargs, |
| 146 | ): |
| 147 | super(QuantizedRowParallelLinear, self).__init__(input_size, output_size, *args, **kwargs) |
| 148 | self.input_size = input_size |
| 149 | self.output_size = output_size |
| 150 | self.weight_bit_width = weight_bit_width |
| 151 | if "skip_bias_add" in kwargs: |
| 152 | self.skip_bias_add = kwargs["skip_bias_add"] |
| 153 | else: |
| 154 | self.skip_bias_add = False |
| 155 | del self.weight |
| 156 | |
| 157 | if weight is None: |
| 158 | self.weight = torch.empty( |
| 159 | self.output_size, self.input_size * weight_bit_width // 8, dtype=torch.int8, device=kwargs["device"] |
| 160 | ) |
| 161 | self.weight_scale = torch.empty(self.output_size, dtype=kwargs["params_dtype"], device=kwargs["device"]) |
| 162 | else: |
| 163 | self.weight_scale = (weight.abs().max(dim=-1).values / ((2 ** (weight_bit_width - 1)) - 1)).half() |
| 164 | self.weight = torch.round(weight / self.weight_scale[:, None]).to(torch.int8) |
| 165 | if weight_bit_width == 4: |
| 166 | self.weight = compress_int4_weight(self.weight) |
| 167 | |
| 168 | if bias is None: |
| 169 | self.register_parameter('bias', None) |
| 170 | else: |
| 171 | del self.bias |
| 172 | self.bias = bias |
| 173 | |
| 174 | self.weight = Parameter(self.weight.to(kwargs["device"]), requires_grad=False) |
| 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 |