nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

llama: inplace amax update (#17064)

Browse Source 
* llama: inplace amax update

* remove amax_out return

* work

* fit

* work

* work

* keep

* diff cleanup
This commit is contained in:
qazal
2026-07-20 15:05:41 +09:00
committed by GitHub
parent ac3f56a1a2
commit 1cf8f2f68c
7 changed files with 87 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/mlperf/model_train.py
+2
View File
@@ -1462,6 +1462,8 @@ def train_llama3():
@TinyJit
def minibatch(tokens:Tensor):
for nxt in fp8_next_amax: nxt.assign(0)
for nxt in fp8_next_grad_amax: nxt.assign(0)
if is_dp: tokens = tokens.to(None).shard(device, 0)
if is_mp: tokens = tokens.shard(device)
if not is_sharding: tokens = tokens.to(None)
examples/mlperf/models/flat_llama.py
+60 -58
View File
@@ -37,8 +37,8 @@ def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax
def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
x_fp8:Tensor|None=None, x_new_amax:Tensor|None=None,
grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None) -> tuple[Tensor,...]:
x_fp8:Tensor|None=None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None,
next_amax_x:Tensor|None=None) -> tuple[Tensor,...]:
if not fp8:
if ASM_GEMM:
from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
@@ -56,13 +56,14 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
else:
x_phys = (x_q.cast(dtypes.bfloat16) * _mx_block_scale(x_e8)).reshape(*l_shape, x_q.shape[-1])
out = x_phys @ (w.cast(dtypes.bfloat16) * _mx_block_scale(w_inv_scale)).T
return out, (amax_x.detach() if amax_x is not None else None), x_q
return out, x_q
if x_fp8 is None:
if FUSED_INPUT_QUANTIZE and amax_x is not None:
if FUSED_INPUT_QUANTIZE:
from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
x_fp8, _, x_new_amax, _ = quantize_fp8_delayed(x, amax_x, FP8_DTYPE)
x_fp8, _ = quantize_fp8_delayed(x, amax_x, next_amax_x, FP8_DTYPE)
else:
x_fp8, _, x_new_amax = quantize_fp8(x, amax_state=amax_x)
x_fp8, _, new_amax_x = quantize_fp8(x, amax_state=amax_x)
next_amax_x.assign(new_amax_x)
if ASM_GEMM:
from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
if can_use_asm_gemm(x_fp8, w.T):
@@ -73,51 +74,51 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
else:
out = asm_gemm(x_fp8, w.T, x_scale=amax_x, w_scale=w_inv_scale, grad_amax_state=grad_amax_state,
next_grad_amax_state=next_grad_amax_state)
return out, x_new_amax, x_fp8
return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_new_amax, x_fp8
return out, x_fp8
return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_fp8
def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
grad_amax_state:Tensor, next_grad_amax_state:Tensor):
next_amax_x:Tensor, grad_amax_state:Tensor, next_grad_amax_state:Tensor):
if FUSED_ADD_NORM_MUL_QUANTIZE:
from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
x_fp8, new_amax, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE)
out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x, x_new_amax=new_amax,
x_fp8, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
return out, x_normed, rrms, ret
x_normed, rrms = rmsnorm(x, eps)
out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
next_grad_amax_state=next_grad_amax_state)
next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
return out, x_normed, rrms, ret
def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
next_amax_x:Tensor, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
if FUSED_ADD_NORM_MUL_QUANTIZE:
from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
x_fp8, new_amax, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE)
out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x, x_new_amax=new_amax,
x_fp8, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
return out, h, x_normed, rrms, ret
h = x + residual
x_normed, rrms = rmsnorm(h, eps)
out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
next_grad_amax_state=next_grad_amax_state)
next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
return out, h, x_normed, rrms, ret
def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
amax_x2:Tensor,
amax_x2:Tensor, next_amax_x2:Tensor,
grad_amax_xw13:Tensor, next_grad_amax_xw13:Tensor,
grad_amax_xout:Tensor, next_grad_amax_xout:Tensor):
if FUSED_SILU_W13:
from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
x2_fp8, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
next_grad_amax_state=next_grad_amax_xw13)
out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2, x_new_amax=new_amax_x2,
x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
next_grad_amax_state=next_grad_amax_xw13, amax_out=next_amax_x2)
out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2,
grad_amax_state=grad_amax_xout, next_grad_amax_state=next_grad_amax_xout)
return out, ret
hidden = x_w13.shape[-1] // 2
x_w1, x_w3 = x_w13[..., :hidden], x_w13[..., hidden:]
out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
next_grad_amax_state=next_grad_amax_xout)
next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2)
return out, ret
class FlatTransformer:
@@ -186,14 +187,14 @@ class FlatTransformer:
def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
amax_xqkv:Tensor, amax_xo:Tensor, s_qkv:Tensor, s_o:Tensor,
next_amax_xqkv:Tensor, next_amax_xo:Tensor,
grad_amax_xqkv:Tensor, grad_amax_xo:Tensor, next_grad_amax_xqkv:Tensor, next_grad_amax_xo:Tensor):
bsz, seqlen, _ = x.shape
amaxs, saves = [], []
saves = []
xqkv, x_normed, rrms, (new_amax, *s) = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
xqkv, x_normed, rrms, s = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv,
next_grad_amax_state=next_grad_amax_xqkv)
amaxs.append(new_amax)
next_grad_amax_state=next_grad_amax_xqkv, next_amax_x=next_amax_xqkv)
saves.extend([x_normed, rrms, *s, xqkv])
if getenv("HK_FLASH_ATTENTION"):
from extra.thunder.amd.fa import flash_attention, fused_qkv_rope
@@ -211,64 +212,62 @@ class FlatTransformer:
attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
attn = attn.reshape(bsz, seqlen, -1)
out, new_amax, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
next_grad_amax_state=next_grad_amax_xo)
amaxs.append(new_amax)
out, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
next_grad_amax_state=next_grad_amax_xo, next_amax_x=next_amax_xo)
saves.extend([*s, out])
return out, amaxs, saves
return out, saves
def feed_forward(self, x:Tensor, residual:Tensor, **kwargs):
amaxs, saves = [], []
saves = []
if SPLIT_W13:
h = x + residual
x_normed, rrms = rmsnorm(h, self.norm_eps)
saves.extend([x_normed, rrms])
inp = x_normed * kwargs["ffn_norm"]
x_w1, new_amax, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"])
amaxs.append(new_amax)
x_w1, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"],
next_amax_x=kwargs["next_amax_x1"])
saves.extend([*s, x_w1])
x_w3, new_amax, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"])
amaxs.append(new_amax)
x_w3, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"],
next_amax_x=kwargs["next_amax_x3"])
saves.extend([*s, x_w3])
if FUSED_SILU_W13 and MXFP8:
from extra.llama_kernels.fused_silu_mul_quantize_mxfp8 import fused_silu_mul_quantize_mxfp8
aq, ae8, asi = fused_silu_mul_quantize_mxfp8(x_w1.reshape(-1, x_w1.shape[-1]), x_w3.reshape(-1, x_w3.shape[-1]))
out, new_amax, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
next_grad_amax_state=kwargs["next_grad_amax_xout"])
out, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
next_grad_amax_state=kwargs["next_grad_amax_xout"], next_amax_x=kwargs["next_amax_x2"])
out = out.reshape(*x_w1.shape[:-1], kwargs["w2"].shape[0])
else:
out, new_amax, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
grad_amax_state=kwargs["grad_amax_xout"], next_grad_amax_state=kwargs["next_grad_amax_xout"])
amaxs.append(new_amax)
out, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
grad_amax_state=kwargs["grad_amax_xout"], next_grad_amax_state=kwargs["next_grad_amax_xout"],
next_amax_x=kwargs["next_amax_x2"])
saves.extend([*s, out])
else:
x_w13, h, x_normed, rrms, (new_amax, *s) = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
x_w13, h, x_normed, rrms, s = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
self.norm_eps, amax_x=kwargs["amax_x13"],
next_amax_x=kwargs["next_amax_x13"],
grad_amax_state=kwargs["grad_amax_xw13"],
next_grad_amax_state=kwargs["next_grad_amax_xw13"])
amaxs.append(new_amax)
saves.extend([x_normed, rrms, *s, x_w13])
out, (new_amax, *s) = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
out, s = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
next_amax_x2=kwargs["next_amax_x2"],
grad_amax_xw13=kwargs["grad_amax_xw13"],
next_grad_amax_xw13=kwargs["next_grad_amax_xw13"],
grad_amax_xout=kwargs["grad_amax_xout"],
next_grad_amax_xout=kwargs["next_grad_amax_xout"])
amaxs.append(new_amax)
saves.extend([*s, out])
return out, h, amaxs, saves
return out, h, saves
@function(precompile=True, precompile_backward=True)
def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
attn, attn_amaxs, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
ffn, h, ffn_amaxs, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
attn, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
ffn, h, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
h = h + ffn
amaxs = tuple(a.detach() for a in (*attn_amaxs, *ffn_amaxs))
if save: return (h, *amaxs, *attn_saves, *ffn_saves)
else: return (h, *amaxs)
if save: return (h, *attn_saves, *ffn_saves)
else: return (h,)
def shard(self, device:tuple[str, ...], mp:bool=False):
from tinygrad.nn.state import get_parameters
@@ -319,21 +318,21 @@ class FlatTransformer:
for i in range(self.n_layers):
attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i],
amax_xqkv=a["xqkv"][i], amax_xo=a["xo"][i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
next_amax_xqkv=na["xqkv"][i], next_amax_xo=na["xo"][i],
grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["xo"][i],
next_grad_amax_xqkv=nga["xqkv"][i], next_grad_amax_xo=nga["xo"][i])
ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i],
amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i], next_grad_amax_xout=nga["xout"][i])
amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i], next_grad_amax_xout=nga["xout"][i],
next_amax_x2=na["x2"][i])
if SPLIT_W13:
ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], amax_x1=a["x1"][i], amax_x3=a["x3"][i],
next_amax_x1=na["x1"][i], next_amax_x3=na["x3"][i],
s_1=s["w1"][i], s_3=s["w3"][i], grad_amax_xw1=ga["xw1"][i], grad_amax_xw3=ga["xw3"][i],
next_grad_amax_xw1=nga["xw1"][i], next_grad_amax_xw3=nga["xw3"][i])
else:
ffn_kwargs.update(w13=self.w13[i], amax_x13=a["x13"][i], s_13=s["w13"][i], grad_amax_xw13=ga["xw13"][i],
next_grad_amax_xw13=nga["xw13"][i])
h, *ret = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
amax_names = ["xqkv", "xo"] + (["x1", "x3"] if SPLIT_W13 else ["x13"]) + ["x2"]
for name, new_val in zip(amax_names, ret[:len(amax_names)]):
na[name][i].assign(new_val)
next_grad_amax_xw13=nga["xw13"][i], next_amax_x13=na["x13"][i])
h, *_ = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
return logits
@@ -416,6 +415,9 @@ if __name__ == "__main__":
@TinyJit
def fwd_bwd(tokens:Tensor):
with Timing("python forward: "):
for amax_dict in (model._fp8_next_amax, model._fp8_next_grad_amax):
for ts in amax_dict.values():
for nxt in ts: nxt.assign(0)
logits = model(tokens[:, :-1], save=llama_size=="8B")
loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
with Timing("python backward: "):
extra/gemm/cdna_asm_gemm.py
+1 -4
View File
@@ -280,10 +280,7 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=
elif getenv("FUSED_GRAD_QUANTIZE", 0):
grad_amax_t = Tensor(grad_amax_state, device=a.device)
g_amax = grad_amax_t
g_fp8, _, new_grad_amax, _ = quantize_fp8_delayed(g_t, g_amax)
store_effect = next_grad_amax_state.store(new_grad_amax.uop)
assert g_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {g_fp8.uop.op}"
g_fp8 = Tensor(g_fp8.uop.replace(src=g_fp8.uop.src + (store_effect,)), device=a.device)
g_fp8, _ = quantize_fp8_delayed(g_t, g_amax, Tensor(next_grad_amax_state, device=a.device))
else:
grad_amax_t = Tensor(grad_amax_state, device=a.device)
g_amax = grad_amax_t
extra/llama_kernels/cast_amax/__init__.py
+5 -8
View File
@@ -43,7 +43,7 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
device = xw13.device
axis = xw13.axis if isinstance(device, tuple) else None
grad_xw13_fp8 = alloc_like(xw13.shape, dtypes.fp8e4m3, device, axis)
grad_amax_next = Tensor.zeros((), dtype=dtypes.float32, device=device).contiguous()
grad_amax_next = Tensor(next_grad_amax_state, device=device)
grad_amax_state_t = Tensor(grad_amax_state, device=device)
fxn = functools.partial(_custom_fused_bwd_w13, dname=dname_of(device))
grad_amax = grad_amax_state_t.empty_like()
@@ -52,16 +52,14 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
Tensor(xw13, device=device), Tensor(gradient, device=device).cast(dtypes.bfloat16),
Tensor(amax_state, device=device), grad_amax_state_t, fxn=fxn)
grad_xw13_uop = grad_xw13_fp8.uop.cast(dtypes.bfloat16)
store_effect = next_grad_amax_state.store(grad_amax_next.uop)
assert grad_xw13_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {grad_xw13_fp8.uop.op}"
grad_xw13_fp8_uop = grad_xw13_fp8.uop.replace(src=grad_xw13_fp8.uop.src + (store_effect,))
# Stash fp8 companion for cdna_asm_gemm's bwd to attach to grad_a.
_grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8_uop, grad_amax_state_t.uop)
_grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8.uop, grad_amax_state_t.uop)
return (None, None, grad_xw13_uop, None, None, None)
def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_state:Tensor,
next_grad_amax_state:Tensor) -> tuple[Tensor, Tensor]:
# NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, new_amax)
next_grad_amax_state:Tensor, amax_out:Tensor) -> Tensor:
# NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns fp8.
# grad_amax_state: delayed amax for grad_xw13 fp8 quantization in the backward.
assert xw13.dtype == dtypes.bfloat16, f"expected bf16, got {xw13.dtype}"
MBS, SEQ, H2 = xw13.shape
@@ -69,8 +67,7 @@ def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_
HIDDEN = H2 // 2
axis = xw13.uop.axis if isinstance(xw13.device, tuple) else None
fp8_out = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype, xw13.device, axis)
amax_out = Tensor.zeros((), dtype=dtypes.float32, device=xw13.device).contiguous()
fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname_of(xw13.device))
fp8_out, amax_out, *_ = Tensor.custom_kernel(fp8_out, amax_out, xw13, amax_state, grad_amax_state, next_grad_amax_state,
fxn=fxn, grad_fxn=_fused_quantize_bwd_w13)
return fp8_out, amax_out
return fp8_out
extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py
+7 -8
View File
@@ -112,8 +112,9 @@ def _fused_add_bwd(*args, **kwargs):
grad_h, grad_w = _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel)
return (None, None, None, None, None, grad_h, grad_h, grad_w, None)
def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, eps:float, fp8_dtype) -> tuple[Tensor, Tensor, Tensor, Tensor]:
# NOTE: rmsnorm(x) * weight -> fp8 + amax. Returns (fp8, new_amax, x_normed, rrms).
def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, eps:float, fp8_dtype,
amax_out:Tensor) -> tuple[Tensor, Tensor, Tensor]:
# NOTE: rmsnorm(x) * weight -> fp8 + amax. Returns (fp8, x_normed, rrms).
# x_normed + rrms are saved for the rmsnorm backward (also recomputed here from x regs).
assert x.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
assert x.shape[-1] == weight.shape[-1], f"HIDDEN mismatch: x={x.shape}, weight={weight.shape}"
@@ -123,16 +124,15 @@ def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, e
fp8_out = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype, x.device, axis)
x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
rrms_out = alloc_like((MBS, SEQ), dtypes.float32, x.device, axis)
amax_out = Tensor.zeros((), dtype=dtypes.float32, device=x.device).contiguous()
fxn = functools.partial(_custom_fwd, dname=dname_of(x.device), eps_val=eps)
fp8_out, x_normed_out, rrms_out, amax_out, *_ = Tensor.custom_kernel(
fp8_out, x_normed_out, rrms_out, amax_out, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_bwd)
return fp8_out, amax_out, x_normed_out, rrms_out
return fp8_out, x_normed_out, rrms_out
def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor, amax_state:Tensor,
eps:float, fp8_dtype) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
eps:float, fp8_dtype, amax_out:Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
# NOTE: h = x + residual; y_normed = rmsnorm(h); fp8 = quantize(y_normed * weight).
# Returns (fp8, new_amax, h, x_normed, rrms). h is also written so downstream can
# Returns (fp8, h, x_normed, rrms). h is also written so downstream can
# reuse it without recomputing x+residual — eliminates the separate residual-add kernel.
assert x.dtype == dtypes.bfloat16 and residual.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
assert x.shape == residual.shape
@@ -143,9 +143,8 @@ def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor,
h_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
rrms_out = alloc_like((MBS, SEQ), dtypes.float32, x.device, axis)
amax_out = Tensor.zeros((), dtype=dtypes.float32, device=x.device).contiguous()
fxn = functools.partial(_custom_fwd_add, dname=dname_of(x.device), eps_val=eps)
fp8_out, h_out, x_normed_out, rrms_out, amax_out, *_ = Tensor.custom_kernel(
fp8_out, h_out, x_normed_out, rrms_out, amax_out, x, residual, weight, amax_state,
fxn=fxn, grad_fxn=_fused_add_bwd)
return fp8_out, amax_out, h_out, x_normed_out, rrms_out
return fp8_out, h_out, x_normed_out, rrms_out
extra/llama_kernels/quantize_fp8_delayed/__init__.py
+3 -8
View File
@@ -73,24 +73,19 @@ def _quantize_fp8_delayed_bwd(gradient:UOp, kernel:UOp):
grad_x = (Tensor(gradient, device=device).float() * scale).cast(dtypes.bfloat16)
return (None, None, grad_x.uop, None)
def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -> tuple[Tensor, Tensor, Tensor, UOp]:
# NOTE: one-pass bf16 -> fp8 quantize with delayed scaling. Returns (fp8, inv_scale, new_amax, store_effect).
def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, amax_out:Tensor, fp8_dtype=dtypes.fp8e4m3) -> tuple[Tensor, Tensor]:
# NOTE: one-pass bf16 -> fp8 quantize with delayed scaling.
# Fused kernel reads x once and writes fp8 + scalar amax via global atomic max.
# store_effect writes new_amax into amax_state's buffer — the caller must thread it into a realized
# output via `.after(store_effect)`. Calling `amax_state.assign(new_amax)` inside a grad_fxn does
# NOT work because .assign mutates only the temp Tensor's .uop, not the original layer-owned buffer.
assert x.dtype == dtypes.bfloat16, f"expected bf16, got {x.dtype}"
axis = x.uop.axis if isinstance(x.device, tuple) else None
fp8_out = alloc_like(x.shape, fp8_dtype, x.device, axis)
n_elems = prod(x.uop.shard_shape)
assert n_elems % NUM_WG == 0, f"{n_elems=} must divide over {NUM_WG=}"
amax_out = Tensor.zeros((), dtype=dtypes.float32, device=x.device).contiguous()
fxn = functools.partial(_custom_quantize_fp8_with_amax, device=x.device)
fp8_out, amax_out, *_ = Tensor.custom_kernel(fp8_out, amax_out, x, amax_state,
fxn=fxn, grad_fxn=_quantize_fp8_delayed_bwd)
inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
store_effect = amax_state.uop.store(amax_out.uop)
return fp8_out, inv_scale, amax_out, store_effect
return fp8_out, inv_scale
def quantize_fp8_scalar(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -> Tensor:
# NOTE: pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation.
test/backend/test_llama_kernels.py
+9 -7
View File
@@ -49,9 +49,10 @@ def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None:
with Context(DEBUG=0): Tensor.realize(x, amax_state)
if delayed:
fp8, inv_scale, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
amax_out = Tensor.zeros((), dtype=dtypes.float32, device=x.device).realize()
fp8, inv_scale = quantize_fp8_delayed(x, amax_state, amax_out, FP8_DTYPE)
ref_fp8, ref_inv_scale, ref_new_amax = quantize_fp8(x, amax_state=amax_state)
Tensor.realize(fp8, inv_scale, new_amax)
Tensor.realize(fp8, inv_scale)
Tensor.realize(ref_fp8, ref_inv_scale, ref_new_amax)
else:
fp8 = quantize_fp8_scalar(x, amax_state, FP8_DTYPE)
@@ -63,8 +64,8 @@ def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None:
assert fp8.cast(dtypes.float).allclose(ref_fp8.cast(dtypes.float), atol=0, rtol=0).item(), "fp8 mismatch"
if delayed:
assert inv_scale.allclose(ref_inv_scale, atol=0, rtol=0).item(), "inv_scale mismatch"
assert new_amax.allclose(ref_new_amax, atol=0, rtol=0).item(), \
f"amax mismatch: got={new_amax.item()} ref={ref_new_amax.item()} diff={abs(new_amax.item()-ref_new_amax.item())}"
assert amax_out.allclose(ref_new_amax, atol=0, rtol=0).item(), \
f"amax mismatch: got={amax_out.item()} ref={ref_new_amax.item()} diff={abs(amax_out.item()-ref_new_amax.item())}"
@unittest.skipUnless(Device.DEFAULT == "AMD", "requires atomic max")
class TestQuantizeFP8(unittest.TestCase):
@@ -82,10 +83,11 @@ class TestQuantizeFP8(unittest.TestCase):
x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.multi(0)
x = Tensor(x, device=devs)
amax_state = Tensor.full((), 2.0, dtype=dtypes.float32, device=devs).contiguous()
fp8, _, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
Tensor.realize(fp8, new_amax)
amax_out = Tensor.zeros((), dtype=dtypes.float32, device=devs).realize()
fp8, _ = quantize_fp8_delayed(x, amax_state, amax_out, FP8_DTYPE)
Tensor.realize(fp8)
assert fp8.uop.shape == x.uop.shape
assert new_amax.shape == ()
assert amax_out.shape == ()
class TestLocalAmax(unittest.TestCase):
def test_multi_tensor_local_shard_amax(self):