diff --git a/.github/actions/setup-tinygrad/action.yml b/.github/actions/setup-tinygrad/action.yml index 2fcdeebb57..b75bdc0f6c 100644 --- a/.github/actions/setup-tinygrad/action.yml +++ b/.github/actions/setup-tinygrad/action.yml @@ -233,7 +233,7 @@ runs: shell: bash run: | sudo mkdir -p /usr/local/lib - curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/nimlgen/amdcomgr_dylib/releases/latest | \ + curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/tinygrad/amdcomgr_dylib/releases/latest | \ jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \ sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib cargo build --release --manifest-path ./extra/remu/Cargo.toml diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index 63da56c26d..76c9c9db09 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -1336,11 +1336,13 @@ def train_llama3(): # vocab_size from the mixtral tokenizer if not SMALL: model_params |= {"vocab_size": 32000} real_vocab_size = model_params['vocab_size'] - if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP) - vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers print(f"model parameters: {model_params}") + # pad vocab + if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP) + vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size + model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True) params = get_parameters(model) # weights are all bfloat16 for now @@ -1409,7 +1411,7 @@ def train_llama3(): device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)) tokens = tokens.shard(device) logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan) - loss = vocab_mask.where(-float("inf"), logits).sparse_categorical_crossentropy(tokens[:, 1:]) + loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:]) loss.backward() assert all(p.grad is g for p,g in zip(optim.params, grads)) Tensor.realize(loss, *grads) @@ -1439,13 +1441,15 @@ def train_llama3(): device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)) tokens = tokens.shard(device) logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan) - loss = vocab_mask.where(-float("inf"), logits).sparse_categorical_crossentropy(tokens[:, 1:]) + loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:]) return loss.flatten().float().to("CPU") # ** data iters ** def fake_data(bs, samples): + import numpy as np for _ in range(samples // bs): - yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT) + fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32) + yield Tensor(fake_data_np, device="NPY") def get_train_iter(): if getenv("FAKEDATA", 0): @@ -1550,7 +1554,7 @@ def train_llama3(): # run eval eval_losses = [] eval_iter = get_eval_iter() - tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences") + tqdm.write(f"evaluating {EVAL_SAMPLES//EVAL_BS} batches of {EVAL_BS} sequences") for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS): eval_losses += eval_step(tokens).tolist() diff --git a/examples/mlperf/optim.py b/examples/mlperf/optim.py index 7961bd8fb8..e76a1f8e73 100644 --- a/examples/mlperf/optim.py +++ b/examples/mlperf/optim.py @@ -12,7 +12,20 @@ class GradAccClipAdamW(Optimizer): self.v = self._new_optim_param() self.grad_acc, self.clip_norm = grad_acc, clip_norm + def fstep(self, grads:list[Tensor]): + if self.fused: + out, extra = self._step([], grads) + updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)] + else: + updates, extra = self._step([], grads) + for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i])) + to_realize = extra+self.params+self.buffers + + Tensor.realize(*to_realize) + def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]: + grads = list(grads) + for i in range(len(grads)): if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device) @@ -33,13 +46,13 @@ class GradAccClipAdamW(Optimizer): ret = [] self.b1_t *= self.b1 self.b2_t *= self.b2 - for i, (t, g) in enumerate(zip(params, grads)): + for i, g in enumerate(grads): self.m[i].assign((self.b1 * self.m[i] + (1.0 - self.b1) * g).cast(self.m[i].dtype)) self.v[i].assign((self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)).cast(self.v[i].dtype)) m_hat = self.m[i] / (1.0 - self.b1_t) v_hat = self.v[i] / (1.0 - self.b2_t) up = m_hat / (v_hat.sqrt() + self.eps) - ret.append((self.lr * up).cast(t.dtype)) + ret.append((self.lr * up).cast(g.dtype)) return ret, [self.b1_t, self.b2_t] + self.m + self.v def _apply_update(self, t:Tensor, up:Tensor) -> Tensor: diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh index c729d1b947..89c48e4d6d 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh @@ -5,6 +5,7 @@ export DEV=${DEV:-AMD} export EMULATE="AMD_CDNA4" export CHECK_OOB=0 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000 +export DEVICE_IN_FUNCTION_BUG=1 export DEBUG=${DEBUG:-0} export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1} diff --git a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh index dab82946a0..188aea51af 100755 --- a/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh +++ b/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh @@ -3,4 +3,4 @@ export BENCHMARK=5 export EVAL_BS=0 export VIZ=${VIZ:--1} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh -PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20 +extra/viz/cli.py --profile --device "AMD" --top 20 diff --git a/extra/torch_backend/test_kernel_fusion.py b/extra/torch_backend/test_kernel_fusion.py index 0546323366..dffcfe067f 100644 --- a/extra/torch_backend/test_kernel_fusion.py +++ b/extra/torch_backend/test_kernel_fusion.py @@ -1,7 +1,6 @@ # simple tests import unittest import torch -import warnings from tinygrad.helpers import getenv, GlobalCounters if getenv("TINY_BACKEND2"): import extra.torch_backend.backend2 @@ -18,9 +17,7 @@ class TestKernelFusionRegression(unittest.TestCase): torch.manual_seed(42) GlobalCounters.reset() fn().detach().cpu().numpy() - expectation = f"{GlobalCounters.kernel_count} vs {expected_kernels} expected." - if GlobalCounters.kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning) - self.assertLessEqual(GlobalCounters.kernel_count, expected_kernels, f"{expectation}") + self.assertEqual(GlobalCounters.kernel_count, expected_kernels) def test_elementwise_fusion(self): def fn(): @@ -34,7 +31,7 @@ class TestKernelFusionRegression(unittest.TestCase): conv = torch.nn.Conv2d(3, 16, 3, padding=1).to(device) with torch.no_grad(): return torch.nn.functional.relu(conv(x)) - self._check_kernel_count(fn, 8) + self._check_kernel_count(fn, 6) def test_batchnorm_fusion(self): def fn(): @@ -44,7 +41,7 @@ class TestKernelFusionRegression(unittest.TestCase): bn.eval() with torch.no_grad(): return torch.nn.functional.relu(bn(conv(x))) - self._check_kernel_count(fn, 16) + self._check_kernel_count(fn, 10) def test_reduce_fusion(self): def fn(): @@ -92,7 +89,7 @@ class TestKernelFusionRegression(unittest.TestCase): out = bn(conv(x)) out += identity return torch.nn.functional.relu(out) - self._check_kernel_count(fn, 17) + self._check_kernel_count(fn, 12) def test_multiple_inplace_ops_fusion(self): def fn(): @@ -117,7 +114,7 @@ class TestKernelFusionRegression(unittest.TestCase): bn.train() with torch.no_grad(): return bn(x) - self._check_kernel_count(fn, 10) + self._check_kernel_count(fn, 8) # this is a minimal extra/other_mnist/beautiful_mnist_torch.py to cover fusion for training with optimizer def test_mnist_training_fusion(self): @@ -138,7 +135,7 @@ class TestKernelFusionRegression(unittest.TestCase): loss.backward() optimizer.step() return loss - self._check_kernel_count(fn, 28) + self._check_kernel_count(fn, 24) if __name__ == "__main__": unittest.main() diff --git a/extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension/TinyGPUDriver.NV.Release.entitlements b/extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension/TinyGPUDriver.NV.Release.entitlements new file mode 100644 index 0000000000..792787d37f --- /dev/null +++ b/extra/usbgpu/tbgpu/installer/TinyGPUDriverExtension/TinyGPUDriver.NV.Release.entitlements @@ -0,0 +1,17 @@ + + + + + com.apple.application-identifier + 9YG3G8543N.org.tinygrad.tinygpu.edriver + com.apple.developer.driverkit + + com.apple.developer.driverkit.transport.pci + + + IOPCIPrimaryMatch + 0x000010de&0x0000FFFF + + + + diff --git a/extra/usbgpu/tbgpu/installer/build_and_sign_nv.sh b/extra/usbgpu/tbgpu/installer/build_and_sign_nv.sh new file mode 100755 index 0000000000..bbaed99df8 --- /dev/null +++ b/extra/usbgpu/tbgpu/installer/build_and_sign_nv.sh @@ -0,0 +1,33 @@ +#!/bin/bash +set -e + +xcodebuild clean build CODE_SIGN_IDENTITY="" CODE_SIGNING_REQUIRED=NO -alltargets -configuration Release build + +cp "../profiles/edriver_rel_2.provisionprofile" "./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext/embedded.provisionprofile" +cp "../profiles/installer_provisioning.provisionprofile" "./build/Release/TinyGPU.app/Contents/embedded.provisionprofile" + +codesign \ + --sign "Developer ID Application: tinygrad, Corp. (9YG3G8543N)" \ + --entitlements ./TinyGPUDriverExtension/TinyGPUDriver.NV.Release.entitlements \ + --verbose \ + --options runtime \ + --timestamp \ + --force \ + ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext + +codesign \ + --sign "Developer ID Application: tinygrad, Corp. (9YG3G8543N)" \ + --entitlements ./macOS/macOS.entitlements \ + --options runtime \ + --verbose \ + --timestamp \ + --force \ + ./build/Release/TinyGPU.app + +codesign --verify --deep --strict --verbose=4 ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext + +codesign --verify --deep --strict --verbose=4 ./build/Release/TinyGPU.app + +spctl -a -vv ./build/Release/TinyGPU.app + +spctl -a -vv ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext diff --git a/extra/viz/README b/extra/viz/README index b1c76bad66..a7db4f9f66 100644 --- a/extra/viz/README +++ b/extra/viz/README @@ -1,17 +1,17 @@ A command line tool for exploring the VIZ trace. -After running with VIZ=-1, use `PYTHONPATH=. extra/viz/cli.py` to explore the saved trace files. +After running with VIZ=-1, use `extra/viz/cli.py` to explore the saved trace files. ## Inspect runtime profiling -Use `PYTHONPATH=. extra/viz/cli.py --profile` to list all traced devices. +Use `extra/viz/cli.py --profile` to list all traced devices. List top slowest kernels on a device: `--profile --device "AMD"` List samples of a kernel on a device: `--profile --device "AMD" --kernel E_3` ## Inspect codegen and PatternMatcher -Use `PYTHONPATH=. extra/viz/cli.py --rewrites` to list all traced kernels. +Use `extra/viz/cli.py --rewrites` to list all traced kernels. List all codegen steps for a kernel: `--rewrites --kernel E_3` Get source code: `--rewrites --kernel E_3 --select "View Source"` diff --git a/extra/viz/cli.py b/extra/viz/cli.py index a8f942dede..e6dfe04d99 100755 --- a/extra/viz/cli.py +++ b/extra/viz/cli.py @@ -1,44 +1,73 @@ #!/usr/bin/env python3 import os os.environ["VIZ"] = "0" -import argparse, pathlib +import argparse, pathlib, sys, struct, json from typing import Iterator from tinygrad.viz import serve as viz from tinygrad.uop.ops import RewriteTrace from tinygrad.helpers import temp, ansistrip, colored, time_to_str, ansilen -from test.null.test_viz import load_profile + +# ** generic helpers def optional_eq(val:dict, arg:str|None) -> bool: return arg is None or ansistrip(val["name"]) == arg def print_data(data:dict) -> None: if isinstance(data.get("value"), Iterator): for m in data["value"]: - if m.get("uop"): - print("Input UOp:") - print(m["uop"]) - if not m["diff"]: continue - print("Rewrites:") - fp = pathlib.Path(m["upat"][0][0]) - print(f"{fp.parent.name}/{fp.name}:{m['upat'][0][1]}") - print(m["upat"][1]) - for line in m["diff"]: - color = "red" if line.startswith("-") else "green" if line.startswith("+") else None - print(colored(line, color)) + if m.get("uop"): print(f"Input UOp:\n{m['uop']}") + if m.get("diff"): + loc = pathlib.Path(m["upat"][0][0]) + print(f"Rewrite at {loc.parent.name}/{loc.name}:{m['upat'][0][1]}\n{m['upat'][1]}") + for line in m["diff"]: print(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None)) if data.get("src") is not None: print(data["src"]) +# ** Profiler trace decoder + +# 0 means None, otherwise it's an enum value +def option(i:int) -> int|None: return None if i == 0 else i-1 + +def decode_profile(data:bytes) -> dict: + ret, off = data, 0 + def u(fmt:str) -> tuple: + nonlocal off + vals = struct.unpack_from(fmt, ret, off) + off += struct.calcsize(fmt) + return vals + total_dur, global_peak, index_len, layout_len = u(" 0.01 else None) if et is not None else "" name = e["name"]+(" " * (46 - ansilen(e["name"]))) print(f"{name} {ptm}/{(et or 0)*1e3:9.2f}ms "+e['fmt'].replace('\n', ' | ')+" ") @@ -81,7 +110,7 @@ if __name__ == "__main__": other_t = total-sum(t for _, (t, _) in sel) table.append([f"Other ({len(other)} unique)", time_to_str(other_t, w=9), sum(c for _,(_,c) in other), f"{other_t/total*100.0:.2f}%"]) print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github")) - exit(0) + sys.exit(0) for k in viz.ctxs: if not optional_eq(k, args.kernel): continue diff --git a/test/backend/test_multitensor.py b/test/backend/test_multitensor.py index ebb1bd9660..7c9c423900 100644 --- a/test/backend/test_multitensor.py +++ b/test/backend/test_multitensor.py @@ -228,17 +228,17 @@ class TestMultiTensor(unittest.TestCase): a,b = _test_allreduce(Tensor.rand(256, 256)) np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5) - def test_multiple_to_single_device_naive(self): - with Context(RING=0): - t = Tensor.arange(32).shard(devices_4, 0).to(Device.DEFAULT).realize() - self.assertEqual(t.device, Device.DEFAULT) - np.testing.assert_equal(t.numpy(), np.arange(32)) - - def test_multiple_to_single_device_ring(self): - with Context(RING=2): - t = Tensor.arange(32).shard(devices_4, 0).to(Device.DEFAULT).realize() - self.assertEqual(t.device, Device.DEFAULT) - np.testing.assert_equal(t.numpy(), np.arange(32)) + def test_multiple_to_single_device(self): + kernel_counts = {} + for ring in (0, 2): + GlobalCounters.reset() + with Context(RING=ring, SCACHE=0): + t = Tensor.arange(32).contiguous().shard(devices_4, 0).to(Device.DEFAULT) + t.realize() + kernel_counts[ring] = GlobalCounters.kernel_count + self.assertEqual(t.device, Device.DEFAULT) + np.testing.assert_equal(t.numpy(), np.arange(32)) + self.assertNotEqual(kernel_counts[0], kernel_counts[2]) def test_allreduce_all2all(self): with Context(ALL2ALL=2): diff --git a/test/backend/test_schedule.py b/test/backend/test_schedule.py index 260a2b33d2..6ad348fe66 100644 --- a/test/backend/test_schedule.py +++ b/test/backend/test_schedule.py @@ -795,6 +795,38 @@ class TestSchedule(unittest.TestCase): self.assertIsNotNone(out.uop.base.realized) self.assertIsInstance(out.uop.base.realized.dtype, ImageDType) + @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL") + @unittest.expectedFailure + def test_image_dot_f16_fusion(self): + with Context(FLOAT16=1): + def cnt(): + x, y, z = Tensor.empty((64, 64), dtype='float'), Tensor.empty((64, 64), dtype='float'), Tensor.empty((64, 64), dtype='float') + a = (x @ y).relu() + sched = ((a @ z).relu() + a).schedule() + for si in sched: si.lower() + return len([si for si in sched if isinstance(si.prg, CompiledRunner)]) + + with Context(IMAGE=1): cnt1 = cnt() + with Context(IMAGE=2): cnt2 = cnt() + + self.assertEqual(cnt1, cnt2) + + @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL") + @unittest.expectedFailure + def test_image_conv_fusion(self): + def cnt(): + x, y, z = Tensor.empty((1, 4, 3, 3)), Tensor.empty((4, 1, 3, 3)), Tensor.empty((4, 1, 7, 7)) + a = x.conv2d(y, Tensor.empty(4), groups=4, padding=1) + b = a.conv2d(z, groups=4, padding=3) + sched = (a + b).schedule() + for si in sched: si.lower() + return len([si for si in sched if isinstance(si.prg, CompiledRunner)]) + + with Context(IMAGE=1): cnt1 = cnt() + with Context(IMAGE=2): cnt2 = cnt() + + self.assertEqual(cnt1, cnt2) + def _test_fusion(self, shapes, f, cnt): with Context(DEBUG=0, TRACK_MATCH_STATS=0): args = [Tensor.randn(s).realize() for s in shapes] run_schedule(check_schedule(compare:=f(*args), cnt)) diff --git a/test/backend/test_symbolic_ops.py b/test/backend/test_symbolic_ops.py index 56f02f297e..856b5e191a 100644 --- a/test/backend/test_symbolic_ops.py +++ b/test/backend/test_symbolic_ops.py @@ -69,6 +69,58 @@ class TestSymbolicOps(unittest.TestCase): # symbolic shape dropout is not supported self.test_attention(dropout_p=0.5) + def test_sdpa_symbolic_seq_len(self): + # symbolic seq_len on all of q/k/v (dim -2 after transpose) + q = Tensor.rand(2, 10, 4, 8) + k = Tensor.rand(2, 10, 4, 8) + v = Tensor.rand(2, 10, 4, 8) + for i in range(1, 5): + vi = Variable("i", 1, 10).bind(i) + Tensor.realize(q, k, v) + symbolic = q[:, :vi].transpose(1, 2).scaled_dot_product_attention( + k[:, :vi].transpose(1, 2), v[:, :vi].transpose(1, 2)).realize()[:2, :4, :i, :8].numpy() + expected = q[:, :i].transpose(1, 2).scaled_dot_product_attention( + k[:, :i].transpose(1, 2), v[:, :i].transpose(1, 2)).realize().numpy() + np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) + + def test_sdpa_symbolic_seq_len_query_only(self): + # symbolic seq_len on query only (dim -2 after transpose) + q = Tensor.rand(2, 10, 4, 8) + k = Tensor.rand(2, 5, 4, 8) + v = Tensor.rand(2, 5, 4, 8) + for i in range(1, 5): + vi = Variable("i", 1, 10).bind(i) + Tensor.realize(q, k, v) + symbolic = q[:, :vi].transpose(1, 2).scaled_dot_product_attention( + k.transpose(1, 2), v.transpose(1, 2)).realize()[:2, :4, :i, :8].numpy() + expected = q[:, :i].transpose(1, 2).scaled_dot_product_attention( + k.transpose(1, 2), v.transpose(1, 2)).realize().numpy() + np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) + + def test_sdpa_symbolic_batch(self): + # symbolic batch dim (dim 0) + q = Tensor.rand(10, 4, 3, 8) + k = Tensor.rand(10, 4, 3, 8) + v = Tensor.rand(10, 4, 3, 8) + for i in range(1, 5): + vi = Variable("i", 1, 10).bind(i) + Tensor.realize(q, k, v) + symbolic = q[:vi].scaled_dot_product_attention(k[:vi], v[:vi]).realize()[:i, :4, :3, :8].numpy() + expected = q[:i].scaled_dot_product_attention(k[:i], v[:i]).realize().numpy() + np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) + + def test_sdpa_symbolic_heads(self): + # symbolic heads dim (dim -3) + q = Tensor.rand(2, 10, 3, 8) + k = Tensor.rand(2, 10, 3, 8) + v = Tensor.rand(2, 10, 3, 8) + for i in range(1, 5): + vi = Variable("i", 1, 10).bind(i) + Tensor.realize(q, k, v) + symbolic = q[:, :vi].scaled_dot_product_attention(k[:, :vi], v[:, :vi]).realize()[:2, :i, :3, :8].numpy() + expected = q[:, :i].scaled_dot_product_attention(k[:, :i], v[:, :i]).realize().numpy() + np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6) + def test_attention_pos_0_sz_0(self): Attention(128, 8)(Tensor.ones(1, 0, 128), Variable("start_pos", 0, 128).bind(0), None) diff --git a/test/backend/test_tensor_variable.py b/test/backend/test_tensor_variable.py index b05529c71c..9e9d26520b 100644 --- a/test/backend/test_tensor_variable.py +++ b/test/backend/test_tensor_variable.py @@ -136,6 +136,30 @@ class TestTensorVariable(unittest.TestCase): with self.assertRaises(AssertionError): t.chunk(2, dim=0) # can't split along symbolic dim + def test_symbolic_var_sum(self, var_name="u"): + t = Variable("t", 1, 10).bind(4) + v = Variable(var_name, 1, 5).bind(1) + mask = (Tensor.full((1, 1, t, v+t), 1) + 1).contiguous() + mask.shrink(((0, 1), (0, 1), (0, 4), (0, 4))).numpy() + def test_symbolic_var_sum_alt_name(self): self.test_symbolic_var_sum("s") + + def test_symbolic_triu(self): + t = Variable("t", 1, 10).bind(4) + for start_pos in (0, 1, 3): + var_start_pos = Variable("start_pos", 0, 5).bind(start_pos) + mask = Tensor.full((1, 1, t, var_start_pos+t), float("-inf")).triu(var_start_pos+1) + out = mask.shrink(((0, 1), (0, 1), (0, 4), (0, start_pos+4))).numpy() + expected = np.triu(np.full((1, 1, 4, start_pos+4), float("-inf")), k=start_pos+1) + np.testing.assert_equal(out, expected) + + def test_symbolic_tril(self): + t = Variable("t", 1, 10).bind(4) + for start_pos in (0, 1, 3): + var_start_pos = Variable("start_pos", 0, 5).bind(start_pos) + mask = Tensor.full((1, 1, t, var_start_pos+t), float("-inf")).tril(var_start_pos+1) + out = mask.shrink(((0, 1), (0, 1), (0, 4), (0, start_pos+4))).numpy() + expected = np.tril(np.full((1, 1, 4, start_pos+4), float("-inf")), k=start_pos+1) + np.testing.assert_equal(out, expected) if __name__ == '__main__': unittest.main() diff --git a/test/external/external_fuzz_beam_timeout_recovery.py b/test/external/external_fuzz_beam_timeout_recovery.py new file mode 100644 index 0000000000..73b67c8551 --- /dev/null +++ b/test/external/external_fuzz_beam_timeout_recovery.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python3 +""" +Stress test for beam timeout + device recovery on AM devices. + +Usage: + AMD=1 python test/external/external_test_beam_timeout_recovery.py +""" +from tinygrad import Tensor, Device +from tinygrad.helpers import Context +from tinygrad.runtime.ops_amd import AMDDevice + +if __name__ == "__main__": + dev = Device["AMD"] + assert isinstance(dev, AMDDevice) and dev.is_am(), "not am" + + N = 10000 + for i in range(N): + with Context(DEBUG=0, BEAM=0): + a = Tensor.rand(4096, 4096, device="AMD").contiguous().realize() + b = Tensor.rand(4096, 4096, device="AMD").contiguous().realize() + c = a.matmul(b) + c.realize() + try: dev.synchronize(timeout=1) + except RuntimeError as e: print(e) + with Context(DEBUG=0, BEAM=0): + a = Tensor.ones(512, 512, device="AMD").contiguous().realize() + b = Tensor.ones(512, 512, device="AMD").contiguous().realize() + result = a.matmul(b).realize()[0, 0].item() + assert result == 512.0, f"iter {i}: got {result}" + print(f" iter {i+1}/{N}: ok") + print(f"=== All {N} iterations passed ===") diff --git a/test/external/external_test_gpu_crash.py b/test/external/external_test_gpu_crash.py index f8df34bc96..143e700954 100644 --- a/test/external/external_test_gpu_crash.py +++ b/test/external/external_test_gpu_crash.py @@ -4,13 +4,19 @@ These tests intentionally cause GPU faults to verify error handling. Run with: AMD=1 python -m pytest test/external/external_test_gpu_crash.py -v """ -import unittest, re +import unittest, re, importlib from tinygrad.device import Device -from tinygrad.runtime.autogen.amd.rdna3.ins import * # noqa: F403 from tinygrad.renderer.amd.dsl import s, v, Inst, NULL -def assemble(code:str, name:str="test") -> str: - kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "wavefront_size32": 1, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8} +RDNA3_CDNA3_MAP = {"v_mov_b32_e32": "v_mov_b32_e32", "s_mov_b32": "s_mov_b32", "s_waitcnt": "s_waitcnt", "s_endpgm": "s_endpgm", + "global_load_b32": "global_load_dword", "global_store_b32": "global_store_dword", + "global_atomic_add_u32": "global_atomic_add", "flat_load_b32": "flat_load_dword", + "flat_store_b32": "flat_store_dword", "flat_atomic_add_u32": "flat_atomic_add", "s_load_b32": "s_load_dword"} + +def assemble(code:str, name:str="test", is_cdna:bool=False) -> str: + kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8} + if is_cdna: kd["accum_offset"] = 8 + else: kd["wavefront_size32"] = 1 return f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{code}\n.rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + \ "\n".join(f".amdhsa_{k} {v}" for k,v in kd.items()) + "\n.end_amdhsa_kernel" @@ -21,6 +27,10 @@ class TestGPUCrash(unittest.TestCase): from tinygrad.runtime.support.compiler_amd import HIPCompiler cls.dev = Device["AMD"] cls.compiler = HIPCompiler(cls.dev.arch) + cls.is_cdna = cls.dev.target[0] < 10 + ins = importlib.import_module('tinygrad.runtime.autogen.amd.' + ('cdna' if cls.is_cdna else 'rdna3') + '.ins') + for rdna3_name, cdna3_name in RDNA3_CDNA3_MAP.items(): + setattr(cls, rdna3_name, getattr(ins, cdna3_name if cls.is_cdna else rdna3_name)) def setUp(self): # Verify device works before each test @@ -33,7 +43,7 @@ class TestGPUCrash(unittest.TestCase): def _run(self, code: str): from tinygrad.runtime.ops_amd import AMDProgram - prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code))) + prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code, is_cdna=self.is_cdna))) prg(self.dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True) def _run_insts(self, insts: list[Inst]): @@ -57,32 +67,32 @@ class TestOutOfBoundsMemoryAccess(TestGPUCrash): def test_global_load_null_ptr(self): """Global load from NULL pointer.""" - insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), - global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()] + insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), + self.global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()] self._assert_gpu_fault(lambda: self._run_insts(insts)) def test_global_store_null_ptr(self): """Global store to NULL pointer.""" - insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF), - global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()] + insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 0xDEADBEEF), + self.global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()] self._assert_gpu_fault(lambda: self._run_insts(insts)) def test_global_load_unmapped_high_address(self): """Global load from high unmapped address (0xDEAD00000000).""" - insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD), - global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()] + insts = [self.v_mov_b32_e32(v[0], 0x00000000), self.v_mov_b32_e32(v[1], 0xDEAD), + self.global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()] self._assert_gpu_fault(lambda: self._run_insts(insts)) def test_global_store_unmapped_high_address(self): """Global store to high unmapped address.""" - insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 0x12345678), - global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()] + insts = [self.v_mov_b32_e32(v[0], 0x00000000), self.v_mov_b32_e32(v[1], 0xDEAD), self.v_mov_b32_e32(v[2], 0x12345678), + self.global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()] self._assert_gpu_fault(lambda: self._run_insts(insts)) def test_global_atomic_unmapped(self): """Atomic operation on unmapped memory.""" - insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 1), - global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()] + insts = [self.v_mov_b32_e32(v[0], 0xBEEF0000), self.v_mov_b32_e32(v[1], 0xDEAD), self.v_mov_b32_e32(v[2], 1), + self.global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()] self._assert_gpu_fault(lambda: self._run_insts(insts)) @@ -91,14 +101,14 @@ class TestSMEMFaults(TestGPUCrash): def test_smem_load_null(self): """SMEM load from NULL base.""" - insts = [s_mov_b32(s[2], 0), s_mov_b32(s[3], 0), - s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()] + insts = [self.s_mov_b32(s[2], 0), self.s_mov_b32(s[3], 0), + self.s_load_b32(s[4], s[2:3], 0, soffset=NULL), self.s_waitcnt(0), self.s_endpgm()] self._assert_gpu_fault(lambda: self._run_insts(insts)) def test_smem_load_unmapped(self): """SMEM load from unmapped address.""" - insts = [s_mov_b32(s[2], 0xBEEF0000), s_mov_b32(s[3], 0xDEAD), - s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()] + insts = [self.s_mov_b32(s[2], 0xBEEF0000), self.s_mov_b32(s[3], 0xDEAD), + self.s_load_b32(s[4], s[2:3], 0, soffset=NULL), self.s_waitcnt(0), self.s_endpgm()] self._assert_gpu_fault(lambda: self._run_insts(insts)) @@ -107,20 +117,20 @@ class TestFlatMemoryFaults(TestGPUCrash): def test_flat_load_null(self): """FLAT load from NULL address.""" - insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), - flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()] + insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), + self.flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()] self._assert_gpu_fault(lambda: self._run_insts(insts)) def test_flat_store_null(self): """FLAT store to NULL address.""" - insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF), - flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()] + insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 0xDEADBEEF), + self.flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()] self._assert_gpu_fault(lambda: self._run_insts(insts)) def test_flat_atomic_null(self): """FLAT atomic on NULL address.""" - insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 1), - flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()] + insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 1), + self.flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()] self._assert_gpu_fault(lambda: self._run_insts(insts)) diff --git a/test/external/external_test_speed_llama.py b/test/external/external_test_speed_llama.py index 7d50fd51af..1234113e77 100644 --- a/test/external/external_test_speed_llama.py +++ b/test/external/external_test_speed_llama.py @@ -10,7 +10,7 @@ from tinygrad.helpers import Profiling class FakeProgram: def __init__(self, name:str, prg:bytes, **kwargs): pass - def __call__(self, *bufs, global_size, local_size, vals=(), wait=False): pass + def __call__(self, *bufs, global_size, local_size, vals=(), wait=False, **kw): pass class FakeAllocator(Allocator[Compiled]): def _alloc(self, sz, options): return None diff --git a/test/models/test_onnx.py b/test/models/test_onnx.py index 691e80de0c..c036ecf51b 100644 --- a/test/models/test_onnx.py +++ b/test/models/test_onnx.py @@ -87,7 +87,7 @@ class TestHuggingFaceOnnxModels(unittest.TestCase): "input_ids": np.random.randint(0, 250002, (1, 11), dtype=np.int64), "attention_mask": np.ones((1, 11), dtype=np.int64), } - self._validate(repo_id, model_file, custom_inputs) + self._validate(repo_id, model_file, custom_inputs, atol=1e-3) if __name__ == "__main__": unittest.main() diff --git a/test/null/test_schedule.py b/test/null/test_schedule.py index 1a14c1d646..f5f63e19da 100644 --- a/test/null/test_schedule.py +++ b/test/null/test_schedule.py @@ -234,6 +234,18 @@ class TestSchedule(unittest.TestCase): d = Tensor.empty(1).assign(c) check_schedule(d, 1) + def test_detach_assign(self): + a = Tensor.ones(4, 4).contiguous().realize() + buf1, buf2 = Tensor.empty(4, 4).contiguous(), Tensor.empty(4, 4).contiguous() + r = buf2.assign(buf1.assign(a + 1.0) * 2.0) + check_schedule(r.detach().contiguous(), 2) + + def test_contiguous_backward_assign(self): + a = Tensor.ones(4, 4).contiguous().realize() + buf1, buf2 = Tensor.empty(4, 4).contiguous(), Tensor.empty(4, 4).contiguous() + r = buf2.assign(buf1.assign(a + 1.0) * 2.0) + check_schedule(r.contiguous_backward().contiguous(), 2) + def test_mulacc_relu_fusion(self): a = Tensor.empty(10) b = Tensor.empty(10) diff --git a/test/null/test_viz.py b/test/null/test_viz.py index 8d8c74e6d9..b564d2eaae 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -1,4 +1,4 @@ -import unittest, decimal, json, struct, sys +import unittest, decimal, sys from dataclasses import dataclass from typing import Generator @@ -357,41 +357,9 @@ class TestVizIntegration(BaseTestViz): from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry from tinygrad.viz.serve import get_profile +from extra.viz.cli import decode_profile -class TinyUnpacker: - def __init__(self, buf): self.buf, self.offset = buf, 0 - def __call__(self, fmt:str) -> tuple: - ret = struct.unpack_from(fmt, self.buf, self.offset) - self.offset += struct.calcsize(fmt) - return ret - -# 0 means None, otherwise it's an enum value -def option(i:int) -> int|None: return None if i == 0 else i-1 - -def load_profile(lst:list[ProfileEvent]) -> dict: - ret = get_profile(lst) - u = TinyUnpacker(ret) - total_dur, global_peak, index_len, layout_len = u(" dict: return decode_profile(get_profile(lst)) class TestVizProfiler(BaseTestViz): def test_transfer_uses_copy_device(self): diff --git a/test/unit/test_function.py b/test/unit/test_function.py index 2a3fa6a1e2..216c5763fc 100644 --- a/test/unit/test_function.py +++ b/test/unit/test_function.py @@ -70,6 +70,14 @@ class TestFunction(unittest.TestCase): b = Tensor([4,5,6]) np.testing.assert_equal(f(a, b).numpy(), [5,7,9]) + def test_contiguous_backward(self): + @function + def f(a:Tensor, b:Tensor) -> Tensor: return (a + b).contiguous_backward() + + a = Tensor([1,2,3]) + b = Tensor([4,5,6]) + np.testing.assert_equal(f(a, b).numpy(), [5,7,9]) + def test_method(self): class Foo: def __init__(self): self.w = Tensor([10,20,30]) @@ -129,6 +137,31 @@ class TestFunction(unittest.TestCase): a = Tensor([1., 2., 3.]) np.testing.assert_allclose(g(f(a)).numpy(), [110., 440., 990.]) + def test_nested_calls_backward(self): + w = Tensor([[1., 2.], [3., 4.]]).contiguous().realize() + @function + def inner(x:Tensor) -> Tensor: return x + w + @function + def outer(a:Tensor, b:Tensor) -> Tensor: return inner(a.reshape(1,2) + b.reshape(1,2)) + + a = Tensor([1., 2.], requires_grad=True) + b = Tensor([3., 4.], requires_grad=True) + outer(a, b).sum().backward() + np.testing.assert_allclose(a.grad.numpy(), [2., 2.]) + np.testing.assert_allclose(b.grad.numpy(), [2., 2.]) + + def test_unused_param_backward(self): + @function + def f(a:Tensor, b:Tensor, c:Tensor) -> Tensor: return a + c # b is unused + + a = Tensor([1., 2., 3.], requires_grad=True) + b = Tensor([4., 5., 6.], requires_grad=True) + c = Tensor([7., 8., 9.], requires_grad=True) + f(a, b, c).sum().backward() + np.testing.assert_allclose(a.grad.numpy(), [1., 1., 1.]) + np.testing.assert_allclose(b.grad.numpy(), [0., 0., 0.]) + np.testing.assert_allclose(c.grad.numpy(), [1., 1., 1.]) + def test_name(self): @function def f(a:Tensor) -> Tensor: return a + 1 @@ -230,5 +263,77 @@ class TestFunctionMulti(unittest.TestCase): f(x).sum().backward() np.testing.assert_allclose(w.grad.numpy(), [4., 5., 6., 7.]) + def test_call_axis(self): + @function + def f(x:Tensor, w:Tensor) -> Tensor: return x @ w + + x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]]).shard(self.devices_2, axis=0) + w = Tensor([[1.,2.],[3.,4.]]).shard(self.devices_2, axis=None) + result = f(x, w) + # CALL output should inherit axis=0 from the sharded input + self.assertEqual(result.uop.axis, 0) + # reduce on the sharded axis should remove it + self.assertIsNone(result.sum().uop.axis) + + def test_call_axis_shard_inside(self): + @function + def f(x:Tensor, w:Tensor) -> Tensor: + return x.shard(self.devices_2, axis=0) @ w.shard(self.devices_2, axis=None) + + x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]]) + w = Tensor([[1.,2.],[3.,4.]]) + result = f(x, w) + self.assertEqual(result.uop.axis, 0) + np.testing.assert_allclose(result.numpy(), x.numpy() @ w.numpy()) + + def test_data_parallel_backward(self): + @function + def f(x:Tensor, w:Tensor) -> Tensor: return x @ w + + x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]], requires_grad=True).shard(self.devices_2, axis=0) + w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(self.devices_2, axis=None) + w.realize() + f(x, w).sum().backward() + # d/dx = ones @ w^T = [[1,3],[1,3],[1,3],[1,3]], but sum so ones(4,2) @ w^T? no: + # L = sum(x @ w), dL/dx = ones(4,2) @ w^T... actually dL/d(xw) = ones(4,2), dL/dx = ones(4,2) @ w^T + np.testing.assert_allclose(x.grad.numpy(), np.ones((4,2)) @ np.array([[1,3],[2,4]])) + + def test_data_parallel_backward_4(self): + devices_4 = tuple(f"CPU:{i}" for i in range(4)) + @function + def f(x:Tensor, w:Tensor) -> Tensor: return x @ w + + x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0) + w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None) + w.realize() + f(x, w).sum().backward() + np.testing.assert_allclose(x.grad.numpy(), np.ones((8,2)) @ np.array([[1,3],[2,4]])) + + def test_data_parallel_backward_implicit(self): + devices_4 = tuple(f"CPU:{i}" for i in range(4)) + w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None) + w.realize() + @function + def f(x:Tensor) -> Tensor: return x @ w + + x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0) + f(x).sum().backward() + np.testing.assert_allclose(x.grad.numpy(), np.ones((8,2)) @ np.array([[1,3],[2,4]])) + + def test_data_parallel_backward_twice(self): + devices_4 = tuple(f"CPU:{i}" for i in range(4)) + w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None) + w.realize() + # pre-init grads like the training loop does + w.grad = w.zeros_like().contiguous().realize() + @function + def f(x:Tensor) -> Tensor: return x @ w + + expected = np.ones((8,2)) @ np.array([[1,3],[2,4]]) + for _ in range(2): + x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0) + f(x).sum().backward() + np.testing.assert_allclose(x.grad.numpy(), expected) + if __name__ == '__main__': unittest.main() diff --git a/test/unit/test_linalg.py b/test/unit/test_linalg.py index 9bdff0b5cf..2e97ad6c9b 100644 --- a/test/unit/test_linalg.py +++ b/test/unit/test_linalg.py @@ -83,6 +83,15 @@ class TestLinAlg(unittest.TestCase): s_diag = (S.unsqueeze(-2) * Tensor.eye(2)) reconstruction_helper([U, s_diag, V], a) + def test_svd_identity_4x4(self): + a = Tensor.eye(4) + U,S,V = a.svd() + assert not np.isnan(U.numpy()).any() + assert not np.isnan(S.numpy()).any() + assert not np.isnan(V.numpy()).any() + s_diag = (S.unsqueeze(-2) * Tensor.eye(4)) + reconstruction_helper([U, s_diag, V], a) + def test_svd_rank1(self): a = Tensor([[1.0, 1.0], [2.0, 2.0]]).realize() U, S, V = a.svd() diff --git a/tinygrad/apps/llm.py b/tinygrad/apps/llm.py index d5f621d8de..a2fb7342e3 100644 --- a/tinygrad/apps/llm.py +++ b/tinygrad/apps/llm.py @@ -143,7 +143,8 @@ class TransformerBlock: #v = self.cache_kv[1, :, :, 0:start_pos+T, :] # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True - mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(int(start_pos)+1) if T > 1 else None + # TODO: this if statement should be removed and it shouldn't generate extra kernels + mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(start_pos+1) if T > 1 else None attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True) # (B,H,T,Hd) attn = attn.transpose(1, 2).reshape(B, T, -1) # back to (B,T,D) attn = self.attn_output(attn) @@ -263,7 +264,7 @@ CHAT_HTML = b'''tinygrad chat