From 69f35c46e063702e6ba4b6c69dd86c08e4967f57 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 14 Aug 2026 12:38:07 +0800 Subject: [PATCH] fix: quantize amax from raw values, not scaled fp8 values - amax for delayed scale was the quantized max (always ~448), so scale collapsed to 1 - this made fp8 gradients diverge (cosine 0.05) and training stall - stop w/x transpose-quantize amax from polluting the grad scale --- astrai/extension/fp8.py | 68 +++++++++++++++++++++++++++++++++++------ csrc/kernels/fp8_mm.cu | 27 ++++++++-------- 2 files changed, 74 insertions(+), 21 deletions(-) diff --git a/astrai/extension/fp8.py b/astrai/extension/fp8.py index d73e859..74d4ec8 100644 --- a/astrai/extension/fp8.py +++ b/astrai/extension/fp8.py @@ -52,8 +52,15 @@ class FP8TensorMeta: "idx", "x_scale", "x_scale_inv", + "x_history", + "x_idx", "g_scale", "g_scale_inv", + "g_history", + "g_idx", + "w_init", + "x_init", + "g_init", ) def __init__(self, device: torch.device, update_interval: int): @@ -65,8 +72,43 @@ class FP8TensorMeta: 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.x_history = torch.ones(update_interval, device=device, dtype=torch.float32) + self.x_idx = 0 self.g_scale = torch.ones(1, device=device, dtype=torch.float32) self.g_scale_inv = torch.ones(1, device=device, dtype=torch.float32) + self.g_history = torch.ones(update_interval, device=device, dtype=torch.float32) + self.g_idx = 0 + self.w_init = False + self.x_init = False + self.g_init = False + + def init_scale(self, t: torch.Tensor) -> None: + """Immediate scale from the current amax; used on the first call. + + A scale of 1 would underflow small activations/gradients (e4m3 min + normal is 2^-6); initialize from the actual amax once, then delayed + updates take over. + """ + amax = t.abs().amax().to(torch.float32).clamp_min(1e-12) + self.scale.copy_(amax / E4M3_MAX) + self.scale_inv.copy_(E4M3_MAX / amax) + self.record(amax) + + def push_x_scale(self, amax: torch.Tensor) -> None: + """Window update for the activation scale (delayed, TE style).""" + self.x_history[self.x_idx] = amax.reshape(()) + self.x_idx = (self.x_idx + 1) % self.x_history.numel() + m = self.x_history.max() + self.x_scale.copy_(m / E4M3_MAX) + self.x_scale_inv.copy_(E4M3_MAX / m) + + def push_g_scale(self, amax: torch.Tensor) -> None: + """Window update for the gradient scale (delayed, TE style).""" + self.g_history[self.g_idx] = amax.reshape(()) + self.g_idx = (self.g_idx + 1) % self.g_history.numel() + m = self.g_history.max() + self.g_scale.copy_(m / E4M3_MAX) + self.g_scale_inv.copy_(E4M3_MAX / m) def record(self, amax: torch.Tensor) -> None: """Push the latest amax into the ring buffer (device-side copy, no sync).""" @@ -155,13 +197,6 @@ def fp8_autocast(enabled: bool = True, update_interval: int = 16): 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). @@ -173,6 +208,15 @@ def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None): bias = torch.empty(0, device=x.device, dtype=x.dtype) state = fp8_state() meta = state.get_weight_meta(w) + if not meta.w_init: + meta.init_scale(w) + meta.w_init = True + if not meta.x_init: + amax = x.abs().amax().to(torch.float32).clamp_min(1e-12) + meta.x_history.fill_(amax) + meta.x_scale.copy_(amax / E4M3_MAX) + meta.x_scale_inv.copy_(E4M3_MAX / amax) + meta.x_init = True amax_x = torch.empty(1, device=x.device, dtype=torch.float32) amax_w = torch.empty(1, device=x.device, dtype=torch.float32) out = linear_forward_scaled( @@ -187,7 +231,7 @@ def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None): amax_w, ) meta.record(amax_w) - _update_delayed_scale(meta.x_scale, meta.x_scale_inv, amax_x) + meta.push_x_scale(amax_x) return out @@ -195,6 +239,12 @@ def fp8_linear_backward(g, x, w, masks): """TE-style scaled fp8 linear backward (called from aten::linear_backward).""" state = fp8_state() meta = state.get_weight_meta(w) + if not meta.g_init: + amax = g.abs().amax().to(torch.float32).clamp_min(1e-12) + meta.g_history.fill_(amax) + meta.g_scale.copy_(amax / E4M3_MAX) + meta.g_scale_inv.copy_(E4M3_MAX / amax) + meta.g_init = True amax_g = torch.empty(1, device=g.device, dtype=torch.float32) out = linear_backward_scaled( g, @@ -209,7 +259,7 @@ def fp8_linear_backward(g, x, w, masks): meta.x_scale_inv, amax_g, ) - _update_delayed_scale(meta.g_scale, meta.g_scale_inv, amax_g) + meta.push_g_scale(amax_g) return out diff --git a/csrc/kernels/fp8_mm.cu b/csrc/kernels/fp8_mm.cu index b6e248d..581f9f5 100644 --- a/csrc/kernels/fp8_mm.cu +++ b/csrc/kernels/fp8_mm.cu @@ -154,9 +154,9 @@ __global__ void quantize_kernel(const __nv_bfloat16* __restrict__ src, 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(v); - amax = fabsf(v); + float raw = __bfloat162float(src[i]); + dst[i] = cast_fp8(raw * *scale_inv); + amax = fabsf(raw); } for (int off = 16; off; off >>= 1) amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, off)); @@ -182,9 +182,9 @@ __global__ void transpose_quantize_kernel( float amax = 0.f; for (int j = 0; j < 32; j += 8) { if (x < cols && y + j < rows) { - float v = __bfloat162float(src[(y + j) * cols + x]) * *scale_inv; - tile[threadIdx.y + j][threadIdx.x] = cast_fp8(v); - amax = fmaxf(amax, fabsf(v)); + float raw = __bfloat162float(src[(y + j) * cols + x]); + tile[threadIdx.y + j][threadIdx.x] = cast_fp8(raw * *scale_inv); + amax = fmaxf(amax, fabsf(raw)); } } __syncthreads(); @@ -382,6 +382,9 @@ std::tuple fp8_linear_backward_scal auto gt8 = masks[1] ? torch::empty({n, m}, fp8_options) : torch::Tensor(); auto wt8 = masks[0] ? torch::empty({k, n}, fp8_options) : torch::Tensor(); auto xt8 = masks[1] ? torch::empty({k, m}, fp8_options) : torch::Tensor(); + // w/x transpose-quantize amax goes to a scratch buffer, NOT amax_g: the + // gradient scale must only see the gradient's own max-abs. + auto amax_t = torch::zeros({1}, g_c.options().dtype(torch::kFloat32)); int64_t block = 256; quantize_kernel<__nv_fp8_e4m3> @@ -394,8 +397,8 @@ std::tuple fp8_linear_backward_scal transpose_quantize_kernel<__nv_fp8_e4m3> <<>>( reinterpret_cast(w_c.data_ptr()), swi_ptr, - reinterpret_cast<__nv_fp8_e4m3*>(wt8.data_ptr()), amax_g_ptr, - n, k); + reinterpret_cast<__nv_fp8_e4m3*>(wt8.data_ptr()), + amax_t.data_ptr(), n, k); fp8_gemm_into(g8, wt8, grad_input.reshape({m, k}), m, n, k, sg_ptr, sw_ptr, stream.stream()); } @@ -405,13 +408,13 @@ std::tuple fp8_linear_backward_scal transpose_quantize_kernel<__nv_fp8_e4m3> <<>>( reinterpret_cast(g_c.data_ptr()), sgi_ptr, - reinterpret_cast<__nv_fp8_e4m3*>(gt8.data_ptr()), amax_g_ptr, - m, n); + reinterpret_cast<__nv_fp8_e4m3*>(gt8.data_ptr()), + amax_t.data_ptr(), m, n); transpose_quantize_kernel<__nv_fp8_e4m3> <<>>( reinterpret_cast(x_c.data_ptr()), sxi_ptr, - reinterpret_cast<__nv_fp8_e4m3*>(xt8.data_ptr()), amax_g_ptr, - m, k); + reinterpret_cast<__nv_fp8_e4m3*>(xt8.data_ptr()), + amax_t.data_ptr(), m, k); fp8_gemm_into(gt8, xt8, grad_weight, n, m, k, sg_ptr, sx_ptr, stream.stream()); }