dc66cd97ef
Co-authored-by: William Lin <SolitaryThinker@users.noreply.github.com> Co-authored-by: Loay Rashid <42599591+loaydatrain@users.noreply.github.com> Co-authored-by: Kaiqin Kong <k1kong@ucsd.edu> Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
131 lines
4.8 KiB
Python
131 lines
4.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""FP8 quantization-aware training for linear layers.
|
|
|
|
Mirror of ``fp4linear.py`` but for FP8 (e4m3). The forward pass quantizes both
|
|
activations and weights to FP8 and runs ``torch._scaled_mm``; the backward pass
|
|
is a bf16 straight-through estimator so the high-precision master weights stay
|
|
trainable. Falls back to a bf16 fake-quant forward on GPUs older than sm89.
|
|
"""
|
|
import torch
|
|
|
|
FP8_DTYPE = torch.float8_e4m3fn
|
|
FP8_MAX = float(torch.finfo(FP8_DTYPE).max) # 448.0
|
|
FP8_MIN_SCALE = 1.0 / (FP8_MAX * 512.0)
|
|
|
|
|
|
def _supports_fp8_compute() -> bool:
|
|
"""Whether the active device supports FP8 ``_scaled_mm`` (sm89+)."""
|
|
if not torch.cuda.is_available():
|
|
return False
|
|
cap = torch.cuda.get_device_capability()
|
|
return cap[0] > 8 or (cap[0] == 8 and cap[1] >= 9)
|
|
|
|
|
|
def _quantize_tensorwise(x_2d: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Returns ``(x_fp8 [M, K], x_scale [1] float32)``."""
|
|
x_absmax = x_2d.abs().amax().float()
|
|
x_scale = (x_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE)
|
|
x_fp8 = (x_2d / x_scale.to(x_2d.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE)
|
|
return x_fp8, x_scale.view(1)
|
|
|
|
|
|
def _quantize_rowwise(x_2d: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Returns ``(x_fp8 [M, K], x_scale [M, 1] float32)``."""
|
|
x_absmax = x_2d.abs().amax(dim=-1, keepdim=True).float()
|
|
x_scale = (x_absmax / FP8_MAX).clamp(min=FP8_MIN_SCALE)
|
|
x_fp8 = (x_2d / x_scale.to(x_2d.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE)
|
|
return x_fp8, x_scale
|
|
|
|
|
|
def _fake_quant(x_2d: torch.Tensor, granularity: str) -> torch.Tensor:
|
|
"""bf16 fake-quant (quantize then dequantize) for pre-sm89 fallback."""
|
|
if granularity == "channel":
|
|
x_fp8, x_scale = _quantize_rowwise(x_2d)
|
|
else:
|
|
x_fp8, x_scale = _quantize_tensorwise(x_2d)
|
|
return x_fp8.to(x_2d.dtype) * x_scale.to(x_2d.dtype)
|
|
|
|
|
|
class _LinearFWD8BWD16Fn(torch.autograd.Function):
|
|
|
|
@staticmethod
|
|
def forward(ctx, x, weight, bias, granularity="tensor"):
|
|
# assert/normalize activation dtype
|
|
if x.dtype not in (torch.float16, torch.bfloat16):
|
|
x = x.to(dtype=torch.bfloat16)
|
|
|
|
# cast params (can be fp32) to activation dtype for quantization
|
|
weight_cast = weight.to(dtype=x.dtype)
|
|
bias_cast = bias.to(dtype=x.dtype) if bias is not None else None
|
|
|
|
orig_shape = x.shape
|
|
k = weight_cast.shape[1]
|
|
n = weight_cast.shape[0]
|
|
x2d = x.reshape(-1, k).contiguous()
|
|
|
|
if not _supports_fp8_compute():
|
|
# bf16 fake-quant fallback: simulate the FP8 rounding error but
|
|
# compute the matmul in bf16.
|
|
x_fq = _fake_quant(x2d, granularity)
|
|
w_fq = _fake_quant(weight_cast, granularity)
|
|
out2d = x_fq.matmul(w_fq.t())
|
|
if bias_cast is not None:
|
|
out2d = out2d + bias_cast
|
|
ctx.save_for_backward(x2d, weight, bias)
|
|
ctx.n = n
|
|
ctx.orig_shape = orig_shape
|
|
return out2d.reshape(*orig_shape[:-1], n)
|
|
|
|
if granularity == "channel":
|
|
x_fp8, x_scale = _quantize_rowwise(x2d)
|
|
w_fp8, w_scale = _quantize_rowwise(weight_cast)
|
|
scale_b = w_scale.view(1, -1)
|
|
else:
|
|
x_fp8, x_scale = _quantize_tensorwise(x2d)
|
|
w_fp8, w_scale = _quantize_tensorwise(weight_cast)
|
|
scale_b = w_scale
|
|
|
|
out2d = torch._scaled_mm(
|
|
x_fp8,
|
|
w_fp8.t(),
|
|
scale_a=x_scale,
|
|
scale_b=scale_b,
|
|
out_dtype=x.dtype,
|
|
)
|
|
if isinstance(out2d, tuple):
|
|
out2d = out2d[0]
|
|
|
|
if bias_cast is not None:
|
|
out2d = out2d + bias_cast
|
|
|
|
# save tensors for backward (keep original dtypes)
|
|
ctx.save_for_backward(x2d, weight, bias)
|
|
ctx.n = n
|
|
ctx.orig_shape = orig_shape
|
|
return out2d.reshape(*orig_shape[:-1], n)
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_out):
|
|
x2d, weight, bias = ctx.saved_tensors
|
|
M = x2d.shape[0]
|
|
n = ctx.n
|
|
|
|
grad_out_2d = grad_out.reshape(M, n).contiguous()
|
|
|
|
# bf16 straight-through estimator: gradients flow through the
|
|
# full-precision master weights, not the FP8 quantized values.
|
|
weight_cast = weight.to(dtype=grad_out.dtype)
|
|
x_cast = x2d.to(dtype=grad_out.dtype)
|
|
|
|
grad_x = grad_out_2d.matmul(weight_cast).reshape(*ctx.orig_shape)
|
|
grad_w = grad_out_2d.t().matmul(x_cast)
|
|
grad_b = grad_out_2d.sum(dim=0) if bias is not None else None
|
|
|
|
# None for the extra forward arg (granularity)
|
|
return grad_x, grad_w, grad_b, None
|
|
|
|
|
|
def fp8_linear_forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
# pass config **positionally**; autograd.Function.apply ignores kwargs
|
|
return _LinearFWD8BWD16Fn.apply(x, self.weight, self.bias, "tensor"), None
|