(ctx, grad_output: torch.Tensor)
| 59 | |
| 60 | @staticmethod |
| 61 | def backward(ctx, grad_output: torch.Tensor): |
| 62 | inp, quant_w, scale_w = ctx.saved_tensors |
| 63 | weight = extract_weight_to_half(quant_w, scale_w, ctx.weight_bit_width) |
| 64 | grad_output = grad_output.contiguous().view(-1, weight.size(0)) |
| 65 | grad_input = grad_output.mm(weight) |
| 66 | grad_weight = grad_output.t().mm(inp) |
| 67 | return grad_input.view(ctx.inp_shape), grad_weight.view(ctx.weight_shape), None, None |
| 68 | |
| 69 | |
| 70 | def compress_int4_weight(weight: torch.Tensor): # (n, m) |