From 876a8be97a04d8b030c877c940893569e1dd739e Mon Sep 17 00:00:00 2001 From: Ignacio Sica Date: Wed, 2 Apr 2025 14:34:07 +0800 Subject: [PATCH 1/3] Debug env var breakdown (#9663) * add debug level breakdown * hotfix * Update env_vars.md --- docs/env_vars.md | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/docs/env_vars.md b/docs/env_vars.md index 2ac4066141..b23b435675 100644 --- a/docs/env_vars.md +++ b/docs/env_vars.md @@ -30,7 +30,7 @@ These control the behavior of core tinygrad even when used as a library. Variable | Possible Value(s) | Description ---|---|--- -DEBUG | [1-6] | enable debugging output, with 4 you get operations, timings, speed, generated code and more +DEBUG | [1-7] | enable debugging output (operations, timings, speed, generated code and more) GPU | [1] | enable the GPU (OpenCL) backend CUDA | [1] | enable CUDA backend AMD | [1] | enable AMD backend @@ -49,4 +49,16 @@ VISIBLE_DEVICES | [list[int]]| restricts the NV/AMD devices that are availab JIT | [0-2] | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled VIZ | [1] | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz) ALLOW_TF32 | [1] | enable TensorFloat-32 tensor cores on Ampere or newer GPUs. -WEBGPU_BACKEND | [WGPUBackendType_Metal, ...] | Force select a backend for WebGPU (Metal, DirectX, OpenGL, Vulkan...) \ No newline at end of file +WEBGPU_BACKEND | [WGPUBackendType_Metal, ...] | Force select a backend for WebGPU (Metal, DirectX, OpenGL, Vulkan...) + +## Debug breakdown + +Variable | Value | Description +---|---|--- +DEBUG | >= 1 | Enables debugging and lists devices being used +DEBUG | >= 2 | Provides performance metrics for operations, including timing, memory usage, bandwidth for each kernel execution +DEBUG | >= 3 | Outputs buffers used for each kernel (shape, dtype and strides) and the applied optimizations at a kernel level +DEBUG | >= 4 | Outputs the generated kernel code +DEBUG | >= 5 | Displays the intermediate representation of the computation UOps (AST) +DEBUG | >= 6 | Displays the intermediate representation of the computation UOps in a linearized manner, detailing the operation sequence +DEBUG | >= 7 | Outputs the assembly code generated for the target hardware From 6a5eacba8b8bddcee185f1362897713821813ebf Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 2 Apr 2025 03:19:09 -0400 Subject: [PATCH 2/3] disable CI red llama 3 4 gpu beam (#9690) device hangs and ci would fail --- .github/workflows/benchmark.yml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 053b66f4d8..a75090bddd 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -401,8 +401,9 @@ jobs: # run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_six_gpu.txt - name: Run LLaMA-3 8B BEAM run: AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt - - name: Run LLaMA-3 8B on 4 GPUs with BEAM - run: AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt + # TODO: device hangs + # - name: Run LLaMA-3 8B on 4 GPUs with BEAM + # run: AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt # - name: Run LLaMA-3 8B on 6 GPUs # run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt - name: Restore amdgpu From 4514fd91c10d9cc7b827b1c45d7eeef8626e1660 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 2 Apr 2025 15:27:48 +0800 Subject: [PATCH 3/3] more stuff from DSP (#9689) * more good stuff from dsp branch * test pkl imagenet --- examples/benchmark_onnx.py | 9 ++------- examples/test_pkl_imagenet.py | 19 +++++++++++++++++++ tinygrad/codegen/symbolic.py | 10 +++++++--- tinygrad/engine/jit.py | 7 ++++--- tinygrad/renderer/cstyle.py | 2 +- 5 files changed, 33 insertions(+), 14 deletions(-) create mode 100644 examples/test_pkl_imagenet.py diff --git a/examples/benchmark_onnx.py b/examples/benchmark_onnx.py index de7ac9122b..e88033bd0a 100644 --- a/examples/benchmark_onnx.py +++ b/examples/benchmark_onnx.py @@ -1,12 +1,12 @@ import sys, onnx, time, pickle -from tinygrad import TinyJit, Device, GlobalCounters, fetch, getenv +from tinygrad import TinyJit, GlobalCounters, fetch, getenv from tinygrad.frontend.onnx import OnnxRunner from extra.onnx_helpers import get_example_inputs, validate def load_onnx_model(onnx_file): onnx_model = onnx.load(onnx_file) run_onnx = OnnxRunner(onnx_model) - run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(Device.DEFAULT) for k,v in kwargs.items()}).values())), prune=True) + run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True, optimize=True) return run_onnx_jit, run_onnx.graph_inputs if __name__ == "__main__": @@ -34,8 +34,3 @@ if __name__ == "__main__": if getenv("ORT"): validate(onnx_file, new_inputs, rtol=1e-3, atol=1e-3) print("model validated") - - if (fn:=getenv("SAVE_PKL", "")) != "": - with open(fn, "wb") as f: - pickle.dump(run_onnx_jit, f) - print(f"pkl saved to {fn}") diff --git a/examples/test_pkl_imagenet.py b/examples/test_pkl_imagenet.py new file mode 100644 index 0000000000..8110abf309 --- /dev/null +++ b/examples/test_pkl_imagenet.py @@ -0,0 +1,19 @@ +import sys, pickle +from tinygrad import GlobalCounters +from tinygrad.helpers import fetch, getenv +from examples.test_onnx_imagenet import imagenet_dataloader + +if __name__ == "__main__": + with open(fetch(sys.argv[1]), "rb") as f: + run_onnx_jit = pickle.load(f) + input_name = run_onnx_jit.captured.expected_names[0] + device = run_onnx_jit.captured.expected_st_vars_dtype_device[0][-1] + print(f"input goes into {input_name=} on {device=}") + hit = 0 + for i,(img,y) in enumerate(imagenet_dataloader(cnt=getenv("CNT", 100))): + GlobalCounters.reset() + p = run_onnx_jit(**{input_name:img.to(device)}) + assert p.shape == (1,1000) + t = p.to('cpu').argmax().item() + hit += y==t + print(f"target: {y:3d} pred: {t:3d} acc: {hit/(i+1)*100:.2f}%") diff --git a/tinygrad/codegen/symbolic.py b/tinygrad/codegen/symbolic.py index d589f35a61..df95f8d8ec 100644 --- a/tinygrad/codegen/symbolic.py +++ b/tinygrad/codegen/symbolic.py @@ -175,7 +175,7 @@ def div_and_mod_folding(x: UOp, y: UOp, which: Literal[Ops.MOD, Ops.IDIV], split gep_pushing = PatternMatcher([ # GEP/VECTORIZE, GEP/GEP, GEP/CONST, GEP/VCONST (UPat(Ops.GEP, src=(UPat(Ops.GEP, name='g2'),), name='g1'), - lambda g1, g2: g2.src[0].gep(tuple(g2.arg[g1.arg[i]] for i in range(g1.dtype.count)))), + lambda g1, g2: g2.src[0].gep(tuple(g2.arg[g1.arg[i]] for i in range(len(g1.arg))))), (UPat(Ops.GEP, src=(UPat(Ops.VECTORIZE, name="vec"),), name="gep"), lambda gep, vec: UOp(Ops.VECTORIZE, gep.dtype, tuple(vec.src[i] for i in gep.arg)) if len(gep.arg) > 1 else vec.src[gep.arg[0]]), (UPat(Ops.GEP, src=(UPat.cvar("c", vec=False),), name="gep"), lambda gep, c: gep.const_like(c.arg)), @@ -191,9 +191,13 @@ gep_pushing = PatternMatcher([ (UPat(Ops.VECTORIZE, name="v", src=UPat(Ops.GEP, src=(UPat.var("x"),))), lambda v,x: x.gep(tuple(get_single_element(i.arg) for i in v.src))), ]) -symbolic = symbolic_simple+PatternMatcher([ +commutative = PatternMatcher([ # ** COMMUTATIVE flipping (only for ints) ** - (UPat(GroupOp.Commutative, dtype=dtypes.ints, name='x'), lambda x: x.replace(src=x.src[::-1]) if x.src[1].tuplize < x.src[0].tuplize else None), + # NOTE: this can break merging vector math by only flipping some of them + (UPat(GroupOp.Commutative, dtype=dtypes.int, name='x'), lambda x: x.replace(src=x.src[::-1]) if x.src[1].tuplize < x.src[0].tuplize else None), +]) + +symbolic = symbolic_simple+commutative+PatternMatcher([ # ** boolean algebra ** (UPat.var("x") | (UPat.var("x") & UPat.var()), lambda x: x), # x|(x&y) -> x # ** combine terms ** diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index ac65ce5c7a..04fbfb1cb9 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -149,8 +149,7 @@ class CapturedJit(Generic[ReturnType]): expected_st_vars_dtype_device: list[tuple[ShapeTracker, tuple[Variable, ...], DType, str]] def __reduce__(self): - # TODO: free_intermediates here? - self.optimize_weights() + # TODO: free_intermediates here? optimize_weights here? return self.__class__, (self.ret, self.jit_cache, self.input_replace, self.extra_view_inputs, self.expected_names, self.expected_st_vars_dtype_device) @@ -218,12 +217,13 @@ def _prepare_jit_inputs(args, kwargs): return input_buffers, var_vals, names, st_vars_dtype_device class TinyJit(Generic[ReturnType]): - def __init__(self, fxn:Optional[Callable[..., ReturnType]], captured:Optional[CapturedJit]=None, prune=False): + def __init__(self, fxn:Optional[Callable[..., ReturnType]], captured:Optional[CapturedJit]=None, prune=False, optimize=False): assert fxn or captured, "need either a function or a CapturedJit" self.fxn = fxn self.captured: Optional[CapturedJit] = captured self.cnt: int = 2 if self.fxn is None else 0 self.prune = prune + self.optimize = optimize def add_buffer(self, b:Buffer) -> Buffer: if found:=self._buffer_replace.get(b, None): return found @@ -314,6 +314,7 @@ class TinyJit(Generic[ReturnType]): # set this for next run self.captured = CapturedJit(ret, jit_cache, input_replace, extra_view_inputs, names, st_vars_dtype_device) + if self.optimize: self.captured.optimize_weights() elif self.cnt >= 2: # jit exec assert self.captured is not None diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index e900fc875b..857683118a 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -62,7 +62,7 @@ base_rewrite = PatternMatcher([ extra_pm = PatternMatcher([ # insert a NOOP before BITCAST to force it to be rendered. not needed on all backends? (UPat(Ops.BITCAST, name="x"), - lambda x: UOp(Ops.BITCAST, x.dtype, (UOp(Ops.NOOP, x.src[0].dtype, x.src),)) if x.src[0].op is not Ops.NOOP else None), + lambda x: UOp(Ops.BITCAST, x.dtype, (UOp(Ops.NOOP, x.src[0].dtype, x.src),)) if x.src[0].op not in {Ops.NOOP, Ops.LOAD, Ops.CUSTOM} else None), # rewrite MAX to CMPLT + WHERE (max function is annoying on many cstyle backends) (UPat(Ops.MAX, name="m"), lambda m: (m.src[0] < m.src[1]).where(m.src[1], m.src[0])), # devectorize any bools