superrob/IQ.Pilot
1
0
Fork 0
forked from IQ.Lvbs/IQ.Pilot
Code Pull Requests Activity

IQ.Pilot Release Commit @ bec7652

Browse Source
This commit is contained in:
IQ.Lvbs history cleanup
2026-08-22 23:42:41 -05:00
commit 58039e647c
4603 changed files with 1236178 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

453
artifacts/package_sources/tinygrad/examples/mlperf/models/flat_llama.py Normal file
View File

@@ -0,0 +1,453 @@
import math, os
if __name__ == "__main__":
os.environ["DEFAULT_FLOAT"] = "bfloat16"
os.environ["OPTIM_DTYPE"] = "bfloat16"
if "DEV" not in os.environ: os.environ["DEV"] = "NULL::gfx950"
# CDNA
os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
os.environ["ALL2ALL"] = "1"
os.environ["USE_ATOMICS"] = "1"
if "HK_FLASH_ATTENTION" not in os.environ:
os.environ["HK_FLASH_ATTENTION"] = "1"
if "ASM_GEMM" not in os.environ:
os.environ["ASM_GEMM"] = "1"
from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker, round_up
from tinygrad.uop.ops import Ops, UOp
from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
from extra.llama_kernels.rmsnorm import rmsnorm
from extra.llama_kernels import FP8_MAX, local_abs_max
ASM_GEMM = getenv("ASM_GEMM", 0)
FUSED_INPUT_QUANTIZE = getenv("FUSED_INPUT_QUANTIZE", 0)
FUSED_ADD_NORM_MUL_QUANTIZE = getenv("FUSED_ADD_NORM_MUL_QUANTIZE", 0)
FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
SPLIT_W13 = getenv("SPLIT_W13", 0)
COLUMNWISE_WEIGHT_SCALE = getenv("COLUMNWISE_WEIGHT_SCALE", 0)
MXFP8 = getenv("MXFP8", 0)
MXFP4 = getenv("MXFP4", 0)
FP8_DTYPE = dtypes.fp8e4m3
FP8_GRAD_DTYPE = dtypes.fp8e5m2
def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
new_amax = (local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach().cast(dtypes.float32)
scale = FP8_MAX / ((amax_state if amax_state is not None else new_amax) + 1e-8)
x_scaled = x * scale
x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach()) # STE
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, 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
if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),)
return (x @ w.T,)
if MXFP4:
assert x is not None, "MXFP4 matmul requires an unquantized input"
from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm
if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T, mxfp4=True),)
return (x @ w.T,)
assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
if MXFP8:
from extra.gemm.cdna_asm_gemm import asm_gemm, quantize_mxfp8, mx_pack, can_use_asm_gemm, _mx_block_scale
if x_prequant_mx is not None: x_q, x_e8, x_si = x_prequant_mx # fused producer already quantized (2d)
else: x_q, x_e8, x_si = quantize_mxfp8(x.reshape(-1, x.shape[-1]))
l_shape = x.shape[:-1] if x is not None else x_q.shape[:-1]
if can_use_asm_gemm(x_q, w.T):
out = asm_gemm(x_q, w.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(w_inv_scale), w_inv_scale),
mx_w_stored=True).reshape(*l_shape, w.shape[0])
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, x_q
if x_fp8 is None:
if FUSED_INPUT_QUANTIZE:
from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
x_fp8, _ = quantize_fp8_delayed(x, amax_x, next_amax_x, FP8_DTYPE)
else:
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):
assert amax_x is not None
if COLUMNWISE_WEIGHT_SCALE:
out = asm_gemm(x_fp8, w.T, x_scale=amax_x, grad_amax_state=grad_amax_state,
next_grad_amax_state=next_grad_amax_state, w_post_scale=w_inv_scale)
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_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|None,
next_amax_x:Tensor|None, grad_amax_state:Tensor|None, next_grad_amax_state:Tensor|None):
if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
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_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|None,
next_amax_x:Tensor|None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
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_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|None, next_amax_x2:Tensor|None,
grad_amax_xw13:Tensor|None, next_grad_amax_xw13:Tensor|None,
grad_amax_xout:Tensor|None, next_grad_amax_xout:Tensor|None):
if FUSED_SILU_W13 and not MXFP4:
from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
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_amax_x=next_amax_x2)
return out, ret
class FlatTransformer:
def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
rope_theta:int=10000, max_context:int=1024):
self.vocab_size = vocab_size
self.n_layers = n_layers
self.n_heads = n_heads
self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
self.head_dim = dim // n_heads
self.n_rep = self.n_heads // self.n_kv_heads
self.hidden_dim = hidden_dim
scaled_std = 0.02 / math.sqrt(2 * n_layers)
# Attention
self.wqkv, s_qkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
self.wo, s_o = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
# FeedForward
if SPLIT_W13:
self.w1, s_1 = self.lin_per_layer(dim, hidden_dim)
self.w3, s_3 = self.lin_per_layer(dim, hidden_dim)
else:
self.w13, s_13 = self.lin_per_layer(dim, hidden_dim * 2)
self.w2, s_2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
self.norm_eps = norm_eps
self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
# output
self.norm = nn.RMSNorm(dim, norm_eps)
self.tok_embeddings = nn.Embedding(vocab_size, dim)
self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).clone().is_param_(False)
def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
n_amax = 0 if MXFP4 else n_layers
names = ["xqkv", "xo", "x2"]
names += ["x1", "x3"] if SPLIT_W13 else ["x13"]
self._fp8_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
self._fp8_next_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
grad_names = ["xqkv", "xo", "xout"]
grad_names += ["xw1", "xw3"] if SPLIT_W13 else ["xw13"]
self._fp8_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names}
self._fp8_next_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names}
w_scales = [("wqkv", s_qkv), ("wo", s_o), ("w2", s_2)]
w_scales += [("w1", s_1), ("w3", s_3)] if SPLIT_W13 else [("w13", s_13)]
self._fp8_inv_scale = {name: (s if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales}
self._fp8_next_inv_scale = {name: (s if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales}
def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02, w:Tensor|None=None):
if w is None:
if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features)
else: w = Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)
if MXFP8:
from extra.gemm.cdna_asm_gemm import quantize_mxfp8
w_q, w_e8, _ = quantize_mxfp8(w.reshape(self.n_layers * out_features, in_features))
return w_q.reshape(self.n_layers, out_features, in_features), w_e8.reshape(self.n_layers, out_features, in_features // 32)
if MXFP4:
# FP4 is produced dynamically so optimizer updates always start from the current BF16 weight.
return w.cast(dtypes.bfloat16), Tensor.ones(self.n_layers)
amax = (w.abs().max(axis=2) if COLUMNWISE_WEIGHT_SCALE else w.abs().flatten(1).max(1)).detach()
scale = FP8_MAX / (amax + 1e-8)
inv_scale = (amax + 1e-8) / FP8_MAX
scale_b = scale.reshape(self.n_layers, out_features, 1) if COLUMNWISE_WEIGHT_SCALE else scale.reshape(-1, 1, 1)
return (w * scale_b).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE), inv_scale
def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
amax_xqkv:Tensor|None, amax_xo:Tensor|None, s_qkv:Tensor, s_o:Tensor,
next_amax_xqkv:Tensor|None, next_amax_xo:Tensor|None,
grad_amax_xqkv:Tensor|None, grad_amax_xo:Tensor|None,
next_grad_amax_xqkv:Tensor|None, next_grad_amax_xo:Tensor|None):
bsz, seqlen, _ = x.shape
saves = []
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, 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
xq, xk, xv = fused_qkv_rope(xqkv, freqs_cis, self.n_heads, self.n_kv_heads, self.head_dim)
attn, *save = flash_attention(xq, xk, xv, is_causal=True, write_flat=True)
saves.extend(save)
else:
xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
attn = attn.reshape(bsz, seqlen, -1)
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, saves
def feed_forward(self, x:Tensor, residual:Tensor, **kwargs):
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, *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, *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, *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, *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, 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"])
saves.extend([x_normed, rrms, *s, x_w13])
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"])
saves.extend([*s, out])
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_saves = self.attention(x, freqs_cis, **attn_kwargs)
ffn, h, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
h = h + ffn
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
if not mp:
for v in get_parameters(self): v.shard_(device, axis=None)
else:
# flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
def _shard_fp8(name:str, axis:int, std:float=0.02):
w = getattr(self, name)
if MXFP8:
from extra.gemm.cdna_asm_gemm import quantize_mxfp8
w_bf16 = Tensor.empty(self.n_layers, w.shape[1], w.shape[2], dtype=dtypes.bfloat16).shard(device, axis=axis).randn_like() * std
w_q, w_e8, _ = quantize_mxfp8(w_bf16)
w.replace(w_q)
self._fp8_inv_scale[name].replace(w_e8.contiguous()).is_param_(False)
self._fp8_next_inv_scale[name].replace(w_e8.contiguous()).is_param_(False)
else:
w.shard_(device, axis=axis)
scale_axis = (1 if axis == 1 else None) if COLUMNWISE_WEIGHT_SCALE else None
self._fp8_inv_scale[name] = self._fp8_inv_scale[name].shard(device, axis=scale_axis).contiguous().is_param_(False)
self._fp8_next_inv_scale[name] = self._fp8_next_inv_scale[name].shard(device, axis=scale_axis).contiguous().is_param_(False)
Tensor.realize(w, self._fp8_inv_scale[name], self._fp8_next_inv_scale[name])
sstd = 0.02 / math.sqrt(2 * self.n_layers)
_shard_fp8("wqkv", 1) # (n_layers, out, dim) shard out
_shard_fp8("wo", 2, sstd) # (n_layers, dim, in) shard in
if SPLIT_W13:
_shard_fp8("w1", 1)
_shard_fp8("w3", 1)
else:
_shard_fp8("w13", 1) # (n_layers, hidden*2, dim) shard out
_shard_fp8("w2", 2, sstd) # (n_layers, dim, hidden) shard in
self.attention_norm.shard_(device, axis=None).realize()
self.ffn_norm.shard_(device, axis=None).realize()
self.norm.weight.shard_(device, axis=None).realize()
self.tok_embeddings.weight.shard_(device, axis=0).realize()
self.output.shard_(device, axis=1).realize()
self.freqs_cis.shard_(device, axis=None).realize()
for amax_dict in (self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax):
for name in amax_dict:
for i in range(len(amax_dict[name])):
amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().is_param_(False)
def reset_amax(self):
for st in (self._fp8_next_amax, self._fp8_next_grad_amax):
for ts in st.values():
for t in ts: t.assign(0)
def update_amax(self):
for cur, nxt in ((self._fp8_amax, self._fp8_next_amax), (self._fp8_grad_amax, self._fp8_next_grad_amax)):
for name in cur:
for c, n in zip(cur[name], nxt[name]): c.assign(n)
def __call__(self, tokens:Tensor, save:bool=True):
h = self.tok_embeddings(tokens)
freqs_cis = self.freqs_cis.cast(h.dtype)
if not getenv("HK_FLASH_ATTENTION"): freqs_cis = freqs_cis[:, :tokens.shape[1], :, :, :]
a, na, ga, nga, s = self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax, self._fp8_inv_scale
def amax_kwargs(i:int, act_names:tuple[str, ...], grad_names:tuple[str, ...]) -> dict[str, Tensor|None]:
specs = (("amax_", a, act_names), ("next_amax_", na, act_names), ("grad_amax_", ga, grad_names), ("next_grad_amax_", nga, grad_names))
if MXFP4: return dict.fromkeys(f"{prefix}{name}" for prefix, _, names in specs for name in names)
return {f"{prefix}{name}":val[name][i] for prefix, val, names in specs for name in names}
for i in range(self.n_layers):
attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
**amax_kwargs(i, ("xqkv", "xo"), ("xqkv", "xo")))
ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i], s_2=s["w2"][i], **amax_kwargs(i, ("x2",), ("xout",)))
if SPLIT_W13:
ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], s_1=s["w1"][i], s_3=s["w3"][i], **amax_kwargs(i, ("x1", "x3"), ("xw1", "xw3")))
else:
ffn_kwargs.update(w13=self.w13[i], s_13=s["w13"][i], **amax_kwargs(i, ("x13",), ("xw13",)))
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
def _get_pads(uop:UOp) -> list[UOp]:
if uop.op == Ops.ADD: return _get_pads(uop.src[0]) + _get_pads(uop.src[1])
return [uop]
def apply_grad(grad_buf:Tensor, new_grad:UOp):
pads = _get_pads(new_grad)
if len(pads) <= 1:
new_grad = new_grad.cast(grad_buf.dtype)
grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
return
cur = grad_buf.uop
for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
if pad.op == Ops.PAD:
grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
buf_slice = cur.shrink(grad_shrink)
cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
else:
cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
grad_buf.uop = cur
if __name__ == "__main__":
config = {}
BS = config["BS"] = getenv("BS", 16)
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
SMALL = config["SMALL"] = getenv("SMALL", 0)
from examples.llama3 import MODEL_PARAMS
model_params = MODEL_PARAMS[llama_size:=getenv("LLAMA3_SIZE", "8B")]["args"]
# vocab_size from mixtral tokenizer
if not SMALL: model_params |= {"vocab_size": 32000}
real_vocab_size = model_params['vocab_size']
if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params["n_layers"] = llama_layers
# pad vocab
if (MP := getenv("MP", 1)) > 1: model_params["vocab_size"] = round_up(model_params["vocab_size"], 256 * MP)
vocab_mask:Tensor = Tensor.arange(model_params["vocab_size"]).reshape(1, 1, -1) >= real_vocab_size
model = FlatTransformer(**model_params, max_context=SEQLEN)
state = nn.state.get_state_dict(model)
print("tensor count:", len(state))
# shard the model
from tinygrad import Device
is_dp = (DP := getenv("DP", 1)) > 1
is_mp = (MP := getenv("MP", 1)) > 1
is_sharding = is_dp or is_mp
device_count = max(DP, MP)
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
model.shard(device, is_mp)
if is_dp: vocab_mask.shard_(device, axis=None).realize()
if is_mp: vocab_mask.shard_(device, axis=2).realize()
# preallocate all the grad buffers and zero them out
grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype
grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param}
fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts]
# print model size
sz = 0
for k,v in state.items():
print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device} {v.nbytes()/1e9:.2f} GB")
sz += v.nbytes()
print(f"total sz: {sz/1e9:.2f} GB")
with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=real_vocab_size, dtype=dtypes.int)
with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
if DP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)), axis=0)
if MP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)))
@TinyJit
def fwd_bwd(tokens:Tensor):
with Timing("python forward: "):
model.reset_amax()
logits = model(tokens[:, :-1], save=llama_size=="8B")
loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
with Timing("python backward: "):
for t,g in zip(grads, loss.gradient(*grads)):
apply_grad(grads[t], g.uop)
with Timing("run fwd_bwd: "): loss.realize(*grads.values(), *fp8_amax, *fp8_grad_amax)
@TinyJit
def optim_step():
for g in grads.values(): g.assign(g.zeros_like())
Tensor.realize(*grads.values())
for i in range(6):
GlobalCounters.reset()
profile_marker(f"step {i}")
with Timing(colored(f"*** step {i}: ", "red")):
fwd_bwd(tokens)
optim_step()
print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))

