31 lines
899 B
Python
31 lines
899 B
Python
from tinygrad import Tensor,nn
|
|
import math
|
|
|
|
class CosineLR:
|
|
def __init__(self,optm,totalSteps,maxlr,minlr):
|
|
self.optm = optm
|
|
self.maxlr = maxlr
|
|
self.minlr = minlr
|
|
self.totalSteps = totalSteps
|
|
self.steps = 0
|
|
|
|
def step(self):
|
|
lr = self.minlr + 0.5 * (self.maxlr - self.minlr) * (1 + math.cos((self.steps / self.totalSteps) * math.pi))
|
|
for o in self.optm:
|
|
o.lr = lr
|
|
self.optm.step()
|
|
self.steps += 1
|
|
|
|
def zero_grad(self):
|
|
self.optm.zero_grad()
|
|
|
|
|
|
def llmOptimizer(params,steps,maxlr,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=maxlr)
|
|
o2 = nn.optim.AdamW(adamw_params, lr=maxlr)
|
|
optimizer = nn.optim.OptimizerGroup(o1,o2)
|
|
return CosineLR(optimizer,steps,maxlr,minlr)
|