From 7690ea62c87133410ffe1596aa222a6e0e9069f7 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 6 Sep 2026 10:51:35 -0600 Subject: [PATCH] Add variable precision adam convrot optimizer with stochastic rounding and accumulation. --- toolkit/optimizer.py | 5 + toolkit/optimizers/adamconvrot.py | 687 ++++++++++++++++++++++++++ toolkit/optimizers/test_optimizers.py | 13 +- 3 files changed, 703 insertions(+), 2 deletions(-) create mode 100644 toolkit/optimizers/adamconvrot.py diff --git a/toolkit/optimizer.py b/toolkit/optimizer.py index f1c86d8..8cb92a1 100644 --- a/toolkit/optimizer.py +++ b/toolkit/optimizer.py @@ -106,6 +106,11 @@ def get_optimizer( elif lower_type == 'automagicexperiment': from toolkit.optimizers.automagicEXPERIMENT import AutomagicEXPERIMENT optimizer = AutomagicEXPERIMENT(params, lr=float(learning_rate), **optimizer_params) + elif lower_type == 'adamconvrot': + from toolkit.optimizers.adamconvrot import AdamConvRot + if 'eps' not in optimizer_params: + optimizer_params['eps'] = 1e-6 + optimizer = AdamConvRot(params, lr=float(learning_rate), **optimizer_params) else: raise ValueError(f'Unknown optimizer type {optimizer_type}') return optimizer diff --git a/toolkit/optimizers/adamconvrot.py b/toolkit/optimizers/adamconvrot.py new file mode 100644 index 0000000..9124858 --- /dev/null +++ b/toolkit/optimizers/adamconvrot.py @@ -0,0 +1,687 @@ +""" +AdamConvRot: Adam / AdamW with both moments stored in ConvRot n-bit codes. + +qtype picks the storage width: "convrot8" or "convrotint2".."convrotint8" +(the convrot_quant names; 8 = plain int8 codes, everything else is the +bit-packed n-bit format of pack_intn_rows). adamw=True is decoupled weight +decay (AdamW), adamw=False is classic L2-in-the-gradient Adam. + +Storage. Each moment is viewed flat and cut into 256-element blocks; a block +is rotated with the ConvRot regular-Hadamard transform (rotate(x, 256) from +convrot_quant -- spreads a block's outliers so one symmetric scale covers the +block without clipping the rest) and stored as n-bit symmetric codes with one +fp32 scale per block. The second moment is kept in the SQRT domain (what the +update divides by; half the dynamic range). Cost: bits/4 bytes per element +for both moments plus 1/32 byte of scales -- 2.03 B/elem at 8 bit, 1.03 at 4, +0.53 at 2. + +Requantization is NOT dequantize -> EMA -> re-round. That scheme is dead on +arrival for the slow second moment: with beta2 = 0.999 the per-step change +of sqrt(v) is ~0.05%, far below half an int8 step, so nearest rounding +freezes the state (it only ever ratchets up on spikes and never decays), and +per-step stochastic rounding is unbiased but injects one code of noise every +step that the EMA only damps at rate (1 - beta2): ~9 codes rms at 8 bit, +more than the whole 4-bit grid. Instead the uniform decay of an EMA is +applied EXACTLY to the block scale (scale *= beta each step, an fp32 +multiply), and only the increment -- (1-beta1)*g for m, the exact +sqrt-domain increment for s -- is rotated and added to the codes with +stochastic rounding, so an element with no gradient decays exactly and the +rounding noise is Poisson in the accumulated increment instead of one code +per step. A block is re-rounded (new scale from its exact fp32 values, +stochastic) only when a code would leave the grid; the new scale leaves +1.25x headroom so growth does not retrigger it every step. A dead block +(scale 0) restarts from its increment on the first non-zero gradient. + +The update uses the exact fp32 moments of the current step (before their +requantization), bias-corrected as in torch Adam; parameters are updated in +fp32 and written back with stochastic rounding for bf16 / fp16. + +Everything for one parameter -- unpack, 2 rotations back, EMA, update, +parameter write, 2 rotations forward, stochastic requant, pack -- is one +triton kernel launch, one pass over memory, no fp32 temporaries (the +rotation is two batched 16x16 tensor-core dots, exact on integer codes). +CPU / non-triton devices run the identical algorithm in torch. Tensors +below min_quant_numel (or 1D) keep plain fp32 moments. + +fused=True updates each parameter inside the backward pass (a +post-accumulate-grad hook) and frees its gradient at once, so the peak never +holds a full set of gradients; .step() then only runs the closure. That +bypasses the trainer's grad clipping / non-finite skip, so non-finite grads +are zeroed in the hook, and it is incompatible with multi-backward gradient +accumulation. Default is the traditional unfused step, where low-precision +(bf16 / fp16) parameters accumulate their micro-batch grads through a +post-accumulate-grad hook that sums in fp32 and stochastically rounds back +into the param's dtype (optimizer_utils.stochastic_grad_accummulation), so +small per-micro-batch grads are not lost to round-to-nearest; .step() +consumes that buffer. Note the trainer's grad clipping does not see it. +""" + +import math + +import torch + +from toolkit.optimizers.optimizer_utils import ( + copy_stochastic, + stochastic_grad_accummulation, +) +from toolkit.util.convrot_quant import ( + _import_triton, + _triton_available, + pack_intn_rows, + regular_hadamard, + rotate, + unpack_intn_rows, +) + +BLOCK = 256 # rotation block == scale group +HEADROOM = 1.25 # rescale puts the block max at qmax / HEADROOM +TINY = 1e-30 + + +def _bits_for_qtype(qtype: str) -> int: + if qtype == "convrot8": + return 8 + if qtype.startswith("convrotint"): + try: + bits = int(qtype[len("convrotint") :]) + except ValueError: + bits = 0 + if 2 <= bits <= 8: + return bits + raise ValueError( + f"AdamConvRot: unsupported qtype {qtype!r}; use convrot8 or convrotint2..convrotint8" + ) + + +# ---------------------------------------------------------------- triton path + +_kernel = None +_h16_cache = {} + + +def _h16(device): + key = str(device) + h = _h16_cache.get(key) + if h is None: + # regular_hadamard(256) == kron(H16, H16): the 256-rotation is H16 over + # the high base-16 digit and H16 over the low one + h = regular_hadamard(16, device, torch.float32).contiguous() + _h16_cache[key] = h + return h + + +def _get_kernel(): + global _kernel + if _kernel is not None: + return _kernel + triton, tl = _import_triton() + + @triton.jit + def _rot256(x, hb, R: tl.constexpr): + # tf32 dots are exact for integer codes: |c| <= 127 and H16 = +-1/4 keep + # every partial sum within 11 mantissa bits + x3 = tl.reshape(x, (R, 16, 16)) + z = tl.dot(x3, hb, input_precision="tf32") + y3 = tl.dot(hb, z, input_precision="tf32") + return tl.reshape(y3, (R, 256)) + + @triton.jit + def _unpack256( + c_ptr, + b, + bm, + BITS: tl.constexpr, + BPOW: tl.constexpr, + QMAX: tl.constexpr, + R: tl.constexpr, + ): + # pack_intn_rows layout: 8 codes -> one BITS-byte little-endian word + gi = tl.arange(0, 32) + bi = tl.arange(0, BPOW) + j = tl.arange(0, 8) + ptrs = ( + c_ptr + + b[:, None, None] * (32 * BITS) + + gi[None, :, None] * BITS + + bi[None, None, :] + ) + msk = bm[:, None, None] & (bi < BITS)[None, None, :] + by = tl.load(ptrs, mask=msk, other=0).to(tl.int64) + word = tl.sum(by << (8 * bi)[None, None, :].to(tl.int64), axis=2) + v = (word[:, :, None] >> (BITS * j)[None, None, :].to(tl.int64)) & ( + (1 << BITS) - 1 + ) + return tl.reshape((v - QMAX).to(tl.float32), (R, 256)) + + @triton.jit + def _pack256( + c_ptr, + b, + bm, + q, + BITS: tl.constexpr, + BPOW: tl.constexpr, + QMAX: tl.constexpr, + R: tl.constexpr, + ): + u = tl.reshape(q + QMAX, (R, 32, 8)).to(tl.int64) + j = tl.arange(0, 8) + word = tl.sum(u << (BITS * j)[None, None, :].to(tl.int64), axis=2) + gi = tl.arange(0, 32) + bi = tl.arange(0, BPOW) + by = ((word[:, :, None] >> (8 * bi)[None, None, :].to(tl.int64)) & 0xFF).to( + tl.uint8 + ) + ptrs = ( + c_ptr + + b[:, None, None] * (32 * BITS) + + gi[None, :, None] * BITS + + bi[None, None, :] + ) + msk = bm[:, None, None] & (bi < BITS)[None, None, :] + tl.store(ptrs, by, mask=msk) + + @triton.jit + def _requant256( + c_ptr, + s_ptr, + b, + bm, + offs, + codes, + scale, + decay, + inc_rot, + rqmax, + tiny, + seed_a, + BITS: tl.constexpr, + BPOW: tl.constexpr, + QMAX: tl.constexpr, + SR: tl.constexpr, + R: tl.constexpr, + ): + s_dec = scale * decay + values = codes * s_dec[:, None] + inc_rot # exact fp32 block values + if SR: + u = tl.rand(seed_a, offs) + else: + u = tl.full((R, 256), 0.5, tl.float32) + q = tl.floor(values / tl.maximum(s_dec, tiny)[:, None] + u) + over = tl.max(tl.abs(q), axis=1) > QMAX + # rescale path: new scale with headroom, same random draw (only the + # triggering element's rounding is conditioned on it) + s_out = tl.where(over, tl.max(tl.abs(values), axis=1) / rqmax, s_dec) + q = tl.floor(values / tl.maximum(s_out, tiny)[:, None] + u) + q = tl.minimum(tl.maximum(q, -1.0 * QMAX), 1.0 * QMAX) + tl.store(s_ptr + b, s_out, mask=bm) + _pack256(c_ptr, b, bm, q.to(tl.int32), BITS, BPOW, QMAX, R) + return s_out + + @triton.jit(do_not_specialize=["numel", "n_blocks", "seed"]) + def adam_convrot_kernel( + p_ptr, + g_ptr, + cm_ptr, + sm_ptr, + cs_ptr, + ss_ptr, + h_ptr, + numel, + n_blocks, + lr, + wd, + beta1, + beta2, + eps, + inv_bc1, + inv_sqrt_bc2, + sqrt_beta2, + rqmax, + tiny, + seed, + BITS: tl.constexpr, + BPOW: tl.constexpr, + QMAX: tl.constexpr, + ADAMW: tl.constexpr, + P_KIND: tl.constexpr, + SR: tl.constexpr, + SANITIZE: tl.constexpr, + R: tl.constexpr, + ): + pid = tl.program_id(0) + b = pid * R + tl.arange(0, R) + bm = b < n_blocks + e = tl.arange(0, 256) + offs = b[:, None] * 256 + e[None, :] + em = bm[:, None] & (offs < numel) + r = tl.arange(0, 16) + hb = tl.load( + h_ptr + + tl.zeros((R, 1, 1), tl.int32) + + r[None, :, None] * 16 + + r[None, None, :] + ) + + g = tl.load(g_ptr + offs, mask=em, other=0.0).to(tl.float32) + if SANITIZE: + g = tl.where((g == g) & (tl.abs(g) < 3.0e38), g, 0.0) + if not ADAMW: + p0 = tl.load(p_ptr + offs, mask=em, other=0.0).to(tl.float32) + g = g + wd * p0 + + # second moment (sqrt domain) first: only its reciprocal stays live + ss = tl.load(ss_ptr + b, mask=bm, other=0.0) + cs = _unpack256(cs_ptr, b, bm, BITS, BPOW, QMAX, R) + s_old = _rot256(cs, hb, R) * ss[:, None] + # increment from the rectified value: quant noise can leave s negative, + # and an increment against a negative s would grow with the noise + s_pos = tl.maximum(s_old, 0.0) + s_new = tl.sqrt(beta2 * s_pos * s_pos + (1.0 - beta2) * g * g) + inc_s = _rot256(s_new - sqrt_beta2 * s_pos, hb, R) + ss_new = _requant256( + cs_ptr, + ss_ptr, + b, + bm, + offs, + cs, + ss, + sqrt_beta2, + inc_s, + rqmax, + tiny, + seed + 2, + BITS, + BPOW, + QMAX, + SR, + R, + ) + # denominator floored at the block's code resolution -- a stored + # sqrt(v) below one code step is unresolved, not small + den = 1.0 / (tl.maximum(s_new, ss_new[:, None]) * inv_sqrt_bc2 + eps) + + # first moment and the parameter update + sm = tl.load(sm_ptr + b, mask=bm, other=0.0) + cm = _unpack256(cm_ptr, b, bm, BITS, BPOW, QMAX, R) + m_new = beta1 * (_rot256(cm, hb, R) * sm[:, None]) + (1.0 - beta1) * g + upd = m_new * inv_bc1 * den + p = tl.load(p_ptr + offs, mask=em, other=0.0).to(tl.float32) + if ADAMW: + p_new = p - lr * (upd + wd * p) + else: + p_new = p - lr * upd + inc_m = _rot256(g * (1.0 - beta1), hb, R) + _requant256( + cm_ptr, + sm_ptr, + b, + bm, + offs, + cm, + sm, + beta1, + inc_m, + rqmax, + tiny, + seed, + BITS, + BPOW, + QMAX, + SR, + R, + ) + + if P_KIND == 0: + tl.store(p_ptr + offs, p_new, mask=em) + elif P_KIND == 1: + if SR: + # copy_stochastic_bf16: random 16-bit add, truncate + xi = p_new.to(tl.int32, bitcast=True) + r16 = (tl.randint(seed + 4, offs) & 0xFFFF).to(tl.int32) + xi = (xi + r16) & -65536 + out = xi.to(tl.float32, bitcast=True).to(tl.bfloat16) + else: + out = p_new.to(tl.bfloat16) + tl.store(p_ptr + offs, out, mask=em) + else: + if SR: + # fp16 ulp from the fp32 exponent (subnormals below 2^-14) + ex = (p_new.to(tl.int32, bitcast=True) >> 23) & 0xFF + ulp = tl.exp2((tl.maximum(ex - 127, -14) - 10).to(tl.float32)) + u = tl.rand(seed + 4, offs) + out = (tl.floor(p_new / ulp + u) * ulp).to(tl.float16) + else: + out = p_new.to(tl.float16) + tl.store(p_ptr + offs, out, mask=em) + + _kernel = adam_convrot_kernel + return _kernel + + +# ------------------------------------------------------------------ optimizer + + +class AdamConvRot(torch.optim.Optimizer): + def __init__( + self, + params, + lr=1e-3, + betas=(0.9, 0.999), + eps=1e-8, + weight_decay=0.0, + qtype="convrot8", + adamw=True, + min_quant_numel=4096, + fused=False, + blocks_per_program=8, # 8x4 warps: no register spills, ~1.1 ms per 64M elems on a 5090 + num_warps=4, + ): + if lr < 0.0: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 <= betas[0] < 1.0 or not 0.0 <= betas[1] < 1.0: + raise ValueError(f"Invalid betas: {betas}") + self.bits = _bits_for_qtype(qtype) + self.qtype = qtype + self.adamw = bool(adamw) + self.min_quant_numel = int(min_quant_numel) + self.blocks_per_program = int(blocks_per_program) + self.num_warps = int(num_warps) + # off only for validation (nearest rounding, kernel == torch path) + self.stochastic_rounding = True + defaults = {"lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay} + super().__init__(params, defaults) + + self.fused = bool(fused) + self._hook_handles = [] + self._rebuild_group_index() + for group in self.param_groups: + for p in group["params"]: + if not p.requires_grad: + continue + if self.fused: + self._hook_handles.append( + p.register_post_accumulate_grad_hook(self._fused_hook) + ) + elif p.dtype != torch.float32: + # stochastic bf16/fp16 grad accumulation across micro-batches + self._hook_handles.append( + p.register_post_accumulate_grad_hook( + stochastic_grad_accummulation + ) + ) + + def _rebuild_group_index(self): + # the hooks cannot rely on group-dict identity: load_state_dict + # replaces the group dicts + self._param_group_index = { + p: gi for gi, group in enumerate(self.param_groups) for p in group["params"] + } + + @torch.no_grad() + def _fused_hook(self, p): + gi = self._param_group_index.get(p) + if gi is None: + self._rebuild_group_index() + gi = self._param_group_index.get(p, 0) + seed = int(torch.randint(0, 2**31 - 1, (1,)).item()) + self._update_param(p, self.param_groups[gi], seed, sanitize=True) + p.grad = None + + @property + def qmax(self) -> int: + return (1 << (self.bits - 1)) - 1 + + @property + def rescale_qmax(self) -> float: + return max(self.qmax / HEADROOM, 1.0) + + def _quantizes(self, p: torch.Tensor) -> bool: + return p.dim() >= 2 and p.numel() >= self.min_quant_numel + + # ------------------------------------------------------------------ state + + def _init_state(self, p: torch.Tensor): + state = self.state[p] + state["step"] = 0 + if self._quantizes(p): + n_blocks = -(-p.numel() // BLOCK) + zero_codes = pack_intn_rows( + torch.zeros(n_blocks, BLOCK, dtype=torch.int8, device=p.device), + self.bits, + ) + for key in ("exp_avg", "exp_avg_sq"): + state[key] = zero_codes.clone() + # fp32 scales as a uint8 view: the parent load_state_dict casts + # floating state to the param dtype + state[key + "_scale"] = torch.zeros( + n_blocks, dtype=torch.float32, device=p.device + ).view(torch.uint8) + else: + state["exp_avg"] = torch.zeros_like(p, dtype=torch.float32) + state["exp_avg_sq"] = torch.zeros_like(p, dtype=torch.float32) + + def _state_ok(self, p: torch.Tensor, state) -> bool: + if "exp_avg" not in state or "step" not in state: + return False + if self._quantizes(p): + n_blocks = -(-p.numel() // BLOCK) + return ( + "exp_avg_scale" in state + and tuple(state["exp_avg"].shape) == (n_blocks, 32 * self.bits) + and state["exp_avg"].dtype == torch.uint8 + ) + return state["exp_avg"].shape == p.shape + + def load_state_dict(self, state_dict): + # the parent casts every state tensor to the param dtype, which would + # turn the uint8 code / scale buffers (and fp32 moments) into bf16 + super().load_state_dict(state_dict) + saved_ids = [pid for g in state_dict["param_groups"] for pid in g["params"]] + cur_params = [p for g in self.param_groups for p in g["params"]] + id_map = dict(zip(saved_ids, cur_params)) + for pid, saved in state_dict["state"].items(): + p = id_map.get(pid) + if p is None: + continue + for k, v in saved.items(): + if k != "step" and torch.is_tensor(v): + self.state[p][k] = v.to(p.device, copy=True) + self._rebuild_group_index() + + # ------------------------------------------------------------------- step + + @torch.no_grad() + def step(self, closure=None): + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + if self.fused: + return loss # every param was updated inside backward + + step_seed = int(torch.randint(0, 2**31 - 1, (1,)).item()) + param_index = 0 + for group in self.param_groups: + for p in group["params"]: + param_index += 1 + accum = getattr(p, "_accum_grad", None) + if accum is not None: + p.grad = accum + del p._accum_grad + if p.grad is None: + continue + seed = (step_seed ^ (param_index * 0x9E3779B1)) & 0x7FFFFFFF + self._update_param(p, group, seed, sanitize=False) + return loss + + def _update_param(self, p, group, seed, sanitize): + if p.grad.is_sparse: + raise RuntimeError("AdamConvRot does not support sparse gradients") + beta1, beta2 = group["betas"] + state = self.state[p] + if not self._state_ok(p, state): + if len(state) > 0: + print( + "WARNING: AdamConvRot state does not match the parameter " + "or qtype; re-initializing its moments." + ) + self._init_state(p) + state["step"] += 1 + t = state["step"] + bc1 = 1.0 - beta1**t + bc2 = 1.0 - beta2**t + if self._quantizes(p): + self._step_quant(p, state, group, bc1, bc2, seed, sanitize) + else: + self._step_fp32(p, state, group, bc1, bc2, sanitize) + + def _step_fp32(self, p, state, group, bc1, bc2, sanitize=False): + beta1, beta2 = group["betas"] + lr, wd, eps = group["lr"], group["weight_decay"], group["eps"] + g = p.grad.to(torch.float32) + if sanitize: + g = g.nan_to_num(nan=0.0, posinf=0.0, neginf=0.0) + p32 = p.to(torch.float32) + if not self.adamw and wd != 0.0: + g = g.add(p32, alpha=wd) + m, v = state["exp_avg"], state["exp_avg_sq"] + m.mul_(beta1).add_(g, alpha=1.0 - beta1) + v.mul_(beta2).addcmul_(g, g, value=1.0 - beta2) + upd = (m / bc1) / ((v / bc2).sqrt_().add_(eps)) + if self.adamw and wd != 0.0: + upd.add_(p32, alpha=wd) + if p.dtype == torch.float32: + p.add_(upd, alpha=-lr) + return + p32.add_(upd, alpha=-lr) + if self.stochastic_rounding and p.is_cuda: + copy_stochastic(p, p32) + else: + p.copy_(p32) + + def _step_quant(self, p, state, group, bc1, bc2, seed, sanitize=False): + use_triton = ( + p.is_cuda + and p.grad.is_cuda + and p.is_contiguous() + and p.dtype in (torch.float32, torch.bfloat16, torch.float16) + and _triton_available() + ) + if use_triton: + self._step_quant_triton(p, state, group, bc1, bc2, seed, sanitize) + else: + self._step_quant_torch(p, state, group, bc1, bc2, sanitize) + + def _step_quant_triton(self, p, state, group, bc1, bc2, seed, sanitize=False): + beta1, beta2 = group["betas"] + g = p.grad + if not g.is_contiguous(): + g = g.contiguous() + n_blocks = state["exp_avg"].shape[0] + R = self.blocks_per_program + kernel = _get_kernel() + p_kind = {torch.float32: 0, torch.bfloat16: 1, torch.float16: 2}[p.dtype] + kernel[(-(-n_blocks // R),)]( + p, + g, + state["exp_avg"], + state["exp_avg_scale"].view(torch.float32), + state["exp_avg_sq"], + state["exp_avg_sq_scale"].view(torch.float32), + _h16(p.device), + p.numel(), + n_blocks, + float(group["lr"]), + float(group["weight_decay"]), + float(beta1), + float(beta2), + float(group["eps"]), + 1.0 / bc1, + 1.0 / math.sqrt(bc2), + math.sqrt(beta2), + float(self.rescale_qmax), + TINY, + int(seed), + BITS=self.bits, + BPOW=max(2, 1 << (self.bits - 1).bit_length()), + QMAX=self.qmax, + ADAMW=self.adamw, + P_KIND=p_kind, + SR=self.stochastic_rounding, + SANITIZE=bool(sanitize), + R=R, + num_warps=self.num_warps, + ) + + # ------------------------------------------------------------ torch path + + def _dequant(self, codes, scale_u8, n_blocks): + c = unpack_intn_rows(codes, self.bits, n_blocks, BLOCK).to(torch.float32) + return (rotate(c, BLOCK) * scale_u8.view(torch.float32)[:, None]).reshape(-1) + + def _requant(self, codes, scale_u8, decay, inc_rot): + qmax = self.qmax + scale = scale_u8.view(torch.float32) + n_blocks = codes.shape[0] + c = unpack_intn_rows(codes, self.bits, n_blocks, BLOCK).to(torch.float32) + s_dec = scale * decay + values = c * s_dec[:, None] + inc_rot + u = torch.rand_like(inc_rot) if self.stochastic_rounding else 0.5 + q = torch.floor(values / s_dec.clamp_min(TINY)[:, None] + u) + over = q.abs().amax(dim=1) > qmax + s_out = torch.where(over, values.abs().amax(dim=1) / self.rescale_qmax, s_dec) + q = torch.floor(values / s_out.clamp_min(TINY)[:, None] + u).clamp_(-qmax, qmax) + codes.copy_(pack_intn_rows(q.to(torch.int8), self.bits)) + scale.copy_(s_out) + return scale + + def _step_quant_torch(self, p, state, group, bc1, bc2, sanitize=False): + beta1, beta2 = group["betas"] + lr, wd, eps = group["lr"], group["weight_decay"], group["eps"] + numel = p.numel() + n_blocks = state["exp_avg"].shape[0] + n_pad = n_blocks * BLOCK + g = torch.zeros(n_pad, dtype=torch.float32, device=p.device) + g[:numel] = p.grad.reshape(-1).to(torch.float32) + if sanitize: + g.nan_to_num_(nan=0.0, posinf=0.0, neginf=0.0) + p32 = torch.zeros(n_pad, dtype=torch.float32, device=p.device) + p32[:numel] = p.reshape(-1).to(torch.float32) + if not self.adamw: + g.add_(p32, alpha=wd) + + m_old = self._dequant(state["exp_avg"], state["exp_avg_scale"], n_blocks) + m_new = beta1 * m_old + (1.0 - beta1) * g + self._requant( + state["exp_avg"], + state["exp_avg_scale"], + beta1, + rotate((g * (1.0 - beta1)).view(n_blocks, BLOCK), BLOCK), + ) + s_old = self._dequant(state["exp_avg_sq"], state["exp_avg_sq_scale"], n_blocks) + s_pos = s_old.clamp_min(0.0) + s_new = torch.sqrt(beta2 * s_pos * s_pos + (1.0 - beta2) * g * g) + sqrt_beta2 = math.sqrt(beta2) + ss_new = self._requant( + state["exp_avg_sq"], + state["exp_avg_sq_scale"], + sqrt_beta2, + rotate((s_new - sqrt_beta2 * s_pos).view(n_blocks, BLOCK), BLOCK), + ) + + # denominator floored at the block's code resolution (see the kernel) + s_eff = torch.maximum(s_new, ss_new.repeat_interleave(BLOCK)) + upd = (m_new / bc1) / (s_eff / math.sqrt(bc2) + eps) + if self.adamw: + upd.add_(p32, alpha=wd) + p32.add_(upd, alpha=-lr) + new = p32[:numel].view(p.shape) + if p.dtype == torch.float32: + p.copy_(new) + elif self.stochastic_rounding and p.is_cuda: + copy_stochastic(p, new) + else: + p.copy_(new) diff --git a/toolkit/optimizers/test_optimizers.py b/toolkit/optimizers/test_optimizers.py index 62edb33..70f6a9c 100644 --- a/toolkit/optimizers/test_optimizers.py +++ b/toolkit/optimizers/test_optimizers.py @@ -10,6 +10,8 @@ Compares speed (ms/step) and peak VRAM across: - Automagic v2 (fused-backward) - Automagic v3 (fused-backward and traditional/unfused) - Automagic AdamW (fused-backward and traditional/unfused) + - AdamConvRot (AdamW with ConvRot-quantized moments; 8, 4 and 2 bit, + traditional/unfused and fused-backward) - Prodigy Pass optimizer keys as CLI args to run a subset (all when none passed), e.g.: @@ -181,6 +183,7 @@ def main(): from toolkit.optimizers.automagic2 import Automagic2 from toolkit.optimizers.automagic3 import Automagic3 from toolkit.optimizers.automagicEXPERIMENT import AutomagicEXPERIMENT + from toolkit.optimizers.adamconvrot import AdamConvRot from toolkit.optimizers.adafactor import Adafactor from prodigyopt import Prodigy import bitsandbytes as bnb @@ -195,18 +198,24 @@ def main(): ("automagic3", "Automagic v3 unfused", lambda p: Automagic3(p, lr=1e-4, fused=False)), ("automagicEXPERIMENT", "AutomagicEXPERIMENT", lambda p: AutomagicEXPERIMENT(p, lr=1e-4, fused=True)), ("automagicEXPERIMENT", "AutomagicEXPERIMENT", lambda p: AutomagicEXPERIMENT(p, lr=1e-4, fused=False)), + ("adamconvrot", "AdamConvRot 8bit", lambda p: AdamConvRot(p, lr=1e-4, eps=1e-6, qtype="convrot8")), + ("adamconvrot", "AdamConvRot 4bit", lambda p: AdamConvRot(p, lr=1e-4, eps=1e-6, qtype="convrotint4")), + ("adamconvrot", "AdamConvRot 2bit", lambda p: AdamConvRot(p, lr=1e-4, eps=1e-6, qtype="convrotint2")), + ("adamconvrot", "AdamConvRot 8bit fused", lambda p: AdamConvRot(p, lr=1e-4, eps=1e-6, qtype="convrot8", fused=True)), + ("adamconvrot", "AdamConvRot 4bit fused", lambda p: AdamConvRot(p, lr=1e-4, eps=1e-6, qtype="convrotint4", fused=True)), + ("adamconvrot", "AdamConvRot 2bit fused", lambda p: AdamConvRot(p, lr=1e-4, eps=1e-6, qtype="convrotint2", fused=True)), ("prodigy", "Prodigy", lambda p: Prodigy(p, lr=1.0, eps=1e-6)), ] selected = [a.lower() for a in sys.argv[1:]] if selected: - valid = {key for key, _, _ in optimizers} + valid = {key.lower() for key, _, _ in optimizers} unknown = [a for a in selected if a not in valid] if unknown: print(f"Unknown optimizer keys: {', '.join(unknown)}") print(f"Valid keys: {', '.join(sorted(valid))}") sys.exit(1) - optimizers = [o for o in optimizers if o[0] in selected] + optimizers = [o for o in optimizers if o[0].lower() in selected] for dtype in DTYPES: dtype_name = str(dtype).replace("torch.", "")