diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index ca3e12f0d5..1e4da1dc11 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -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) diff --git a/examples/mlperf/models/flat_llama.py b/examples/mlperf/models/flat_llama.py index 86f27304ba..f93a752ddb 100644 --- a/examples/mlperf/models/flat_llama.py +++ b/examples/mlperf/models/flat_llama.py @@ -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: "): diff --git a/extra/gemm/cdna_asm_gemm.py b/extra/gemm/cdna_asm_gemm.py index 168f451983..c206b70ce0 100644 --- a/extra/gemm/cdna_asm_gemm.py +++ b/extra/gemm/cdna_asm_gemm.py @@ -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 diff --git a/extra/llama_kernels/cast_amax/__init__.py b/extra/llama_kernels/cast_amax/__init__.py index 7eed8e628a..07e98b3ce6 100644 --- a/extra/llama_kernels/cast_amax/__init__.py +++ b/extra/llama_kernels/cast_amax/__init__.py @@ -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 diff --git a/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py b/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py index 8a4a4e84f6..6267bc30b3 100644 --- a/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py +++ b/extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py @@ -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 diff --git a/extra/llama_kernels/quantize_fp8_delayed/__init__.py b/extra/llama_kernels/quantize_fp8_delayed/__init__.py index 015b8e1ec2..b3a824f915 100644 --- a/extra/llama_kernels/quantize_fp8_delayed/__init__.py +++ b/extra/llama_kernels/quantize_fp8_delayed/__init__.py @@ -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. diff --git a/test/backend/test_llama_kernels.py b/test/backend/test_llama_kernels.py index e0cd4c5559..878e5302a7 100644 --- a/test/backend/test_llama_kernels.py +++ b/test/backend/test_llama_kernels.py @@ -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):