From 0378e62e17c62b0b88ad4782c7553490a3d11740 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 14 Aug 2026 12:22:54 +0800 Subject: [PATCH] refactor: split fp8 into fp8_ops adapter and fp8 policy module - fp8_ops is the only module touching the pybind (kernel interface) - fp8.py keeps scaling state, delayed amax and aten::linear dispatch - remove circular imports between old fp8_ops/fp8_state/fp8_dispatch --- astrai/extension/{fp8_state.py => fp8.py} | 120 ++++++++++++++++++---- astrai/extension/fp8_dispatch.py | 84 --------------- astrai/extension/fp8_ops.py | 109 +++++++++----------- 3 files changed, 150 insertions(+), 163 deletions(-) rename astrai/extension/{fp8_state.py => fp8.py} (61%) delete mode 100644 astrai/extension/fp8_dispatch.py diff --git a/astrai/extension/fp8_state.py b/astrai/extension/fp8.py similarity index 61% rename from astrai/extension/fp8_state.py rename to astrai/extension/fp8.py index bdd4f4e..d73e859 100644 --- a/astrai/extension/fp8_state.py +++ b/astrai/extension/fp8.py @@ -1,24 +1,41 @@ -"""FP8 training state: per-tensor scales, amax history, delayed scaling. +"""FP8 training: scaling state and aten::linear dispatch. -TE-style (TransformerEngine) delayed scaling: -- weight tensors carry an ``FP8TensorMeta`` keyed by (data_ptr, shape) with a - fixed scale derived from a 16-step amax history window; -- activations/gradients reuse the quantize kernel's free atomic amax, delayed - one step (scale updated after each call, used by the next call); -- ``fp8_autocast()`` context manager toggles fp8 dispatch (like - ``torch.autocast``) and advances the scale-update counter once per step. - Entering it also ensures the aten::linear CUDA impl is registered, so - ``import astrai.extension.fp8_dispatch`` is not required by callers. +Layered (see also ``fp8_ops.py`` for the CUDA interface adapter): + +1. Kernel interface: "fp8_ops" — the only module touching the pybind. +2. Training state (this module): per-tensor scales, amax history, delayed + scaling, and the ``fp8_autocast`` context (TE-style, like + ``torch.autocast``). +3. aten::linear integration (this module): registers the CUDA impl and the + M/N alignment guard. + +Usage:: + + from astrai.extension.fp8 import fp8_autocast + + with fp8_autocast(enabled=True): + logits = model(input_ids) + loss.backward() + +Importing this module registers the aten::linear CUDA implementation. """ from contextlib import contextmanager import torch +from torch.library import Library + +from astrai.extension.fp8_ops import ( + linear_backward_scaled, + linear_forward_scaled, +) E4M3_MAX = 448.0 -# FP8 GEMM layout: D = A_SCALE * B_SCALE * A * B, so the per-tensor scales are -# amax/448 (e4m3) and the quantization divides by scale (multiplies by 1/scale). + +# --------------------------------------------------------------------------- +# Layer 2: training state (scales, amax history, delayed scaling, autocast) +# --------------------------------------------------------------------------- class FP8TensorMeta: @@ -155,11 +172,10 @@ def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None): if bias is None: bias = torch.empty(0, device=x.device, dtype=x.dtype) state = fp8_state() - mod = _mod() meta = state.get_weight_meta(w) amax_x = torch.empty(1, device=x.device, dtype=torch.float32) amax_w = torch.empty(1, device=x.device, dtype=torch.float32) - out = mod.fp8_linear_forward_scaled( + out = linear_forward_scaled( x, w, bias, @@ -178,14 +194,13 @@ def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None): def fp8_linear_backward(g, x, w, masks): """TE-style scaled fp8 linear backward (called from aten::linear_backward).""" state = fp8_state() - mod = _mod() meta = state.get_weight_meta(w) amax_g = torch.empty(1, device=g.device, dtype=torch.float32) - out = mod.fp8_linear_backward_scaled( + out = linear_backward_scaled( g, x, w, - list(masks), + masks, meta.g_scale, meta.scale, meta.x_scale, @@ -198,7 +213,72 @@ def fp8_linear_backward(g, x, w, masks): return out -def _mod(): - from astrai.extension.loader import get_module +# --------------------------------------------------------------------------- +# Layer 3: aten::linear integration +# --------------------------------------------------------------------------- - return get_module("fp8_mm") + +def fp8_linear_enable(enabled: bool = True) -> None: + """Toggle fp8 dispatch for aten::linear (global; backward runs on engine + worker threads, so a thread-local flag would be lost during backward).""" + fp8_state().enabled = enabled + + +def fp8_linear_enabled() -> bool: + return fp8_state().enabled + + +def _fp8_supported(x: torch.Tensor, w: torch.Tensor) -> bool: + """cuBLASLt fp8 requires M % 16 == 0 and N % 16 == 0 (K is padded).""" + m = x.numel() // x.size(-1) + return m % 16 == 0 and w.size(0) % 16 == 0 + + +def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None): + if ( + fp8_linear_enabled() + and x.dtype == torch.bfloat16 + and w.dtype == torch.bfloat16 + and _fp8_supported(x, w) + ): + return fp8_linear_forward(x, w, bias) + return torch.ops.aten.linear.default.redispatch( + torch._C.DispatchKeySet(torch._C.DispatchKey.CompositeImplicitAutograd), + x, + w, + bias, + ) + + +def _linear_backward_cuda_impl(input_tensor, grad_output, weight, output_mask): + if ( + fp8_linear_enabled() + and weight.dtype == torch.bfloat16 + and _fp8_supported(grad_output, weight) + ): + return fp8_linear_backward(grad_output, input_tensor, weight, list(output_mask)) + compute_dtype = weight.dtype + grad = grad_output.to(compute_dtype) + grad_2d = grad.reshape(-1, weight.size(0)) + input_2d = input_tensor.reshape(-1, input_tensor.size(-1)).to(compute_dtype) + grad_input = ( + torch.mm(grad_2d, weight) + if output_mask[0] + else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype) + ) + grad_weight = ( + torch.mm(grad_2d.t(), input_2d) + if output_mask[1] + else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype) + ) + grad_bias = ( + grad.sum(dim=0) + if output_mask[2] + else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype) + ) + return grad_input.reshape_as(input_tensor), grad_weight, grad_bias + + +_lib = Library("aten", "IMPL", "CUDA") +_lib.impl("linear", _linear_cuda_impl) +_lib.impl("linear_backward", _linear_backward_cuda_impl) diff --git a/astrai/extension/fp8_dispatch.py b/astrai/extension/fp8_dispatch.py deleted file mode 100644 index 5153f61..0000000 --- a/astrai/extension/fp8_dispatch.py +++ /dev/null @@ -1,84 +0,0 @@ -"""FP8 linear dispatch: replace aten::linear on the CUDA key, no model changes. - -``F.linear`` -> ``aten::linear`` -> dispatcher -> this CUDA impl (fp8 when -enabled) or the original composite implementation via ``redispatch``. -Enabling is per-thread; model code stays untouched. -""" - -import threading - -import torch -from torch.library import Library - -from astrai.extension.fp8_ops import fp8_linear_backward, fp8_linear_forward -from astrai.extension.fp8_state import fp8_autocast, fp8_state - - -def fp8_linear_enable(enabled: bool = True) -> None: - """Toggle fp8 dispatch for aten::linear on this thread.""" - fp8_state().enabled = enabled - - -def fp8_linear_enabled() -> bool: - return fp8_state().enabled - - -def _fp8_supported(x: torch.Tensor, w: torch.Tensor) -> bool: - """cuBLASLt fp8 requires M % 16 == 0 and N % 16 == 0 (K is padded); else fall back.""" - m = x.numel() // x.size(-1) - return m % 16 == 0 and w.size(0) % 16 == 0 - - -def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None): - if ( - fp8_linear_enabled() - and x.dtype == torch.bfloat16 - and w.dtype == torch.bfloat16 - and _fp8_supported(x, w) - ): - return fp8_linear_forward(x, w, bias) - return torch.ops.aten.linear.default.redispatch( - torch._C.DispatchKeySet(torch._C.DispatchKey.CompositeImplicitAutograd), - x, - w, - bias, - ) - - -def _linear_backward_cuda_impl(input_tensor, grad_output, weight, output_mask): - # VariableType wraps aten::linear; its backward runs aten::linear_backward - # with schema (self, grad_output, weight, mask). When fp8 is enabled the - # fused CUDA backward runs in one call (scale-corrected); otherwise the - # plain bf16/fp32 math, dtype aligned to the leaf weight: - # grad_input = g @ W, grad_weight = g^T @ X, grad_bias = sum(g, dim=0) - if ( - fp8_linear_enabled() - and weight.dtype == torch.bfloat16 - and _fp8_supported(grad_output, weight) - ): - return fp8_linear_backward(grad_output, input_tensor, weight, list(output_mask)) - compute_dtype = weight.dtype - grad = grad_output.to(compute_dtype) - grad_2d = grad.reshape(-1, weight.size(0)) - input_2d = input_tensor.reshape(-1, input_tensor.size(-1)).to(compute_dtype) - grad_input = ( - torch.mm(grad_2d, weight) - if output_mask[0] - else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype) - ) - grad_weight = ( - torch.mm(grad_2d.t(), input_2d) - if output_mask[1] - else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype) - ) - grad_bias = ( - grad.sum(dim=0) - if output_mask[2] - else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype) - ) - return grad_input.reshape_as(input_tensor), grad_weight, grad_bias - - -_lib = Library("aten", "IMPL", "CUDA") -_lib.impl("linear", _linear_cuda_impl) -_lib.impl("linear_backward", _linear_backward_cuda_impl) diff --git a/astrai/extension/fp8_ops.py b/astrai/extension/fp8_ops.py index ace30c5..de0f13b 100644 --- a/astrai/extension/fp8_ops.py +++ b/astrai/extension/fp8_ops.py @@ -1,10 +1,12 @@ -"""FP8 matrix-multiply op (torch.library custom_op) and FP8 linear replacement. +"""FP8 CUDA kernel interface adapter (the only module touching the pybind. -Dispatch table: -- Meta (register_fake): shapes only, for torch.compile / dynamic shapes -- CUDA: csrc fp8_mm kernel (cuBLASLt TN fp8 GEMM, e4m3 in, fp32 acc/out) -- CPU: fp32 fallback (testing) -- AutogradCUDA (register_autograd): bf16 backward, scale-corrected +Isolates the ``fp8_mm`` CUDA extension behind stable Python functions: +- availability / dtype checks and clear errors +- torch.library ``custom::fp8_mm`` registration (meta + CPU fallback) +- quantize-in-GEMM primitives used by ``fp8.py`` training state + +Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``; +this module is stateless. """ import torch @@ -13,70 +15,59 @@ from torch.library import custom_op from astrai.extension.loader import get_module, is_available -@custom_op("custom::fp8_mm", mutates_args=()) -def fp8_mm( - a: torch.Tensor, b: torch.Tensor, sx: torch.Tensor, sw: torch.Tensor -) -> torch.Tensor: - """FP8 e4m3 GEMM: a[M,K] x b[N,K] -> fp32[M,N], scales applied by the caller. - - a/b arrive pre-scaled (divided by sx/sw) fp8 tensors; the op returns the - unscaled fp32 result so scale math stays in autograd-land. - """ - - -@fp8_mm.register_fake -def _fp8_mm_fake(a, b, sx, sw): - return torch.empty((a.size(0), b.size(1)), device=a.device, dtype=torch.float32) - - -@fp8_mm.register_kernel("cuda") -def _fp8_mm_cuda(a, b, sx, sw): +def _mod(): if not is_available("fp8_mm"): raise RuntimeError( "CUDA kernel 'fp8_mm' is not available. Build with CSRC_KERNELS=true." ) - return get_module("fp8_mm").fp8_mm(a, b) + return get_module("fp8_mm") + + +@custom_op("custom::fp8_mm", mutates_args=()) +def fp8_mm( + a: torch.Tensor, b: torch.Tensor, sx: torch.Tensor, sw: torch.Tensor +) -> torch.Tensor: + """FP8 e4m3 GEMM: a[M,K] x b[N,K] -> bf16[M,N] (pre-scaled inputs).""" + + +@fp8_mm.register_fake +def _fp8_mm_fake(a, b, sx, sw): + return torch.empty((a.size(0), b.size(1)), device=a.device, dtype=torch.bfloat16) + + +@fp8_mm.register_kernel("cuda") +def _fp8_mm_cuda(a, b, sx, sw): + return _mod().fp8_mm(a, b) @fp8_mm.register_kernel("cpu") def _fp8_mm_cpu(a, b, sx, sw): - return torch.mm(a.float(), b.float().t()) + return torch.mm(a.float(), b.float().t()).to(torch.bfloat16) -def _fp8_mm_setup_context(ctx, inputs, output): - ctx.save_for_backward(*inputs) +def linear_forward_scaled(x, w, bias, sx, sw, sx_inv, sw_inv, amax_x, amax_w): + """Quantize x/w with per-tensor scales + cuBLASLt GEMM + bias -> bf16. - -def _fp8_mm_backward(ctx, g): - """Scale-corrected straight-through gradients. - - out = F(a, b) * (sx * sw) with F(a, b) = a @ b^T, a = x/sx, b = w/sw: - dx = g * sw @ b (dout/dx = dF/da * 1/sx * sx*sw) - dW = (g * sx)^T @ a (dout/dw = dF/db * 1/sw * sx*sw) - bf16 GEMMs keep gradients in range (e4m3 saturates at 448). + x/w: [..., K] / [N, K] bf16; sx/sw: f32 scale tensors (device scalars); + sx_inv/sw_inv: 1/scale; amax_x/amax_w: f32 buffers receiving max-abs. """ - a, b, sx, sw = ctx.saved_tensors - ga = torch.mm(g * sw, b.float()) - gb = torch.mm((g * sx).t(), a.float()) - return ga.to(torch.bfloat16), gb.to(torch.bfloat16), None, None + if not (x.dtype == torch.bfloat16 and w.dtype == torch.bfloat16): + raise TypeError(f"fp8 forward requires bf16 inputs, got {x.dtype}/{w.dtype}") + return _mod().fp8_linear_forward_scaled( + x, w, bias, sx, sw, sx_inv, sw_inv, amax_x, amax_w + ) -fp8_mm.register_autograd(_fp8_mm_backward, setup_context=_fp8_mm_setup_context) - - -def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None): - """TE-style scaled fp8 linear forward (delegates to fp8_state).""" - from astrai.extension.fp8_state import fp8_linear_forward as _f - - return _f(x, w, bias) - - -def fp8_linear_backward(g, x, w, masks): - """TE-style scaled fp8 linear backward (delegates to fp8_state).""" - from astrai.extension.fp8_state import fp8_linear_backward as _b - - return _b(g, x, w, masks) - - -def fp8_available() -> bool: - return is_available("fp8_mm") +def linear_backward_scaled(g, x, w, masks, sg, sw, sx, sg_inv, sw_inv, sx_inv, amax_g): + """dX = g @ W, dW = g^T @ X, dB = sum(g) with per-tensor scales.""" + 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}" + ) + return _mod().fp8_linear_backward_scaled( + g, x, w, masks, sg, sw, sx, sg_inv, sw_inv, sx_inv, amax_g + )