fix stupid OPENCL=1 OOM
MODIFY
README.md
README.md
@@ -164,6 +164,8 @@ PROPROTIP: Set "DEBUG=1" environment variable if you want to see why it's slow.
Run `TORCH=1 python3 examples/stable_diffusion.py`
+(or without torch: `OPT=2 OPENCL=1 python3 examples/stable_diffusion.py`)
+
<p align="center">
<img src="https://raw.githubusercontent.com/geohot/tinygrad/master/docs/stable_diffusion_by_tinygrad.jpg">
</p>

MODIFY
accel/opencl/ops_opencl.py
accel/opencl/ops_opencl.py
@@ -195,7 +195,8 @@ class OpenCLBuffer(GPUBuffer):
#ewtypes.append(f"read_only image2d_t {name}_g")
return super().contiguous_view_constant_fold(name)
- seen = set()
+ # THIS WAS STUPIDLY OOMING THE WHOLE THING!
+ #seen = set()
def _processing_op(ret, bufs: List[Tuple[str, OpenCLBuffer]]=[], code:str="acc", C=None, op=ReduceOps.SUM, reduce_shape=None, earlybufs:Set[str]=set(), earlycode:str="acc"):
if C is None or earlycode != "acc":
# TODO: handle an opencl conv without the conv part
@@ -206,9 +207,9 @@ class OpenCLBuffer(GPUBuffer):
w = [x for x in bufs if x[0] == "weight"][0][1]
ewbufs = [x for x in bufs if x[0] not in ["input", "weight"]]
- if (x,w) in OpenCLBuffer.seen:
- print("WARNING: recomputing CONV with", x, w)
- OpenCLBuffer.seen.add((x,w))
+ #if (x,w) in OpenCLBuffer.seen:
+ # print("WARNING: recomputing CONV with", x, w)
+ #OpenCLBuffer.seen.add((x,w))
# remove fakebufs
fakebufs, ewtypes, getters = get_getters(ewbufs, ret)

MODIFY
examples/stable_diffusion.py
examples/stable_diffusion.py
@@ -11,6 +11,7 @@ from extra.utils import fake_torch_load_zipped, get_child
from tinygrad.nn import Conv2d
from tinygrad.tensor import Tensor
from tinygrad.helpers import prod
+from extra.introspection import print_objects
# TODO: rename to GroupNorm and put in nn.py
class Normalize:
@@ -363,7 +364,7 @@ class UNetModel:
saved_inputs = []
for i,b in enumerate(self.input_blocks):
- print("input block", i)
+ #print("input block", i)
for bb in b:
x = run(x, bb)
saved_inputs.append(x)
@@ -371,7 +372,7 @@ class UNetModel:
for bb in self.middle_block:
x = run(x, bb)
for i,b in enumerate(self.output_blocks):
- print("output block", i)
+ #print("output block", i)
x = x.cat(saved_inputs.pop(), dim=1)
for bb in b:
x = run(x, bb)
@@ -559,30 +560,15 @@ if __name__ == "__main__":
def get_model_output(latent, t):
# put into diffuser
timesteps = Tensor([t])
- from tinygrad.llops.ops_gpu import CL
- from tinygrad.llops.ops_gpu import CLBuffer
- from tinygrad.llops.ops_opencl import CLImage, OpenCLBuffer
- import gc
-
- def print_ram():
- print(CL.mem_used/1e9, sum([prod(x.shape)*4 for x in gc.get_objects() if isinstance(x, Tensor)])/1e9)
- img_count = sum([x.is_image() for x in gc.get_objects() if isinstance(x, OpenCLBuffer)])
- print("img_count", img_count)
- buffer_bytes = sum([x.cl.size for x in gc.get_objects() if isinstance(x, CLBuffer)])
- image_bytes = sum([x.cl.row_pitch*x.cl.height for x in gc.get_objects() if isinstance(x, CLImage)])
- print("buffer bytes", buffer_bytes/1e9, "image bytes", image_bytes/1e9, "sum", (buffer_bytes+image_bytes)/1e9)
-
- print_ram()
+ print_objects()
unconditional_latent = model.model.diffusion_model(latent, timesteps, unconditional_context).realize()
- print_ram()
latent = model.model.diffusion_model(latent, timesteps, context).realize()
- print_ram()
unconditional_guidance_scale = 7.5
e_t = unconditional_latent + unconditional_guidance_scale * (latent - unconditional_latent)
return e_t
- TIMESTEPS = 5
+ TIMESTEPS = int(os.getenv("TIMESTEPS", "5"))
timesteps = list(np.arange(1, 1000, 1000//TIMESTEPS))
print(f"running for {timesteps} timesteps")
alphas = [model.alphas_cumprod.numpy()[t] for t in timesteps]

ADD
extra/introspection.py
+# TODO: move the GRAPH and DEBUG stuff to here
+import gc
+from tinygrad.helpers import prod
+from tinygrad.tensor import Tensor
+from tinygrad.ops import LazyBuffer
+from tinygrad.llops.ops_gpu import CL, GPUBuffer
+#from tinygrad.llops.ops_opencl import CLImage, OpenCLBuffer
+
+def print_objects():
+ gc.collect()
+ tensors = [x for x in gc.get_objects() if isinstance(x, Tensor)]
+ tensor_ram_used = sum([prod(x.shape)*4 for x in tensors])
+ lazybuffers = [x for x in gc.get_objects() if isinstance(x, LazyBuffer)]
+ gpubuffers = [x for x in gc.get_objects() if isinstance(x, GPUBuffer)]
+ realized_buffers = [x.realized for x in lazybuffers if x.realized]
+ gpubuffers_orphaned = [x for x in gpubuffers if x not in realized_buffers]
+
+ print(f"{len(tensors)} tensors allocated in {tensor_ram_used/1e9:.2f} GB, GPU using {CL.mem_used/1e9:.2f} GB")
+ print(f"{len(lazybuffers)} lazybuffers {len(realized_buffers)} realized, {len(gpubuffers)} GPU buffers")
+ print(f"{len(gpubuffers_orphaned)} GPU buffers are orphaned")
+
+ cnt = 0
+ for tb in gpubuffers_orphaned:
+ bb = gc.get_referrers(tb)
+ for b in bb:
+ if b is not gpubuffers and b is not gpubuffers_orphaned:
+ print(tb, "reference", type(b),len(b), str(b)[0:150])
+ for x in gc.get_referrers(b):
+ print(str(x)[0:100])
+ if cnt == 10:
+ break
+ cnt += 1
+
+ for x in gpubuffers_orphaned:
+ if getattr(x, '_buf', None): del x._buf
+ if getattr(x, '_image', None): del x._image
+
+"""
+import gc
+
+def print_ram():
+ print(CL.mem_used/1e9, sum([prod(x.shape)*4 for x in gc.get_objects() if isinstance(x, Tensor)])/1e9)
+ img_count = sum([x.is_image() for x in gc.get_objects() if isinstance(x, OpenCLBuffer)])
+ print("img_count", img_count)
+ buffer_bytes = sum([x.cl.size for x in gc.get_objects() if isinstance(x, CLBuffer)])
+ image_bytes = sum([x.cl.row_pitch*x.cl.height for x in gc.get_objects() if isinstance(x, CLImage)])
+ print("buffer bytes", buffer_bytes/1e9, "image bytes", image_bytes/1e9, "sum", (buffer_bytes+image_bytes)/1e9)
+"""

fix stupid OPENCL=1 OOM