adv-optm 1.2.0__py3-none-any.whl

adv_optm/optim/Simplified_AdEMAMix.py

@@ -0,0 +1,298 @@
+import torch
+from typing import Optional, Callable
+
+import math
+
+from ..util.BF16_Stochastic_Rounding import add_stochastic_
+from ..util.Effective_Shape import _get_effective_shape
+from ..util.NNMF import _nnmf,_unnmf
+from ..util.OrthoGrad import _orthogonalize_gradient
+from ..util.One_Bit_Boolean import _pack_bools, _unpack_bools
+from ..util.Kourkoutas import KourkoutasHelper
+
+# A little helper from the original simplified_AdEMAMix
+def linear_hl_warmup_scheduler(step, beta_end, beta_start=0, warmup=1):
+
+    def f(beta, eps=1e-8):
+        return math.log(0.5)/math.log(beta+eps)-1
+
+    def f_inv(t):
+        return math.pow(0.5, 1/(t+1))
+
+    if step < warmup:
+        a = step / float(warmup)
+        return f_inv((1.0-a) * f(beta_start) + a * f(beta_end))
+    return beta_end
+
+class Simplified_AdEMAMix(torch.optim.Optimizer):
+    """
+    Implements the Simplified AdEMAMix algorithm.
+    Refactored from:
+    https://github.com/DepenM/Simplified-AdEMAMix/blob/main/simplified_AdEMAMix.py
+
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float): learning rate (default: 1e-5)
+        betas (tuple[float, float]): coefficients used for computing running
+            averages of gradient and its square (default: (0.99, 0.999))
+        eps (float): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float): weight decay (L2 penalty) (default: 0).
+        alpha_grad (float): Coeficient for mixing the current gradient and EMA. for small batch
+            sizes set it to high values, up to 100. And for large batch sized set it to small
+            value, down to 0. (default: 100)
+        beta1_warmup (int, optional): number of warmup steps used to increase beta1 (default: None)
+        min_beta1 (float, optional): minimum value of beta1 to start from (default 0.9)
+        vector_reshape (bool): whether to reshape 1D vectors into 2D
+            matrices to apply low-rank compression (default: True).
+        stochastic_rounding (bool): whether to use stochastic
+            rounding for BF16 parameter updates (default: True).
+        orthogonal_gradient (bool): whether to use OrthoGrad. (default: False)
+        kourkoutas_beta (bool): whether to enable the layer-wise dynamic β₂ logic.
+            If `False`, the optimizer behaves as standard Simplified_AdEMAMix. (default: False)
+        beta2_min (float): The minimum value for dynamic β₂, used during periods of
+            high gradient variance ("sunspikes"). Must be less than `betas[1]`.
+            (default: 0.88)
+        ema_alpha (float): The decay rate for the Exponential Moving Average (EMA) of
+            the pooled gradient norms. Corresponds to `α` in the paper.
+            (default: 0.93)
+        tiny_spike (float): A small constant added to the denominator of the
+            "sunspike" ratio calculation to prevent division by zero. Corresponds
+            to `ε_spike` in the paper. (default: 1e-9)
+        k_warmup_steps (int): The number of initial steps during which β₂ is held
+            at a fixed beta2 value before the
+            dynamic logic activates. (default: 0)
+        k_logging (int): if > 0 and kourkoutas_beta=True, enables periodic console
+            logging of Kourkoutas-β statistics (min, max, mean of `β₂` across layers)
+            every logging steps. Useful for debugging and tuning. Set to 0 to disable
+            logging (default: 0). 
+        layer_key_fn (Optional[Callable]): A function that takes a parameter `p`
+            and returns a unique, hashable key representing its "layer" or "bucket".
+            If `None`, parameters are bucketed by their memory ID (tensor-wise).
+            (default: None)
+        nnmf_factor (bool): whether to use the factorization or disable it to use
+            the uncompressed optimizer. (default: False)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-5,
+        betas: tuple[float, float] = (0.99, 0.999),
+        eps: float = 1e-8,
+        weight_decay: float = 0.0,
+        alpha_grad: float = 100.0,
+        beta1_warmup: int | None = None,
+        min_beta1: float | None = 0.9,
+        use_bias_correction: bool = True,
+        vector_reshape: bool = True,
+        stochastic_rounding: bool = True,
+        orthogonal_gradient: bool = False,
+        kourkoutas_beta: bool = False,
+        beta2_min: float = 0.9,
+        ema_alpha: float = 0.95,
+        tiny_spike: float = 1e-9,
+        k_warmup_steps: int = 0,
+        k_logging: int = 0,
+        layer_key_fn: Optional[Callable] = None,
+        nnmf_factor: bool = False,
+    ):
+        if not (lr >= 0.0):
+            raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
+        if not (0.0 <= betas[0] < 1.0 and 0.0 <= betas[1] < 1.0):
+            raise ValueError(f"Betas should be in [0.0, 1.0). Got {betas}")
+        if not (eps >= 0.0):
+            raise ValueError(f"Epsilon should be >= 0.0. Got {eps}")
+        if not (weight_decay >= 0.0):
+            raise ValueError(f"Weight-decay should be >= 0.0. Got {weight_decay}")
+        if not 0.0 <= alpha_grad:
+            raise ValueError("Invalid alpha value: {}".format(alpha_grad))
+        if kourkoutas_beta and not (betas[1] > beta2_min): raise ValueError(f"For Kourkoutas-β, betas[1] (as beta2_max) must be > beta2_min. Got {betas[1]} and {beta2_min}")
+
+        defaults = {
+            "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay,
+            "alpha_grad": alpha_grad, "beta1_warmup": beta1_warmup, "min_beta1": min_beta1,
+            "vector_reshape": vector_reshape,
+            "orthogonal_gradient": orthogonal_gradient, "use_bias_correction": use_bias_correction,
+            "kourkoutas_beta": kourkoutas_beta, "beta2_min": beta2_min, "ema_alpha": ema_alpha,
+            "tiny_spike": tiny_spike, "k_warmup_steps": k_warmup_steps, "k_logging": k_logging,
+        }
+        self.stochastic_rounding = stochastic_rounding
+        self.factored = nnmf_factor
+        self.kourkoutas_beta = kourkoutas_beta
+        self.layer_key_fn = layer_key_fn
+        super().__init__(params, defaults)
+
+        if self.kourkoutas_beta:
+            self.kourkoutas_helper = KourkoutasHelper(self)
+
+    @property
+    def supports_fused_back_pass(self):
+        return True
+
+    @property
+    def supports_memory_efficient_fp16(self):
+        return True
+
+    @property
+    def supports_flat_params(self):
+        return False
+
+    @torch.no_grad()
+    def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):
+        if p.grad is None:
+            return
+
+        grad = p.grad
+        if grad.dtype != torch.float32 and self.factored:
+            grad = grad.float()
+        if group["orthogonal_gradient"]:
+            grad = _orthogonalize_gradient(p, grad)
+        state = self.state[p]
+
+        # State Initialization
+        if 'step' not in state:
+            state['step'] = 0
+
+            should_factor = (
+                self.factored and
+                not (len(p.shape) == 1 and not group['vector_reshape'])
+            )
+
+            state['factored'] = should_factor
+
+            dtype = torch.float32 if self.factored else p.dtype
+            device = p.device
+
+            if state['factored']:
+                state['effective_shape'] = _get_effective_shape(p.numel())
+                d1, d2 = state['effective_shape']
+
+                # First moment (m)
+                state['mu_m_nmf'] = torch.zeros(d1, device=device, dtype=dtype) 
+                state['mv_m_nmf'] = torch.zeros(d2, device=device, dtype=dtype)
+                packed_d2 = (d2 + 7) // 8
+                state['sign'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=device)
+                # Second moment (v)
+                state['mu_v_nmf'] = torch.zeros(d1, device=device, dtype=dtype) 
+                state['mv_v_nmf'] = torch.zeros(d2, device=device, dtype=dtype)
+            else:  # Fallback to standard optimizer for non-factored tensors
+                state['exp_avg'] = torch.zeros_like(p, device=device, dtype=dtype)
+                state['exp_avg_sq'] = torch.zeros_like(p, device=device, dtype=dtype)
+            
+            if group['use_bias_correction']:
+                state['num_sum'] = 0.0
+                state['den_sum'] = 0.0
+            else:
+                state['num_sum'] = 1.0
+                state['den_sum'] = 1.0
+
+        beta1_final, beta2 = group["betas"]
+
+        current_step = state['step']
+        if group.get('kourkoutas_beta', False):
+            # Call prepare_step() once at the beginning of the step for all params
+            self.kourkoutas_helper.maybe_prepare_step(current_step)
+            # Accumulate current grad's norm for the *next* step
+            self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
+            # Get the dynamic beta2 calculated in prepare_step()
+            beta2 = self.kourkoutas_helper.get_beta2(p, group, current_step)
+
+        beta1_warmup = group["beta1_warmup"]
+        alpha_grad = group["alpha_grad"]
+
+        if beta1_warmup is not None:
+            step = state['step'] + 1
+            beta1 = linear_hl_warmup_scheduler(step, beta_end=beta1_final, beta_start=group['min_beta1'], warmup=beta1_warmup)
+        else:
+            beta1 = beta1_final
+
+        if group['use_bias_correction']:
+            state['num_sum'] = beta1 * state['num_sum'] + 1.0
+            if group.get('kourkoutas_beta', False):
+                state['den_sum'] = group['betas'][1] * state['den_sum'] + (1.0 - group['betas'][1])
+            else:
+                state['den_sum'] = beta2 * state['den_sum'] + (1.0 - beta2)
+
+        if state['factored']:
+            d1, d2 = state['effective_shape']
+
+            # Reconstruct momentum from previous step's factors
+            mt = _unnmf((state['mu_m_nmf'], state['mv_m_nmf']))
+            unpacked_sign = _unpack_bools(state['sign'], original_m=d2)
+            torch.where(unpacked_sign, mt, -mt, out=mt)
+            del unpacked_sign
+            # Update momentum in full-size
+            grad_reshaped = grad.view(d1, d2)
+            mt.mul_(beta1).add_(grad_reshaped, alpha=1.0)
+
+            vt = _unnmf((state['mu_v_nmf'], state['mv_v_nmf']))
+            vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped, value=1.0 - beta2)
+
+            update = torch.add(mt, grad_reshaped, alpha=alpha_grad)
+            del grad_reshaped
+
+            denom = vt.sqrt().add_(group['eps'] * math.sqrt(state['den_sum']))
+            update.div_(denom)
+            del denom
+
+            if group['use_bias_correction']:
+                update = (update / state['num_sum']) * math.sqrt(state['den_sum'])
+
+            update = update.view(p.shape).mul_(group['lr'])
+
+            # Compress updated moments and store new factors
+            state['sign'] = _pack_bools(mt > 0)
+            _nnmf(mt.abs(), out=(state['mu_m_nmf'], state['mv_m_nmf']))
+            del mt
+            _nnmf(vt, out=(state['mu_v_nmf'], state['mv_v_nmf']))
+            del vt
+
+        else:  # Standard optimizer logic for non-factored tensors
+            exp_avg_sq = state['exp_avg_sq']
+
+            exp_avg = state['exp_avg']
+            exp_avg.mul_(beta1).add_(grad, alpha=1.0)
+
+            update = torch.add(exp_avg, grad, alpha=alpha_grad)
+
+            exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+            denom = exp_avg_sq.sqrt().add_(group['eps'] * math.sqrt(state['den_sum']))
+            update.div_(denom)
+            del denom
+
+            if group['use_bias_correction']:
+                update = (update / state['num_sum']) * math.sqrt(state['den_sum'])
+
+            update.mul_(group['lr'])
+
+        # Decoupled weight decay
+        if group["weight_decay"] != 0:
+            if p.dtype == torch.bfloat16 and self.stochastic_rounding:
+                add_stochastic_(p.data, p.data, alpha=-group["weight_decay"] * group["lr"])
+            else:
+                p.data.add_(p.data, alpha=-group["weight_decay"] * group["lr"])
+
+        if p.dtype == torch.bfloat16 and self.stochastic_rounding:
+            add_stochastic_(p.data, -update)
+        else:
+            p.data.add_(-update)
+        del update
+
+        state['step'] += 1
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step."""
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for i, p in enumerate(group['params']):
+                self.step_parameter(p, group, i)
+
+        return loss

