diff --git a/astrai/extension/fp8.py b/astrai/extension/fp8.py index 8120cb4..fccc326 100644 --- a/astrai/extension/fp8.py +++ b/astrai/extension/fp8.py @@ -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, diff --git a/astrai/extension/ops/fp8.py b/astrai/extension/ops/fp8.py index c399d2e..549f2a1 100644 --- a/astrai/extension/ops/fp8.py +++ b/astrai/extension/ops/fp8.py @@ -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, diff --git a/csrc/kernels/fp8/ops.cu b/csrc/kernels/fp8/ops.cu index 790f220..dad2072 100644 --- a/csrc/kernels/fp8/ops.cu +++ b/csrc/kernels/fp8/ops.cu @@ -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}; } diff --git a/tests/extension/test_fp8_mma.py b/tests/extension/test_fp8_mma.py index dc75799..61175a0 100644 --- a/tests/extension/test_fp8_mma.py +++ b/tests/extension/test_fp8_mma.py @@ -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)