DeepSpeed

1
# Copyright (c) Microsoft Corporation.
2
# SPDX-License-Identifier: Apache-2.0
3

4
# DeepSpeed Team
5

6
import torch
7

8
from deepspeed.ops.op_builder import QuantizerBuilder
9

10
# Cuda modules will be imported if needed
11
quantizer_cuda_module = None
12

13

14
def ds_quantizer(input, groups=1, bit_num=8, sr=False, asym=False):
15
    # Load cuda modules if needed
16
    global quantizer_cuda_module
17
    if quantizer_cuda_module is None:
18
        quantizer_cuda_module = QuantizerBuilder().load()
19
    if sr:
20
        if asym:
21
            quantize_func = quantizer_cuda_module.ds_sr_quantize_asym_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_sr_quantize_asym_fp32
22
        else:
23
            quantize_func = quantizer_cuda_module.ds_sr_quantize_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_sr_quantize_fp32
24
    else:
25
        if asym:
26
            quantize_func = quantizer_cuda_module.ds_quantize_asym_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_quantize_asym_fp32
27
        else:
28
            quantize_func = quantizer_cuda_module.ds_quantize_fp16 if input.dtype == torch.half else quantizer_cuda_module.ds_quantize_fp32
29
    return quantize_func(input, groups, bit_num)
30