tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 22:46:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

gptoss optim fixes (#17356)

Browse Source
This commit is contained in:
wozeparrot
2026-08-01 09:10:06 -07:00
committed by GitHub
parent 8e524ca467
commit 98b700bad1
1 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/mlperf/optim.py
+6 -3
View File
@@ -2,7 +2,7 @@ from tinygrad.tensor import Tensor
from tinygrad.dtype import dtypes
from tinygrad.nn.optim import Optimizer, OptimizerGroup
from tinygrad.helpers import FUSE_OPTIM, getenv
from tinygrad.uop.ops import UOp, Ops
from tinygrad.uop.ops import UOp, Ops, AxisType
STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
@@ -42,8 +42,8 @@ class GradAccClipAdamW(Optimizer):
self.master_params = None
def _zero_shard(self, t:Tensor) -> Tensor:
if not self.zero or (t.shape[0] % len(self.device)) != 0: return t
return Tensor(t.uop._shard(0, len(self.device)).unshard(0)).clone()
if not self.zero or t.ndim < 2 or (t.shape[0] % len(self.device)) != 0: return t
return Tensor(t.uop._shard(0, UOp.range(len(self.device), -1, AxisType.DEVICE)).unshard(0)).clone()
def _zero_gather(self, t:Tensor) -> Tensor:
if not isinstance(t.device, tuple) or t.uop.axis != 0: return t
@@ -123,6 +123,9 @@ class GradAccClipAdamW(Optimizer):
return out.shard_like(t) if offloaded else out
class GradAccClipAdamWGroup(OptimizerGroup):
def __init__(self, *optimizers:GradAccClipAdamW):
super().__init__(*optimizers)
for o in self.optimizers[1:]: o.lr = self.optimizers[0].lr
def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
offset = 0
to_realize = []