From ee12e801a35d39181ad8f4da0e2b3a1c84cb6adb Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 28 May 2025 13:50:30 -0700 Subject: [PATCH] optional fused optimizers (#10549) * enumerate cases of Tensors in the JIT * optional fused optimizers * add fused optimizer test * move that there * ugh --- .github/workflows/test.yml | 2 ++ tinygrad/helpers.py | 2 +- tinygrad/nn/optim.py | 68 +++++++++++++++++++++++--------------- 3 files changed, 45 insertions(+), 27 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 8bb7044053..9a3c9ccbef 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -415,6 +415,8 @@ jobs: run: | PYTHONPATH="." GPU=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20 PYTHONPATH="." GPU=1 IMAGE=2 python3 test/models/test_end2end.py TestEnd2End.test_linear_mnist + - name: Run fused optimizer tests + run: PYTHONPATH="." GPU=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py - name: Run process replay tests uses: ./.github/actions/process-replay diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 71ee2be8e5..fb4e6e9ab9 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -118,7 +118,7 @@ CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), Contex DISABLE_COMPILER_CACHE = ContextVar("DISABLE_COMPILER_CACHE", 0) DONT_REALIZE_EXPAND, DONT_GROUP_REDUCES = ContextVar("DONT_REALIZE_EXPAND", 0), ContextVar("DONT_GROUP_REDUCES", 0) QUANTIZE, VALIDATE_WITH_CPU, IGNORE_OOB = ContextVar("QUANTIZE", 0), ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("IGNORE_OOB", 1) -CORRECT_DIVMOD_FOLDING = ContextVar("CORRECT_DIVMOD_FOLDING", 0) +CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0) @dataclass(frozen=True) class Metadata: diff --git a/tinygrad/nn/optim.py b/tinygrad/nn/optim.py index e6736bca70..9b469febf5 100644 --- a/tinygrad/nn/optim.py +++ b/tinygrad/nn/optim.py @@ -1,5 +1,6 @@ # sorted in order of increasing complexity -from tinygrad.helpers import dedup, flatten, getenv, unwrap +import itertools +from tinygrad.helpers import dedup, flatten, getenv, unwrap, FUSE_OPTIM from tinygrad.tensor import Tensor from tinygrad.dtype import dtypes, least_upper_dtype @@ -7,7 +8,7 @@ class Optimizer: """ Base class for all optimizers. """ - def __init__(self, params: list[Tensor], lr: float): + def __init__(self, params: list[Tensor], lr: float, fused=FUSE_OPTIM): # if it's None, but being put into an optimizer, set it to True for x in params: if x.requires_grad is None: x.requires_grad = True @@ -16,9 +17,15 @@ class Optimizer: assert len(self.params) != 0, "optimizer must have at least one param" self.device = self.params[0].device self.buffers: list[Tensor] = dedup([x for x in params if not x.requires_grad]) # buffers are still realized + self.fused = fused # store lr in at least float32 precision self.lr = Tensor(lr if getenv("CONST_LR") else [lr], requires_grad=False, device=self.device, dtype=least_upper_dtype(dtypes.default_float, dtypes.float32)) + if self.fused: self.pos_params = list(itertools.accumulate(self.params, lambda x,y: x+y.flatten().shape[0], initial=0)) + + def _new_optim_param(self) -> list[Tensor]: + if self.fused: return [Tensor.zeros(self.pos_params[-1], dtype=dtypes.float32, device=self.device, requires_grad=False).contiguous()] + return [Tensor.zeros(*t.shape, dtype=dtypes.float32, device=t.device, requires_grad=False).contiguous() for t in self.params] def zero_grad(self): """ @@ -39,9 +46,17 @@ class Optimizer: if not Tensor.training: raise RuntimeError( 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.schedule_step_with_grads([unwrap(t.grad) for t in self.params])+self.params+self.buffers + if self.fused: + # optimizer fusion just concatentates all the buffers, runs the _step, then splits them back up + out, extra = self._step([Tensor.cat(*[t.flatten() for t in self.params], dim=0)], + [Tensor.cat(*[unwrap(t.grad).flatten() for t in self.params], dim=0)]) + updated_params = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)] + else: + updated_params, extra = self._step(self.params, [unwrap(t.grad) for t in self.params]) + for i, tt in enumerate(self.params): tt.assign(updated_params[i]) + return extra+self.params+self.buffers - def schedule_step_with_grads(self, grads:list[Tensor]) -> list[Tensor]: raise NotImplementedError + def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]: raise NotImplementedError class OptimizerGroup(Optimizer): """ @@ -55,7 +70,7 @@ class OptimizerGroup(Optimizer): def schedule_step(self) -> list[Tensor]: return [x for o in self.optimizers for x in o.schedule_step()] # LARS is essentially just trust ratio to SGD so if we just set the trust coeff 0.0 it's just standard SGD. -def SGD(params: list[Tensor], lr=0.001, momentum=0.0, weight_decay=0.0, nesterov=False, classic=False): +def SGD(params: list[Tensor], lr=0.001, momentum=0.0, weight_decay=0.0, nesterov=False, classic=False, fused=FUSE_OPTIM): """ Stochastic Gradient Descent (SGD) optimizer with optional momentum and weight decay. @@ -63,7 +78,7 @@ def SGD(params: list[Tensor], lr=0.001, momentum=0.0, weight_decay=0.0, nesterov - Described: https://paperswithcode.com/method/sgd """ - return LARS(params, lr, momentum, weight_decay, nesterov, classic, tcoef=0.0) + return LARS(params, lr, momentum, weight_decay, nesterov, classic, tcoef=0.0, fused=fused) class LARS(Optimizer): """ @@ -72,14 +87,14 @@ class LARS(Optimizer): - Described: https://paperswithcode.com/method/lars - Paper: https://arxiv.org/abs/1708.03888v3 """ - def __init__(self, params:list[Tensor], lr=0.001, momentum=0.9, weight_decay=1e-4, nesterov=False, classic=True, tcoef=0.001): - super().__init__(params, lr) + def __init__(self, params:list[Tensor], lr=0.001, momentum=0.9, weight_decay=1e-4, nesterov=False, classic=True, tcoef=0.001, fused=FUSE_OPTIM): + super().__init__(params, lr, fused) self.momentum, self.wd, self.nesterov, self.classic, self.tcoef = momentum, weight_decay, nesterov, classic, tcoef - self.b = [Tensor.zeros(*t.shape, dtype=dtypes.float32, device=t.device, requires_grad=False).contiguous() for t in self.params] \ - if self.momentum else [] + self.b = self._new_optim_param() if self.momentum else [] - def schedule_step_with_grads(self, grads:list[Tensor]) -> list[Tensor]: - for i, (t, g) in enumerate(zip(self.params, grads)): + def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]: + ret = [] + for i, (t, g) in enumerate(zip(params, grads)): if self.tcoef != 0: r1 = t.detach().square().sum().sqrt() r2 = g.square().sum().sqrt() @@ -95,26 +110,26 @@ class LARS(Optimizer): g = (g + self.momentum * self.b[i]) if self.nesterov else self.b[i] # popular momentum does pre learning rate update if not self.classic: g = g * r * self.lr - t.assign((t.detach() - g).cast(t.dtype)) - return self.b + ret.append((t.detach() - g).cast(t.dtype)) + return ret, self.b # LAMB is essentially just the trust ratio part of LARS applied to Adam/W so if we just set the trust ratio to 1.0 it's just Adam/W. -def AdamW(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, weight_decay=0.01): +def AdamW(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, weight_decay=0.01, fused=FUSE_OPTIM): """ AdamW optimizer with optional weight decay. - Described: https://paperswithcode.com/method/adamw - Paper: https://arxiv.org/abs/1711.05101v3 """ - return LAMB(params, lr, b1, b2, eps, weight_decay, adam=True) -def Adam(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8): + return LAMB(params, lr, b1, b2, eps, weight_decay, adam=True, fused=fused) +def Adam(params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, fused=FUSE_OPTIM): """ Adam optimizer. - Described: https://paperswithcode.com/method/adam - Paper: https://arxiv.org/abs/1412.6980 """ - return LAMB(params, lr, b1, b2, eps, 0.0, adam=True) + return LAMB(params, lr, b1, b2, eps, 0.0, adam=True, fused=fused) class LAMB(Optimizer): """ @@ -123,17 +138,18 @@ class LAMB(Optimizer): - Described: https://paperswithcode.com/method/lamb - Paper: https://arxiv.org/abs/1904.00962 """ - def __init__(self, params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, adam=False): - super().__init__(params, lr) + def __init__(self, params: list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, adam=False, fused=FUSE_OPTIM): + super().__init__(params, lr, fused) self.b1, self.b2, self.eps, self.wd, self.adam = b1, b2, eps, weight_decay, adam self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False).contiguous() for _ in [b1, b2]) - 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] + self.m = self._new_optim_param() + self.v = self._new_optim_param() - def schedule_step_with_grads(self, grads:list[Tensor]) -> list[Tensor]: + def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]: + ret = [] self.b1_t *= self.b1 self.b2_t *= self.b2 - for i, (t, g) in enumerate(zip(self.params, grads)): + for i, (t, g) in enumerate(zip(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) @@ -145,5 +161,5 @@ class LAMB(Optimizer): r: Tensor|float = Tensor.where(r1 > 0, Tensor.where(r2 > 0, r1 / r2, 1.0), 1.0) else: r = 1.0 - t.assign((t.detach() - self.lr * r * up).cast(t.dtype)) - return [self.b1_t, self.b2_t] + self.m + self.v + ret.append((t.detach() - self.lr * r * up).cast(t.dtype)) + return ret, [self.b1_t, self.b2_t] + self.m + self.v