|
|
|
|
|
|
|
|
import mlx.core as mx |
|
|
import mlx.nn as nn |
|
|
from time_utils import time_fn |
|
|
|
|
|
|
|
|
def rms_norm(x, w, eps): |
|
|
ot = x.dtype |
|
|
x = x.astype(mx.float32) |
|
|
n = mx.rsqrt(x.square().mean(-1, keepdims=True) + eps) |
|
|
y = (x * n).astype(ot) |
|
|
if w is not None: |
|
|
y = y * w |
|
|
return y |
|
|
|
|
|
|
|
|
def time_rms_norm(): |
|
|
f1 = lambda x, w, y: (rms_norm(x, w, 1e-5) * y).sum() |
|
|
f2 = lambda x, w, y: (mx.fast.rms_norm(x, w, 1e-5) * y).sum() |
|
|
g1 = mx.grad(f1, argnums=(0, 1)) |
|
|
g2 = mx.grad(f2, argnums=(0, 1)) |
|
|
|
|
|
x = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16) |
|
|
w = mx.random.uniform(shape=(4096,)).astype(mx.float16) |
|
|
y = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16) |
|
|
mx.eval(x, w, y) |
|
|
|
|
|
def rms_norm_loop(g, x, w): |
|
|
gx, gw = x, w |
|
|
for _ in range(32): |
|
|
gx, gw = g(gx, gw, y) |
|
|
return gx, gw |
|
|
|
|
|
time_fn(rms_norm_loop, g1, x, w) |
|
|
time_fn(rms_norm_loop, g2, x, w) |
|
|
time_fn(rms_norm_loop, mx.compile(g1), x, w) |
|
|
time_fn(rms_norm_loop, mx.compile(g2), x, w) |
|
|
|
|
|
f1 = lambda x, y: (rms_norm(x, None, 1e-5) * y).sum() |
|
|
f2 = lambda x, y: (mx.fast.rms_norm(x, None, 1e-5) * y).sum() |
|
|
g1 = mx.grad(f1, argnums=(0,)) |
|
|
g2 = mx.grad(f2, argnums=(0,)) |
|
|
|
|
|
x = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16) |
|
|
w = mx.random.uniform(shape=(4096,)).astype(mx.float16) |
|
|
y = mx.random.uniform(shape=(8, 1024, 4096)).astype(mx.float16) |
|
|
mx.eval(x, w, y) |
|
|
|
|
|
def rms_norm_loop(g, x): |
|
|
gx = x |
|
|
for _ in range(32): |
|
|
gx = g(gx, y) |
|
|
return gx |
|
|
|
|
|
time_fn(rms_norm_loop, g1, x) |
|
|
time_fn(rms_norm_loop, g2, x) |
|
|
time_fn(rms_norm_loop, mx.compile(g1), x) |
|
|
time_fn(rms_norm_loop, mx.compile(g2), x) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
time_rms_norm() |
|
|
|
