ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: quantize amax from raw values, not scaled fp8 values

Browse Source 
- 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
This commit is contained in:
ViperEkura
2026-08-14 14:26:16 +08:00
parent 0378e62e17
commit 69f35c46e0
2 changed files with 74 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/fp8.py
+59 -9
View File
@@ -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
csrc/kernels/fp8_mm.cu
+15 -12
View File
@@ -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<T8>(v);
amax = fabsf(v);
float raw = __bfloat162float(src[i]);
dst[i] = cast_fp8<T8>(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<T8>(v);
amax = fmaxf(amax, fabsf(v));
float raw = __bfloat162float(src[(y + j) * cols + x]);
tile[threadIdx.y + j][threadIdx.x] = cast_fp8<T8>(raw * *scale_inv);
amax = fmaxf(amax, fabsf(raw));
}
}
__syncthreads();
@@ -382,6 +382,9 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> 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<torch::Tensor, torch::Tensor, torch::Tensor> fp8_linear_backward_scal
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);
reinterpret_cast<__nv_fp8_e4m3*>(wt8.data_ptr()),
amax_t.data_ptr<float>(), 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<torch::Tensor, torch::Tensor, torch::Tensor> fp8_linear_backward_scal
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);
reinterpret_cast<__nv_fp8_e4m3*>(gt8.data_ptr()),
amax_t.data_ptr<float>(), 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);
reinterpret_cast<__nv_fp8_e4m3*>(xt8.data_ptr()),
amax_t.data_ptr<float>(), m, k);
fp8_gemm_into(gt8, xt8, grad_weight, n, m, k, sg_ptr, sx_ptr,
stream.stream());
}