adv_optm/optim/__init__.py

@@ -0,0 +1,19 @@
+from .AdamW_adv import AdamW_adv
+from .Prodigy_adv import Prodigy_adv
+from .Adopt_adv import Adopt_adv
+from .Simplified_AdEMAMix import Simplified_AdEMAMix
+from .Lion_adv import Lion_adv
+from .Lion_Prodigy_adv import Lion_Prodigy_adv
+from .Muon_adv import Muon_adv
+from .AdaMuon_adv import AdaMuon_adv
+
+__all__ = [
+    "AdamW_adv",
+    "Prodigy_adv",
+    "Adopt_adv",
+    "Simplified_AdEMAMix",
+    "Lion_adv",
+    "Lion_Prodigy_adv",
+    "Muon_adv",
+    "AdaMuon_adv",
+]

adv_optm/util/BF16_Stochastic_Rounding.py

@@ -0,0 +1,47 @@
+import torch
+from torch import Tensor
+
+def copy_stochastic_(target: Tensor, source: Tensor):
+    """
+    Nerogar's implementation of stochastic rounding in the paper "Revisiting BFloat16 Training"
+    (https://arxiv.org/abs/2010.06192).
+    see:
+    https://github.com/pytorch/pytorch/issues/120376
+    https://github.com/Nerogar/OneTrainer/blob/daae18eaed8c0fa39289b2ff79cc2c1e08577fcb/modules/util/bf16_stochastic_rounding.py
+
+    Args:
+        target: the target tensor with dtype=bfloat16
+        source: the target tensor with dtype=float32
+    """
+    # create a random 16 bit integer
+    result = torch.randint_like(
+        source,
+        dtype=torch.int32,
+        low=0,
+        high=(1 << 16),
+    )
+
+    # add the random number to the lower 16 bit of the mantissa
+    result.add_(source.view(dtype=torch.int32))
+
+    # mask off the lower 16 bit of the mantissa
+    result.bitwise_and_(-65536)  # -65536 = FFFF0000 as a signed int32
+
+    # copy the higher 16 bit into the target tensor
+    target.copy_(result.view(dtype=torch.float32))
+
+    del result
+
+def add_stochastic_(input: Tensor, other: Tensor, alpha: float = 1.0):
+    """
+    adds other to input using stochastic rounding
+
+    Args:
+        input: the input tensor with dtype=bfloat16
+        other: the other tensor
+        alpha: a multiplier for other
+    """
+    result = other.clone() if other.dtype == torch.float32 else other.to(dtype=torch.float32)
+
+    result.add_(input, alpha=alpha)
+    copy_stochastic_(input, result)

