| 52 | LayerCache = Tuple[torch.Tensor, torch.Tensor] |
| 53 | |
| 54 | class BitLinearKernel(nn.Module): |
| 55 | in_features: int |
| 56 | out_features: int |
| 57 | weight: torch.Tensor |
| 58 | weight_scale: torch.Tensor |
| 59 | |
| 60 | def __init__(self, in_features: int, out_features: int, bias: bool = False): |
| 61 | super().__init__() |
| 62 | self.in_features = in_features |
| 63 | self.out_features = out_features |
| 64 | |
| 65 | self.weight = torch.nn.Parameter(torch.zeros(out_features, in_features//4, dtype=torch.int8), requires_grad=False) |
| 66 | self.weight_scale = torch.nn.Parameter(torch.zeros(4, dtype=torch.bfloat16), requires_grad=False) |
| 67 | |
| 68 | @torch.compile |
| 69 | def quant_input(self, input): |
| 70 | s = 127 / input.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5) |
| 71 | return (input * s).round().clamp(-128, 127).to(torch.int8), s |
| 72 | |
| 73 | def forward(self, input): |
| 74 | input, s = self.quant_input(input) |
| 75 | return bitnet_int8xint2_linear(input, self.weight, s, self.weight_scale) |
| 76 | |
| 77 | class BitLinear(nn.Linear): |
| 78 | @torch.compile |
nothing calls this directly
no outgoing calls
no test coverage detected