github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

update benchmark bert scripts with BS and ACC_DTYPE (#9826)

Browse Source 
BS=16, ACC_DTYPE=half for tinybox, BS=128, ACC_DTYPE=float for mi300x
This commit is contained in:
chenyu
2025-04-10 02:06:02 -04:00
committed by GitHub
parent 498a2bf738
commit c462162db8
2 changed files with 10 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
test/external/external_benchmark_bert_matmuls.py vendored
View File

@@ -1,9 +1,12 @@
from tinygrad import Tensor, dtypes
dtypes.default_float = dtypes.float16
from tinygrad.dtype import to_dtype
from tinygrad.helpers import getenv
if __name__ == "__main__":
# matmuls in bert layers
BS = 96//6
BS = getenv("BS", 96//6)
acc_dtype = to_dtype(getenv("ACC_DTYPE", "half"))
tensors = [
(Tensor.empty(BS, 512, 1024), Tensor.empty(1024, 1024).T), # linear to get qkv
(Tensor.empty(BS, 512, 16, 64).permute(0,2,1,3), Tensor.empty(BS, 512, 16, 64).permute(0,2,3,1)), # q@k
@@ -12,4 +15,4 @@ if __name__ == "__main__":
for t0, t1 in tensors:
print(f"{t0.shape=}, {t0.lazydata.st.real_strides()=}, {t1.shape=}, {t1.lazydata.st.real_strides()=}")
for _ in range(5):
t0.dot(t1, dtype="half").realize()
t0.dot(t1, dtype=acc_dtype).realize()

7
test/external/external_benchmark_bert_softmax.py vendored
View File

@@ -1,15 +1,18 @@
from tinygrad import Tensor, dtypes, Context, GlobalCounters
dtypes.default_float = dtypes.float16
from tinygrad.dtype import to_dtype
from tinygrad.helpers import getenv
from test.test_softmax_fusion import single_kernel_softmax
if __name__ == "__main__":
# softmax in bert layers
BS = 96//6
BS = getenv("BS", 96//6)
acc_dtype = to_dtype(getenv("ACC_DTYPE", "half"))
t = Tensor.empty(BS, 16, 512, 512)
t.softmax(-1, dtype="half").realize()
# test single kernel softmax
GlobalCounters.reset()
with Context(DONT_GROUP_REDUCES=1):
single_kernel_softmax(t, -1, "half").realize()
single_kernel_softmax(t, -1, acc_dtype).realize()