29 lines
864 B
Python
29 lines
864 B
Python
from tinygrad import Tensor
|
|
import math
|
|
|
|
class CosineLR:
|
|
def __init__(self,optm,totalSteps,minlr):
|
|
self.optm = optm
|
|
self.maxlr = optm.lr
|
|
self.minlr = minlr
|
|
self.totalSteps = totalSteps
|
|
self.steps = 0
|
|
|
|
def step(self):
|
|
self.optm.lr = self.minlr + 0.5 * (self.maxlr - self.minlr) * (1 + math.cos((step / self.totalSteps) * math.pi))
|
|
self.optm.step()
|
|
self.steps += 1
|
|
|
|
def zero_grad(self):
|
|
self.optm.zero_grad()
|
|
|
|
|
|
def llmOptimizer(params,steps,minlr):
|
|
muon_params = [p for p in params if len(p.shape) >= 2]
|
|
adamw_params = [p for p in params if len(p.shape) < 2]
|
|
|
|
o1 = nn.optim.Muon(muon_params, lr=hypr["starting_lr"])
|
|
o2 = nn.optim.AdamW(adamw_params, lr=hypr["starting_lr"])
|
|
optimizer = nn.optim.OptimizerGroup([o1,o2])
|
|
return CosineLR(optimizer,steps,minlr)
|