diff --git a/test/test_tensor.py b/test/test_tensor.py index d34556792e..e3a9ceea2a 100644 --- a/test/test_tensor.py +++ b/test/test_tensor.py @@ -77,38 +77,23 @@ class TestTinygrad(unittest.TestCase): np.testing.assert_allclose(xgrad2, xgrad * 2., atol=1e-6) np.testing.assert_allclose(wgrad2, wgrad * 2., atol=1e-6) - @unittest.expectedFailure def test_second_order_backward_pass(self): def test_pytorch(): - x = torch.tensor(x_init) - m = torch.tensor(m_init, requires_grad=True) - out = x.mul(m).sum() - # use retain graph so we can compute second order derivatives later - out.backward(retain_graph=True) - # save first-order gradient (dO/dm). they still contain graph information on how they were constructed wrt x and W - grad_m = m.grad - # zero gradients so second-order gradients are correct - m.grad = None - # compute second-order gradients - grad_m.sum().backward(retain_graph=True) - - # d2O/dm2 - second_grad_m = m.grad - return second_grad_m.numpy() + x_val = torch.tensor([2.0], requires_grad=True) + f = x_val**3 + first_derivative = torch.autograd.grad(outputs=f, inputs=x_val, create_graph=True)[0] + second_derivative = torch.autograd.grad(outputs=first_derivative, inputs=x_val)[0] + # d^2f/dx^2 = 6x = 6*2 = 12 + return second_derivative.numpy() def test_tinygrad(): - x = Tensor(x_init) - m = Tensor(m_init, requires_grad=True) - out = x.mul(m).sum() - out.backward() - grad_m = m.grad - m.grad = None - grad_m.sum().backward() - second_grad_m = m.grad # currently, this will be None (incorrect) - return second_grad_m.numpy() + x_val = Tensor(2.0) + f = x_val**3 + first_derivative = f.gradient(x_val)[0] + second_derivative = first_derivative.gradient(x_val)[0] + return second_derivative.numpy() - for x,y in zip(test_tinygrad(), test_pytorch()): - np.testing.assert_allclose(x, y, atol=1e-5) + np.testing.assert_allclose(test_tinygrad(), test_pytorch(), atol=1e-5) # passing `gradient` to backward def test_backward_pass_vjp(self): diff --git a/test/unit/test_gradient.py b/test/unit/test_gradient.py index 36d5a8947a..4e6f092e4b 100644 --- a/test/unit/test_gradient.py +++ b/test/unit/test_gradient.py @@ -5,7 +5,7 @@ import jax.numpy as jnp from tinygrad import Tensor from tinygrad.dtype import dtypes from tinygrad.ops import UOp -from tinygrad.gradient import gradient +from tinygrad.gradient import compute_gradient class TestGradient(unittest.TestCase): def _cmp_nan_okay(self, x, y): @@ -14,7 +14,7 @@ class TestGradient(unittest.TestCase): def _test_one_input_function(self, f:Callable, jf:Callable|None=None): x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float) - gx = gradient(f(x), [x])[0] + gx = compute_gradient(f(x), UOp.const(dtypes.float, 1.0), [x])[x] gf = jax.grad(f if jf is None else jf) for val in [-5., -2.0, 0.0, 2.0, 5.]: @@ -24,7 +24,8 @@ class TestGradient(unittest.TestCase): def _test_two_input_function(self, f:Callable, jf:Callable|None=None): x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float) y = UOp.variable('y', -math.inf, math.inf, dtype=dtypes.float) - gx, gy = gradient(f(x, y), [x, y]) + grads = compute_gradient(f(x, y), UOp.const(dtypes.float, 1.0), [x, y]) + gx, gy = grads[x], grads[y] gf = jax.grad(f if jf is None else jf, argnums=(0, 1)) for valx in [-5., -2.0, 0.0, 2.0, 5.]: @@ -54,9 +55,7 @@ class TestGradient(unittest.TestCase): def test_chain_binop(self): self._test_two_input_function(lambda x,y: (x*y)+x*y) def test_big_add_sin(self): self._test_two_input_function(lambda x,y: x.sin()+3.0/y, lambda x,y: jnp.sin(x)+3.0/y) def test_big_chain(self): self._test_two_input_function(lambda x,y: (1.0/x*y)+x*y) - - # TODO: this isn't working - #def test_where(self): self._test_two_input_function(lambda x,y: (x list[UOp]: - grads = {root: root.const_like(1.0)} +def compute_gradient(root:UOp, root_grad:UOp, targets:list[UOp]) -> dict[UOp, UOp]: + grads = {root: root_grad} for t0 in reversed(_deepwalk(root, targets)): if t0 not in grads: continue lgrads: tuple[UOp|None, ...]|None = cast(tuple[UOp, ...]|None, pm_gradient.rewrite(t0, ctx=grads[t0])) @@ -65,8 +65,4 @@ def gradient(root:UOp, targets:list[UOp]) -> list[UOp]: if v is None: continue if k in grads: grads[k] = grads[k] + v else: grads[k] = v - ret = [grads.get(x, None) for x in targets] - for i,x in enumerate(ret): - if x is None: raise RuntimeError(f"{targets[i]}\n\nnot found in\n\n{root}") - return ret - + return grads diff --git a/tinygrad/nn/optim.py b/tinygrad/nn/optim.py index 5a189c450b..bc3b8dbfa8 100644 --- a/tinygrad/nn/optim.py +++ b/tinygrad/nn/optim.py @@ -1,6 +1,6 @@ # sorted in order of increasing complexity from typing import List -from tinygrad.helpers import dedup, flatten, getenv +from tinygrad.helpers import dedup, flatten, getenv, unwrap from tinygrad.tensor import Tensor from tinygrad.dtype import dtypes, least_upper_dtype @@ -39,8 +39,8 @@ class Optimizer: assert Tensor.training, ( f"""Tensor.training={Tensor.training}, Tensor.training must be enabled to use the optimizer. - help: Consider setting Tensor.training=True before calling Optimizer.step().""") - return self._step()+self.params+self.buffers - def _step(self) -> List[Tensor]: raise NotImplementedError + return self.schedule_step_with_grads([unwrap(t.grad) for t in self.params])+self.params+self.buffers + def schedule_step_with_grads(self, grads:List[Tensor]) -> List[Tensor]: raise NotImplementedError class OptimizerGroup(Optimizer): """ @@ -76,12 +76,11 @@ class LARS(Optimizer): self.momentum, self.wd, self.nesterov, self.classic, self.tcoef = momentum, weight_decay, nesterov, classic, tcoef self.b = [Tensor.zeros(*t.shape, dtype=t.dtype, device=t.device, requires_grad=False) for t in self.params] if self.momentum else [] - def _step(self) -> List[Tensor]: - for i, t in enumerate(self.params): - assert t.grad is not None + def schedule_step_with_grads(self, grads:List[Tensor]) -> List[Tensor]: + for i, (t, g) in enumerate(zip(self.params, grads)): # contiguous is needed since the grads can allegedly form a "diamond" # TODO: fix this in lazy.py - g = t.grad.contiguous() + g = g.contiguous() if self.tcoef != 0: r1 = t.detach().square().sum().sqrt() r2 = g.square().sum().sqrt() @@ -130,13 +129,12 @@ class LAMB(Optimizer): self.m = [Tensor.zeros(*t.shape, dtype=dtypes.float32, device=t.device, requires_grad=False).contiguous() for t in self.params] self.v = [Tensor.zeros(*t.shape, dtype=dtypes.float32, device=t.device, requires_grad=False).contiguous() for t in self.params] - def _step(self) -> List[Tensor]: + def schedule_step_with_grads(self, grads:List[Tensor]) -> List[Tensor]: self.b1_t *= self.b1 self.b2_t *= self.b2 - for i, t in enumerate(self.params): - assert t.grad is not None - self.m[i].assign(self.b1 * self.m[i] + (1.0 - self.b1) * t.grad) - self.v[i].assign(self.b2 * self.v[i] + (1.0 - self.b2) * (t.grad * t.grad)) + for i, (t, g) in enumerate(zip(self.params, grads)): + self.m[i].assign(self.b1 * self.m[i] + (1.0 - self.b1) * g) + self.v[i].assign(self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)) m_hat = self.m[i] / (1.0 - self.b1_t) v_hat = self.v[i] / (1.0 - self.b2_t) up = (m_hat / (v_hat.sqrt() + self.eps)) + self.wd * t.detach() diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index b2bce208bf..88feee46ce 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -7,7 +7,7 @@ from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, leas from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten, dedup from tinygrad.helpers import IMAGE, DEBUG, WINO, _METADATA, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap from tinygrad.multi import MultiLazyBuffer -from tinygrad.gradient import gradient +from tinygrad.gradient import compute_gradient from tinygrad.ops import smax, smin, resolve, UOp, Ops, sint, Variable, SimpleMathTrait, identity_element from tinygrad.device import Device, Buffer, BufferSpec from tinygrad.engine.realize import run_schedule @@ -865,7 +865,7 @@ class Tensor(SimpleMathTrait): # ***** toposort and backward pass ***** - def gradient(self, *targets:Tensor) -> list[Tensor]: + def gradient(self, *targets:Tensor, gradient:Optional[Tensor]=None) -> list[Tensor]: """ Compute the gradient of the targets with respect to self. @@ -881,7 +881,13 @@ class Tensor(SimpleMathTrait): """ assert isinstance(self.lazydata, UOp), "multi isn't supported yet" target_uops: List[UOp] = [x.lazydata for x in targets if isinstance(x.lazydata, UOp)] - return [Tensor(y) for y in gradient(self.lazydata, target_uops)] + assert gradient is not None or self.shape == tuple(), "when no gradient is provided, backward must be called on a scalar tensor" + grads = compute_gradient(self.lazydata, self.lazydata.const_like(1) if gradient is None else cast(UOp, gradient.lazydata), target_uops) + ret = [] + for x in target_uops: + if (y:=grads.get(x)) is None: raise RuntimeError(f"{x}\n\nnot found in\n\n{self.lazydata}") + ret.append(Tensor(y, device=x.device)) + return ret def _deepwalk(self): def _walk(node, visited):