forked from tinygrad/tinygrad
qazalandGitHub
44f1f45cd5
* start by copying the C * uop kernel * cleanup tests * estimates is part of SPEC
50 lines
2.5 KiB
Python
50 lines
2.5 KiB
Python
import functools, math
|
|
from tinygrad import Tensor, dtypes
|
|
from tinygrad.uop.ops import UOp, KernelInfo
|
|
from tinygrad.renderer import Estimates
|
|
from extra.llama_kernels import alloc_like
|
|
|
|
LOG2E = 1.4426950408889634
|
|
|
|
@functools.cache
|
|
def _custom_swiglu(out:UOp, x_w13:UOp) -> UOp:
|
|
rows, hidden = math.prod(x_w13.shape[:-1]), x_w13.shape[-1]//2
|
|
n_elems = rows * hidden
|
|
out, x_w13 = out.reshape(n_elems), x_w13.reshape(rows, 2*hidden)
|
|
i = UOp.range(n_elems, 0)
|
|
row, col = i // hidden, i % hidden
|
|
act, gate = x_w13[row, col].cast(dtypes.float), x_w13[row, hidden+col].cast(dtypes.float)
|
|
sigmoid = (1.0 + (-LOG2E * act).exp2()).reciprocal()
|
|
store = out[i].store((act * sigmoid * gate).cast(out.dtype))
|
|
return store.end(i).sink(arg=KernelInfo(f"swiglu_fwd_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=6*n_elems)))
|
|
|
|
@functools.cache
|
|
def _custom_swiglu_bwd(grad_out:UOp, x_w13:UOp, grad_act:UOp) -> UOp:
|
|
rows, hidden = math.prod(x_w13.shape[:-1]), x_w13.shape[-1]//2
|
|
n_elems = rows * hidden
|
|
grad_out, x_w13, grad_act = grad_out.reshape(rows, 2*hidden), x_w13.reshape(rows, 2*hidden), grad_act.reshape(n_elems)
|
|
i = UOp.range(n_elems, 0)
|
|
row, col = i // hidden, i % hidden
|
|
act, gate = x_w13[row, col].cast(dtypes.float), x_w13[row, hidden+col].cast(dtypes.float)
|
|
grad = grad_act[i].cast(dtypes.float)
|
|
sigmoid = (1.0 + (-LOG2E * act).exp2()).reciprocal()
|
|
silu = act * sigmoid
|
|
dact = grad_out[row, col].store((grad * (sigmoid + silu * (1.0 - sigmoid)) * gate).cast(grad_out.dtype))
|
|
dgate = grad_out.after(dact)[row, hidden+col].store((grad * silu).cast(grad_out.dtype))
|
|
return dgate.end(i).sink(arg=KernelInfo(f"swiglu_bwd_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=10*n_elems)))
|
|
|
|
def _swiglu_bwd(gradient:UOp, kernel:UOp):
|
|
_, x_w13 = kernel.src[1:]
|
|
axis = x_w13.axis if isinstance(x_w13.device, tuple) else None
|
|
grad_out = alloc_like(x_w13.shape, dtypes.bfloat16, x_w13.device, axis)
|
|
grad_out, *_ = Tensor.custom_kernel(grad_out, Tensor(x_w13, device=x_w13.device), Tensor(gradient, device=x_w13.device),
|
|
fxn=_custom_swiglu_bwd)
|
|
return (None, grad_out.uop)
|
|
|
|
def swiglu(x_w13:Tensor) -> Tensor:
|
|
assert x_w13.dtype == dtypes.bfloat16 and x_w13.ndim >= 2 and x_w13.shape[-1] % 32 == 0
|
|
*prefix, two_k = x_w13.shape
|
|
axis = x_w13.uop.axis if isinstance(x_w13.device, tuple) else None
|
|
out = alloc_like((*prefix, two_k//2), dtypes.bfloat16, x_w13.device, axis)
|
|
return Tensor.custom_kernel(out, x_w13, fxn=_custom_swiglu, grad_fxn=_swiglu_bwd)[0]
|