from typing import List import torch from optimum.quanto import QBytesTensor from toolkit.optimizers.optimizer_utils import ( copy_stochastic, stochastic_grad_accummulation, ) from toolkit.util.convrot_quant import ( largest_pow4_divisor, quantize_int8_rows_fused, rotate, ) class _ConvRot8State: """ ConvRot int8 storage for optimizer moment state: values are rotated with the block regular-Hadamard transform along the trailing dims (the rotation spreads outliers so one symmetric per-row scale is safe -- same scheme as the convrot8 weight backend) and stored as int8 codes with one fp32 scale per row. Quantization runs through the fused triton kernel (quantize_int8_rows_fused): a single pass instead of a chain of eager kernels; dequantization is one scale multiply plus the (self-inverse) rotation matmul. Tensors are viewed as (shape[0], numel // shape[0]) -- per-output-channel rows, matching the convrot8 weight layout. """ def __init__(self, source, rot_size: int = 256): if isinstance(source, dict): # constructor from state dict self._load_from_state_dict(source) return self.shape = source.shape self.numel = source.numel() self.rows = source.shape[0] self.K = self.numel // self.rows rot = min(rot_size, largest_pow4_divisor(self.K)) self.rot = rot if rot >= 16 else 1 self.quantized = None self.scale = None self.quantize_(source) def quantize_(self, values_fp32: torch.Tensor): x = rotate(values_fp32.reshape(self.rows, self.K), self.rot) self.quantized, self.scale = quantize_int8_rows_fused(x) def dequantize(self) -> torch.Tensor: w = self.quantized.to(torch.float32).mul_(self.scale.unsqueeze(1)) return rotate(w, self.rot).reshape(self.shape) # self-inverse def state_dict(self): return { 'quantized': self.quantized, 'scale': self.scale, 'shape': tuple(self.shape), 'numel': self.numel, 'rot': self.rot, } def _load_from_state_dict(self, state_dict): self.quantized = state_dict['quantized'] self.scale = state_dict['scale'].to(torch.float32) self.shape = torch.Size(state_dict['shape']) self.numel = state_dict['numel'] self.rows = self.shape[0] self.K = self.numel // self.rows self.rot = state_dict['rot'] def to_(self, device): self.quantized = self.quantized.to(device) self.scale = self.scale.to(device) return self class AutomagicEXPERIMENT(torch.optim.Optimizer): """ Automagic3's learning-rate controller with a FULL per-element second moment (AdamW-style, stored in ConvRot int8) instead of Adafactor's factored approximation. Deliberately the proven v3 design, ported faithfully; validated to settle where actual Automagic3 settles. NO MOMENTUM -- beta1 must be 0, and this is a theorem of the control law, not a preference: v3's overshoot detector is perfect period-2 gradient alternation, which only exists when the applied update tracks the instantaneous gradient. Momentum smooths the step into long-period heavy-ball orbits, so an overshoot bounce produces RUNS of same-sign gradients that the sign window reads as consistency -- the controller then votes the lr UP during a bounce (verified: every acceptance scenario exploded with beta1 = 0.9). This is why automagic has always been momentum-free. The update: v = beta2 EMA of grad^2 (full per-element, bias-corrected), update = grad / (sqrt(v_hat) + eps), passed through v3's trust region (RMS-scaled then clamped elementwise to clip_threshold -- required with 8-bit state, whose under-resolved second-moment elements dequantize to zero and would otherwise produce exploding elements). Applied step: weights -= group_lr * update (+ decoupled weight decay). THE CONTROLLER (v3, verbatim): each element records the SIGN of its gradient into an H-step 1-bit ring buffer (H = polarity_history). Only the two perfectly decisive window states vote -- all H signs agreeing (either direction) votes up ("step too small"), all H-1 transitions flipping (the period-2 overshoot bounce) votes down ("step too large"); everything else is noise and votes 0. Two patterns each, identical pure-noise probability, exact balance. Votes are weighted by |update| and pooled across every element of every tensor in the group; ONE lr per group moves as lr *= exp(pooled vote), clamped only by the min_lr/max_lr failsafes (parked decades outside any operating range -- if they are ever touched, the math failed). Pooling at group level is the load-bearing choice: coupled tensors (q/k pairs) fight per-tensor lrs; one shared lr makes their opposing votes cancel in the pool. The controller is low-gain, integral, and bounded -- there is no estimator whose failure can become a training failure, the lesson of five destroyed finetune runs on a measurement-based secant controller that this file replaces. With fused=True (default) the step is fused into the backward pass via register_post_accumulate_grad_hook: each parameter is updated and its grad freed as soon as autograd finishes accumulating into it, so .step() only applies the pooled group votes. Note this bypasses the trainer's grad clipping / nan-skip (they run after backward) -- non-finite grads are neutralized in the hook instead -- and is not compatible with multi-backward gradient accumulation. With fused=False it behaves like a traditional optimizer: grads accumulate across backward passes and the update happens in .step(); low-precision grads are accumulated with stochastic rounding. State per element: the second moment in ConvRot int8 for >=2D params (stored in SQRT domain -- halves the dynamic range the linear code has to cover and is the quantity the update divides by; fp32 for 1D params) plus the H/8-byte packed sign history -- ~2 bytes/element total at the default H=8. Updates to low-precision parameters are applied in fp32 and stochastically rounded on write-back. The reported lr (get_avg_learning_rate) is the parameter-count-weighted average of the group lrs, in ABSOLUTE units, directly comparable to a classic optimizer lr. """ def __init__( self, params, lr=1e-6, # start lr; the controller adapts away from it min_lr=1e-30, # FAILSAFES only, parked decades outside any operating max_lr=1e3, # range -- if either is ever touched, the math failed betas=(0.0, 0.999), # beta1 MUST be 0: momentum breaks the sign-window eps=1e-8, # overshoot detector (see the class docstring) clip_threshold=1.0, weight_decay=0.0, polarity_history=8, # sign-window length H (2-64); H/8 bytes/element fused=True, ): self.lr = lr self.min_lr = min_lr self.max_lr = max_lr defaults = { "lr": lr, "betas": betas, "eps": eps, "clip_threshold": clip_threshold, "weight_decay": weight_decay, "polarity_history": max(2, min(64, int(polarity_history))), } super().__init__(params, defaults) self.base_lrs: List[float] = [ lr for group in self.param_groups ] self.fused = fused self.is_stochastic_rounding_accumulation = False self._hook_handles = [] self._rebuild_group_index() for group in self.param_groups: for param in group['params']: if not param.requires_grad: continue if self.fused: # Fused: update each param the moment its grad is ready. self._hook_handles.append( param.register_post_accumulate_grad_hook( self._make_backward_hook() ) ) elif param.dtype != torch.float32: # Non-fused: accumulate low-precision grads across # micro-batches with stochastic rounding; the update # happens in .step(). self.is_stochastic_rounding_accumulation = True self._hook_handles.append( param.register_post_accumulate_grad_hook( stochastic_grad_accummulation ) ) total = 0 for group in self.param_groups: for param in group['params']: total += torch.numel(param) print(f"Total training paramiters: {total:,}") # ------------------------------------------------------------------ utils def _rebuild_group_index(self): # param -> index of its param group, plus per-group vote accumulators # (gathered across every tensor in the group during the step and # applied once in .step()). The map exists because the fused hooks # cannot rely on group-dict identity: the parent's 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'] } self._group_num: List = [None] * len(self.param_groups) self._group_den: List = [None] * len(self.param_groups) def _make_backward_hook(self): def _hook(p: torch.Tensor): gi = self._param_group_index.get(p) if gi is None: self._rebuild_group_index() gi = self._param_group_index.get(p, 0) self._update_param(p, self.param_groups[gi], gi) return _hook @staticmethod def _rms(t: torch.Tensor) -> torch.Tensor: return t.norm(2) / (t.numel() ** 0.5) # Per-device cached constants for 1-bit pack/unpack (from automagic3). _PACK_CONSTS: dict = {} @classmethod def _pack_consts(cls, device): consts = cls._PACK_CONSTS.get(device) if consts is None: consts = ( torch.tensor( [1, 2, 4, 8, 16, 32, 64, 128], device=device, dtype=torch.uint8 ), torch.tensor( [0, 1, 2, 3, 4, 5, 6, 7], device=device, dtype=torch.uint8 ), ) cls._PACK_CONSTS[device] = consts return consts @classmethod def _pack_bits(cls, bits: torch.Tensor) -> torch.Tensor: # Pack sign bits (bool / {0, 1}) 8 per byte (uint8). weights, _ = cls._pack_consts(bits.device) flat = bits.reshape(-1).to(torch.uint8) pad = (-flat.numel()) % 8 if pad: flat = torch.cat([flat, flat.new_zeros(pad)]) return (flat.view(-1, 8) * weights).sum(-1, dtype=torch.uint8) # ------------------------------------------------------------ lr reporting @staticmethod def _get_lr(param_group, param_state): if 'lr' in param_state: return param_state['lr'] return 0.0 def _get_group_lr(self, group): # average weighted by parameter count total = 0.0 count = 0 for p in group["params"]: n = torch.numel(p) total = total + self._get_lr(group, self.state[p]) * n count += n if count == 0: return self.lr return total / count def step_hook(self): if not self.is_stochastic_rounding_accumulation: return # copy over stochastically rounded grads for group in self.param_groups: for param in group['params']: if param.requires_grad and hasattr(param, "_accum_grad"): param.grad = param._accum_grad del param._accum_grad # automagic manages its own lr def get_learning_rates(self): lrs = [ self._get_group_lr(group) for group in self.param_groups ] if len(lrs) == 0: lrs = self.base_lrs # if called before stepping return lrs def get_avg_learning_rate(self): # weighted by parameter count across all groups total = 0.0 count = 0 for group in self.param_groups: for p in group["params"]: n = torch.numel(p) total = total + self._get_lr(group, self.state[p]) * n count += n if count == 0: return self.lr return total / count # -------------------------------------------------------------- per-param @torch.no_grad() def _update_param(self, p, group, group_index): if p.grad is None: return grad = p.grad if grad.is_sparse: raise RuntimeError( "AutomagicAdamW does not support sparse gradients.") if grad.dtype != torch.float32: grad = grad.to(torch.float32) # In fused mode this runs inside backward, so the trainer's grad # clipping and nan/inf-skip come too late to protect us. A single # non-finite gradient would poison the moment EMAs (NaN stays NaN # forever), so neutralize non-finite grads; those elements contribute # nothing this step. grad = grad.nan_to_num_(nan=0.0, posinf=0.0, neginf=0.0) state = self.state[p] if len(state) == 0: self.initialize_state(p, group) if ('exp_avg_sq' not in state or 'sign_history' not in state or 'lr' not in state): self.initialize_state(p, group) state["step"] += 1 step = state["step"] beta1, beta2 = group["betas"] bias_correction2 = 1.0 - beta2 ** step if beta1 != 0.0: raise ValueError( "AutomagicAdamW requires beta1 = 0: momentum smooths the " "applied step into long-period heavy-ball orbits, so an " "overshoot bounce produces RUNS of same-sign gradients " "instead of period-2 alternation -- the sign-window " "controller then reads the bounce as consistency and votes " "the lr UP. v3's control law is only valid when the update " "tracks the instantaneous gradient." ) quantized = isinstance(state["exp_avg_sq"], _ConvRot8State) if quantized: # 8-bit second moment (see _ConvRot8State); the EMA math runs # on an fp32 dequantized copy which is requantized right after. # The store holds sqrt(v), so square it back for the EMA -- # clamping first: rotation quant noise can leave tiny negatives, # and squaring them would bias v upward. v = state["exp_avg_sq"].dequantize().clamp_(min=0.0).square_() v.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) v_sqrt = v.sqrt_() # v is not needed past here; sqrt in place state["exp_avg_sq"].quantize_(v_sqrt) denom = v_sqrt.div_(bias_correction2 ** 0.5).add_(group["eps"]) else: # 1D params: plain fp32 second moment v = state["exp_avg_sq"] v.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) denom = (v / bias_correction2).sqrt_().add_(group["eps"]) update = grad / denom # v3's trust region: scale the update so its RMS is <= clip_threshold, # then clamp each element to +/-clip_threshold. NOT optional with # 8-bit moments: elements whose sqrt(v) falls below the int8 # resolution dequantize to zero, so their denominator collapses to # eps and the raw update explodes by orders of magnitude. update.div_( (self._rms(update) / group["clip_threshold"]).clamp_(min=1.0)) update.clamp_(-group["clip_threshold"], group["clip_threshold"]) # ---- the vote bits: v3's rule, verbatim, on GRADIENT signs ---- # Adafactor's update sign IS the gradient sign (no momentum), so # v3's polarity bits were always gradient-sign bits; porting them # unchanged keeps every property that made v3 robust. A window of H # consistent signs (all ones or all zeros -- direction-agnostic) # votes up; perfect alternation (the overshoot bounce: the gradient # flips every step regardless of momentum smoothing in the applied # update) votes down; anything mixed is noise. Two patterns each, # identical pure-noise probability, exact balance. Weighted by # |update| and pooled into the group accumulators; the single group # lr is nudged once per step in .step(). (A hypergradient bit # against the momentum direction was tried here and exploded every # scenario: momentum holds the reference direction through a bounce, # so the overshoot signature never fires. Gauge immunity comes from # POOLING -- opposing votes cancel -- not from the bit definition.) H = group["polarity_history"] hist = state["sign_history"] # (H, numel/8) 1-bit packed uint8 cur_bits = grad > 0 idx = state["hist_idx"] hist[idx].copy_(self._pack_bits(cur_bits)) state["hist_idx"] = (idx + 1) % H fill = min(H, state["hist_fill"] + 1) state["hist_fill"] = fill if fill == H: _, shifts = self._pack_consts(hist.device) chron = torch.roll(hist, -state["hist_idx"], dims=0) bits = ( (chron.unsqueeze(-1) >> shifts) .bitwise_and_(1) .view(H, -1)[:, : update.numel()] ) s1 = bits.sum(0, dtype=torch.int16) flips = (bits[1:] ^ bits[:-1]).sum(0, dtype=torch.int16) up = s1.eq(H).logical_or_(s1.eq(0)) down = flips.eq(H - 1) w = update.abs().view(-1) num = (w * up).sum().sub_((w * down).sum()) den = w.sum() if self._group_num[group_index] is None: self._group_num[group_index] = num self._group_den[group_index] = den else: acc = self._group_num[group_index] if num.device != acc.device: num = num.to(acc.device) den = den.to(acc.device) acc.add_(num) self._group_den[group_index].add_(den) lr_t = state['lr'] # this param's mirror of the shared group lr p_data_fp32 = p if isinstance(p_data_fp32, QBytesTensor): p_data_fp32 = p_data_fp32.dequantize() if p.dtype != torch.float32: p_data_fp32 = p_data_fp32.clone().float() if group["weight_decay"] != 0: # Decoupled weight decay folded into the direction: # p -= lr * (update + weight_decay * p) update.add_(p_data_fp32, alpha=group["weight_decay"]) p_data_fp32.addcmul_(update, lr_t, value=-1.0) if p.dtype != torch.float32: # apply stochastic rounding copy_stochastic(p, p_data_fp32) p.grad = None # ----------------------------------------------------------- optimizer API @torch.no_grad() def step(self, closure=None): loss = None if closure is not None: with torch.enable_grad(): loss = closure() # Fused mode already updated every param inside the backward pass; # only the pooled group votes remain. if not self.fused: self.step_hook() for group in self.param_groups: for p in group["params"]: if p.grad is None or not p.requires_grad: continue gidx = self._param_group_index.get(p) if gidx is None: self._rebuild_group_index() gidx = self._param_group_index.get(p, 0) self._update_param(p, group, gidx) self._apply_group_votes() return loss def _apply_group_votes(self): # ONE lr nudge per group per step, from the pooled vote of every # element of every tensor in the group. Each param's lr tensor # receives the same multiplicative factor, so they stay identical -- # effectively a single group lr, stored per param only so it rides # the normal state_dict machinery. All tensor ops: no GPU sync. for gi, group in enumerate(self.param_groups): num = self._group_num[gi] if num is None: continue den = self._group_den[gi] signal = num.div_(den.clamp_(min=1e-30)).clamp_(-1.0, 1.0) factor = torch.exp(signal) for p in group["params"]: st = self.state.get(p) if st is None or 'lr' not in st: continue lr_t = st['lr'] f = factor if factor.device == lr_t.device \ else factor.to(lr_t.device) lr_t.mul_(f).clamp_(min=self.min_lr, max=self.max_lr) self._group_num[gi] = None self._group_den[gi] = None # ------------------------------------------------------------------- state def initialize_state(self, p, group=None): if group is None: gi = self._param_group_index.get(p, 0) group = self.param_groups[gi] state = self.state[p] if "step" not in state: state["step"] = 0 if 'lr' not in state: # the group lr, mirrored per param (identical multiplicative # nudges keep them equal) state['lr'] = torch.tensor( min(max(float(self.lr), self.min_lr), self.max_lr), dtype=torch.float32, device=p.device) H = group["polarity_history"] width = (p.numel() + 7) // 8 if 'sign_history' not in state: state['sign_history'] = torch.zeros( (H, width), dtype=torch.uint8, device=p.device) state['hist_idx'] = 0 state['hist_fill'] = 0 if 'exp_avg_sq' not in state: zeros = torch.zeros(p.shape, dtype=torch.float32, device=p.device) if p.dim() >= 2: # holds sqrt(exp_avg_sq) -- see the class docstring state["exp_avg_sq"] = _ConvRot8State(zeros) else: # 1D params: plain fp32 buffer, negligible memory state["exp_avg_sq"] = zeros # keys stored as quantized-tensor objects, serialized via their own # state_dicts and restored by hand in load_state_dict _QUANT_KEYS = ('exp_avg_sq',) def state_dict(self, *args, **kwargs): orig_state_dict = super().state_dict(*args, **kwargs) new_save_state = {} for p, state in orig_state_dict['state'].items(): save_state = {k: v for k, v in state.items() if k not in self._QUANT_KEYS} for key in self._QUANT_KEYS: if key in state: val = state[key] save_state[key] = ( val if isinstance(val, torch.Tensor) else val.state_dict() ) new_save_state[p] = save_state orig_state_dict['state'] = new_save_state return orig_state_dict def load_state_dict(self, state_dict, strict=True): # Validate the format; older formats start fresh. is_valid_state = False if 'state' in state_dict and isinstance(state_dict['state'], dict): for param_id, param_state in state_dict['state'].items(): if isinstance(param_state, dict) and 'sign_history' in param_state: is_valid_state = True break if not is_valid_state: return # Parent load without the quantized buffers (its recursive cast would # convert their fp32 scales to the param dtype). state_dict_copy = { 'state': {}, 'param_groups': state_dict['param_groups'] } for param_id, param_state in state_dict['state'].items(): state_dict_copy['state'][param_id] = { k: v for k, v in param_state.items() if k not in self._QUANT_KEYS } super().load_state_dict(state_dict_copy) # Hyperparameters are NOT loaded from the checkpoint: constructor # args always win (any setting can be changed mid-run by resuming # with a different value). Only adaptive state is restored. for group in self.param_groups: for k, v in self.defaults.items(): group[k] = v self._rebuild_group_index() current_params = [ p for group in self.param_groups for p in group['params'] ] saved_param_count = sum( len(g['params']) for g in state_dict['param_groups']) if len(current_params) != saved_param_count: print(f"WARNING: Number of parameters doesn't match between saved state ({saved_param_count}) " f"and current model ({len(current_params)}). Optimizer state may not be correctly loaded.") # One lr per group: unify restored lrs to their geometric median # (identical already for checkpoints from this version). for group in self.param_groups: lrs = [ st['lr'] for p in group['params'] if (st := self.state.get(p)) is not None and isinstance(st.get('lr'), torch.Tensor) ] med = None if lrs: dev = lrs[0].device med = ( torch.stack([t.to(torch.float32).to(dev) for t in lrs]) .log_() .median() .exp_() ) H = group["polarity_history"] for p in group['params']: st = self.state.get(p) if st is None: continue if isinstance(st.get('lr'), torch.Tensor): st['lr'] = st['lr'].to(torch.float32) if med is not None: st['lr'].copy_(med.to(st['lr'].device)) # Sign history: keep when its geometry matches the current # config; otherwise start fresh (one re-warmup of H steps). numel = p.numel() width = (numel + 7) // 8 sh = st.get('sign_history') hist_ok = ( isinstance(sh, torch.Tensor) and sh.shape == (H, width) and isinstance(st.get('hist_idx'), int) and 0 <= st['hist_idx'] < H and isinstance(st.get('hist_fill'), int) and 0 <= st['hist_fill'] <= H ) if hist_ok: st['sign_history'] = sh.to(torch.uint8) else: st['sign_history'] = torch.zeros( (H, width), dtype=torch.uint8, device=p.device) st['hist_idx'] = 0 st['hist_fill'] = 0 for saved_param_id, saved_state in state_dict['state'].items(): if 'sign_history' not in saved_state: continue if not isinstance(saved_param_id, int) or not (0 <= saved_param_id < len(current_params)): continue i = saved_param_id current_param = current_params[i] if current_param not in self.state: self.initialize_state(current_param) current_state = self.state[current_param] # Reconstruct the quantized buffers: 8-bit ConvRot dicts for >=2D # params, plain fp32 tensors for 1D params for key in self._QUANT_KEYS: saved_buf = saved_state.get(key) restored = None if saved_buf is not None: try: if ( isinstance(saved_buf, dict) and saved_buf.get('numel') == current_param.numel() ): restored = _ConvRot8State( saved_buf).to_(current_param.device) elif ( isinstance(saved_buf, torch.Tensor) and saved_buf.shape == current_param.shape ): restored = saved_buf.to( device=current_param.device, dtype=torch.float32) else: print(f"WARNING: Could not restore {key} for parameter {i}. " f"Initializing fresh.") except Exception as e: print(f"ERROR: Failed to load {key} for parameter {i}: {e}") if restored is None: zeros = torch.zeros( current_param.shape, dtype=torch.float32, device=current_param.device) restored = ( _ConvRot8State(zeros) if current_param.dim() >= 2 else zeros ) current_state[key] = restored