351
artifacts/package_sources/tinygrad/examples/mlperf/models/gpt_oss.py Normal file
View File

@@ -0,0 +1,351 @@
import math, os, functools
if __name__ == "__main__":
os.environ["DEFAULT_FLOAT"] = "bfloat16"
os.environ["OPTIM_DTYPE"] = "bfloat16"
if "DEV" not in os.environ: os.environ["DEV"] = "NULL::gfx950"
# CDNA
os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
os.environ["ALL2ALL"] = "1"
os.environ["USE_ATOMICS"] = "1"
from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
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.02
ASM_GEMM = getenv("ASM_GEMM", 0)
def _quant_dequant_fwd(x:Tensor) -> Tensor:
# x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis)
M, K = x.shape
scale_K = K // 32
amax = x.float().reshape(M, scale_K, 32).abs().max(axis=-1)
e8 = (amax.maximum(1e-38).log2().floor() + 127).clamp(0, 254).cast(dtypes.uint8)
qscale = (127.0 - e8.cast(dtypes.float32)).exp2().reshape(M, scale_K, 1).expand(M, scale_K, 32).reshape(M, K)
x_fp8 = (x.float() * qscale).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE).cast(dtypes.float32)
return (x_fp8 * _mx_block_scale(e8)).cast(dtypes.bfloat16)
@functools.cache
def _quant_dequant_fwd_fxn(x_p, device):
return _quant_dequant_fwd(Tensor(x_p, device=device))
def _quant_dequant_bwd(grad:UOp, call:UOp) -> tuple:
return (Tensor(grad).cast(dtypes.bfloat16).uop,)
def quant_dequant_mx(x:Tensor) -> Tensor:
fxn = _quant_dequant_fwd_fxn(x.as_param(0).uop, x.device)
return Tensor(UOp.maketuple(fxn.uop).call(x.uop, grad_fxn=_quant_dequant_bwd).gettuple(0))
def _mx_scale(e8:Tensor) -> Tensor:
return _mx_block_scale(e8) if e8.ndim == 2 else _mx_block_scale_3d(e8)
def _dequant_fwd(w_q:Tensor, w_scale:Tensor) -> Tensor:
return w_q.cast(dtypes.bfloat16) * _mx_scale(w_scale)
@functools.cache
def _dequant_fwd_fxn(wq_p, ws_p, device):
return _dequant_fwd(Tensor(wq_p, device=device), Tensor(ws_p, device=device))
def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
return (Tensor(grad).cast(dtypes.bfloat16).uop, None)
def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)
call = UOp.maketuple(fxn.uop).call(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd)
return Tensor(call.gettuple(0))
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)
x_linear = x_linear.clamp(-limit, limit)
return (x_glu * (alpha * x_glu).sigmoid()) * (x_linear + 1)
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, dtype=dtypes.float32)[:(dim // 2)] / dim))
freqs = Tensor.arange(end, dtype=dtypes.float32).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).cast(dtypes.default_float).reshape(1, end, 1, dim//2, 2)
class GPTOSS:
def __init__(self, dim:int, n_layers:int, n_heads:int, n_kv_heads:int, head_dim:int, n_experts:int, experts_per_tok:int,
intermediate_size:int, vocab_size:int, norm_eps:float=1e-5, rope_theta:int=150000, sliding_window:int=128,
swiglu_limit:float=7.0, max_context:int=8192):
self.dim, self.n_layers, self.n_heads, self.n_kv_heads, self.head_dim = dim, n_layers, n_heads, n_kv_heads, head_dim
self.n_rep = n_heads // n_kv_heads
self.n_experts, self.experts_per_tok, self.intermediate_size = n_experts, experts_per_tok, intermediate_size
self.vocab_size, self.norm_eps, self.sliding_window, self.swiglu_limit = vocab_size, norm_eps, sliding_window, swiglu_limit
self.sm_scale = 1.0 / math.sqrt(head_dim)
scaled_std = INIT_STD / math.sqrt(2 * n_layers)
q_dim, qkv_dim = n_heads * head_dim, head_dim * (n_heads + 2 * n_kv_heads)
# attn
self.wqkv, self.wqkv_scale = self._quant_weight(n_layers, qkv_dim, dim)
self.wqkv_bias = Tensor.zeros(n_layers, qkv_dim, dtype=dtypes.bfloat16).contiguous()
self.wo, self.wo_scale = self._quant_weight(n_layers, dim, q_dim, std=scaled_std)
self.wo_bias = Tensor.zeros(n_layers, dim, dtype=dtypes.bfloat16).contiguous()
self.sinks = Tensor.zeros(n_layers, n_heads, dtype=dtypes.bfloat16).contiguous()
self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
# moe ffn
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, 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, moe=True)
self.w_down_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous()
# output
self.norm = nn.RMSNorm(dim, norm_eps)
self.tok_embeddings = nn.Embedding(vocab_size, dim)
self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
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, 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)
return (j <= i).where(0.0, -1e30).cast(dtype).contiguous()
def _sliding_attention(self, xq:Tensor, xk:Tensor, xv:Tensor, sinks:Tensor) -> Tensor:
bsz, seqlen, H, hd = xq.shape
KV, R, W = self.n_kv_heads, self.n_rep, self.sliding_window
assert seqlen % W == 0, f"seqlen {seqlen} must be a multiple of sliding_window {W} for banded attention"
nb = seqlen // W
q = xq.reshape(bsz, seqlen, KV, R, hd).permute(0, 2, 3, 1, 4).reshape(bsz, KV, R, nb, W, hd).float()
k, v = (x.permute(0, 2, 1, 3).reshape(bsz, KV, 1, nb, W, hd).float() for x in (xk, xv))
kk, vv = (x.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb].cat(x, dim=-2) for x in (k, v))
sc = (q @ kk.transpose(-1, -2)) * self.sm_scale # (B,KV,R,nb,W,2W)
i, j, pv = Tensor.arange(W).reshape(W, 1), Tensor.arange(2 * W).reshape(1, 2 * W), Tensor.arange(nb).reshape(nb, 1, 1) >= 1
sc = ((j > i) & (j <= i + W) & (pv | (j >= W))).where(sc, -float("inf"))
sink = sinks.reshape(1, KV, R, 1, 1, 1).float()
m = sc.max(-1, keepdim=True).maximum(sink)
e = (sc - m).exp()
p = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
attn = p @ vv.cast(dtypes.bfloat16)
return attn.reshape(bsz, KV, R, seqlen, hd).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, H * hd)
def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, *, attention_norm:Tensor, wqkv:Tensor,
wqkv_scale:Tensor, wqkv_bias:Tensor, wo:Tensor, wo_scale:Tensor, wo_bias:Tensor, sinks:Tensor):
bsz, seqlen, _ = x.shape
x_normed, rrms = rmsnorm(x, self.norm_eps)
qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
qkv = qkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
xq = qkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
xk, xv = qkv[:, :, :, self.n_rep], qkv[:, :, :, self.n_rep + 1]
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16) # (B,N,H,D)/(B,N,KV,D)
if sliding:
attn = self._sliding_attention(xq, xk, xv, sinks)
elif getenv("HK_FLASH_ATTENTION"):
from extra.thunder.amd.fa import flash_attention
attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
else:
xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)
scores = (xqm @ xkm.transpose(-2, -1)).float() * self.sm_scale + mask
sink = sinks.reshape(1, self.n_kv_heads, self.n_rep, 1, 1).float()
m = scores.max(-1, keepdim=True).maximum(sink)
e = (scores - m).exp()
w = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
attn = (w @ xvm).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)
out = matmul_mx(attn, wo, wo_scale) + wo_bias
return out, [x_normed, rrms, attn]
def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
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()
dim, inter = self.dim, self.intermediate_size
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)
def run_layer(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
attn, attn_saves = self.attention(x, freqs_cis, mask, sliding, **attn_kwargs)
h = x + attn
ffn, ffn_saves = self.feed_forward(h, **ffn_kwargs)
h = h + ffn
if save: return (h, *attn_saves, *ffn_saves)
return (h,)
def shard(self, device:tuple[str, ...], mp:bool=False):
assert not mp, "MP not supported"
from tinygrad.nn.state import get_parameters
for v in get_parameters(self): v.shard_(device, axis=None)
Tensor.realize(*get_parameters(self))
def __call__(self, tokens:Tensor, save:bool=True):
h = self.tok_embeddings(tokens)
bsz, seqlen = tokens.shape
freqs_cis = self.freqs_cis.cast(h.dtype)[:, :seqlen, :, :, :]
mask_full = None if getenv("HK_FLASH_ATTENTION") else self._attn_mask(seqlen, dtypes.float32)
for i in range(self.n_layers):
attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wqkv_scale=self.wqkv_scale[i],
wqkv_bias=self.wqkv_bias[i], wo=self.wo[i], wo_scale=self.wo_scale[i], wo_bias=self.wo_bias[i],
sinks=self.sinks[i])
ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], gate=self.gate[i], gate_bias=self.gate_bias[i],
w_gate_up=self.w_gate_up[i], w_gate_up_scale=self.w_gate_up_scale[i], w_gate_up_bias=self.w_gate_up_bias[i],
w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)
logits = self.norm(h) @ self.output.T
return logits
def _get_pads(uop:UOp) -> list[UOp]:
if uop.op == Ops.ADD: return _get_pads(uop.src[0]) + _get_pads(uop.src[1])
return [uop]
def apply_grad(grad_buf:Tensor, new_grad:UOp):
pads = _get_pads(new_grad)
if len(pads) <= 1:
new_grad = new_grad.cast(grad_buf.dtype)
grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
return
cur = grad_buf.uop
for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
if pad.op == Ops.PAD:
grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
buf_slice = cur.shrink(grad_shrink)
cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
else:
cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
grad_buf.uop = cur
GPT_OSS_20B = dict(dim=2880, n_layers=24, n_heads=64, n_kv_heads=8, head_dim=64, n_experts=32, experts_per_tok=4,
intermediate_size=2880, vocab_size=128256, norm_eps=1e-5, rope_theta=150000, sliding_window=128,
swiglu_limit=7.0)
if __name__ == "__main__":
config = {}
BS = config["BS"] = getenv("BS", 16)
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
model_params = GPT_OSS_20B
real_vocab_size = model_params["vocab_size"]
if (layers := getenv("LAYERS")) != 0: model_params["n_layers"] = layers
model = GPTOSS(**model_params, max_context=SEQLEN)
state = nn.state.get_state_dict(model)
print("tensor count:", len(state))
from tinygrad import Device
is_dp = (DP := getenv("DP", 1)) > 1
device_count = DP
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
if is_dp: model.shard(device)
# preallocate all the grad buffers and zero them out
grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype
grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param}
# print model size
sz = 0
for k,v in state.items():
print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device} {v.nbytes()/1e9:.2f} GB")
sz += v.nbytes()
print(f"total sz: {sz/1e9:.2f} GB")
with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=real_vocab_size, dtype=dtypes.int)
with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
if is_dp: tokens = tokens.shard(device, axis=0)
@TinyJit
def fwd_bwd(tokens:Tensor):
with Timing("python forward: "):
logits = model(tokens[:, :-1], save=True)
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
with Timing("python backward: "):
for t,g in zip(grads, loss.gradient(*grads)):
apply_grad(grads[t], g.uop)
with Timing("run fwd_bwd: "): loss.realize(*grads.values())
@TinyJit
def optim_step():
for g in grads.values(): g.assign(g.zeros_like())
Tensor.realize(*grads.values())
for i in range(6):
GlobalCounters.reset()
profile_marker(f"step {i}")
with Timing(colored(f"*** step {i}: ", "red")):
fwd_bwd(tokens)
optim_step()
print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))

