ZeroPoint v0 research prototype: new from-scratch NumPy framework, memory-slot architecture, audited Transformer comparison
5c6c847 verified | """优化器。""" | |
| import numpy as np | |
| class Optimizer: | |
| def __init__(self, params, lr): | |
| self.params = list(params) | |
| self.lr = lr | |
| def zero_grad(self): | |
| for p in self.params: | |
| p.grad = None | |
| def step(self): | |
| raise NotImplementedError | |
| class SGD(Optimizer): | |
| def __init__(self, params, lr=0.01, momentum=0.0): | |
| super().__init__(params, lr) | |
| self.momentum = momentum | |
| self.v = [np.zeros_like(p.data) for p in self.params] | |
| def step(self): | |
| for p, v in zip(self.params, self.v): | |
| if p.grad is None: | |
| continue | |
| v *= self.momentum | |
| v += p.grad | |
| p.data -= self.lr * v | |
| class AdamW(Optimizer): | |
| """Adam + 解耦权重衰减。weight_decay=0 时就是普通 Adam。""" | |
| def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0): | |
| super().__init__(params, lr) | |
| self.b1, self.b2 = betas | |
| self.eps = eps | |
| self.weight_decay = weight_decay | |
| self.t = 0 | |
| self.m = [np.zeros_like(p.data) for p in self.params] | |
| self.v = [np.zeros_like(p.data) for p in self.params] | |
| def step(self): | |
| self.t += 1 | |
| c1, c2 = 1 - self.b1 ** self.t, 1 - self.b2 ** self.t | |
| for p, m, v in zip(self.params, self.m, self.v): | |
| if p.grad is None: | |
| continue | |
| g = p.grad | |
| m *= self.b1 | |
| m += (1 - self.b1) * g | |
| v *= self.b2 | |
| v += (1 - self.b2) * g * g | |
| if self.weight_decay: | |
| p.data -= self.lr * self.weight_decay * p.data | |
| p.data -= self.lr * (m / c1) / (np.sqrt(v / c2) + self.eps) | |
| def state_dict(self): | |
| """优化器状态（步数 + 一阶/二阶动量）。和 model.state_dict() 一起存，才能精确续训。""" | |
| state = {"t": np.array(self.t), "lr": np.array(self.lr)} | |
| for i, (m, v) in enumerate(zip(self.m, self.v)): | |
| state[f"m.{i}"], state[f"v.{i}"] = m.copy(), v.copy() | |
| return state | |
| def load_state_dict(self, state): | |
| assert sum(k.startswith("m.") for k in state) == len(self.params), "参数个数和存档对不上" | |
| self.t, self.lr = int(state["t"]), float(state["lr"]) | |
| for i in range(len(self.params)): | |
| self.m[i][...] = state[f"m.{i}"] | |
| self.v[i][...] = state[f"v.{i}"] | |
| Adam = AdamW | |
| def clip_grad_norm(params, max_norm): | |
| """按全局 L2 范数裁剪梯度，返回裁剪前的范数。""" | |
| grads = [p.grad for p in params if p.grad is not None] | |
| total = float(np.sqrt(sum((g.astype(np.float64) ** 2).sum() for g in grads))) | |
| if total > max_norm: | |
| scale = max_norm / (total + 1e-6) | |
| for p in params: | |
| if p.grad is not None: | |
| p.grad = p.grad * scale | |
| return total | |