diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 00f0a0e5ec..2c3ac71e68 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -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 diff --git a/examples/mlperf/dataloader.py b/examples/mlperf/dataloader.py index 121698c97b..c82c42410d 100644 --- a/examples/mlperf/dataloader.py +++ b/examples/mlperf/dataloader.py @@ -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([ diff --git a/test/test_schedule.py b/test/test_schedule.py index 87530de80e..59aba3ee74 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -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) diff --git a/test/unit/test_llm_tokenizer.py b/test/unit/test_llm_tokenizer.py index fca60bd7bb..7b65818a6f 100644 --- a/test/unit/test_llm_tokenizer.py +++ b/test/unit/test_llm_tokenizer.py @@ -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 }, { "": 5, "": 6, "": 7 }) diff --git a/tinygrad/schedule/grouper.py b/tinygrad/schedule/grouper.py index af62971b44..70682ad685 100644 --- a/tinygrad/schedule/grouper.py +++ b/tinygrad/schedule/grouper.py @@ -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