forked from tinygrad/tinygrad
69 lines
2.1 KiB
Python
69 lines
2.1 KiB
Python
from tinygrad import Tensor, dtypes, Device, TinyJit
|
|
from tinygrad.helpers import getenv, Timing, round_up
|
|
|
|
# llama 70b shapes
|
|
DIM, HIDDEN, N_HEADS, N_KV_HEADS, HEAD_DIM = 8192, 28672, 64, 8, 128
|
|
QKV_OUT = N_HEADS*HEAD_DIM + 2*N_KV_HEADS*HEAD_DIM
|
|
LAYER_SHAPES = {
|
|
"wqkv": ((QKV_OUT, DIM), 0),
|
|
"wo": ((DIM, N_HEADS*HEAD_DIM), 1),
|
|
"w1": ((HIDDEN, DIM), 0),
|
|
"w3": ((HIDDEN, DIM), 0),
|
|
"w2": ((DIM, HIDDEN), 1),
|
|
}
|
|
|
|
def build_params(layers:int, devs:tuple[str,...], dtype):
|
|
ts, nbytes = [], 0
|
|
ndev = len(devs)
|
|
for _ in range(layers):
|
|
for (a,b), ax in LAYER_SHAPES.values():
|
|
shape = (round_up(a, ndev) if ax==0 else a, round_up(b, ndev) if ax==1 else b)
|
|
t = Tensor.empty(*shape, dtype=dtype).shard(devs, axis=ax).realize()
|
|
ts.append(t)
|
|
nbytes += t.numel()*dtype.itemsize
|
|
return ts, nbytes
|
|
|
|
def sync():
|
|
Device["CPU"].synchronize()
|
|
Device[Device.DEFAULT].synchronize()
|
|
|
|
if __name__ == "__main__":
|
|
ngpus = getenv("GPUS", 8)
|
|
layers = getenv("LAYERS", 8)
|
|
devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(ngpus))
|
|
|
|
# to cpu
|
|
grads, nbytes = build_params(layers, devs, dtypes.bfloat16)
|
|
cpu_dst = [Tensor.rand(*g.shape, dtype=dtypes.bfloat16, device="CPU").contiguous().realize() for g in grads]
|
|
sync()
|
|
def to_cpu():
|
|
for d,g in zip(cpu_dst, grads):
|
|
d.assign(g.to("CPU"))
|
|
return Tensor.realize(*cpu_dst)
|
|
to_cpu_jit = TinyJit(to_cpu)
|
|
for _ in range(3):
|
|
to_cpu_jit()
|
|
sync()
|
|
for _ in range(5):
|
|
with Timing("to cpu ", on_exit=lambda ns: f" @ {nbytes/ns:.2f} GB/s"):
|
|
to_cpu_jit()
|
|
sync()
|
|
|
|
# to gpu
|
|
from extra.gemm.cdna_asm_gemm import FP8_DTYPE
|
|
w_gpu, nbytes = build_params(layers, devs, FP8_DTYPE)
|
|
cpu_src = [Tensor.rand(*w.shape, dtype=FP8_DTYPE, device="CPU").contiguous().realize() for w in w_gpu]
|
|
sync()
|
|
def to_gpu():
|
|
for w,c in zip(w_gpu, cpu_src):
|
|
w.assign(c.shard_like(w))
|
|
return Tensor.realize(*w_gpu)
|
|
to_gpu_jit = TinyJit(to_gpu)
|
|
for _ in range(3):
|
|
to_gpu_jit()
|
|
sync()
|
|
for _ in range(5):
|
|
with Timing("to gpu ", on_exit=lambda ns: f" @ {nbytes/ns:.2f} GB/s"):
|
|
to_gpu_jit()
|
|
sync()
|