68
artifacts/package_sources/tinygrad/examples/mlperf/models/test_apply_grad.py Normal file
View File

@@ -0,0 +1,68 @@
import unittest
from tinygrad import Tensor, TinyJit
from tinygrad.nn.state import get_parameters
from examples.mlperf.models.flat_llama import apply_grad
class FlatModel:
def __init__(self, n_layers:int, dim:int, hidden:int):
self.n_layers = n_layers
self.w1 = Tensor.uniform(n_layers, dim, hidden, low=-0.1, high=0.1)
self.w2 = Tensor.uniform(n_layers, hidden, dim, low=-0.1, high=0.1)
self.scale = Tensor.uniform(dim, low=0.9, high=1.1)
self.bias = Tensor.zeros(dim).contiguous()
def __call__(self, x:Tensor) -> Tensor:
h = x
for i in range(self.n_layers):
h = (h @ self.w1[i]).relu() @ self.w2[i] + h
return (h * self.scale + self.bias).sum()
class TestApplyGradE2E(unittest.TestCase):
def _run_with_apply_grad(self, model, xs):
grads = {p: Tensor.zeros(p.shape, dtype=p.dtype).contiguous().realize() for p in get_parameters(model)}
for x in xs:
loss = model(x)
for p, g in zip(grads, loss.gradient(*grads)):
apply_grad(grads[p], g.uop)
Tensor.realize(loss, *grads.values())
return [grads[p] for p in get_parameters(model)]
def _run_reference(self, model, xs):
for x in xs: model(x).backward()
return [p.grad for p in get_parameters(model)]
def _assert_close(self, got, expected, atol, rtol):
for g, e in zip(got, expected):
self.assertTrue(g.allclose(e, atol=atol, rtol=rtol).item(), f"grad mismatch (max abs diff {(g - e).abs().max().item()})")
def _assert_match(self, model, xs, atol, rtol):
self._assert_close(self._run_with_apply_grad(model, xs), self._run_reference(model, xs), atol, rtol)
def test_e2e_single_step(self):
model = FlatModel(n_layers=3, dim=8, hidden=16)
Tensor.realize(*get_parameters(model))
self._assert_match(model, [Tensor.randn(2, 8).realize()], atol=1e-4, rtol=1e-4)
def test_e2e_multi_step_accumulation(self):
model = FlatModel(n_layers=4, dim=8, hidden=16)
Tensor.realize(*get_parameters(model))
self._assert_match(model, [Tensor.randn(2, 8).realize() for _ in range(3)], atol=1e-4, rtol=1e-4)
def test_e2e_jit(self):
model = FlatModel(n_layers=3, dim=8, hidden=16)
Tensor.realize(*get_parameters(model))
grads = {p: Tensor.zeros(p.shape, dtype=p.dtype).contiguous().realize() for p in get_parameters(model)}
@TinyJit
def fwd_bwd(x:Tensor):
loss = model(x)
for p, g in zip(grads, loss.gradient(*grads)): apply_grad(grads[p], g.uop)
Tensor.realize(loss, *grads.values())
xs = [Tensor.randn(2, 8).realize() for _ in range(3)]
for x in xs: fwd_bwd(x)
self._assert_close([grads[p] for p in get_parameters(model)], self._run_reference(model, xs), atol=1e-3, rtol=1e-3)
if __name__ == "__main__":
unittest.main()

