Method compute_quantization

deepspeed/runtime/quantize.py:129–173

(self, input, index=0, factor=1)

        return input_q

    def compute_quantization(self, input, index=0, factor=1):
        # fixing the quantization bits based on the training steps
        # when reducing 1 bit at each period, we increase the period
        # to go slowly toward the target quantization bits
        # the period and starting bit can be configured

        if input.start_bits != input.target_bits:
            if self.qsteps >= input.q_period:
                self.quantize_real_ratio = 1.0
                input.q_period <<= 1
                input.q_period *= factor
                input.start_bits -= 1
                if self.q_verbose:
                    logger.info(
                        f'Quantization settings: current bit-precision = {input.start_bits}, step = {self.qsteps}, quantization period = {input.q_period}, index = {index}'
                    )
        assert (input.start_bits >= input.target_bits), \
            'Quantization bit is lower than target precision bits!'

        if self.use_quantizer_kernel:
            if input.start_bits <= 2:
                raise ValueError('Quantization bit is too low, please do it without quantization kernel!')
            input_q = ds_quantizer(input.data.clone(),
                                   self.q_groups,
                                   input.start_bits,
                                   asym=False if self.q_type == 'symmetric' else True,
                                   sr=False if self.q_rounding == 'nearest_neighbor' else True)
        else:
            if input.start_bits >= 3:
                input_flat = self.quantize_highbit(input.data, input.start_bits)
            elif input.start_bits == 2:
                assert self.q_type == 'symmetric', 'Quantization type is not symmetric!'
                assert self.q_rounding == 'nearest', 'Quantization rounding is not nearest_neighbor!'
                input_flat = self.quantize_tenary(input.data)
            elif input.start_bits == 1:
                assert self.q_type == 'symmetric', 'Quantization type is not symmetric!'
                assert self.q_rounding == 'nearest', 'Quantization rounding is not nearest_neighbor!'
                input_flat = self.quantize_binary(input.data)
        if self.use_quantizer_kernel:
            return self.mixed_fp16_quantize(input.data, input_q, index)
        else:
            if self.q_mixed_fp16 and input.start_bits >= input.target_bits - 1:
                input_flat = self.quantize_real_ratio * input.data + \
                              (1 - self.quantize_real_ratio) * input_flat
            return input_flat

    def update_fp16_ratio(self):
        if self.q_mixed_fp16: