feat: add te-style scaled fp8 training via fp8_autocast
- per-tensor scales applied inside cublasLt via A_SCALE/B_SCALE - delayed scaling: weight amax history ring, refresh every 16 steps - quantize kernels emit atomic amax, device-side scale updates - fp8_autocast context toggles aten::linear dispatch like torch.autocast - fallback to bf16 when M/N not 16-aligned (fp8 gemm constraint) - x/g scales delayed one step, reuse free atomic amax (no abs/max reduce)
This commit is contained in:
@@ -11,21 +11,31 @@ import torch
|
||||
from torch.library import Library
|
||||
|
||||
from astrai.extension.fp8_ops import fp8_linear_backward, fp8_linear_forward
|
||||
|
||||
_state = threading.local()
|
||||
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."""
|
||||
_state.enabled = enabled
|
||||
fp8_state().enabled = enabled
|
||||
|
||||
|
||||
def fp8_linear_enabled() -> bool:
|
||||
return getattr(_state, "enabled", False)
|
||||
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:
|
||||
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),
|
||||
@@ -41,7 +51,11 @@ def _linear_backward_cuda_impl(input_tensor, grad_output, weight, output_mask):
|
||||
# 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:
|
||||
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)
|
||||
|
||||
@@ -65,21 +65,17 @@ 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):
|
||||
"""FP8 replacement for F.linear(x, w, bias), fused in one CUDA call.
|
||||
"""TE-style scaled fp8 linear forward (delegates to fp8_state)."""
|
||||
from astrai.extension.fp8_state import fp8_linear_forward as _f
|
||||
|
||||
x: [..., K] bf16 (any leading dims), w: [N,K] bf16 (in_dim=K).
|
||||
The kernel pipeline (scale cast -> cublasLt fp8 GEMM -> unscale + bias ->
|
||||
transpose -> bf16) runs inside a single extension call, so Python-side
|
||||
dispatch overhead is paid once per linear instead of per operator.
|
||||
"""
|
||||
if bias is None:
|
||||
bias = torch.empty(0, device=x.device, dtype=x.dtype)
|
||||
return get_module("fp8_mm").fp8_linear_forward(x, w, bias)
|
||||
return _f(x, w, bias)
|
||||
|
||||
|
||||
def fp8_linear_backward(g, x, w, masks):
|
||||
"""Fused linear backward (dX/dW/dB in one CUDA call, scale-corrected)."""
|
||||
return get_module("fp8_mm").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:
|
||||
|
||||
@@ -0,0 +1,204 @@
|
||||
"""FP8 training state: per-tensor scales, amax history, delayed scaling.
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
from contextlib import contextmanager
|
||||
|
||||
import torch
|
||||
|
||||
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).
|
||||
|
||||
|
||||
class FP8TensorMeta:
|
||||
"""Scales + amax state for one weight tensor and its paired activations.
|
||||
|
||||
- weight: delayed scale from a 16-step amax history window (TE style)
|
||||
- x/g: delayed one step, reuse the quantize kernel's free atomic amax
|
||||
"""
|
||||
|
||||
__slots__ = (
|
||||
"scale",
|
||||
"scale_inv",
|
||||
"amax_history",
|
||||
"idx",
|
||||
"x_scale",
|
||||
"x_scale_inv",
|
||||
"g_scale",
|
||||
"g_scale_inv",
|
||||
)
|
||||
|
||||
def __init__(self, device: torch.device, update_interval: int):
|
||||
self.scale = torch.ones(1, device=device, dtype=torch.float32)
|
||||
self.scale_inv = torch.ones(1, device=device, dtype=torch.float32)
|
||||
self.amax_history = torch.ones(
|
||||
update_interval, device=device, dtype=torch.float32
|
||||
)
|
||||
self.idx = 0
|
||||
self.x_scale = torch.ones(1, device=device, dtype=torch.float32)
|
||||
self.x_scale_inv = torch.ones(1, device=device, dtype=torch.float32)
|
||||
self.g_scale = torch.ones(1, device=device, dtype=torch.float32)
|
||||
self.g_scale_inv = torch.ones(1, device=device, dtype=torch.float32)
|
||||
|
||||
def record(self, amax: torch.Tensor) -> None:
|
||||
"""Push the latest amax into the ring buffer (device-side copy, no sync)."""
|
||||
self.amax_history[self.idx] = amax.reshape(())
|
||||
self.idx = (self.idx + 1) % self.amax_history.numel()
|
||||
|
||||
def refresh(self) -> None:
|
||||
"""Recompute scale from the amax history window (delayed scaling)."""
|
||||
amax = self.amax_history.max()
|
||||
if amax > 0:
|
||||
self.scale.copy_(amax / E4M3_MAX)
|
||||
self.scale_inv.copy_(E4M3_MAX / amax)
|
||||
|
||||
|
||||
class FP8State:
|
||||
"""Global fp8 training state, TE-style."""
|
||||
|
||||
def __init__(self, update_interval: int = 16):
|
||||
self.enabled = False
|
||||
self.update_interval = update_interval
|
||||
self.step_count = 0
|
||||
self._metas: dict[tuple, FP8TensorMeta] = {}
|
||||
self._last_device: torch.device | None = None
|
||||
|
||||
def _get_device(self, t: torch.Tensor) -> torch.device:
|
||||
if self._last_device is None:
|
||||
self._last_device = t.device
|
||||
return t.device
|
||||
|
||||
def get_weight_meta(self, w: torch.Tensor) -> FP8TensorMeta:
|
||||
key = (w.data_ptr(), w.shape, w.dtype)
|
||||
meta = self._metas.get(key)
|
||||
if meta is None:
|
||||
meta = FP8TensorMeta(self._get_device(w), self.update_interval)
|
||||
self._metas[key] = meta
|
||||
return meta
|
||||
|
||||
def step(self) -> None:
|
||||
"""Advance the counter and refresh all weight scales every N steps."""
|
||||
self.step_count += 1
|
||||
if self.step_count % self.update_interval == 0:
|
||||
for meta in self._metas.values():
|
||||
meta.refresh()
|
||||
|
||||
def reset(self) -> None:
|
||||
self.enabled = False
|
||||
self.step_count = 0
|
||||
self._metas.clear()
|
||||
self._last_device = None
|
||||
|
||||
|
||||
# Global singleton: autograd backward runs on the engine worker threads, so
|
||||
# thread-local state would lose the fp8 flag during loss.backward(). The GIL
|
||||
# protects Python-side mutation; the CUDA kernels take their own mutex.
|
||||
_state = FP8State()
|
||||
|
||||
|
||||
def fp8_state() -> FP8State:
|
||||
return _state
|
||||
|
||||
|
||||
@contextmanager
|
||||
def fp8_autocast(enabled: bool = True, update_interval: int = 16):
|
||||
"""Autocast-style context: fp8 linear dispatch on this thread.
|
||||
|
||||
Usage::
|
||||
|
||||
with fp8_autocast(enabled=True):
|
||||
logits = model(input_ids) # aten::linear -> fp8 path
|
||||
loss.backward()
|
||||
|
||||
The scale-update counter advances once per ``enter`` (one training step),
|
||||
refreshing weight scales from their amax history every ``update_interval``.
|
||||
"""
|
||||
state = fp8_state()
|
||||
prev_enabled = state.enabled
|
||||
prev_interval = state.update_interval
|
||||
state.enabled = enabled
|
||||
state.update_interval = update_interval
|
||||
try:
|
||||
if enabled:
|
||||
state.step()
|
||||
yield
|
||||
finally:
|
||||
state.enabled = prev_enabled
|
||||
state.update_interval = prev_interval
|
||||
|
||||
|
||||
def _update_delayed_scale(scale, scale_inv, amax) -> None:
|
||||
"""scale = amax / 448 for the *next* call (device-side, no sync)."""
|
||||
amax_f = amax.reshape(()).to(torch.float32).clamp_min(1e-12)
|
||||
scale.copy_(amax_f / E4M3_MAX)
|
||||
scale_inv.copy_(E4M3_MAX / amax_f)
|
||||
|
||||
|
||||
def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None):
|
||||
"""TE-style scaled fp8 linear forward (called from the aten::linear impl).
|
||||
|
||||
x uses the delayed scale of its paired weight meta (amax from the previous
|
||||
forward of this linear); the quantize kernel emits the current amax for the
|
||||
next step. No extra abs/max reduce.
|
||||
"""
|
||||
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(
|
||||
x,
|
||||
w,
|
||||
bias,
|
||||
meta.x_scale,
|
||||
meta.scale,
|
||||
meta.x_scale_inv,
|
||||
meta.scale_inv,
|
||||
amax_x,
|
||||
amax_w,
|
||||
)
|
||||
meta.record(amax_w)
|
||||
_update_delayed_scale(meta.x_scale, meta.x_scale_inv, amax_x)
|
||||
return out
|
||||
|
||||
|
||||
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(
|
||||
g,
|
||||
x,
|
||||
w,
|
||||
list(masks),
|
||||
meta.g_scale,
|
||||
meta.scale,
|
||||
meta.x_scale,
|
||||
meta.g_scale_inv,
|
||||
meta.scale_inv,
|
||||
meta.x_scale_inv,
|
||||
amax_g,
|
||||
)
|
||||
_update_delayed_scale(meta.g_scale, meta.g_scale_inv, amax_g)
|
||||
return out
|
||||
|
||||
|
||||
def _mod():
|
||||
from astrai.extension.loader import get_module
|
||||
|
||||
return get_module("fp8_mm")
|
||||
+159
-56
@@ -87,7 +87,11 @@ static cublasStatus_t get_algo_cached(int64_t m, int64_t k, int64_t n,
|
||||
cublasLtMatmulAlgo_t* algo);
|
||||
|
||||
static void fp8_gemm_into(torch::Tensor lhs, torch::Tensor rhs, torch::Tensor out,
|
||||
int64_t m, int64_t k, int64_t n, cudaStream_t stream);
|
||||
int64_t m, int64_t k, int64_t n,
|
||||
const float* a_scale, const float* b_scale,
|
||||
cudaStream_t stream);
|
||||
|
||||
static const float k_scale_one = 1.0f;
|
||||
|
||||
static void set_layout(cublasLtMatrixLayout_t layout, int64_t rows, int64_t cols,
|
||||
int64_t ld) {
|
||||
@@ -117,34 +121,70 @@ torch::Tensor fp8_mm(torch::Tensor a, torch::Tensor b) {
|
||||
|
||||
auto buf = torch::empty({m, n}, a_c.options().dtype(torch::kBFloat16));
|
||||
ensure_cublas_lt();
|
||||
fp8_gemm_into(a_c, b_c, buf, m, k, n, stream.stream());
|
||||
fp8_gemm_into(a_c, b_c, buf, m, k, n, &k_scale_one, &k_scale_one,
|
||||
stream.stream());
|
||||
return buf;
|
||||
}
|
||||
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Fused FP8 linear forward: one call = scale cast x8/w8 -> cublasLt GEMM
|
||||
// (bf16 output) -> bias in-place -> bf16 [..., N].
|
||||
// Quantize: bf16 * scale_inv -> fp8, one atomicMax amax per kernel call.
|
||||
// amax_ptr must be zeroed before launch; float-bits atomicMax works because
|
||||
// |v| >= 0 has a monotonic IEEE bit pattern.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
__global__ void cast_bf16_to_fp8_kernel(
|
||||
const __nv_bfloat16* __restrict__ src, __nv_fp8_e4m3* __restrict__ dst,
|
||||
int64_t n) {
|
||||
int64_t i = blockIdx.x * (int64_t)blockDim.x + threadIdx.x;
|
||||
if (i >= n) return;
|
||||
dst[i] = __nv_fp8_e4m3(__bfloat162float(src[i]));
|
||||
template <typename T8>
|
||||
__device__ __forceinline__ T8 cast_fp8(float v);
|
||||
|
||||
template <>
|
||||
__device__ __forceinline__ __nv_fp8_e4m3 cast_fp8<__nv_fp8_e4m3>(float v) {
|
||||
return __nv_fp8_e4m3(v);
|
||||
}
|
||||
|
||||
__global__ void transpose_cast_bf16_to_fp8_kernel(
|
||||
const __nv_bfloat16* __restrict__ src, __nv_fp8_e4m3* __restrict__ dst,
|
||||
int64_t rows, int64_t cols) {
|
||||
__shared__ __nv_fp8_e4m3 tile[32][33];
|
||||
template <>
|
||||
__device__ __forceinline__ __nv_fp8_e5m2 cast_fp8<__nv_fp8_e5m2>(float v) {
|
||||
return __nv_fp8_e5m2(v);
|
||||
}
|
||||
|
||||
template <typename T8>
|
||||
__global__ void quantize_kernel(const __nv_bfloat16* __restrict__ src,
|
||||
const float* __restrict__ scale_inv,
|
||||
T8* __restrict__ dst,
|
||||
float* __restrict__ amax_ptr, int64_t n) {
|
||||
int64_t i = blockIdx.x * (int64_t)blockDim.x + threadIdx.x;
|
||||
float amax = 0.f;
|
||||
if (i < n) {
|
||||
float v = __bfloat162float(src[i]) * *scale_inv;
|
||||
dst[i] = cast_fp8<T8>(v);
|
||||
amax = fabsf(v);
|
||||
}
|
||||
for (int off = 16; off; off >>= 1)
|
||||
amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, off));
|
||||
__shared__ float sm[8];
|
||||
if ((threadIdx.x & 31) == 0) sm[threadIdx.x >> 5] = amax;
|
||||
__syncthreads();
|
||||
if (threadIdx.x == 0) {
|
||||
float m = 0.f;
|
||||
for (int w = 0; w < blockDim.x / 32; ++w) m = fmaxf(m, sm[w]);
|
||||
atomicMax(reinterpret_cast<unsigned*>(amax_ptr), __float_as_uint(m));
|
||||
}
|
||||
}
|
||||
|
||||
// Same but with a transpose (rows x cols bf16 row-major -> fp8 [cols, rows]).
|
||||
template <typename T8>
|
||||
__global__ void transpose_quantize_kernel(
|
||||
const __nv_bfloat16* __restrict__ src, const float* __restrict__ scale_inv,
|
||||
T8* __restrict__ dst, float* __restrict__ amax_ptr, int64_t rows,
|
||||
int64_t cols) {
|
||||
__shared__ T8 tile[32][33];
|
||||
int64_t x = blockIdx.x * 32 + threadIdx.x;
|
||||
int64_t y = blockIdx.y * 32 + threadIdx.y;
|
||||
float amax = 0.f;
|
||||
for (int j = 0; j < 32; j += 8) {
|
||||
if (x < cols && y + j < rows) {
|
||||
tile[threadIdx.y + j][threadIdx.x] =
|
||||
__nv_fp8_e4m3(__bfloat162float(src[(y + j) * cols + x]));
|
||||
float v = __bfloat162float(src[(y + j) * cols + x]) * *scale_inv;
|
||||
tile[threadIdx.y + j][threadIdx.x] = cast_fp8<T8>(v);
|
||||
amax = fmaxf(amax, fabsf(v));
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
@@ -156,6 +196,16 @@ __global__ void transpose_cast_bf16_to_fp8_kernel(
|
||||
dst[(y + j) * rows + x] = tile[threadIdx.x][threadIdx.y + j];
|
||||
}
|
||||
}
|
||||
for (int off = 16; off; off >>= 1)
|
||||
amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, off));
|
||||
__shared__ float sm[8];
|
||||
if ((threadIdx.x & 31) == 0) sm[threadIdx.x >> 5] = amax;
|
||||
__syncthreads();
|
||||
if (threadIdx.x == 0) {
|
||||
float m = 0.f;
|
||||
for (int w = 0; w < blockDim.x / 32; ++w) m = fmaxf(m, sm[w]);
|
||||
atomicMax(reinterpret_cast<unsigned*>(amax_ptr), __float_as_uint(m));
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void bias_add_bf16_kernel(
|
||||
@@ -196,11 +246,21 @@ static cublasStatus_t get_algo_cached(int64_t m, int64_t k, int64_t n,
|
||||
}
|
||||
|
||||
static void fp8_gemm_into(torch::Tensor lhs, torch::Tensor rhs, torch::Tensor out,
|
||||
int64_t m, int64_t k, int64_t n, cudaStream_t stream) {
|
||||
int64_t m, int64_t k, int64_t n,
|
||||
const float* a_scale, const float* b_scale,
|
||||
cudaStream_t stream) {
|
||||
std::lock_guard<std::recursive_mutex> lock(g_mutex);
|
||||
set_layout(g_layout_a, k, n, k); // param A = rhs (op=T -> [N,K])
|
||||
set_layout(g_layout_b, k, m, k); // param B = lhs (op=N -> [K,M])
|
||||
set_layout(g_layout_c, n, m, n); // col-major [N,M] == row-major [M,N]
|
||||
// Per-tensor FP32 scales applied inside the GEMM:
|
||||
// D = alpha * A_SCALE * B_SCALE * A * B (alpha = 1).
|
||||
TORCH_CHECK(cublasLtMatmulDescSetAttribute(
|
||||
g_desc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &a_scale,
|
||||
sizeof(a_scale)) == CUBLAS_STATUS_SUCCESS);
|
||||
TORCH_CHECK(cublasLtMatmulDescSetAttribute(
|
||||
g_desc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &b_scale,
|
||||
sizeof(b_scale)) == CUBLAS_STATUS_SUCCESS);
|
||||
float alpha = 1.0f, beta = 0.0f;
|
||||
static AlgoCache cache;
|
||||
cublasLtMatmulAlgo_t algo;
|
||||
@@ -215,11 +275,22 @@ static void fp8_gemm_into(torch::Tensor lhs, torch::Tensor rhs, torch::Tensor ou
|
||||
"cublasLtMatmul failed: ", cublasLtGetStatusName(st));
|
||||
}
|
||||
|
||||
torch::Tensor fp8_linear_forward(torch::Tensor x, torch::Tensor w,
|
||||
torch::Tensor bias) {
|
||||
// ---------------------------------------------------------------------------
|
||||
// Scaled FP8 linear forward: quantize x/w with per-tensor scales -> cublasLt
|
||||
// GEMM (scales applied inside) -> bias in-place -> bf16 [..., N].
|
||||
// sx/sw: f32 scale tensors (device scalars); sx_inv/sw_inv: 1/scale.
|
||||
// amax_x/amax_w: f32 buffers receiving max-abs of the quantized tensors.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
torch::Tensor fp8_linear_forward_scaled(torch::Tensor x, torch::Tensor w,
|
||||
torch::Tensor bias, torch::Tensor sx,
|
||||
torch::Tensor sw, torch::Tensor sx_inv,
|
||||
torch::Tensor sw_inv,
|
||||
torch::Tensor amax_x,
|
||||
torch::Tensor amax_w) {
|
||||
TORCH_CHECK(x.is_cuda() && w.is_cuda(), "CUDA tensors required");
|
||||
TORCH_CHECK(x.dtype() == torch::kBFloat16, "x must be bf16");
|
||||
TORCH_CHECK(w.dtype() == torch::kBFloat16, "w must be bf16");
|
||||
TORCH_CHECK(x.dtype() == torch::kBFloat16 && w.dtype() == torch::kBFloat16,
|
||||
"x and w must be bf16");
|
||||
const at::cuda::OptionalCUDAGuard guard(x.device());
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
@@ -229,22 +300,30 @@ torch::Tensor fp8_linear_forward(torch::Tensor x, torch::Tensor w,
|
||||
TORCH_CHECK(w_c.size(1) == k, "inner dim mismatch");
|
||||
ensure_cublas_lt();
|
||||
|
||||
const float* sx_ptr = sx.data_ptr<float>();
|
||||
const float* sw_ptr = sw.data_ptr<float>();
|
||||
const float* sxi_ptr = sx_inv.data_ptr<float>();
|
||||
const float* swi_ptr = sw_inv.data_ptr<float>();
|
||||
float* amax_x_ptr = amax_x.data_ptr<float>();
|
||||
float* amax_w_ptr = amax_w.data_ptr<float>();
|
||||
C10_CUDA_CHECK(cudaMemsetAsync(amax_x_ptr, 0, sizeof(float), stream.stream()));
|
||||
C10_CUDA_CHECK(cudaMemsetAsync(amax_w_ptr, 0, sizeof(float), stream.stream()));
|
||||
|
||||
auto x8 = torch::empty({m, k}, x_c.options().dtype(torch::kFloat8_e4m3fn));
|
||||
auto w8 = torch::empty({n, k}, w_c.options().dtype(torch::kFloat8_e4m3fn));
|
||||
int64_t block = 256;
|
||||
cast_bf16_to_fp8_kernel<<<(unsigned)((m * k + block - 1) / block), block, 0, stream>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(x_c.data_ptr()),
|
||||
reinterpret_cast<__nv_fp8_e4m3*>(x8.data_ptr()), m * k);
|
||||
cast_bf16_to_fp8_kernel<<<(unsigned)((n * k + block - 1) / block), block, 0, stream>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(w_c.data_ptr()),
|
||||
reinterpret_cast<__nv_fp8_e4m3*>(w8.data_ptr()), n * k);
|
||||
quantize_kernel<__nv_fp8_e4m3>
|
||||
<<<(unsigned)((m * k + block - 1) / block), block, 0, stream.stream()>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(x_c.data_ptr()), sxi_ptr,
|
||||
reinterpret_cast<__nv_fp8_e4m3*>(x8.data_ptr()), amax_x_ptr, m * k);
|
||||
quantize_kernel<__nv_fp8_e4m3>
|
||||
<<<(unsigned)((n * k + block - 1) / block), block, 0, stream.stream()>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(w_c.data_ptr()), swi_ptr,
|
||||
reinterpret_cast<__nv_fp8_e4m3*>(w8.data_ptr()), amax_w_ptr, n * k);
|
||||
C10_CUDA_CHECK(cudaGetLastError());
|
||||
|
||||
// A/B swap makes the col-major [N,M] storage directly represent the
|
||||
// row-major output [M,N]. Keep this buffer as the public output so the
|
||||
// bias path does not need a second allocation or a copy kernel.
|
||||
auto out = torch::empty({m, n}, x_c.options());
|
||||
fp8_gemm_into(x8, w8, out, m, k, n, stream.stream());
|
||||
fp8_gemm_into(x8, w8, out, m, k, n, sw_ptr, sx_ptr, stream.stream());
|
||||
|
||||
if (bias.defined() && bias.numel() > 0) {
|
||||
TORCH_CHECK(bias.scalar_type() == torch::kBFloat16 && bias.numel() == n,
|
||||
@@ -261,12 +340,15 @@ torch::Tensor fp8_linear_forward(torch::Tensor x, torch::Tensor w,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Fused FP8 linear backward: dX = g @ W, dW = g^T @ X, dB = sum(g).
|
||||
// Scaled FP8 linear backward: dX = g @ W, dW = g^T @ X, dB = sum(g).
|
||||
// Scales: g uses sg (immediate), w/x reuse the forward scales.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> fp8_linear_backward(
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> fp8_linear_backward_scaled(
|
||||
torch::Tensor g, torch::Tensor x, torch::Tensor w,
|
||||
std::vector<int64_t> masks) {
|
||||
std::vector<int64_t> masks, torch::Tensor sg, torch::Tensor sw,
|
||||
torch::Tensor sx, torch::Tensor sg_inv, torch::Tensor sw_inv,
|
||||
torch::Tensor sx_inv, torch::Tensor amax_g) {
|
||||
const at::cuda::OptionalCUDAGuard guard(g.device());
|
||||
TORCH_CHECK(g.dtype() == torch::kBFloat16 && x.dtype() == torch::kBFloat16 &&
|
||||
w.dtype() == torch::kBFloat16,
|
||||
@@ -286,6 +368,15 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> fp8_linear_backward(
|
||||
auto grad_bias = torch::empty({0}, g_c.options().dtype(g.dtype()));
|
||||
ensure_cublas_lt();
|
||||
|
||||
const float* sg_ptr = sg.data_ptr<float>();
|
||||
const float* sw_ptr = sw.data_ptr<float>();
|
||||
const float* sx_ptr = sx.data_ptr<float>();
|
||||
const float* sgi_ptr = sg_inv.data_ptr<float>();
|
||||
const float* swi_ptr = sw_inv.data_ptr<float>();
|
||||
const float* sxi_ptr = sx_inv.data_ptr<float>();
|
||||
float* amax_g_ptr = amax_g.data_ptr<float>();
|
||||
C10_CUDA_CHECK(cudaMemsetAsync(amax_g_ptr, 0, sizeof(float), stream.stream()));
|
||||
|
||||
auto fp8_options = g_c.options().dtype(torch::kFloat8_e4m3fn);
|
||||
auto g8 = torch::empty({m, n}, fp8_options);
|
||||
auto gt8 = masks[1] ? torch::empty({n, m}, fp8_options) : torch::Tensor();
|
||||
@@ -293,28 +384,36 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> fp8_linear_backward(
|
||||
auto xt8 = masks[1] ? torch::empty({k, m}, fp8_options) : torch::Tensor();
|
||||
|
||||
int64_t block = 256;
|
||||
cast_bf16_to_fp8_kernel<<<(unsigned)((m * n + block - 1) / block), block, 0, stream>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(g_c.data_ptr()),
|
||||
reinterpret_cast<__nv_fp8_e4m3*>(g8.data_ptr()), m * n);
|
||||
quantize_kernel<__nv_fp8_e4m3>
|
||||
<<<(unsigned)((m * n + block - 1) / block), block, 0, stream.stream()>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(g_c.data_ptr()), sgi_ptr,
|
||||
reinterpret_cast<__nv_fp8_e4m3*>(g8.data_ptr()), amax_g_ptr, m * n);
|
||||
dim3 threads(32, 8);
|
||||
if (masks[0]) {
|
||||
dim3 blocks((k + 31) / 32, (n + 31) / 32);
|
||||
transpose_cast_bf16_to_fp8_kernel<<<blocks, threads, 0, stream>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(w_c.data_ptr()),
|
||||
reinterpret_cast<__nv_fp8_e4m3*>(wt8.data_ptr()), n, k);
|
||||
fp8_gemm_into(g8, wt8, grad_input.reshape({m, k}), m, n, k,
|
||||
stream.stream());
|
||||
transpose_quantize_kernel<__nv_fp8_e4m3>
|
||||
<<<blocks, threads, 0, stream.stream()>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(w_c.data_ptr()), swi_ptr,
|
||||
reinterpret_cast<__nv_fp8_e4m3*>(wt8.data_ptr()), amax_g_ptr,
|
||||
n, k);
|
||||
fp8_gemm_into(g8, wt8, grad_input.reshape({m, k}), m, n, k, sg_ptr,
|
||||
sw_ptr, stream.stream());
|
||||
}
|
||||
if (masks[1]) {
|
||||
dim3 g_blocks((n + 31) / 32, (m + 31) / 32);
|
||||
dim3 x_blocks((k + 31) / 32, (m + 31) / 32);
|
||||
transpose_cast_bf16_to_fp8_kernel<<<g_blocks, threads, 0, stream>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(g_c.data_ptr()),
|
||||
reinterpret_cast<__nv_fp8_e4m3*>(gt8.data_ptr()), m, n);
|
||||
transpose_cast_bf16_to_fp8_kernel<<<x_blocks, threads, 0, stream>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(x_c.data_ptr()),
|
||||
reinterpret_cast<__nv_fp8_e4m3*>(xt8.data_ptr()), m, k);
|
||||
fp8_gemm_into(gt8, xt8, grad_weight, n, m, k, stream.stream());
|
||||
transpose_quantize_kernel<__nv_fp8_e4m3>
|
||||
<<<g_blocks, threads, 0, stream.stream()>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(g_c.data_ptr()), sgi_ptr,
|
||||
reinterpret_cast<__nv_fp8_e4m3*>(gt8.data_ptr()), amax_g_ptr,
|
||||
m, n);
|
||||
transpose_quantize_kernel<__nv_fp8_e4m3>
|
||||
<<<x_blocks, threads, 0, stream.stream()>>>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(x_c.data_ptr()), sxi_ptr,
|
||||
reinterpret_cast<__nv_fp8_e4m3*>(xt8.data_ptr()), amax_g_ptr,
|
||||
m, k);
|
||||
fp8_gemm_into(gt8, xt8, grad_weight, n, m, k, sg_ptr, sx_ptr,
|
||||
stream.stream());
|
||||
}
|
||||
C10_CUDA_CHECK(cudaGetLastError());
|
||||
if (masks[2]) {
|
||||
@@ -327,12 +426,16 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> fp8_linear_backward(
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("fp8_mm", &fp8_mm, py::arg("a"), py::arg("b"),
|
||||
"FP8 e4m3 GEMM: a[M,K] x b[N,K] -> bf16[M,N] (pre-scaled inputs)");
|
||||
m.def("fp8_linear_forward", &fp8_linear_forward,
|
||||
py::arg("x"), py::arg("w"), py::arg("bias"),
|
||||
"Fused FP8 linear forward: scale cast + cublasLt GEMM + bias "
|
||||
"-> bf16, single call");
|
||||
m.def("fp8_linear_backward", &fp8_linear_backward,
|
||||
m.def("fp8_linear_forward_scaled", &fp8_linear_forward_scaled,
|
||||
py::arg("x"), py::arg("w"), py::arg("bias"), py::arg("sx"),
|
||||
py::arg("sw"), py::arg("sx_inv"), py::arg("sw_inv"),
|
||||
py::arg("amax_x"), py::arg("amax_w"),
|
||||
"Scaled FP8 linear forward: quantize with per-tensor scales + "
|
||||
"cublasLt GEMM (scales applied inside) + bias -> bf16");
|
||||
m.def("fp8_linear_backward_scaled", &fp8_linear_backward_scaled,
|
||||
py::arg("g"), py::arg("x"), py::arg("w"), py::arg("masks"),
|
||||
"Fused linear backward: dX = g*sw @ W, dW = (g*sx)^T @ X, "
|
||||
"dB = sum(g), single call");
|
||||
py::arg("sg"), py::arg("sw"), py::arg("sx"), py::arg("sg_inv"),
|
||||
py::arg("sw_inv"), py::arg("sx_inv"), py::arg("amax_g"),
|
||||
"Scaled FP8 linear backward: dX = g*sw @ W, dW = (g*sx)^T @ X, "
|
||||
"dB = sum(g)");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user