| 717 | |
| 718 | |
| 719 | class Int8Params(torch.nn.Parameter): |
| 720 | def __new__( |
| 721 | cls, |
| 722 | data: Optional[torch.Tensor] = None, |
| 723 | requires_grad=True, |
| 724 | has_fp16_weights=False, |
| 725 | CB: Optional[torch.Tensor] = None, |
| 726 | SCB: Optional[torch.Tensor] = None, |
| 727 | **kwargs, |
| 728 | ): |
| 729 | if data is None: |
| 730 | data = torch.empty(0) |
| 731 | obj = torch.Tensor._make_subclass(cls, data, requires_grad) |
| 732 | obj.CB = CB |
| 733 | obj.SCB = SCB |
| 734 | obj.has_fp16_weights = has_fp16_weights |
| 735 | return obj |
| 736 | |
| 737 | def _quantize(self, device): |
| 738 | if self.has_fp16_weights: |
| 739 | return super().to(device) |
| 740 | |
| 741 | # We quantize the weight and store in 8bit row-major |
| 742 | B = self.data.contiguous().to(device=device, dtype=torch.float16) |
| 743 | CB, SCB, _ = bnb.functional.int8_vectorwise_quant(B) |
| 744 | self.data = CB |
| 745 | self.CB = CB |
| 746 | self.SCB = SCB |
| 747 | |
| 748 | return self |
| 749 | |
| 750 | def cpu(self): |
| 751 | return self.to(device="cpu") |
| 752 | |
| 753 | def cuda(self, device: Optional[int | device | str] = None, non_blocking: bool = False): |
| 754 | return self.to(device="cuda" if device is None else device, non_blocking=non_blocking) |
| 755 | |
| 756 | def xpu(self, device: Optional[int | device | str] = None, non_blocking: bool = False): |
| 757 | return self.to(device="xpu" if device is None else device, non_blocking=non_blocking) |
| 758 | |
| 759 | def __deepcopy__(self, memo): |
| 760 | # adjust this if new arguments are added to the constructor |
| 761 | new_instance = type(self).__new__( |
| 762 | type(self), |
| 763 | data=copy.deepcopy(self.data, memo), |
| 764 | requires_grad=self.requires_grad, |
| 765 | has_fp16_weights=self.has_fp16_weights, |
| 766 | CB=copy.deepcopy(self.CB, memo), |
| 767 | SCB=copy.deepcopy(self.SCB, memo), |
| 768 | ) |
| 769 | return new_instance |
| 770 | |
| 771 | @overload |
| 772 | def to( |
| 773 | self: T, |
| 774 | device: Optional[int | device] = ..., |
| 775 | dtype: Optional[dtype | str] = ..., |
| 776 | non_blocking: bool = ..., |