|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +import math |
| 8 | +import struct |
| 9 | +from typing import Callable |
| 10 | + |
| 11 | +import fbgemm_gpu |
| 12 | +import numpy as np |
| 13 | +import torch |
| 14 | + |
| 15 | +# pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`. |
| 16 | +open_source: bool = getattr(fbgemm_gpu, "open_source", False) |
| 17 | + |
| 18 | +try: |
| 19 | + # pyre-ignore[21] |
| 20 | + from fbgemm_gpu import open_source # noqa: F401 |
| 21 | + |
| 22 | +except Exception: |
| 23 | + if torch.version.hip: |
| 24 | + torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_hip") |
| 25 | + else: |
| 26 | + torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") |
| 27 | + |
| 28 | + torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu") |
| 29 | + |
| 30 | +# Eigen/Python round 0.5 away from 0, Numpy rounds to even |
| 31 | +round_to_nearest: Callable[[np.ndarray], np.ndarray] = np.vectorize(round) |
| 32 | + |
| 33 | + |
| 34 | +def bytes_to_floats(byte_matrix: np.ndarray) -> np.ndarray: |
| 35 | + floats = np.empty([np.shape(byte_matrix)[0], 1], dtype=np.float32) |
| 36 | + for i, byte_values in enumerate(byte_matrix): |
| 37 | + (floats[i],) = struct.unpack("f", bytearray(byte_values)) |
| 38 | + return floats |
| 39 | + |
| 40 | + |
| 41 | +def floats_to_bytes(floats: np.ndarray) -> np.ndarray: |
| 42 | + byte_matrix = np.empty([np.shape(floats)[0], 4], dtype=np.uint8) |
| 43 | + for i, value in enumerate(floats): |
| 44 | + assert isinstance(value, np.float32), (value, floats) |
| 45 | + as_bytes = struct.pack("f", value) |
| 46 | + # In Python3 bytes will be a list of int, in Python2 a list of string |
| 47 | + if isinstance(as_bytes[0], int): |
| 48 | + byte_matrix[i] = list(as_bytes) |
| 49 | + else: |
| 50 | + byte_matrix[i] = list(map(ord, as_bytes)) |
| 51 | + return byte_matrix |
| 52 | + |
| 53 | + |
| 54 | +def bytes_to_half_floats(byte_matrix: np.ndarray) -> np.ndarray: |
| 55 | + floats = np.empty([np.shape(byte_matrix)[0], 1], dtype=np.float16) |
| 56 | + for i, byte_values in enumerate(byte_matrix): |
| 57 | + (floats[i],) = np.frombuffer( |
| 58 | + memoryview(byte_values).tobytes(), dtype=np.float16 |
| 59 | + ) |
| 60 | + return floats |
| 61 | + |
| 62 | + |
| 63 | +def half_floats_to_bytes(floats: np.ndarray) -> np.ndarray: |
| 64 | + byte_matrix = np.empty([np.shape(floats)[0], 2], dtype=np.uint8) |
| 65 | + for i, value in enumerate(floats): |
| 66 | + assert isinstance(value, np.float16), (value, floats) |
| 67 | + byte_matrix[i] = np.frombuffer( |
| 68 | + memoryview(value.tobytes()).tobytes(), dtype=np.uint8 |
| 69 | + ) |
| 70 | + return byte_matrix |
| 71 | + |
| 72 | + |
| 73 | +def fused_rowwise_8bit_quantize_reference(data: np.ndarray) -> np.ndarray: |
| 74 | + minimum = np.min(data, axis=-1, keepdims=True) |
| 75 | + maximum = np.max(data, axis=-1, keepdims=True) |
| 76 | + span = maximum - minimum |
| 77 | + bias = minimum |
| 78 | + scale = span / 255.0 |
| 79 | + inverse_scale = 255.0 / (span + 1e-8) |
| 80 | + quantized_data = round_to_nearest((data - bias) * inverse_scale) |
| 81 | + scale_bytes = floats_to_bytes(scale.reshape(-1)) |
| 82 | + scale_bytes = scale_bytes.reshape(data.shape[:-1] + (scale_bytes.shape[-1],)) |
| 83 | + bias_bytes = floats_to_bytes(bias.reshape(-1)) |
| 84 | + bias_bytes = bias_bytes.reshape(data.shape[:-1] + (bias_bytes.shape[-1],)) |
| 85 | + return np.concatenate([quantized_data, scale_bytes, bias_bytes], axis=-1) |
| 86 | + |
| 87 | + |
| 88 | +def fused_rowwise_8bit_dequantize_reference(fused_quantized: np.ndarray) -> np.ndarray: |
| 89 | + scale = bytes_to_floats(fused_quantized[..., -8:-4].astype(np.uint8).reshape(-1, 4)) |
| 90 | + scale = scale.reshape(fused_quantized.shape[:-1] + (scale.shape[-1],)) |
| 91 | + bias = bytes_to_floats(fused_quantized[..., -4:].astype(np.uint8).reshape(-1, 4)) |
| 92 | + bias = bias.reshape(fused_quantized.shape[:-1] + (bias.shape[-1],)) |
| 93 | + quantized_data = fused_quantized[..., :-8] |
| 94 | + return quantized_data * scale + bias |
| 95 | + |
| 96 | + |
| 97 | +def fused_rowwise_8bit_dequantize_reference_half( |
| 98 | + fused_quantized: np.ndarray, |
| 99 | +) -> np.ndarray: |
| 100 | + scale = bytes_to_half_floats( |
| 101 | + fused_quantized[..., -8:-4].astype(np.uint8).reshape(-1, 4) |
| 102 | + ) |
| 103 | + scale = scale.reshape(fused_quantized.shape[:-1] + (scale.shape[-1],)) |
| 104 | + bias = bytes_to_half_floats( |
| 105 | + fused_quantized[..., -4:].astype(np.uint8).reshape(-1, 4) |
| 106 | + ) |
| 107 | + bias = bias.reshape(fused_quantized.shape[:-1] + (bias.shape[-1],)) |
| 108 | + quantized_data = fused_quantized[..., :-8] |
| 109 | + return quantized_data * scale + bias |
| 110 | + |
| 111 | + |
| 112 | +def fused_rowwise_nbit_quantize_reference(data: np.ndarray, bit: int) -> np.ndarray: |
| 113 | + minimum = np.min(data, axis=1).astype(np.float16).astype(np.float32) |
| 114 | + maximum = np.max(data, axis=1) |
| 115 | + span = maximum - minimum |
| 116 | + qmax = (1 << bit) - 1 |
| 117 | + scale = (span / qmax).astype(np.float16).astype(np.float32) |
| 118 | + bias = np.zeros(data.shape[0]) |
| 119 | + quantized_data = np.zeros(data.shape).astype(np.uint8) |
| 120 | + |
| 121 | + for i in range(data.shape[0]): |
| 122 | + bias[i] = minimum[i] |
| 123 | + inverse_scale = 1.0 if scale[i] == 0.0 else 1 / scale[i] |
| 124 | + if scale[i] == 0.0 or math.isinf(inverse_scale): |
| 125 | + scale[i] = 1.0 |
| 126 | + inverse_scale = 1.0 |
| 127 | + quantized_data[i] = np.clip( |
| 128 | + np.round((data[i, :] - minimum[i]) * inverse_scale), 0, qmax |
| 129 | + ) |
| 130 | + |
| 131 | + # pack |
| 132 | + assert 8 % bit == 0 |
| 133 | + num_elem_per_byte = 8 // bit |
| 134 | + packed_dim = (data.shape[1] + num_elem_per_byte - 1) // num_elem_per_byte |
| 135 | + packed_data = np.zeros([data.shape[0], packed_dim]).astype(np.uint8) |
| 136 | + for i in range(data.shape[0]): |
| 137 | + for j in range(data.shape[1]): |
| 138 | + if j % num_elem_per_byte == 0: |
| 139 | + packed_data[i, j // num_elem_per_byte] = quantized_data[i, j] |
| 140 | + else: |
| 141 | + packed_data[i, j // num_elem_per_byte] += quantized_data[i, j] << ( |
| 142 | + (j % num_elem_per_byte) * bit |
| 143 | + ) |
| 144 | + |
| 145 | + scale_bytes = half_floats_to_bytes(scale.astype(np.float16)) |
| 146 | + bias_bytes = half_floats_to_bytes(bias.astype(np.float16)) |
| 147 | + return np.concatenate([packed_data, scale_bytes, bias_bytes], axis=1) |
| 148 | + |
| 149 | + |
| 150 | +def fused_rowwise_nbit_quantize_dequantize_reference( |
| 151 | + data: np.ndarray, bit: int |
| 152 | +) -> np.ndarray: |
| 153 | + fused_quantized = fused_rowwise_nbit_quantize_reference(data, bit) |
| 154 | + scale = bytes_to_half_floats(fused_quantized[:, -4:-2].astype(np.uint8)).astype( |
| 155 | + np.float32 |
| 156 | + ) |
| 157 | + bias = bytes_to_half_floats(fused_quantized[:, -2:].astype(np.uint8)).astype( |
| 158 | + np.float32 |
| 159 | + ) |
| 160 | + quantized_data = fused_quantized[:, :-4] |
| 161 | + |
| 162 | + # unpack |
| 163 | + packed_dim = fused_quantized.shape[1] - 4 |
| 164 | + assert 8 % bit == 0 |
| 165 | + num_elem_per_byte = 8 // bit |
| 166 | + assert packed_dim == ((data.shape[1] + num_elem_per_byte - 1) // num_elem_per_byte) |
| 167 | + unpacked_data = np.zeros(data.shape).astype(np.uint8) |
| 168 | + for i in range(data.shape[0]): |
| 169 | + for j in range(data.shape[1]): |
| 170 | + unpacked_data[i, j] = ( |
| 171 | + quantized_data[i, j // num_elem_per_byte] |
| 172 | + >> ((j % num_elem_per_byte) * bit) |
| 173 | + ) & ((1 << bit) - 1) |
| 174 | + |
| 175 | + return scale * unpacked_data + bias |
0 commit comments