wehub-resource-sync
87 lines
2.9 KiB
Python
87 lines
2.9 KiB
Python
# Copyright (c) 2024 Habana Labs, Ltd. an Intel Company
|
|
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
try:
|
|
# is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed
|
|
# if successful this also means we're doing a local install and not JIT compile path
|
|
from op_builder import __deepspeed__ # noqa: F401 # type: ignore
|
|
from op_builder.builder import OpBuilder
|
|
except ImportError:
|
|
from deepspeed.ops.op_builder.builder import OpBuilder
|
|
|
|
|
|
class FPQuantizerBuilder(OpBuilder):
|
|
BUILD_VAR = "DS_BUILD_FP_QUANTIZER"
|
|
NAME = "fp_quantizer"
|
|
|
|
def __init__(self, name=None):
|
|
name = self.NAME if name is None else name
|
|
super().__init__(name=name)
|
|
|
|
def absolute_name(self):
|
|
return f'deepspeed.ops.fp_quantizer.{self.NAME}_op'
|
|
|
|
def sources(self):
|
|
return []
|
|
|
|
def load(self, verbose=True):
|
|
return FPQuantizer
|
|
|
|
@staticmethod
|
|
def get_default_quant_dtype():
|
|
return torch.float8_e4m3fn
|
|
|
|
@staticmethod
|
|
def get_quant_range(q_bits=None):
|
|
import habana_frameworks.torch.utils.experimental as htexp
|
|
if htexp._get_device_type() == htexp.synDeviceType.synDeviceGaudi2:
|
|
dtype = torch.float8_e4m3fnuz
|
|
else:
|
|
dtype = torch.float8_e4m3fn
|
|
return torch.finfo(dtype).max
|
|
|
|
|
|
class FPQuantizer:
|
|
CUDA_IMPL = False
|
|
|
|
@classmethod
|
|
def selective_dequantize(cls, val_q, scales, indexes, group_size, q_mantisa_bits, q_exponent_bits):
|
|
assert False, "Selective dequantize isn't implemented for HPU!"
|
|
|
|
@classmethod
|
|
def dequantize(cls, fp_out, input_q, scale, group_size, q_mantisa_bits, q_exponent_bits):
|
|
orig_shape = fp_out.shape
|
|
orig_dtype = fp_out.dtype
|
|
dequant_out = torch.ops.hpu.cast_from_fp8(input_q, (1.0 / scale), orig_dtype).view(orig_shape)
|
|
fp_out.copy_(dequant_out)
|
|
return fp_out
|
|
|
|
@classmethod
|
|
def quantize(cls, out, val, scale, group_size, stochastic_rounding, q_bits, q_mantisa_bits):
|
|
assert q_bits == 8, "Quantize on HPU only supports quantization to FP8"
|
|
assert q_mantisa_bits == 3, "Quantize on HPU only supports q_mantissa_bits = 3"
|
|
assert out.dtype.is_floating_point, "Quantization on HPU is only to float dtypes"
|
|
|
|
num_groups, group_size = out.shape
|
|
|
|
# Reshape the tensor
|
|
val_reshaped = val.view(num_groups, group_size).float()
|
|
# Calculate the scale
|
|
max_vals = val_reshaped.abs().max(dim=1, keepdim=True)[0]
|
|
q_range = torch.finfo(out.dtype).max
|
|
tmp_scale = q_range / max_vals
|
|
scale.copy_(tmp_scale)
|
|
# Copy quantized
|
|
quant, _ = torch.ops.hpu.cast_to_fp8_v2(val_reshaped, scale, stochastic_rounding, dtype=out.dtype)
|
|
out.copy_(quant)
|
|
|
|
return out
|
|
|
|
@classmethod
|
|
def get_scales(cls, out, num_groups):
|
|
return out
|