| |
|
|
| 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() |
|
|
