(self, weight_bit_width: int, weight, bias=None, device="cuda", dtype=None, empty_init=False)
| 157 | |
| 158 | class QuantizedLinear(torch.nn.Module): |
| 159 | def __init__(self, weight_bit_width: int, weight, bias=None, device="cuda", dtype=None, empty_init=False): |
| 160 | super().__init__() |
| 161 | weight = weight.to(device) # ensure the weight is on the cuda device |
| 162 | assert str(weight.device).startswith( |
| 163 | 'cuda'), 'The weights that need to be quantified should be on the CUDA device' |
| 164 | self.weight_bit_width = weight_bit_width |
| 165 | shape = weight.shape |
| 166 | |
| 167 | if weight is None or empty_init: |
| 168 | self.weight = torch.empty(shape[0], shape[1] * weight_bit_width // 8, dtype=torch.int8, device=device) |
| 169 | self.weight_scale = torch.empty(shape[0], dtype=dtype, device=device) |
| 170 | else: |
| 171 | self.weight_scale = weight.abs().max(dim=-1).values / ((2 ** (weight_bit_width - 1)) - 1) |
| 172 | self.weight = torch.round(weight / self.weight_scale[:, None]).to(torch.int8) |
| 173 | if weight_bit_width == 4: |
| 174 | self.weight = compress_int4_weight(self.weight) |
| 175 | |
| 176 | self.weight = Parameter(self.weight.to(device), requires_grad=False) |
| 177 | self.weight_scale = Parameter(self.weight_scale.to(device), requires_grad=False) |
| 178 | self.bias = Parameter(bias.to(device), requires_grad=False) if bias is not None else None |
| 179 | |
| 180 | def forward(self, input): |
| 181 | output = W8A16Linear.apply(input, self.weight, self.weight_scale, self.weight_bit_width) |
nothing calls this directly
no test coverage detected