diff --git a/adv_optm/__init__.py b/adv_optm/__init__.py index d786625..42a6598 100644 --- a/adv_optm/__init__.py +++ b/adv_optm/__init__.py @@ -20,4 +20,4 @@ "SinkSGD_adv", ] -__version__ = "2.5.11" +__version__ = "2.6.1.dev3" diff --git a/adv_optm/optim/AdaMuon_adv.py b/adv_optm/optim/AdaMuon_adv.py index 8bd8160..764b384 100644 --- a/adv_optm/optim/AdaMuon_adv.py +++ b/adv_optm/optim/AdaMuon_adv.py @@ -137,6 +137,7 @@ def __init__( # Decoupled/cautious weight decay weight_decay: float = 0, cautious_wd: bool = False, + scaled_wd: bool = False, # Nesterov momentum nesterov: bool = True, nesterov_coef: float | None = None, @@ -227,7 +228,7 @@ def __init__( defaults = { "lr": lr, "betas": betas, "weight_decay": weight_decay, "cautious_wd": cautious_wd, - "eps": eps, "rms_rescaling": rms_rescaling, "ns_steps": ns_steps, + "eps": eps, "rms_rescaling": rms_rescaling, "ns_steps": ns_steps, "scaled_wd": scaled_wd, "ns_eps": ns_eps, "ns_coeffs": ns_coeffs, "nnmf_factor": nnmf_factor, "vector_reshape": vector_reshape, "nesterov":nesterov, "nesterov_coef": nesterov_coef, "use_atan2":use_atan2, diff --git a/adv_optm/optim/AdamW_adv.py b/adv_optm/optim/AdamW_adv.py index 3457550..8d1550c 100644 --- a/adv_optm/optim/AdamW_adv.py +++ b/adv_optm/optim/AdamW_adv.py @@ -98,6 +98,7 @@ def __init__( weight_decay: float = 0.0, fisher_wd: bool = False, cautious_wd: bool = False, + scaled_wd: bool = False, # Adam's Bias Correction use_bias_correction: bool = True, # Stochastic Rounding for BF16 @@ -156,7 +157,7 @@ def __init__( defaults = { "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, - "fisher_wd": fisher_wd, "cautious_wd": cautious_wd, + "fisher_wd": fisher_wd, "cautious_wd": cautious_wd, "scaled_wd": scaled_wd, "use_atan2": use_atan2, "nesterov": nesterov, "nesterov_coef": nesterov_coef, "normed_momentum": normed_momentum, "orthogonal_gradient": orthogonal_gradient, "use_bias_correction": use_bias_correction, diff --git a/adv_optm/optim/Adopt_adv.py b/adv_optm/optim/Adopt_adv.py index f798cd3..954a469 100644 --- a/adv_optm/optim/Adopt_adv.py +++ b/adv_optm/optim/Adopt_adv.py @@ -101,6 +101,7 @@ def __init__( weight_decay: float = 0.0, fisher_wd: bool = False, cautious_wd: bool = False, + scaled_wd: bool = False, # ADOPT clipping clip_lambda: Optional[Callable[[int], float]] = lambda step: step**0.25, # Adam_atan2 (scale invariant) @@ -157,7 +158,7 @@ def __init__( state_precision = "factored" defaults = { - "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, + "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "scaled_wd": scaled_wd, "fisher_wd": fisher_wd, "cautious_wd": cautious_wd, "orthogonal_gradient": orthogonal_gradient, "nesterov": nesterov, "nesterov_coef": nesterov_coef, "kourkoutas_beta": kourkoutas_beta, "beta2_min": beta2_min, "ema_alpha": ema_alpha, diff --git a/adv_optm/optim/Lion_adv.py b/adv_optm/optim/Lion_adv.py index 6810c03..c41668b 100644 --- a/adv_optm/optim/Lion_adv.py +++ b/adv_optm/optim/Lion_adv.py @@ -64,6 +64,7 @@ def __init__( # Decoupled/cautious weight decay weight_decay: float = 0.0, cautious_wd: bool = False, + scaled_wd: bool = False, # Stochastic Rounding for BF16 stochastic_rounding: bool = True, # OrthoGrad @@ -96,6 +97,7 @@ def __init__( betas=betas, weight_decay=weight_decay, cautious_wd=cautious_wd, + scaled_wd=scaled_wd, vector_reshape=vector_reshape, orthogonal_gradient=orthogonal_gradient, kappa_p=kappa_p, diff --git a/adv_optm/optim/Muon_adv.py b/adv_optm/optim/Muon_adv.py index e8c8e5a..1d20754 100644 --- a/adv_optm/optim/Muon_adv.py +++ b/adv_optm/optim/Muon_adv.py @@ -111,6 +111,7 @@ def __init__( # Decoupled/cautious weight decay weight_decay: float = 0.0, cautious_wd: bool = False, + scaled_wd: bool = False, # Nesterov momentum nesterov: bool = True, nesterov_coef: float | None = None, @@ -201,7 +202,7 @@ def __init__( defaults = { "lr": lr, "beta1": beta1, "weight_decay": weight_decay, "cautious_wd": cautious_wd, "nesterov": nesterov, "nesterov_coef": nesterov_coef, "ns_steps": ns_steps, "ns_eps": ns_eps, - "ns_coeffs": ns_coeffs, "nnmf_factor": nnmf_factor, + "ns_coeffs": ns_coeffs, "nnmf_factor": nnmf_factor, "scaled_wd": scaled_wd, "vector_reshape": vector_reshape, "rms_rescaling": rms_rescaling, "orthogonal_gradient": orthogonal_gradient, 'compiled_optimizer': compiled_optimizer, diff --git a/adv_optm/optim/Prodigy_adv.py b/adv_optm/optim/Prodigy_adv.py index d06afa9..5ef4088 100644 --- a/adv_optm/optim/Prodigy_adv.py +++ b/adv_optm/optim/Prodigy_adv.py @@ -115,6 +115,7 @@ def __init__( weight_decay: float = 0.0, fisher_wd: bool = False, cautious_wd: bool = False, + scaled_wd: bool = False, # Stochastic Rounding for BF16 stochastic_rounding: bool = True, # Adam_atan2 (scale invariant) @@ -181,7 +182,7 @@ def __init__( defaults = { "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, - "fisher_wd": fisher_wd, "cautious_wd": cautious_wd, + "fisher_wd": fisher_wd, "cautious_wd": cautious_wd, "scaled_wd": scaled_wd, "use_atan2": use_atan2, "orthogonal_gradient": orthogonal_gradient, "compiled_optimizer": compiled_optimizer, diff --git a/adv_optm/optim/SignSGD_adv.py b/adv_optm/optim/SignSGD_adv.py index c58eed6..398b0ce 100644 --- a/adv_optm/optim/SignSGD_adv.py +++ b/adv_optm/optim/SignSGD_adv.py @@ -59,6 +59,7 @@ def __init__( # weight decay features geometric_wd: bool = False, cautious_wd: bool = False, + scaled_wd: bool = False, # Stochastic Rounding for BF16 stochastic_rounding: bool = True, # OrthoGrad @@ -108,6 +109,7 @@ def __init__( momentum=momentum, weight_decay=weight_decay, cautious_wd=cautious_wd, + scaled_wd=scaled_wd, geometric_wd=geometric_wd, vector_reshape=vector_reshape, orthogonal_gradient=orthogonal_gradient, diff --git a/adv_optm/optim/SinkSGD_adv.py b/adv_optm/optim/SinkSGD_adv.py index 9ac321d..52dbb90 100644 --- a/adv_optm/optim/SinkSGD_adv.py +++ b/adv_optm/optim/SinkSGD_adv.py @@ -66,6 +66,7 @@ def __init__( # weight decay features geometric_wd: bool = False, cautious_wd: bool = False, + scaled_wd: bool = False, # Stochastic Rounding for BF16 stochastic_rounding: bool = True, # OrthoGrad @@ -103,7 +104,7 @@ def __init__( defaults = { "lr": lr, "momentum": momentum, "weight_decay": weight_decay, "nesterov": nesterov, "nesterov_coef": nesterov_coef, "normed_momentum": normed_momentum, "snr_cond": snr_cond, - "geometric_wd": geometric_wd, "cautious_wd": cautious_wd, + "geometric_wd": geometric_wd, "cautious_wd": cautious_wd, "scaled_wd": scaled_wd, "orthogonal_gradient": orthogonal_gradient, "compiled_optimizer": compiled_optimizer, "sinkhorn_iterations": sinkhorn_iterations, diff --git a/adv_optm/util/param_update.py b/adv_optm/util/param_update.py index 047f4bc..d8eaeab 100644 --- a/adv_optm/util/param_update.py +++ b/adv_optm/util/param_update.py @@ -6,7 +6,7 @@ from typing import Dict, Any -from .scaled_optm import adjust_wds +from .scaled_optm import adjust_wds, scale_wd from .centered_decay import dequantize_anchor _generators: Dict[torch.device, torch.Generator] = {} @@ -18,8 +18,8 @@ def _apply_weight_decay( p: Tensor, state: Dict[str, Any], group: Dict[str, Any], - scaled_wd: float | Tensor | None, - scaled_cwd: float | Tensor | None, + eff_wd: float | Tensor | None, + eff_cwd: float | Tensor | None, wd_target: Tensor | None = None, cwd_target: Tensor | None = None, ) -> None: @@ -29,26 +29,26 @@ def _apply_weight_decay( cautious = group.get('cautious_wd', False) # Standard Weight Decay (pulls toward zero) - if scaled_wd is not None: + if eff_wd is not None: if wd_target is None: wd_target = p_calc # Cautious Weight Decay: only decay if the update pushes in the same direction as the decay if cautious: mask = (update_calc * p_calc >= 0).to(p_calc.dtype) - if isinstance(scaled_wd, Tensor): - p_calc.addcmul_(wd_target, mask * scaled_wd, value=-1.0) + if isinstance(eff_wd, Tensor): + p_calc.addcmul_(wd_target, mask * eff_wd, value=-1.0) else: - p_calc.addcmul_(wd_target, mask, value=-scaled_wd) + p_calc.addcmul_(wd_target, mask, value=-eff_wd) del mask else: # Standard decoupled weight decay - if isinstance(scaled_wd, Tensor): - p_calc.addcmul_(wd_target, scaled_wd, value=-1.0) + if isinstance(eff_wd, Tensor): + p_calc.addcmul_(wd_target, eff_wd, value=-1.0) else: - p_calc.add_(wd_target, alpha=-scaled_wd) + p_calc.add_(wd_target, alpha=-eff_wd) # Centered Weight Decay (pulls toward anchor) - if scaled_cwd is not None and 'anchor_data' in state: + if eff_cwd is not None and 'anchor_data' in state: if cwd_target is not None: decay_target = cwd_target else: @@ -59,17 +59,17 @@ def _apply_weight_decay( if cautious: # Cautious Weight Decay: only decay if the update pushes in the same direction as the decay mask = (update_calc * decay_target >= 0).to(p_calc.dtype) - if isinstance(scaled_cwd, Tensor): - p_calc.addcmul_(decay_target, mask * scaled_cwd, value=-1.0) + if isinstance(eff_cwd, Tensor): + p_calc.addcmul_(decay_target, mask * eff_cwd, value=-1.0) else: - p_calc.addcmul_(decay_target, mask, value=-scaled_cwd) + p_calc.addcmul_(decay_target, mask, value=-eff_cwd) del mask else: # Standard decoupled weight decay - if isinstance(scaled_cwd, Tensor): - p_calc.addcmul_(decay_target, scaled_cwd, value=-1.0) + if isinstance(eff_cwd, Tensor): + p_calc.addcmul_(decay_target, eff_cwd, value=-1.0) else: - p_calc.add_(decay_target, alpha=-scaled_cwd) + p_calc.add_(decay_target, alpha=-eff_cwd) if cwd_target is None: del decay_target @@ -105,18 +105,24 @@ def apply_parameter_update( wd = group["weight_decay"] if wd is None else wd cwd = group.get("centered_wd", 0.0) wd, cwd = adjust_wds(wd, cwd, p) + scaled_wd = group.get("scaled_wd", False) + decoupled = scaled_wd or decoupled # Calculate global decay factor for decoupled vs standard decay_factor = (lr / self._init_lr) if decoupled else lr - scaled_wd = (wd * decay_factor) if wd != 0 else None - scaled_cwd = (cwd * decay_factor) if cwd != 0 else None + eff_wd = (wd * decay_factor) if wd != 0 else None + eff_cwd = (cwd * decay_factor) if cwd != 0 else None + + if scaled_wd: + eff_wd = scale_wd(eff_wd, p, skip_vectors=True) + eff_cwd = scale_wd(eff_cwd, p, skip_vectors=False) if wd_scaler is not None: - if scaled_wd is not None: - scaled_wd = scaled_wd * wd_scaler - if scaled_cwd is not None: - scaled_cwd = scaled_cwd * wd_scaler + if eff_wd is not None: + eff_wd = eff_wd * wd_scaler + if eff_cwd is not None: + eff_cwd = eff_cwd * wd_scaler state = self.state[p] @@ -129,8 +135,8 @@ def apply_parameter_update( cwd_t = cwd_target.float() if cwd_target is not None else None # Apply weight decay if needed - if scaled_wd is not None or scaled_cwd is not None: - _apply_weight_decay(p_fp32, update_fp32, p, state, group, scaled_wd, scaled_cwd, wd_t, cwd_t) + if eff_wd is not None or eff_cwd is not None: + _apply_weight_decay(p_fp32, update_fp32, p, state, group, eff_wd, eff_cwd, wd_t, cwd_t) # Apply main update p_fp32.add_(-update_fp32) @@ -147,8 +153,8 @@ def apply_parameter_update( else: # Standard path for non-bfloat16 or without stochastic rounding - if scaled_wd is not None or scaled_cwd is not None: - _apply_weight_decay(p, update, p, state, group, scaled_wd, scaled_cwd, wd_target, cwd_target) + if eff_wd is not None or eff_cwd is not None: + _apply_weight_decay(p, update, p, state, group, eff_wd, eff_cwd, wd_target, cwd_target) # Apply main update p.add_(-update) diff --git a/adv_optm/util/scaled_optm.py b/adv_optm/util/scaled_optm.py index 6df90c4..39dfe17 100644 --- a/adv_optm/util/scaled_optm.py +++ b/adv_optm/util/scaled_optm.py @@ -97,6 +97,25 @@ def adjust_wds(wd: float, cwd: float, p: torch.Tensor) -> tuple[float, float]: # Centered WD safely regularizes the delta without collapsing base feature variance. return wd, cwd +def scale_wd(wd: float | None, p: torch.Tensor, skip_vectors: bool = False) -> float: + """ + Scale-invariant, dimension-scaled weight decay. + """ + if wd is None: + return wd + if getattr(p, '_is_oft', False): + n_el = p.shape[-1] + b = (1.0 + math.sqrt(1.0 + 8.0 * n_el)) / 2.0 + wd = (2 * wd) / (b - 1) + return wd + + is_vector = p.ndim < 2 or getattr(p, '_is_dora_scale', False) or getattr(p, 'is_vector', False) + if is_vector: + return 0.0 if skip_vectors else wd + + if p.ndim >= 2: + width = p.numel() // p.shape[0] + return wd / width def is_spectral(p: torch.Tensor) -> bool: """Determines if a parameter should undergo spectral normalization updates.""" diff --git a/setup.py b/setup.py index 671563e..9d7b91e 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name="adv_optm", - version="2.5.11", + version="2.6.1.dev3", author="Koratahiu", author_email="hiuhonor@gmail.com", license='Apache 2.0',