the big memory gradient didn't even need to be computed
MODIFY
test/test_train.py
test/test_train.py
@@ -32,7 +32,6 @@ class TestTrain(unittest.TestCase):
Y = np.zeros((BS), dtype=np.int32)
train_one_step(model,X,Y)
- @unittest.skip("OOM in GPU test")
def test_vit(self):
model = ViT()
X = np.zeros((BS,3,224,224), dtype=np.float32)

MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
@@ -174,9 +174,11 @@ class Conv2D(Function):
def backward(ctx, grad_output):
x, w, C = ctx.saved_tensors
+ dx, dw = None, None
+ if ctx.needs_input_grad[0]:
#dx = ctx.processing_op(ProcessingOps.CONVT, grad_output, w, x.shape, C) if ctx.needs_input_grad[0] else None
xt = grad_output
- if C.xs > 1 or C.ys > 1: # unstride
+ if C.xs > 1 or C.ys > 1: # unstride. note, this is really memory intensive for big strides.
xt = ctx.movement_op(MovementOps.RESHAPE, xt, (grad_output.shape[0], grad_output.shape[1], grad_output.shape[2], 1, grad_output.shape[3], 1))
xt = ctx.movement_op(MovementOps.SLICE, xt, ((0,xt.shape[0]), (0,xt.shape[1]), (0,xt.shape[2]), (0,C.ys), (0,xt.shape[4]), (0,C.xs)))
xt = ctx.movement_op(MovementOps.RESHAPE, xt, (xt.shape[0], xt.shape[1], xt.shape[2]*C.ys, xt.shape[4]*C.xs))
@@ -189,6 +191,7 @@ class Conv2D(Function):
dx = ctx.processing_op(ProcessingOps.CONV, xt, wt, (Cdx.bs, Cdx.cout, Cdx.oy, Cdx.ox), Cdx)
dx = ctx.movement_op(MovementOps.SLICE, dx, [(0,s) for s in x.shape])
+ if ctx.needs_input_grad[1]:
# compute derivative of weights using ProcessingOps.CONV
xdw = ctx.movement_op(MovementOps.RESHAPE, x, (C.bs, C.groups, C.cin, C.iy, C.ix))
xdw = ctx.movement_op(MovementOps.PERMUTE, xdw, (2,1,0,3,4))

the big memory gradient didn't even need to be computed