nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

Merge branch 'master' into simpler_fusion

Browse Source
This commit is contained in:
George Hotz
2025-08-07 13:46:42 -07:00
committed by GitHub
parent 174747efb3 6ed2dfd187
commit 3969d8574d
5 changed files with 53 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/benchmark.yml
+3 -2
View File
@@ -692,8 +692,9 @@ jobs:
run: AMD=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
- name: Run full CIFAR training w 1 GPU
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
# TODO: enable
# - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
# run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
- name: Run 10 MLPerf Bert training steps (1 gpu)
# TODO: remove BERT_LAYERS once scheduler is fast
run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
examples/mlperf/dataloader.py
+41 -1
View File
@@ -564,10 +564,13 @@ class GPTDataset:
# check for cache
cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
cache_path = base_path.with_name(f"{base_path.name}.{cache_hash}.index_cache")
print(f"try loading GPTDataset from {cache_path}...")
if cache_path.exists():
print("cache found, loading...")
with open(cache_path, "rb") as f:
self.doc_idx, self.sample_idx, self.shuffle_idx = pickle.load(f)
else:
print("cache not found, building index...")
self.doc_idx = self._build_doc_idx()
self.sample_idx = self._build_sample_idx()
self.shuffle_idx = self._build_shuffle_idx()
@@ -692,10 +695,47 @@ class BlendedGPTDataset:
self.datasets = [GPTDataset(path, samples_per_blend[i], seqlen, seed + i, shuffle) for i,path in enumerate(paths)]
# check for cache
cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
cache_path = paths[0].with_name(f"{paths[0].name}.{cache_hash}.blend_cache")
print(f"try loading BlendedGPTDataset from {cache_path}...")
if cache_path.exists():
print("cache found, loading...")
with open(cache_path, "rb") as f:
self.dataset_idx, self.dataset_sample_idx = pickle.load(f)
else:
print("cache not found, building index...")
self.dataset_idx, self.dataset_sample_idx = self._build_blend_idx()
# save cache
with open(cache_path, "wb") as f:
pickle.dump((self.dataset_idx, self.dataset_sample_idx), f)
def get(self, idx:int):
tokens = self.datasets[0][idx]
tokens = self.datasets[self.dataset_idx[idx]][self.dataset_sample_idx[idx]]
return tokens
def _build_blend_idx(self):
dataset_idx = np.zeros(self.samples, dtype=np.int16)
dataset_sample_idx = np.zeros(self.samples, dtype=np.int64)
unspent_datasets = set(range(len(self.datasets)))
dataset_sample_counts = [0] * len(self.datasets)
for i in tqdm(range(self.samples)):
error_argmax, error_max = 0, 0.0
for di in unspent_datasets:
error = self.weights[di] * max(i, 1) - dataset_sample_counts[di]
if error > error_max:
error_max = error
error_argmax = di
dataset_idx[i] = error_argmax
dataset_sample_idx[i] = dataset_sample_counts[error_argmax]
dataset_sample_counts[error_argmax] += 1
return dataset_idx, dataset_sample_idx
def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
if val:
dataset = BlendedGPTDataset([
test/test_schedule.py
+5 -4
View File
@@ -1727,7 +1727,8 @@ class TestIndexing(unittest.TestCase):
s = Tensor.schedule(*lst)
lowered = [x[1] for x in lower_schedule(s.copy())]
kernels = [ei for ei in list(lowered) if isinstance(ei.prg, CompiledRunner)]
if FUSE_ARANGE: self.assertEqual(len(kernels), cnt)
if FUSE_ARANGE and len(kernels) != cnt:
raise KernelCountException(f"{len(kernels)} != {cnt}")
for ei in lowered: ei.run(do_update_stats=True)
return s
@@ -1741,7 +1742,7 @@ class TestIndexing(unittest.TestCase):
def test_simple_indexing_alt(self):
X = Tensor.arange(16).reshape(4, 4)
xt = X[[1, 2], [1, 2]]
self.check_schedule(xt, 5)
self.check_schedule(xt, 3)
np.testing.assert_equal(xt.numpy(), (np.arange(16).reshape(4, 4))[[1, 2], [1, 2]])
def test_advanced_indexing(self):
@@ -1753,13 +1754,13 @@ class TestIndexing(unittest.TestCase):
def test_advanced_indexing_alt(self):
X = Tensor.arange(6).reshape(3, 2)+1
xt = X[[Tensor([2]), Tensor([1])]]
self.check_schedule(xt, 6)
self.check_schedule(xt, 3)
np.testing.assert_equal(xt.numpy(), 6)
def test_advanced_simple_indexing_combined(self):
X = Tensor.arange(16).reshape(4, 4)
xt = X[1:2, [1, 2]]
self.check_schedule(xt, 4)
self.check_schedule(xt, 2)
def test_push_through_reshape(self):
Tensor.manual_seed(0)
test/unit/test_llm_tokenizer.py
+2 -1
View File
@@ -1,7 +1,8 @@
import unittest, base64, functools
import unittest, base64, functools, sys
from tinygrad.apps.llm import SimpleTokenizer, get_llama_re
from tinygrad.helpers import fetch
@unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows")
class TestLLMTokenizer(unittest.TestCase):
@functools.cached_property
def basic_tok(self): return SimpleTokenizer(".*", { b"a": 0, b"b": 1, b"c": 2, b"ab": 3, b"bc": 4 }, { "<x>": 5, "<y>": 6, "<z>": 7 })
tinygrad/schedule/grouper.py
+2 -2
View File
@@ -26,8 +26,8 @@ def realize_before_view(ctx:dict[UOp, None], view:UOp, tr:UOp) -> None:
do_realize = PatternMatcher([
# always realize SINK parents
(UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
# always realize ASSIGN/CONTIGUOUS/GroupOp.Meta
(UPat({Ops.ASSIGN, Ops.CONTIGUOUS, *GroupOp.Meta}, name="tr"), realize),
# always realize ASSIGN/CONTIGUOUS/COPY/BUFFER_VIEW
(UPat({Ops.ASSIGN, Ops.CONTIGUOUS, Ops.COPY, Ops.BUFFER_VIEW}, name="tr"), realize),
# realize before expand or unsafe pad ops
(UPat(Ops.VIEW, src=(UPat(GroupOp.All-ALWAYS_CONTIGUOUS, name="tr"),), name="view"), realize_before_view),
# realize parents of COPY, MSELECT, MSTACK