137
artifacts/package_sources/tinygrad/examples/mlperf/models/test_flat_llama.py Normal file
View File

@@ -0,0 +1,137 @@
import os
os.environ["WQKV"] = "1"
import unittest
import numpy as np
from tinygrad import Tensor, nn, dtypes
from tinygrad.device import Device
from examples.mlperf.models.llama import Transformer
from examples.mlperf.models.flat_llama import FlatTransformer
def copy_weights(flat:FlatTransformer, ref:Transformer):
n_layers = flat.n_layers
Tensor.realize(*nn.state.get_state_dict(ref).values())
flat.wqkv.assign(Tensor(np.stack([ref.layers[i].attention.wqkv.weight.numpy() for i in range(n_layers)])))
flat.wo.assign(Tensor(np.stack([ref.layers[i].attention.wo.weight.numpy() for i in range(n_layers)])))
flat.w1.assign(Tensor(np.stack([ref.layers[i].feed_forward.w1.weight.numpy() for i in range(n_layers)])))
flat.w2.assign(Tensor(np.stack([ref.layers[i].feed_forward.w2.weight.numpy() for i in range(n_layers)])))
flat.w3.assign(Tensor(np.stack([ref.layers[i].feed_forward.w3.weight.numpy() for i in range(n_layers)])))
flat.attention_norm.assign(Tensor(np.stack([ref.layers[i].attention_norm.weight.numpy() for i in range(n_layers)])))
flat.ffn_norm.assign(Tensor(np.stack([ref.layers[i].ffn_norm.weight.numpy() for i in range(n_layers)])))
flat.norm.weight.assign(Tensor(ref.norm.weight.numpy()))
flat.tok_embeddings.weight.assign(Tensor(ref.tok_embeddings.weight.numpy()))
flat.output.weight.assign(Tensor(ref.output.weight.numpy()))
class TestFlatLlama(unittest.TestCase):
def test_forward_match(self):
Tensor.manual_seed(42)
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
ref = Transformer(**params)
flat = FlatTransformer(**params)
copy_weights(flat, ref)
Tensor.realize(*nn.state.get_state_dict(flat).values())
tokens = Tensor([[1, 50, 100, 999, 2]])
ref_logits = ref(tokens).realize()
flat_logits = flat(tokens).realize()
self.assertEqual(ref_logits.shape, flat_logits.shape)
diff = (ref_logits - flat_logits).abs().max().item()
self.assertLess(diff, 1e-5, f"forward mismatch: max abs diff {diff}")
def test_backward_match(self):
Tensor.manual_seed(42)
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
ref = Transformer(**params)
flat = FlatTransformer(**params)
copy_weights(flat, ref)
Tensor.realize(*nn.state.get_state_dict(flat).values())
tokens = Tensor([[1, 50, 100, 999, 2, 10]])
ref_loss = ref(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
ref_loss.backward()
ref_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(ref).items() if v.grad is not None}
flat_loss = flat(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
flat_loss.backward()
flat_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(flat).items() if v.grad is not None}
# check loss matches
self.assertAlmostEqual(ref_loss.item(), flat_loss.item(), places=4)
# check output weight grad matches
diff = abs(ref_grads["output.weight"] - flat_grads["output.weight"]).max()
self.assertLess(diff, 1e-4, f"output.weight grad mismatch: max abs diff {diff}")
# check per-layer weight grads match
for i in range(params["n_layers"]):
for flat_key, ref_key in [
("wqkv", f"layers.{i}.attention.wqkv.weight"),
("wo", f"layers.{i}.attention.wo.weight"),
("w1", f"layers.{i}.feed_forward.w1.weight"),
("w2", f"layers.{i}.feed_forward.w2.weight"),
("w3", f"layers.{i}.feed_forward.w3.weight"),
]:
diff = abs(ref_grads[ref_key] - flat_grads[flat_key][i]).max()
self.assertLess(diff, 1e-4, f"layer {i} {flat_key} grad mismatch: max abs diff {diff}")
@unittest.skipUnless(Device.DEFAULT == "CPU", "multi-device CPU test")
def test_forward_match_mp(self):
Tensor.manual_seed(42)
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
from tinygrad import Device
devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
ref = Transformer(**params)
flat = FlatTransformer(**params)
copy_weights(flat, ref)
Tensor.realize(*nn.state.get_state_dict(flat).values())
flat.shard(devices, mp=True)
tokens = Tensor([[1, 50, 100, 999, 2]], device=devices[0])
ref_logits = ref(tokens.to(devices[0])).numpy()
flat_logits = flat(tokens.shard(devices)).numpy()
self.assertEqual(ref_logits.shape, flat_logits.shape)
np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
@unittest.skipUnless(Device.DEFAULT == "CPU", "multi-device CPU test")
def test_forward_match_dp(self):
Tensor.manual_seed(42)
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
from tinygrad import Device
devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
ref = Transformer(**params)
flat = FlatTransformer(**params)
copy_weights(flat, ref)
Tensor.realize(*nn.state.get_state_dict(flat).values())
flat.shard(devices)
tokens = Tensor([[1, 50, 100, 999, 2], [2, 100, 50, 1, 999]], device=devices[0])
ref_logits = ref(tokens.to(devices[0])).numpy()
flat_logits = flat(tokens.shard(devices, axis=0)).numpy()
self.assertEqual(ref_logits.shape, flat_logits.shape)
np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
@unittest.skipUnless(dtypes.fp8e4m3 in Device[Device.DEFAULT].renderer.supported_dtypes(), "fp8 not supported on this device")
def test_forward_fp8(self):
import examples.mlperf.models.flat_llama as flat_llama_mod
old_fp8 = flat_llama_mod.FP8
try:
flat_llama_mod.FP8 = 1
Tensor.manual_seed(42)
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
ref = Transformer(**params)
flat = FlatTransformer(**params)
copy_weights(flat, ref)
Tensor.realize(*nn.state.get_state_dict(flat).values())
tokens = Tensor([[1, 50, 100, 999, 2]])
ref_logits = ref(tokens).numpy()
flat_logits = flat(tokens).numpy()
self.assertEqual(ref_logits.shape, flat_logits.shape)
# FP8 has lower precision, allow larger tolerance
np.testing.assert_allclose(flat_logits, ref_logits, atol=1.0, rtol=0.1)
finally:
flat_llama_mod.FP8 = old_fp8
if __name__ == "__main__":
unittest.main()