adv_optm/util/Effective_Shape.py

@@ -0,0 +1,8 @@
+def _get_effective_shape(numel: int) -> tuple[int, int]:
+    """Finds two factors of numel that are closest to its square root."""
+    if numel <= 0:
+        return (0, 0)
+    for i in reversed(range(1, int(numel ** 0.5) + 1)):
+        if numel % i == 0:
+            return (numel // i, i)
+    return (numel, 1)

adv_optm/util/Kourkoutas.py

@@ -0,0 +1,159 @@
+import torch
+from torch.optim import Optimizer
+from typing import Callable
+
+class KourkoutasHelper:
+    """
+    A helper class to add layer-wise Kourkoutas-β functionality to a PyTorch optimizer.
+    """
+    def __init__(self, optimizer: Optimizer):
+        # We need a reference to the optimizer to access its param_groups and state
+        if not hasattr(optimizer, 'param_groups'):
+            raise TypeError("optimizer must be a valid torch.optim.Optimizer instance.")
+        self.optimizer = optimizer
+        self.layer_state = {}
+
+        self.layer_info = {}
+        self._layer_info_built = False
+        self._current_step_prepared = -1
+
+        # Store stats for external logging (e.g., TensorBoard)
+        self.last_beta2_stats = {}
+
+        # This ensures the map is complete before the first backward pass,
+        # making it compatible with fused back pass mechanisms.
+        self._build_layer_info_if_needed()
+
+    def _build_layer_info_if_needed(self):
+        """Builds a map of layers and the parameters they contain."""
+        if self._layer_info_built:
+            return
+
+        if hasattr(self.optimizer, 'layer_key_fn') and self.optimizer.layer_key_fn is not None:
+            # A custom key function was provided by the user. We will use it.
+            pass
+        else:
+            # No key function was provided. Default to coarse, shape-based bucketing.
+            self.optimizer.layer_key_fn = lambda p: \
+                (id(p),) if p.dim() == 2 and 1 <= p.shape[0] <= 10 and p.shape[1] in {768, 1280, 4096} \
+                else tuple(p.shape)
+            # This ensures that we won't mix embeddings with tokens (1 to 10)
+            # TODO find a better way to safeguard the embeddings
+
+        for group in self.optimizer.param_groups:
+            if not group.get('kourkoutas_beta', False) and not group.get('adam_kourkoutas_beta', False):
+                continue
+
+            for p in group['params']:
+                # The mapping is static and should not depend on the presence of a gradient.
+                layer_key = self.optimizer.layer_key_fn(p)
+                if layer_key not in self.layer_info:
+                    self.layer_info[layer_key] = {'params': [], 'group_ref': group}
+                self.layer_info[layer_key]['params'].append(p)
+
+        self._layer_info_built = True
+
+    def prepare_step(self, current_step: int):
+        """
+        Calculates dynamic beta2 for all layers using the completed scalar accumulators
+        from the PREVIOUS step. Should be called once at the start of an optimizer step.
+        """
+        
+        beta2_log = []
+        # These are just for the sample log, initialize them
+        sun, pooled_grad_norm, r_ema_tensor = (torch.tensor(0.0),)*3
+
+        # The optimizer that owns this helper holds the master defaults for K-b.
+        # This is crucial in hybrid optimizers where some param_groups might not
+        # have all K-b keys populated, preventing KeyErrors.
+        master_defaults = self.optimizer.defaults
+
+        for layer_key, info in self.layer_info.items():
+            params, group = info['params'], info['group_ref']
+
+            if not group.get('kourkoutas_beta', False) and not group.get('adam_kourkoutas_beta', False):
+                continue
+
+            first_param_in_layer = info['params'][0]
+            param_state = self.optimizer.state[first_param_in_layer]
+
+            if layer_key not in self.layer_state:
+                self.layer_state[layer_key] = {
+                    'sum_sq_accumulator': torch.tensor(0.0, device=first_param_in_layer.device, dtype=torch.float32)
+                }
+            
+            if 'kourkoutas_r_ema' not in param_state:
+                param_state['kourkoutas_r_ema'] = torch.tensor(0.0, device=first_param_in_layer.device, dtype=torch.float32)
+
+            # Use group-specific K-b settings, falling back to the optimizer's master defaults.
+            # This makes the helper robust against param groups that enable kourkoutas_beta
+            # but are missing the other required hyperparameters.
+            ema_alpha = group.get('ema_alpha', master_defaults['ema_alpha'])
+            beta2_max = group.get('betas', master_defaults['betas'])[1]
+            beta2_min = group.get('beta2_min', master_defaults['beta2_min'])
+            tiny_spike = group.get('tiny_spike', master_defaults['tiny_spike'])
+            k_warmup_steps = group.get('k_warmup_steps', master_defaults['k_warmup_steps'])
+
+            r_ema_tensor = param_state['kourkoutas_r_ema']
+            accumulator = self.layer_state[layer_key]['sum_sq_accumulator']
+            
+            pooled_grad_norm = torch.sqrt(accumulator)
+            
+            # Update the persistent EMA tensor in-place.
+            r_ema_tensor.mul_(ema_alpha).add_(pooled_grad_norm, alpha=1.0 - ema_alpha)
+            
+            sun = torch.tensor(0.0, device=r_ema_tensor.device) # Default sun to 0 for warmup
+            
+            if current_step < k_warmup_steps:
+                beta2 = beta2_max
+            else:
+                raw = pooled_grad_norm / (r_ema_tensor + tiny_spike)
+                sun = raw / (1.0 + raw)
+                beta2 = beta2_max - (beta2_max - beta2_min) * sun
+
+            # Store the final calculated beta2 in the helper's transient state for this step.
+            self.layer_state[layer_key]['dynamic_beta2'] = beta2.item() if isinstance(beta2, torch.Tensor) else beta2
+            
+            # Reset the accumulator for the next optimizer step.
+            accumulator.zero_()
+
+            beta2_log.append(self.layer_state[layer_key]['dynamic_beta2'])
+
+        # Always compute stats for TensorBoard
+        if beta2_log:
+            beta2_tensor = torch.tensor(beta2_log, device='cpu')
+            self.last_beta2_stats = {
+                'mean': beta2_tensor.mean().item(),
+            }
+
+    def maybe_prepare_step(self, current_step: int):
+        """
+        A universal guard that calls prepare_step() exactly once per training step.
+        """
+        if self._current_step_prepared < current_step:
+            self.prepare_step(current_step)
+            self._current_step_prepared = current_step
+
+    def accumulate_gradient_sq_norm(self, p: torch.Tensor, grad: torch.Tensor):
+        """
+        Accumulates the squared L2 norm of a single gradient for the next step's calculation.
+        """
+        layer_key = self.optimizer.layer_key_fn(p)
+
+        if layer_key in self.layer_info:
+            # Initialize the transient state for this layer if it's the first time in the step.
+            if layer_key not in self.layer_state:
+                    self.layer_state[layer_key] = {
+                    'sum_sq_accumulator': torch.tensor(0.0, device=p.device, dtype=torch.float32)
+                }
+            # Accumulate for the *next* step's prepare_step call
+            self.layer_state[layer_key]['sum_sq_accumulator'] += torch.sum(grad.detach().pow(2)).float()
+
+    def get_beta2(self, p: torch.Tensor, group: dict, current_step: int) -> float:
+        """
+        Gets the appropriate beta2 for the current parameter, handling warmup and dynamic value fetching.
+        """
+        layer_key = self.optimizer.layer_key_fn(p)
+        # The default is the max value, which is correct for unmapped params or edge cases
+        beta2_default = group.get('betas', group.get('adam_betas'))[1] if group.get('betas', group.get('adam_betas')) else 0.999
+        return self.layer_state.get(layer_key, {}).get('dynamic_beta2', beta2_default)

adv_optm/util/NNMF.py

@@ -0,0 +1,18 @@
+import torch
+
+def _unnmf(row_col: tuple) -> torch.Tensor:
+    """Reconstructs a matrix from its rank-1 factors (outer product)."""
+    return torch.outer(row_col[0], row_col[1])
+
+def _nnmf(matrix: torch.Tensor, out: tuple):
+    """Performs a rank-1 non-negative matrix factorization."""
+    shape = matrix.shape
+    torch.sum(matrix, dim=1, out=out[0])
+    torch.sum(matrix, dim=0, out=out[1])
+    # Normalize one of the factors for stability
+    if shape[0] < shape[1]:
+        scale = out[0].sum()
+        if scale != 0: out[0].div_(scale)
+    else:
+        scale = out[1].sum()
+        if scale != 0: out[1].div_(scale)

adv_optm/util/Newton_Schulz.py

@@ -0,0 +1,87 @@
+import torch
+
+@torch.no_grad()
+def _newton_schulz_iteration(
+    G: torch.Tensor,
+    steps: int = 5,
+    eps: float = 1e-7,
+    coeffs: tuple[float, float, float] = (3.4445, -4.7750, 2.0315),
+    cns: bool = False,
+    cns_a_bound: float = 1e-4,
+) -> torch.Tensor:
+    """
+    Performs the Newton-Schulz iteration to find the nearest orthogonal matrix.
+    This is the core computation of the Muon optimizer.
+
+    Args:
+        G (torch.Tensor): The 2D input matrix (momentum-accumulated gradient).
+        steps (int): The number of iterations to run.
+        eps (float): Small constant for numerical stability during normalization.
+        coeffs (tuple[float, float, float]): The (a, b, c) coefficients for the
+            quintic polynomial update.
+        cns (bool): If True, enables Chebyshev-accelerated Newton-Schulz (CANS)
+            using an iterative 3rd-order polynomial with optimal coefficients
+            derived at each step.
+        cns_a_bound (float): The initial lower bound for singular values when
+            using CANS. The upper bound is assumed to be 1.0 after normalization.
+    Returns:
+        torch.Tensor: The orthogonalized matrix.
+    """
+    assert G.ndim >= 2
+
+    a, b, c = coeffs
+
+    X = G.to(torch.bfloat16)
+
+    transposed = G.size(-2) > G.size(-1)
+    if transposed:
+        X = X.mT 
+
+    X = X / (X.norm(dim=(-2, -1), keepdim=True) + eps)
+
+    if cns:
+        # Chebyshev-accelerated Newton-Schulz (CANS) from
+        # "Accelerating Newton-Schulz Iteration for Orthogonalization via Chebyshev-type Polynomials"
+        # This implements the iterative scheme from Algorithm 1, using the
+        # closed-form 3rd-order polynomial from Proposition 2.
+        lower_bound = cns_a_bound
+        upper_bound = 1.0  # Matrix is normalized, so largest singular value is approx 1.
+
+        for _ in range(steps):
+            # Calculate optimal 3rd-order coefficients c1, c3 for p(x) = c1*x + c3*x^3
+            # based on the current singular value bounds [lower_bound, upper_bound].
+            # Formulas are derived from Proposition 2 and its proof in Appendix B of the paper.
+            a_bound, b_bound = lower_bound, upper_bound
+            term = a_bound*a_bound + a_bound*b_bound + b_bound*b_bound
+            e_sq = term / 3.0
+
+            # Calculate alpha, which scales the polynomial
+            common_den_part = 2.0 * (e_sq**1.5)
+            ab_part = a_bound*a_bound*b_bound + b_bound*b_bound*a_bound
+            alpha_den = common_den_part + ab_part
+            alpha = 6.0 / alpha_den
+
+            c1 = alpha * e_sq
+            c3 = -alpha / 3.0
+
+            # Apply the 3rd-order Newton-Schulz update
+            A = X @ X.mT
+            X = c1 * X + c3 * (A @ X)
+
+            # Update the singular value bounds for the next iteration based on the error
+            eps_num = common_den_part - ab_part
+            eps_val = eps_num / alpha_den
+            lower_bound = 1.0 - eps_val
+            upper_bound = 1.0 + eps_val
+    else:
+        # Perform the iterative updates
+        for _ in range(steps):
+            A = X @ X.mT
+            B = b * A + c * (A @ A)
+            X = a * X + B @ X
+
+    # Transpose back if necessary
+    if transposed:
+        X = X.mT
+
+    return X.to(G.dtype)

adv_optm/util/One_Bit_Boolean.py

@@ -0,0 +1,22 @@
+import torch
+
+@torch.no_grad()
+def _pack_bools(tensor: torch.Tensor) -> torch.Tensor:
+    """Packs a boolean tensor into a uint8 tensor to achieve 1-bit storage."""
+    n, m = tensor.shape
+    packed_m = (m + 7) // 8
+    padded_tensor = torch.nn.functional.pad(tensor, (0, packed_m * 8 - m), 'constant', 0)
+    reshaped = padded_tensor.view(n, packed_m, 8)
+    shifter = torch.arange(8, device=tensor.device, dtype=torch.uint8)
+    packed = (reshaped.to(torch.uint8) * (2**shifter)).sum(dim=2).to(torch.uint8)
+    return packed
+
+@torch.no_grad()
+def _unpack_bools(packed_tensor: torch.Tensor, original_m: int) -> torch.Tensor:
+    """Unpacks a uint8 tensor back into a boolean tensor."""
+    if packed_tensor.dtype != torch.uint8:
+        packed_tensor = packed_tensor.to(torch.uint8)
+    shifter = (2**torch.arange(8, device=packed_tensor.device, dtype=torch.uint8)).view(1, 1, 8)
+    unpacked_padded = (packed_tensor.unsqueeze(2) & shifter) != 0
+    unpacked = unpacked_padded.view(packed_tensor.shape[0], -1)[:, :original_m]
+    return unpacked

adv_optm/util/OrthoGrad.py

@@ -0,0 +1,16 @@
+import torch
+
+def _orthogonalize_gradient(p: torch.Tensor, grad: torch.Tensor) -> torch.Tensor:
+    """Projects the gradient `grad` to be orthogonal to the parameter `p`."""
+    if grad.is_sparse: raise RuntimeError("OrthoGrad logic does not support sparse gradients.")
+    original_shape = grad.shape
+    original_dtype = grad.dtype
+    w = p.view(-1).float()
+    g = grad.view(-1).float()
+    w_norm_sq = torch.dot(w, w).add_(1e-30)
+    proj = torch.dot(w, g) / w_norm_sq
+    g_orth = g.sub(w, alpha=proj)
+    g_norm = g.norm(2)
+    g_orth_norm = g_orth.norm(2).add_(1e-30)
+    g_orth_scaled = g_orth * (g_norm / g_orth_norm)
+    return g_orth_scaled.view(original_shape).to(original_dtype)