mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 10:36:07 +00:00
add test for flat llama (#15327)
* add test for flat llama * simpler * back to split w1/w3 * env * still too much ram * invalid
This commit is contained in:
George Hotz
@@ -3,6 +3,15 @@ if __name__ == "__main__":
|
||||
os.environ["DEFAULT_FLOAT"] = "bfloat16"
|
||||
os.environ["OPTIM_DTYPE"] = "bfloat16"
|
||||
os.environ["DEV"] = "NULL"
|
||||
# CDNA
|
||||
os.environ["EMULATE"] = "AMD_CDNA4"
|
||||
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
|
||||
from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
|
||||
@@ -50,18 +59,14 @@ class FlatTransformer:
|
||||
x = rmsnorm(x, self.norm_eps) * attention_norm
|
||||
xqkv = x @ wqkv.T
|
||||
|
||||
# reshapes
|
||||
xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
|
||||
xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
|
||||
xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
|
||||
xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
|
||||
xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
|
||||
xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
|
||||
xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
|
||||
bsz, seqlen, _ = xqkv.shape
|
||||
# interleaved layout: each kv group has [n_rep q heads, 1 k head, 1 v head] for clean MP sharding
|
||||
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)
|
||||
bsz, seqlen, _, _ = xq.shape
|
||||
|
||||
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)
|
||||
@@ -80,6 +85,24 @@ class FlatTransformer:
|
||||
h = x + self.attention(x, freqs_cis, attention_norm, wqkv, wo)
|
||||
return h + self.feed_forward(h, ffn_norm, w1, w2, w3)
|
||||
|
||||
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
|
||||
self.wqkv.shard_(device, axis=1).realize() # (n_layers, out, dim) shard out
|
||||
self.wo.shard_(device, axis=2).realize() # (n_layers, dim, in) shard in
|
||||
self.w1.shard_(device, axis=1).realize() # (n_layers, hidden, dim) shard out
|
||||
self.w2.shard_(device, axis=2).realize() # (n_layers, dim, hidden) shard in
|
||||
self.w3.shard_(device, axis=1).realize() # (n_layers, hidden, dim) shard out
|
||||
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.weight.shard_(device, axis=0).realize()
|
||||
self.freqs_cis.shard_(device, axis=None).realize()
|
||||
|
||||
def __call__(self, tokens:Tensor):
|
||||
h = self.tok_embeddings(tokens)
|
||||
freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
|
||||
@@ -108,8 +131,16 @@ if __name__ == "__main__":
|
||||
sz += v.nbytes()
|
||||
print(f"total sz: {sz/1e9:.2f} GB")
|
||||
|
||||
from tinygrad import Device
|
||||
if (DP := getenv("DP", 1)) > 1:
|
||||
model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)))
|
||||
if (MP := getenv("MP", 1)) > 1:
|
||||
model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)), mp=True)
|
||||
|
||||
with Timing("realize weights: "): Tensor.realize(*state.values())
|
||||
with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=model.vocab_size, dtype=dtypes.int).realize()
|
||||
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 jit_step(tokens:Tensor):
|
||||
@@ -122,4 +153,4 @@ if __name__ == "__main__":
|
||||
jit_step(tokens)
|
||||
jit_step(tokens)
|
||||
jit_step(tokens)
|
||||
print(f"mem used: {GlobalCounters.mem_used/1e9:.2f} GB")
|
||||
print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
import os
|
||||
os.environ["WQKV"] = "1"
|
||||
import unittest
|
||||
import numpy as np
|
||||
from tinygrad import Tensor, nn
|
||||
from tinygrad.nn.state import get_parameters
|
||||
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)
|
||||
|
||||
for p in get_parameters(ref): p.requires_grad_(True)
|
||||
for p in get_parameters(flat): p.requires_grad_(True)
|
||||
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}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -2702,13 +2702,14 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
|
||||
|
||||
if is_multi:
|
||||
if n_sharded:
|
||||
out = Tensor(Tensor.empty(batch, M, N//len(a.device), dtype=a.dtype, device=a.device).uop.multi(2), device=a.device)
|
||||
out = Tensor(Tensor.invalid(batch, M, N//len(a.device), dtype=a.dtype, device=a.device).uop.multi(2), device=a.device)
|
||||
elif m_sharded:
|
||||
out = Tensor(Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device).uop.multi(1), device=a.device)
|
||||
out = Tensor(Tensor.invalid(batch, M, N, dtype=a.dtype, device=a.device).uop.multi(1), device=a.device)
|
||||
else:
|
||||
out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
|
||||
out = Tensor(Tensor.invalid(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0),
|
||||
device=a.device)
|
||||
else:
|
||||
out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
|
||||
out = Tensor.invalid(batch, M, N, dtype=a.dtype, device=a.device)
|
||||
|
||||
renderer = Device[a.device[0] if is_multi else a.device].renderer
|
||||
dname, arch = renderer.device, getattr(renderer, "arch", "")
|
||||
|
||||
@@ -10,11 +10,11 @@ from tinygrad.uop.ops import UOp, Ops, KernelInfo
|
||||
|
||||
def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None=None) -> Tensor:
|
||||
dtype = dtype or ref.dtype
|
||||
if not isinstance(ref.device, tuple): return Tensor.empty(*shape, dtype=dtype, device=ref.device)
|
||||
if not isinstance(ref.device, tuple): return Tensor.invalid(*shape, dtype=dtype, device=ref.device)
|
||||
shard_axis = ref.uop.axis if axis is None else axis
|
||||
shape = tuple(s // len(ref.device) if i == shard_axis else s for i, s in enumerate(shape))
|
||||
axis = ref.uop.axis if axis is None else axis
|
||||
return Tensor(Tensor.empty(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)
|
||||
return Tensor(Tensor.invalid(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)
|
||||
|
||||
def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
|
||||
return _sharded_empty(ref.shape, ref, axis)
|
||||
|
||||
@@ -10,15 +10,18 @@ def add_to_ctx(ctx, x:UOp):
|
||||
ctx[0].append(x)
|
||||
return ret
|
||||
|
||||
pm_ctx = PatternMatcher([
|
||||
(UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx),
|
||||
(UPat((Ops.AFTER, Ops.CONTIGUOUS), name="x"),
|
||||
lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) and x.op_in_backward_slice_with_self(Ops.BUFFER) else None),
|
||||
pm_transform_unique_const = PatternMatcher([
|
||||
# transform unique consts to LUNIQUE
|
||||
(UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="x"),
|
||||
lambda ctx,x: x.replace(src=(UOp(Ops.LUNIQUE, arg=next(ctx[1])), x.src[1]))),
|
||||
])
|
||||
|
||||
pm_ctx = PatternMatcher([
|
||||
(UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx),
|
||||
(UPat((Ops.AFTER, Ops.CONTIGUOUS), name="x"),
|
||||
lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) and x.op_in_backward_slice_with_self(Ops.BUFFER) else None),
|
||||
])+pm_transform_unique_const
|
||||
|
||||
ReturnType = TypeVar('ReturnType')
|
||||
class _function(Generic[ReturnType]):
|
||||
def __init__(self, fxn:Callable[..., ReturnType], *, precompile:bool, precompile_backward:bool, allow_implicit:bool, grad_fxn:Callable|None):
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
from typing import cast
|
||||
import math, dataclasses
|
||||
from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata
|
||||
import math, dataclasses, itertools
|
||||
from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, graph_rewrite
|
||||
from tinygrad.helpers import argsort
|
||||
|
||||
def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
|
||||
@@ -37,6 +37,9 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
|
||||
grad_bodies = [(i, grads[p]) for i in needed if (p:=params.get(i)) is not None and p in grads]
|
||||
bwd_body = UOp.maketuple(*(gb for _, gb in grad_bodies)).substitute(fwd_subs, walk=True)
|
||||
bwd_body, compact_args = _compact_params(bwd_body, (*args, *grad_args, *fwd_outs))
|
||||
# TODO: is this okay here?
|
||||
from tinygrad.function import pm_transform_unique_const
|
||||
bwd_body = graph_rewrite(bwd_body, pm_transform_unique_const, ctx=(None, itertools.count(0)))
|
||||
bwd_call = bwd_body.call(*compact_args, name=(k.arg.name or "")+"_backward", precompile=k.arg.precompile_backward)
|
||||
gb_map = {i: idx for idx, (i, _) in enumerate(grad_bodies)}
|
||||
return (None,) + tuple(bwd_call.gettuple(gb_map[i]) if i in gb_map else None for i in range(len(args)))
|
||||
|
||||
Reference in New Issue
Block a user