ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: fp8 backward reuses pre-quantized operands

Browse Source 
- g/x/w may each be bf16 or pre-quantized fp8 matching fmt; a pre-quantized operand skips its quantize kernel
- snapshot sx/sw/sg before the ring finalize overwrites the aliased scale slot so the gemm dequantizes with the quantize scale
- forward carries its scale to backward so gradients reuse the forward's scale
- grad_input/grad_weight forced bf16; a pre-quantized g dequantizes before the bias-sum
- regression test: two delayed steps with a changing amax must not leak the scale ratio
This commit is contained in:
ViperEkura
2026-08-25 17:24:20 +08:00
parent 2eeac02d70
commit 3e57cc8069
4 changed files with 123 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/fp8.py
+21 -11
View File
@@ -344,25 +344,31 @@ def fp8_linear_forward(
sx = _dynamic_scale(x.reshape(-1, w.size(1)), cfg.recipe, fmt)
sw = _dynamic_scale(w, cfg.recipe, fmt)
out, *_ = linear_forward_fp8(x, w, bias, sx, sw, fmt)
return out
return out, sx, sw
meta = state.get_weight_meta(w)
if not meta.w.initialized:
meta.w.seed(w, fmt)
if not meta.x.initialized:
meta.x.seed(x, fmt)
# The kernel finalizes each ring in-kernel (overwriting the scale slot), so
# the w/x scales are taken from the ring before the quantize.
# The quantize kernels finalize each ring in-kernel and overwrite the ring's
# scale slot, which ALIASES meta.*.scale (a view into the state buffer).
# Snapshot the scales first so the GEMM dequantizes with the SAME scale the
# operands were quantized with, and so the backward can reuse this step's
# scale (gradient consistency with the forward). The ring finalize still
# publishes the next step's scale into the original slot.
if w.dtype is not torch.bfloat16: # static pre-quantized weight
w_arg, sw_arg, w_ring = w, meta.w.scale, None
else:
w_arg, sw_arg, w_ring = w, meta.w.scale, meta.w.state
sx = meta.x.scale.clone()
sw = sw_arg.clone()
out, _x8, _w8, _ax, _aw = linear_forward_fp8(
x,
w_arg,
bias,
meta.x.scale,
sw_arg,
sx,
sw,
fmt,
None,
meta.x.state,
@@ -375,7 +381,7 @@ def fp8_linear_forward(
meta.x.advance()
if w_ring is not None:
meta.w.advance()
return out
return out, sx, sw
class _LinearFp8(torch.autograd.Function):
@@ -391,8 +397,8 @@ class _LinearFp8(torch.autograd.Function):
@staticmethod
def forward(ctx, x, w, bias):
cfg = _current_config()
out = fp8_linear_forward(x, w, bias, cfg)
ctx.save_for_backward(x, w)
out, sx, sw = fp8_linear_forward(x, w, bias, cfg)
ctx.save_for_backward(x, w, sx, sw)
ctx.fmt_bwd = cfg.fp8_format.bwd()
ctx.recipe = cfg.recipe
ctx.is_dynamic = isinstance(cfg.recipe, DynamicScaling)
@@ -402,7 +408,7 @@ class _LinearFp8(torch.autograd.Function):
@staticmethod
@torch.autograd.function.once_differentiable
def backward(ctx, g):
x, w = ctx.saved_tensors
x, w, _sx_fwd, _sw_fwd = ctx.saved_tensors
fmt = ctx.fmt_bwd
# Per-recipe scale/ring selection; both branches share one call below.
if ctx.is_dynamic:
@@ -414,8 +420,12 @@ class _LinearFp8(torch.autograd.Function):
meta = ctx.meta
if not meta.g.initialized:
meta.g.seed(g, fmt)
sg, ring, idx = meta.g.scale, meta.g.state, meta.g.idx
sw, sx = meta.w.scale, meta.x.scale
# Snapshot the g scale before its ring finalize overwrites the slot
# (same aliasing as the forward); reuse the forward's w/x scales so
# the backward quantizes with the scale the forward actually used.
sg = meta.g.scale.clone()
ring, idx = meta.g.state, meta.g.idx
sw, sx = _sw_fwd, _sx_fwd
grad_x, grad_w, grad_b, _amax_g = linear_backward_fp8(
g,
x,
astrai/extension/ops/fp8.py
+17 -13
View File
@@ -212,20 +212,24 @@ def linear_backward_fp8(
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""FP8 linear backward; returns ``(grad_input, grad_weight, grad_bias, amax_g)``.
The gradient (and the transposed w/x operands) are quantized to ``fmt``
(default E5M2 — larger dynamic range for gradients) and the two GEMMs run
as FP8 tensor-core products sharing a single gradient quantization.
``g_ring`` (delayed scaling) is a ``[hist | scale | counter]`` buffer the
g quantize kernel finalizes in-kernel (see :func:`linear_forward_fp8`).
``g``/``x``/``w`` may each be bf16 (quantized to ``fmt`` here) or already
pre-quantized fp8 matching ``fmt`` — a pre-quantized operand skips its
quantize kernel and is read directly by the GEMM (the ``cached_cast``
analog for the backward, symmetric with :func:`linear_forward_fp8`'s
pre-quantized weight path). ``fmt`` defaults to E5M2 (larger dynamic range
for gradients); the two GEMMs run as FP8 tensor-core products sharing a
single gradient quantization. ``g_ring`` (delayed scaling) is a
``[hist | scale | counter]`` buffer the g quantize kernel finalizes
in-kernel (see :func:`linear_forward_fp8`); a pre-quantized ``g`` does not
finalize it and reports ``amax_g = 0``.
"""
if not (
g.dtype == torch.bfloat16
and x.dtype == torch.bfloat16
and w.dtype == torch.bfloat16
):
raise TypeError(
f"fp8 backward requires bf16 inputs, got {g.dtype}/{x.dtype}/{w.dtype}"
)
f8 = _fmt_dtype(fmt)
for name, t in (("g", g), ("x", x), ("w", w)):
if t.dtype not in (torch.bfloat16, f8):
raise TypeError(
f"fp8 backward requires bf16 or pre-quantized {fmt} inputs, "
f"got {name}={t.dtype}"
)
return get_module("fp8_ops").linear_backward_fp8(
g,
x,
csrc/kernels/fp8/ops.cu
+46 -16
View File
@@ -357,10 +357,19 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
// the w/x quantizes for dX / dW never touch rings — each operand's ring
// is finalized exactly once per step (by the forward or this kernel).
TORCH_CHECK(g.is_cuda() && x.is_cuda() && w.is_cuda(), "CUDA tensors required");
TORCH_CHECK(g.scalar_type() == torch::kBFloat16 &&
x.scalar_type() == torch::kBFloat16 &&
w.scalar_type() == torch::kBFloat16,
"g, x, and w must be bf16");
// Each operand may be bf16 (quantized here) or already fp8 matching fmt
// (reused from the forward — the fp8 cached_cast analog, symmetric with
// the forward's pre-quantized w path). Pre-quantized operands skip their
// quantize kernel; their scale is still passed for the GEMM dequant.
const auto f8opt = fmt ? torch::kFloat8_e5m2 : torch::kFloat8_e4m3fn;
const bool g_prequant = g.scalar_type() == f8opt;
const bool x_prequant = x.scalar_type() == f8opt;
const bool w_prequant = w.scalar_type() == f8opt;
TORCH_CHECK(
(g.scalar_type() == torch::kBFloat16 || g_prequant) &&
(x.scalar_type() == torch::kBFloat16 || x_prequant) &&
(w.scalar_type() == torch::kBFloat16 || w_prequant),
"g, x, and w must be bf16 or pre-quantized fp8 matching fmt");
TORCH_CHECK(g.device() == x.device() && g.device() == w.device(),
"g, x, and w must be on the same device");
TORCH_CHECK(masks.size() == 3, "masks must contain three values");
@@ -375,8 +384,10 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
TORCH_CHECK(x_c.size(0) == m && x_c.size(1) == k && g_c.size(1) == n,
"backward shape mismatch");
auto grad_input = torch::empty_like(x);
auto grad_weight = torch::empty_like(w);
auto grad_input =
torch::empty_like(x, x.options().dtype(torch::kBFloat16));
auto grad_weight =
torch::empty_like(w, w.options().dtype(torch::kBFloat16));
auto grad_bias = torch::empty({0}, g.options());
// amax_g feeds a cross-block atomic_max in the g quantize kernel; zero it
// on the stream (empty + memset, not torch::zeros — see the forward).
@@ -385,8 +396,6 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
// is harmless.
auto amax_g = torch::empty({1}, g.options().dtype(torch::kFloat32));
cudaMemsetAsync(amax_g.data_ptr(), 0, sizeof(float), stream.stream());
auto f8opt = fmt ? g.options().dtype(torch::kFloat8_e5m2)
: g.options().dtype(torch::kFloat8_e4m3fn);
auto quantize = [&](const torch::Tensor& src, torch::Tensor& dst,
const torch::Tensor& scale, torch::Tensor* amax,
@@ -420,14 +429,23 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
torch::Tensor g8;
if (masks[0] || masks[1]) {
g8 = torch::empty({m, n}, f8opt);
quantize(g_c, g8, sg, &amax_g, g_ring, g_ring_idx, g_ring_margin);
if (g_prequant) {
g8 = g_c;
} else {
g8 = torch::empty({m, n}, g.options().dtype(f8opt));
quantize(g_c, g8, sg, &amax_g, g_ring, g_ring_idx, g_ring_margin);
}
}
// dX = g @ w: A = g8 [M,N] (contract over N), B = w8 [N,K] read transposed
// (b[p*b_ld + n] = w[p,n]); out = [M,K], a_ld = N, b_ld = K, contract = N.
if (masks[0]) {
auto w8 = torch::empty({n, k}, f8opt);
quantize(w_c, w8, sw, nullptr, c10::nullopt, 0, 0);
torch::Tensor w8;
if (w_prequant) {
w8 = w_c;
} else {
w8 = torch::empty({n, k}, g.options().dtype(f8opt));
quantize(w_c, w8, sw, nullptr, c10::nullopt, 0, 0);
}
auto grad_input_2d = grad_input.reshape({m, k});
FP8Params gp;
pack_gemm_params(gp, g8.data_ptr(), w8.data_ptr(),
@@ -439,19 +457,31 @@ linear_backward_fp8(torch::Tensor g, torch::Tensor x, torch::Tensor w,
// x8 [M,K] read transposed (b[p*b_ld + n] = x[p,n]); out = [N,K], a_ld = N,
// b_ld = K, contract = M.
if (masks[1]) {
auto x8 = torch::empty({m, k}, f8opt);
quantize(x_c, x8, sx, nullptr, c10::nullopt, 0, 0);
torch::Tensor x8;
if (x_prequant) {
x8 = x_c;
} else {
x8 = torch::empty({m, k}, g.options().dtype(f8opt));
quantize(x_c, x8, sx, nullptr, c10::nullopt, 0, 0);
}
FP8Params gp;
pack_gemm_params(gp, g8.data_ptr(), x8.data_ptr(),
grad_weight.data_ptr(), sg, sx, nullptr, nullptr,
nullptr, n, k, m, n, k);
run_bwd_gemm(gp, true, false);
}
if (!masks[0] && !masks[1]) {
if (!masks[0] && !masks[1] && !g_prequant) {
amax_g.copy_(g_c.abs().amax().to(torch::kFloat32));
}
C10_CUDA_CHECK(cudaGetLastError());
if (masks[2]) grad_bias = g_c.sum(0).to(g.scalar_type());
if (masks[2]) {
// A pre-quantized g has no bf16 source to reduce; dequantize with its
// scale before the batch-sum so grad_bias stays in the true gradient
// domain (sg * sum(g8)).
grad_bias = g_prequant
? (g_c.to(torch::kFloat32) * sg).sum(0).to(torch::kBFloat16)
: g_c.sum(0).to(g.scalar_type());
}
return {grad_input, grad_weight, grad_bias, amax_g};
}
tests/extension/test_fp8_mma.py
+39
View File
@@ -136,6 +136,45 @@ def test_quantize_ring_in_kernel_finalize():
idx = (idx + 1) % 4
@skip_no_fp8
def test_delayed_scaling_forward_uses_snapshot_scale():
"""Regression: the in-kernel ring finalize overwrites the scale slot, which
aliases the scale the GEMM must dequantize with. The forward must snapshot
the delayed scale first, so a changing amax across steps does not leak the
next-step scale into the output (otherwise out is off by
scale_next / scale_current)."""
torch.manual_seed(11)
dev = torch.device("cuda")
state = f8mod.fp8_state()
state.reset()
state.default_recipe = DelayedScaling(history_len=1, margin=0)
state.default_format = FP8Format.E4M3
try:
m, n, k = 32, 16, 64
x1 = torch.randn(m, k, device=dev, dtype=torch.bfloat16) * 0.5
# Smaller amax than x1: the delayed scale (amax(x1)/448) still covers
# x2 without fp8 saturation, while the next-step scale would differ —
# exactly the condition that exposed the overwrite bug.
x2 = torch.randn(m, k, device=dev, dtype=torch.bfloat16) * 0.35
w = torch.randn(n, k, device=dev, dtype=torch.bfloat16) * 0.5
bias = torch.zeros(n, device=dev, dtype=torch.bfloat16)
f8mod.fp8_linear_forward(x1, w, bias) # step 1: seeds the rings
out2, _, _ = f8mod.fp8_linear_forward(x2, w, bias) # amax changes
torch.cuda.synchronize()
# The delayed scale for step 2 is amax(x1)/448 (history_len=1); the
# GEMM must use that same scale for dequant as the quantize used.
sx = _scale(x1)
sw = _scale(w)
qx = _quantize(x2, sx)
qw = _quantize(w, sw)
expected = (qx @ qw.t() * sx * sw + bias).to(torch.bfloat16)
torch.testing.assert_close(out2, expected, atol=0.125, rtol=0.01)
finally:
state.reset()
@skip_no_fp8
def test_fp8_linear_forward_and_backward():
torch.manual_seed(7)