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
This commit is contained in:
@@ -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)
|
||||
@@ -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)
|
||||
+50
-59
@@ -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
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user