Repository: Lightning-AI/pytorch-lightning
File: src/lightning/fabric/plugins/precision/bitsandbytes.py (Python)

Code before:

        def quantize_(self, weight: Optional[torch.Tensor] = None, device: Optional[torch.device] = None) -> None:
            """Inplace quantize."""
            if weight is None:
                weight = self.weight.data
            if weight.data.dtype == torch.int8:
                # already quantized
                return
            assert isinstance(self.weight, bnb.nn.Int8Params)
            self.weight = self.quantize(self.weight, weight, device)

        @staticmethod
        def quantize(
            int8params: bnb.nn.Int8Params, weight: torch.Tensor, device: Optional[torch.device]
        ) -> bnb.nn.Int8Params:
            device = device or torch.device("cuda")
            if device.type != "cuda":
                raise RuntimeError(f"Unexpected device type: {device.type}")
            # https://github.com/TimDettmers/bitsandbytes/blob/0.41.0/bitsandbytes/nn/modules.py#L291-L302
            B = weight.contiguous().to(device=device, dtype=torch.float16)
            if int8params.has_fp16_weights:
                int8params.data = B
            else:
                if hasattr(bnb.functional, "double_quant"):
                    CB, _, SCB, _, _ = bnb.functional.double_quant(B)
                else:  # for versions 0.46+
                    CB, SCB = bnb.functional.int8_double_quant(B)
                int8params.data = CB
                setattr(int8params, "CB", CB)
                setattr(int8params, "SCB", SCB)
            return int8params

        def to_empty(self, *, device: _DEVICE, recurse: bool = True) -> Self:
            if self.weight.device.type == "meta":
                # need custom logic if int8params is on meta device
                raise NotImplementedError
            if self.weight.dtype == torch.uint8:  # was quantized
                # need the original shape here
                raise NotImplementedError
            device = torch.device(device)
            weight = torch.empty_like(self.weight.data, device=device)
            if device.type == "cuda":  # re-quantize
                self.quantize_(weight, device)
            else:
                self.weight = _replace_param(self.weight, weight)
            if self.bias is not None:
                self.bias = _replace_param(self.bias, torch.empty_like(self.bias, device=device))
            return self

        def reset_parameters(self) -> None:
            # from `torch.nn.Linear.reset_parameters`

Review comment (suggestion):
[nitpick] Clarify the comment to specify "for bitsandbytes ≥0.46" rather than just "0.46+" to avoid ambiguity about what the version number refers to.
```suggestion
                else:  # for bitsandbytes versions ≥0.46
```

Code after:

        def quantize_(self, weight: Optional[torch.Tensor] = None, device: Optional[torch.device] = None) -> None:
            """Inplace quantize."""
            if weight is None:
                weight = self.weight.data
            if weight.data.dtype == torch.int8:
                # already quantized
                return
            assert isinstance(self.weight, bnb.nn.Int8Params)
            self.weight = self.quantize(self.weight, weight, device)

        @staticmethod
        def quantize(
            int8params: bnb.nn.Int8Params, weight: torch.Tensor, device: Optional[torch.device]
        ) -> bnb.nn.Int8Params:
            device = device or torch.device("cuda")
            if device.type != "cuda":
                raise RuntimeError(f"Unexpected device type: {device.type}")
            # https://github.com/TimDettmers/bitsandbytes/blob/0.41.0/bitsandbytes/nn/modules.py#L291-L302
            B = weight.contiguous().to(device=device, dtype=torch.float16)
            if int8params.has_fp16_weights:
                int8params.data = B
            else:
                if hasattr(bnb.functional, "double_quant"):
                    CB, _, SCB, _, _ = bnb.functional.double_quant(B)
                else:  # for bitsandbytes versions ≥0.46
                    CB, SCB = bnb.functional.int8_double_quant(B)
                int8params.data = CB
                setattr(int8params, "CB", CB)
                setattr(int8params, "SCB", SCB)
            return int8params

        def to_empty(self, *, device: _DEVICE, recurse: bool = True) -> Self:
            if self.weight.device.type == "meta":
                # need custom logic if int8params is on meta device
                raise NotImplementedError
            if self.weight.dtype == torch.uint8:  # was quantized
                # need the original shape here
                raise NotImplementedError
            device = torch.device(device)
            weight = torch.empty_like(self.weight.data, device=device)
            if device.type == "cuda":  # re-quantize
                self.quantize_(weight, device)
            else:
                self.weight = _replace_param(self.weight, weight)
            if self.bias is not None:
                self.bias = _replace_param(self.bias, torch.empty_like(self.bias, device=device))
            return self

        def reset_parameters(self) -> None:
            # from `torch.nn.Linear.reset_parameters`