Repository: microsoft/DeepSpeed
File: deepspeed/ops/fp_quantizer/quantize.py (Python)

Code before:
                 return_meta_tensor=False) -> torch.Tensor:
        assert input.dtype == torch.bfloat16, "only support bf16 for now"
        if return_meta_tensor:
            assert q_bits == 8, "meta tensor is only supported with q_bit=8"

        self.orig_dtype = input.dtype
        self.orig_shape = input.shape

        if q_bits == 8:
            pass
        elif q_bits == 12:
            q_mantisa_bits = 4
        elif q_bits == 6:
            q_mantisa_bits = 2
        elif q_bits == 4:
            q_mantisa_bits = 1
        else:
            assert (0), \
                f"Missing {q_bits}-quantization, please add the template arguments for the kernel to support this precision!"

        # Adding (group_size - 1) is for padding
        self.num_groups = (input.numel() + self.q_config.group_size - 1) // self.q_config.group_size
        # group_size should be the minimal number between the defined group size and number of elements in tensor.
        group_size = int(min(self.q_config.group_size, input.numel()) * q_bits) // 8
        # CUDA quantization kernel saves the scale as (fp32) inside the quantized tensor for each group
        if not self.is_python_impl:
            group_size += 4
        # CUDA quantization kernel allocates tensors as uint8, but handles them as fp8 inside the kernel.
        self.input_q = torch.ones(self.num_groups, group_size, dtype=self.q_config.q_dtype, device=input.device)
        # CUDA quantization kernel attaches scales to quantized result, in python implementation it can't be done
        # because they are of different types.
        self.scale = torch.ones(self.num_groups, 1, device=input.device)
        out = fp_quant_module.quantize(self.input_q, input, self.scale, group_size, stochastic_mode, q_bits,
                                       q_mantisa_bits)
        if return_meta_tensor:
            if not self.is_python_impl:
                data, self.scale = out.split(group_size, dim=-1)
                data = data.contiguous().reshape(input.shape)
            else:
                data = out.contiguous().reshape(input.shape)
            self.scale = self.scale.contiguous()
            del self.input_q
            del out
            gc.collect()
            get_accelerator().empty_cache()
            return data, self.scale

        return out

    def to(self, *args, **kwargs):
        # Intermediate tensors may need to be moved to different devices

Review comment (suggestion):
Can be make explicit as special handling for cuda
```suggestion
        if self.is_cuda_impl:
```

Code after:
                 return_meta_tensor=False) -> torch.Tensor:
        assert input.dtype == torch.bfloat16, "only support bf16 for now"
        if return_meta_tensor:
            assert q_bits == 8, "meta tensor is only supported with q_bit=8"

        self.orig_dtype = input.dtype
        self.orig_shape = input.shape

        if q_bits == 8:
            pass
        elif q_bits == 12:
            q_mantisa_bits = 4
        elif q_bits == 6:
            q_mantisa_bits = 2
        elif q_bits == 4:
            q_mantisa_bits = 1
        else:
            assert (0), \
                f"Missing {q_bits}-quantization, please add the template arguments for the kernel to support this precision!"

        # Adding (group_size - 1) is for padding
        self.num_groups = (input.numel() + self.q_config.group_size - 1) // self.q_config.group_size
        # group_size should be the minimal number between the defined group size and number of elements in tensor.
        group_size = int(min(self.q_config.group_size, input.numel()) * q_bits) // 8
        # CUDA quantization kernel saves the scale as (fp32) inside the quantized tensor for each group
        if self.cuda_impl:
            group_size += 4
        # CUDA quantization kernel allocates tensors as uint8, but handles them as fp8 inside the kernel.
        self.input_q = torch.ones(self.num_groups, group_size, dtype=self.q_config.q_dtype, device=input.device)
        # CUDA quantization kernel attaches scales to quantized result, in python implementation it can't be done
        # because they are of different types.
        self.scale = torch.ones(self.num_groups, 1, device=input.device)
        out = fp_quant_module.quantize(self.input_q, input, self.scale, group_size, stochastic_mode, q_bits,
                                       q_mantisa_bits)
        if return_meta_tensor:
            if self.cuda_impl:
                data, self.scale = out.split(group_size, dim=-1)
                data = data.contiguous().reshape(input.shape)
            else:
                data = out.contiguous().reshape(input.shape)
            self.scale = self.scale.contiguous()
            del self.input_q
            del out
            gc.collect()
            get_accelerator().empty_cache()
            return data, self.scale

        return out

    def to(self, *args, **kwargs):
        # Intermediate tensors may need to be moved to different devices