[Feat] FP8 per tensor quant support (#4043)
* FP8 per tensor quant support * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Datta Nimmaturi
GitHub
parent
0212f7f7df
commit
f3b5090f24
1 changed files with 54 additions and 52 deletions
|
|
@ -95,8 +95,11 @@ def weight_dequant_block(
|
|||
|
||||
|
||||
def weight_dequant(x: torch.Tensor, s: torch.Tensor, dtype = torch.bfloat16):
|
||||
if s.shape[1] == 1:
|
||||
# this is row quantized weight, just simple multiplication suffices
|
||||
# Per-tensor scale: single value for entire weight matrix
|
||||
if s.numel() == 1:
|
||||
return x.to(dtype) * s.view(1, 1).to(dtype)
|
||||
# Row quantized weight: scale shape is (m, 1) or (n, 1)
|
||||
elif s.ndim == 2 and s.shape[1] == 1:
|
||||
if x.shape[0] == s.shape[0]:
|
||||
y = x.to(dtype) * s.to(dtype)
|
||||
elif x.shape[1] == s.shape[0]:
|
||||
|
|
@ -106,8 +109,8 @@ def weight_dequant(x: torch.Tensor, s: torch.Tensor, dtype = torch.bfloat16):
|
|||
else:
|
||||
raise ValueError(f"Incompatible shapes {x.shape = }, {s.shape = }")
|
||||
return y
|
||||
# Block quantized weight: scale shape is (ceil(m/block_m), ceil(n/block_n))
|
||||
else:
|
||||
# this is block quantized weight
|
||||
return weight_dequant_block(x, s, dtype = dtype)
|
||||
|
||||
|
||||
|
|
@ -238,44 +241,29 @@ def w8a8_block_fp8_matmul_triton(
|
|||
block_size: list[int],
|
||||
output_dtype: torch.dtype = torch.float32,
|
||||
) -> torch.Tensor:
|
||||
"""This function performs matrix multiplication with block-wise
|
||||
quantization.
|
||||
It takes two input tensors `A` and `B` with scales `As` and `Bs`.
|
||||
The output is returned in the specified `output_dtype`.
|
||||
Args:
|
||||
A: The input tensor, e.g., activation.
|
||||
B: The input tensor, e.g., weight.
|
||||
As: The per-token-group quantization scale for `A`.
|
||||
Bs: The per-block quantization scale for `B`.
|
||||
block_size: The block size for per-block quantization. It should
|
||||
be 2-dim, e.g., [128, 128].
|
||||
output_dytpe: The dtype of the returned tensor.
|
||||
Returns:
|
||||
torch.Tensor: The result of matmul.
|
||||
"""
|
||||
assert len(block_size) == 2
|
||||
block_n, block_k = block_size[0], block_size[1]
|
||||
"""Block-wise FP8 matmul."""
|
||||
if block_size is None:
|
||||
block_n, block_k = 128, 128
|
||||
else:
|
||||
assert len(block_size) == 2
|
||||
block_n, block_k = block_size[0], block_size[1]
|
||||
|
||||
N, K = B.shape
|
||||
assert A.shape[-1] == B.shape[-1]
|
||||
assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
|
||||
assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
|
||||
M = A.numel() // A.shape[-1]
|
||||
|
||||
assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
|
||||
N, K = B.shape
|
||||
assert triton.cdiv(N, block_n) == Bs.shape[0]
|
||||
assert triton.cdiv(K, block_k) == Bs.shape[1]
|
||||
|
||||
M = A.numel() // A.shape[-1]
|
||||
C_shape = A.shape[:-1] + (N,)
|
||||
C = A.new_empty(C_shape, dtype = output_dtype)
|
||||
|
||||
BLOCK_SIZE_M = 128
|
||||
if M < BLOCK_SIZE_M:
|
||||
BLOCK_SIZE_M = triton.next_power_of_2(M)
|
||||
BLOCK_SIZE_M = max(BLOCK_SIZE_M, 16)
|
||||
BLOCK_SIZE_K = block_k
|
||||
assert block_k % BLOCK_SIZE_K == 0
|
||||
BLOCK_SIZE_N = block_n
|
||||
BLOCK_SIZE_M = max(triton.next_power_of_2(M), 16)
|
||||
BLOCK_SIZE_K, BLOCK_SIZE_N = block_k, block_n
|
||||
|
||||
def grid(META):
|
||||
return (
|
||||
|
|
@ -342,29 +330,41 @@ fp8_block_matmul = (
|
|||
class FP8BlockQuantLinear(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, X, weight, weight_scale):
|
||||
# block_size = getattr(weight, 'block_size', [128,128])
|
||||
m, n = weight.shape
|
||||
p, q = weight_scale.shape
|
||||
block_size = getattr(weight, "block_size", None) or getattr(
|
||||
weight_scale, "block_size", [128, 128]
|
||||
)
|
||||
assert block_size is not None, "block_size is not set"
|
||||
if triton.cdiv(m, block_size[0]) != p or triton.cdiv(n, block_size[1]) != q:
|
||||
if (
|
||||
triton.cdiv(m, block_size[0]) == q
|
||||
and triton.cdiv(n, block_size[1]) == p
|
||||
):
|
||||
# weights are transposed during backward pass for training :)
|
||||
# We transpose weight scale to counter that. Note that transposing weight would cause issues with matmul with input X
|
||||
weight_scale = weight_scale.T
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Weight shape {weight.shape} and scales shape {weight_scale.shape} is not compatible with block size {block_size}"
|
||||
)
|
||||
|
||||
# Save original scale for backward (before any transformation)
|
||||
original_weight_scale = weight_scale
|
||||
|
||||
# Handle per-tensor quantization: expand scalar to block scale shape
|
||||
if weight_scale.numel() == 1:
|
||||
block_size = [128, 128]
|
||||
# Expand scalar to (ceil(m/128), ceil(n/128)) - same value for all blocks
|
||||
num_blocks_m = triton.cdiv(m, block_size[0])
|
||||
num_blocks_n = triton.cdiv(n, block_size[1])
|
||||
weight_scale = weight_scale.expand(num_blocks_m, num_blocks_n).contiguous()
|
||||
else:
|
||||
# Block quantization path
|
||||
p, q = weight_scale.shape
|
||||
block_size = getattr(weight, "block_size", None) or getattr(
|
||||
weight_scale, "block_size", [128, 128]
|
||||
)
|
||||
assert block_size is not None, "block_size is not set"
|
||||
if triton.cdiv(m, block_size[0]) != p or triton.cdiv(n, block_size[1]) != q:
|
||||
if (
|
||||
triton.cdiv(m, block_size[0]) == q
|
||||
and triton.cdiv(n, block_size[1]) == p
|
||||
):
|
||||
weight_scale = weight_scale.T
|
||||
original_weight_scale = weight_scale # Update for transposed case
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Weight shape {weight.shape} and scales shape {weight_scale.shape} is not compatible with block size {block_size}"
|
||||
)
|
||||
|
||||
if not weight.is_contiguous():
|
||||
weight = weight.contiguous()
|
||||
# this is replica of https://github.com/huggingface/transformers/blob/01c9e1ba683b3e50d7c76bf92f2d470759fd5e81/src/transformers/integrations/finegrained_fp8.py#L331-L353
|
||||
|
||||
# Quantize input and run FP8 matmul
|
||||
qinput, scale = act_quant(X, block_size[1])
|
||||
output = fp8_block_matmul(
|
||||
qinput,
|
||||
|
|
@ -375,8 +375,7 @@ class FP8BlockQuantLinear(torch.autograd.Function):
|
|||
output_dtype = X.dtype,
|
||||
)
|
||||
ctx.weight = weight
|
||||
ctx.weight_scale = weight_scale
|
||||
ctx.block_size = block_size
|
||||
ctx.weight_scale = original_weight_scale # Save original for backward
|
||||
return output.to(X.dtype)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -592,11 +591,14 @@ except:
|
|||
|
||||
@torch_compile
|
||||
def fp8_linear(X, weight, weight_scale, bias = None):
|
||||
if weight_scale.ndim == 2 and weight_scale.shape[1] > 1:
|
||||
# This is block quantized FP8 matmul
|
||||
# Per-tensor quantization: single scalar scale for entire weight
|
||||
# Block quantized FP8: 2D scale tensor with multiple columns
|
||||
if weight_scale.numel() == 1 or (
|
||||
weight_scale.ndim == 2 and weight_scale.shape[1] > 1
|
||||
):
|
||||
out = fp8_block_quant_linear(X, weight, weight_scale)
|
||||
# Row/channel quantized FP8: 2D scale with shape (n, 1)
|
||||
else:
|
||||
# Row quantized FP8
|
||||
out = fbgemm_fp8_linear(X, weight, weight_scale, bias)
|
||||
return out
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue