From 93b74c75fca0664fa0df8891c05ea5b9ff0865c5 Mon Sep 17 00:00:00 2001 From: wozeparrot Date: Fri, 31 Jul 2026 18:25:32 +0800 Subject: [PATCH] gptoss: grouped moe (#17322) --- examples/mlperf/model_train.py | 5 +- examples/mlperf/models/gpt_oss.py | 80 ++++++++++++++++++++++++------- 2 files changed, 68 insertions(+), 17 deletions(-) diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index 6b70032a13..c1345c3f00 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -1748,7 +1748,10 @@ def train_gptoss(): from extra.gemm.cdna_asm_gemm import _mx_block_scale model_state = get_state_dict(model) - fp8_scale_names = {n: f"{n}_scale" for n, t in model_state.items() if t.dtype == FP8_DTYPE} + def _scale_key(n): + if "." in n and (c:=f"{(b:=n.rsplit('.',1))[0]}_scale.{b[1]}") in model_state: return c + return f"{n}_scale" + fp8_scale_names = {n: _scale_key(n) for n, t in model_state.items() if t.dtype == FP8_DTYPE} fp8_inv_scales = [model_state[sname] for sname in fp8_scale_names.values()] for wname, sname in fp8_scale_names.items(): w, scale = model_state[wname], model_state[sname] diff --git a/examples/mlperf/models/gpt_oss.py b/examples/mlperf/models/gpt_oss.py index 8e34e7b65f..4bb8613b94 100644 --- a/examples/mlperf/models/gpt_oss.py +++ b/examples/mlperf/models/gpt_oss.py @@ -13,10 +13,14 @@ from tinygrad.uop.ops import Ops, UOp from extra.models.llama import apply_rotary_emb from extra.llama_kernels.rmsnorm import rmsnorm from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8 +from extra.gemm.moe_gemm import grouped_mx_gemm +from extra.gemm.moe_routing import route, dispatch, combine FP8_DTYPE = dtypes.fp8e4m3 FP8_MAX = 448.0 -INIT_STD = 0.008 +INIT_STD = 0.02 +ASM_GEMM = getenv("ASM_GEMM", 1) + def _quant_dequant_fwd(x:Tensor) -> Tensor: # x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis) @@ -59,10 +63,34 @@ def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor: def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor: l_shape = x.shape[:-1] + if ASM_GEMM: + from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack + x2, K, N = x.reshape(-1, x.shape[-1]), x.shape[-1], w_q.shape[0] + wq, ws = w_q, w_scale + if (pad := (-K) % 256): + x2 = x2.pad(((0, 0), (0, pad))) + wq = wq.pad(((0, 0), (0, pad))) + ws = ws.pad(((0, 0), (0, pad // 32)), value=127).cast(dtypes.uint8) + if (npad := (-N) % 256): + wq = wq.pad(((0, npad), (0, 0))) + ws = ws.pad(((0, npad), (0, 0)), value=127).cast(dtypes.uint8) + x_q, x_e8, x_si = quantize_mxfp8(x2) + if x_si is not None and can_use_asm_gemm(x_q, wq.T): + out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(ws), ws), mx_w_stored=True) + return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16) x_phys = quant_dequant_mx(x.reshape(-1, x.shape[-1])).reshape(*l_shape, x.shape[-1]) w_phys = dequant_weight(w_q, w_scale) return (x_phys @ w_phys.T).cast(dtypes.bfloat16) +def _pad_to_mult(t:Tensor, axis:int, mult:int=256) -> Tensor: + if (r := (-t.shape[axis]) % mult) == 0: return t + pads = [(0, 0)] * t.ndim + pads[axis] = (0, r) + return t.pad(tuple(pads)) + +def _pad_cols(t:Tensor) -> Tensor: return _pad_to_mult(t, -1) +def _pad_rows(t:Tensor) -> Tensor: return _pad_to_mult(t, -2) + def swiglu(x:Tensor, limit:float=7.0, alpha:float=1.702) -> Tensor: x_glu, x_linear = x[..., ::2], x[..., 1::2] x_glu = x_glu.clamp(max_=limit) @@ -99,9 +127,9 @@ class GPTOSS: self.ffn_norm = Tensor.ones(n_layers, dim).contiguous() self.gate = Tensor.normal(n_layers, n_experts, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16) self.gate_bias = Tensor.zeros(n_layers, n_experts, dtype=dtypes.bfloat16).contiguous() - self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim) + self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim, moe=True) self.w_gate_up_bias = Tensor.zeros(n_layers, n_experts, intermediate_size * 2, dtype=dtypes.bfloat16).contiguous() - self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std) + self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std, moe=True) self.w_down_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous() # output @@ -111,10 +139,15 @@ class GPTOSS: self.output = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16) self.freqs_cis = precompute_freqs_cis(head_dim, max_context * 2, rope_theta).contiguous().is_param_(False) - def _quant_weight(self, *shape:int, std:float=INIT_STD): - w = Tensor.zeros(*shape) if getenv("ZEROS") else Tensor.normal(*shape, mean=0.0, std=std) - w_q, w_e8, _ = quantize_mxfp8(w) - return w_q, w_e8.is_param_(False) + def _quant_weight(self, *shape:int, std:float=INIT_STD, moe:bool=False): + def _one(*s:int): + w = Tensor.zeros(*s) if getenv("ZEROS") else Tensor.normal(*s, mean=0.0, std=std) + w_q, w_e8, _ = quantize_mxfp8(_pad_cols(_pad_rows(w)) if moe else w) + return w_q, w_e8.is_param_(False) + if moe: + qs = [_one(*shape[1:]) for _ in range(shape[0])] + return [q[0] for q in qs], [q[1] for q in qs] + return _one(*shape) def _attn_mask(self, seqlen:int, dtype) -> Tensor: i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen) @@ -173,17 +206,32 @@ class GPTOSS: w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor): x_normed, rrms = rmsnorm(x, self.norm_eps) inp = x_normed * ffn_norm - logits = inp.float() @ gate.float().T + gate_bias.float() - thresh = logits.topk(self.experts_per_tok)[0][..., -1:] - weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1) + dim, inter = self.dim, self.intermediate_size - out = None - for e in range(self.n_experts): - gate_up = matmul_mx(inp, w_gate_up[e], w_gate_up_scale[e]) + w_gate_up_bias[e] - y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), w_down[e], w_down_scale[e]) + w_down_bias[e]).contiguous() - contrib = weights[..., e:e+1].cast(y.dtype) * y - out = contrib if out is None else out + contrib + if getenv("GROUPED_MOE", 0): + bsz, seqlen = x.shape[:2] + inp, logits = inp.reshape(-1, dim), logits.reshape(-1, self.n_experts) + r = route(logits, self.experts_per_tok, self.n_experts) + onehot = r.rows_e.one_hot(self.n_experts).float() + xg = dispatch(_pad_cols(inp.cast(dtypes.bfloat16)), r) + h = grouped_mx_gemm(xg, (w_gate_up, w_gate_up_scale), r.off)[:, :2*inter] + (onehot @ w_gate_up_bias.float()).cast(dtypes.bfloat16) + y = swiglu(h, self.swiglu_limit) + z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \ + + (onehot @ w_down_bias.float()).cast(dtypes.bfloat16) + out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim) + else: + thresh = logits.topk(self.experts_per_tok)[0][..., -1:] + weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1) + + out = None + for e in range(self.n_experts): + gu_q, gu_s = w_gate_up[e][:2*inter, :dim].contiguous(), w_gate_up_scale[e][:2*inter, :dim//32].contiguous() + dn_q, dn_s = w_down[e][:dim, :inter].contiguous(), w_down_scale[e][:dim, :inter//32].contiguous() + gate_up = matmul_mx(inp, gu_q, gu_s) + w_gate_up_bias[e] + y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), dn_q, dn_s) + w_down_bias[e]).contiguous() + contrib = weights[..., e:e+1].cast(y.dtype) * y + out = contrib if out is None else out + contrib return out, [x_normed, rrms] @function(precompile=True, precompile_backward=True)