diff --git a/test/external/external_benchmark_bert_matmuls.py b/test/external/external_benchmark_bert_matmuls.py index 526b2745c3..b83e6d4809 100644 --- a/test/external/external_benchmark_bert_matmuls.py +++ b/test/external/external_benchmark_bert_matmuls.py @@ -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() diff --git a/test/external/external_benchmark_bert_softmax.py b/test/external/external_benchmark_bert_softmax.py index e17dd0b84f..176e2061a4 100644 --- a/test/external/external_benchmark_bert_softmax.py +++ b/test/external/external_benchmark_bert_softmax.py @@ -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()