From 009c77f7412ce7644afa0a58e2afff7e0be350bf Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 27 Jun 2026 11:17:22 +0300 Subject: [PATCH 1/9] rename to eff_wd --- adv_optm/util/param_update.py | 52 +++++++++++++++++------------------ 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/adv_optm/util/param_update.py b/adv_optm/util/param_update.py index 047f4bc..ddacd30 100644 --- a/adv_optm/util/param_update.py +++ b/adv_optm/util/param_update.py @@ -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 @@ -109,14 +109,14 @@ def apply_parameter_update( # 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 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 +129,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 +147,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) From 05e83c55b0e05b23c7b74bd50abce8874162babc Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 27 Jun 2026 11:19:19 +0300 Subject: [PATCH 2/9] initial --- adv_optm/util/param_update.py | 7 ++++++- adv_optm/util/scaled_optm.py | 14 ++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/adv_optm/util/param_update.py b/adv_optm/util/param_update.py index ddacd30..a19f7e3 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] = {} @@ -105,6 +105,7 @@ 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) # Calculate global decay factor for decoupled vs standard decay_factor = (lr / self._init_lr) if decoupled else lr @@ -114,8 +115,12 @@ def apply_parameter_update( if wd_scaler is not None: if eff_wd is not None: + if scaled_wd: + eff_wd = scale_wd(eff_wd, p) eff_wd = eff_wd * wd_scaler if eff_cwd is not None: + if scaled_wd: + eff_cwd = scale_wd(eff_cwd, p) eff_cwd = eff_cwd * wd_scaler state = self.state[p] diff --git a/adv_optm/util/scaled_optm.py b/adv_optm/util/scaled_optm.py index 0958157..2d1b8e7 100644 --- a/adv_optm/util/scaled_optm.py +++ b/adv_optm/util/scaled_optm.py @@ -97,6 +97,20 @@ 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, p: torch.Tensor) -> float: + """ + Scale-invariant, dimension-scaled weight decay. + """ + 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 + + 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.""" From fe3bb47ece623f8f9d4cacddf1928d622d56493f Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 27 Jun 2026 11:22:55 +0300 Subject: [PATCH 3/9] initial scaled_wd optm --- adv_optm/optim/AdaMuon_adv.py | 3 ++- adv_optm/optim/AdamW_adv.py | 3 ++- adv_optm/optim/Adopt_adv.py | 3 ++- adv_optm/optim/Lion_adv.py | 2 ++ adv_optm/optim/Muon_adv.py | 3 ++- adv_optm/optim/Prodigy_adv.py | 3 ++- adv_optm/optim/SignSGD_adv.py | 2 ++ adv_optm/optim/SinkSGD_adv.py | 3 ++- 8 files changed, 16 insertions(+), 6 deletions(-) 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, From b60e61934951d3ff1a43115d29d6f1f36cc7568c Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 27 Jun 2026 11:24:32 +0300 Subject: [PATCH 4/9] 2.6.1.dev1 --- adv_optm/__init__.py | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/adv_optm/__init__.py b/adv_optm/__init__.py index b90a66a..d8f998f 100644 --- a/adv_optm/__init__.py +++ b/adv_optm/__init__.py @@ -20,4 +20,4 @@ "SinkSGD_adv", ] -__version__ = "2.5.10" +__version__ = "2.6.1.dev1" diff --git a/setup.py b/setup.py index 2a193d3..e6368f2 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name="adv_optm", - version="2.5.10", + version="2.6.1.dev1", author="Koratahiu", author_email="hiuhonor@gmail.com", license='Apache 2.0', From e400053cf9761f5b54c7766efdd5c713f824bef9 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sat, 27 Jun 2026 11:37:51 +0300 Subject: [PATCH 5/9] dev2 decoupled --- adv_optm/__init__.py | 2 +- adv_optm/util/param_update.py | 1 + setup.py | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/adv_optm/__init__.py b/adv_optm/__init__.py index d8f998f..d947ee8 100644 --- a/adv_optm/__init__.py +++ b/adv_optm/__init__.py @@ -20,4 +20,4 @@ "SinkSGD_adv", ] -__version__ = "2.6.1.dev1" +__version__ = "2.6.1.dev2" diff --git a/adv_optm/util/param_update.py b/adv_optm/util/param_update.py index a19f7e3..f0b484b 100644 --- a/adv_optm/util/param_update.py +++ b/adv_optm/util/param_update.py @@ -106,6 +106,7 @@ def apply_parameter_update( cwd = group.get("centered_wd", 0.0) wd, cwd = adjust_wds(wd, cwd, p) scaled_wd = group.get("scaled_wd", False) + decoupled = scaled_wd # Calculate global decay factor for decoupled vs standard decay_factor = (lr / self._init_lr) if decoupled else lr diff --git a/setup.py b/setup.py index e6368f2..305b2c0 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name="adv_optm", - version="2.6.1.dev1", + version="2.6.1.dev2", author="Koratahiu", author_email="hiuhonor@gmail.com", license='Apache 2.0', From 4b06e338e405abd35d8bcfb70c138e5f0d348c1a Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 28 Jun 2026 04:16:40 +0300 Subject: [PATCH 6/9] fix --- adv_optm/util/param_update.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/adv_optm/util/param_update.py b/adv_optm/util/param_update.py index f0b484b..63d5c7d 100644 --- a/adv_optm/util/param_update.py +++ b/adv_optm/util/param_update.py @@ -106,7 +106,7 @@ def apply_parameter_update( cwd = group.get("centered_wd", 0.0) wd, cwd = adjust_wds(wd, cwd, p) scaled_wd = group.get("scaled_wd", False) - decoupled = scaled_wd + decoupled = scaled_wd or decoupled # Calculate global decay factor for decoupled vs standard decay_factor = (lr / self._init_lr) if decoupled else lr From cd80b7a668a5b97c1c13ba52fbee662fe7f7c8a8 Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 28 Jun 2026 06:05:05 +0300 Subject: [PATCH 7/9] skip_vectors --- adv_optm/util/param_update.py | 4 ++-- adv_optm/util/scaled_optm.py | 7 +++++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/adv_optm/util/param_update.py b/adv_optm/util/param_update.py index 63d5c7d..4b863c3 100644 --- a/adv_optm/util/param_update.py +++ b/adv_optm/util/param_update.py @@ -117,11 +117,11 @@ def apply_parameter_update( if wd_scaler is not None: if eff_wd is not None: if scaled_wd: - eff_wd = scale_wd(eff_wd, p) + eff_wd = scale_wd(eff_wd, p, skip_vectors=True) eff_wd = eff_wd * wd_scaler if eff_cwd is not None: if scaled_wd: - eff_cwd = scale_wd(eff_cwd, p) + eff_cwd = scale_wd(eff_cwd, p, skip_vectors=False) eff_cwd = eff_cwd * wd_scaler state = self.state[p] diff --git a/adv_optm/util/scaled_optm.py b/adv_optm/util/scaled_optm.py index a7dac04..dff7adc 100644 --- a/adv_optm/util/scaled_optm.py +++ b/adv_optm/util/scaled_optm.py @@ -97,7 +97,7 @@ 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, p: torch.Tensor) -> float: +def scale_wd(wd: float, p: torch.Tensor, skip_vectors: bool = False) -> float: """ Scale-invariant, dimension-scaled weight decay. """ @@ -107,11 +107,14 @@ def scale_wd(wd: float, p: torch.Tensor) -> float: 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.""" if getattr(p, '_is_oft', False): From 3a9b8747628b9f49369c1d39c500294825eaf95b Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 28 Jun 2026 13:38:00 +0300 Subject: [PATCH 8/9] fix --- adv_optm/util/param_update.py | 8 ++++---- adv_optm/util/scaled_optm.py | 4 +++- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/adv_optm/util/param_update.py b/adv_optm/util/param_update.py index 4b863c3..d8eaeab 100644 --- a/adv_optm/util/param_update.py +++ b/adv_optm/util/param_update.py @@ -114,14 +114,14 @@ def apply_parameter_update( 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 eff_wd is not None: - if scaled_wd: - eff_wd = scale_wd(eff_wd, p, skip_vectors=True) eff_wd = eff_wd * wd_scaler if eff_cwd is not None: - if scaled_wd: - eff_cwd = scale_wd(eff_cwd, p, skip_vectors=False) eff_cwd = eff_cwd * wd_scaler state = self.state[p] diff --git a/adv_optm/util/scaled_optm.py b/adv_optm/util/scaled_optm.py index dff7adc..39dfe17 100644 --- a/adv_optm/util/scaled_optm.py +++ b/adv_optm/util/scaled_optm.py @@ -97,10 +97,12 @@ 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, p: torch.Tensor, skip_vectors: bool = False) -> float: +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 From b4c49afc7e96f57af49b14afe5f14e54bfe0977c Mon Sep 17 00:00:00 2001 From: Koratahiu~ Date: Sun, 28 Jun 2026 13:38:15 +0300 Subject: [PATCH 9/9] dev3 --- adv_optm/__init__.py | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/adv_optm/__init__.py b/adv_optm/__init__.py index d947ee8..42a6598 100644 --- a/adv_optm/__init__.py +++ b/adv_optm/__init__.py @@ -20,4 +20,4 @@ "SinkSGD_adv", ] -__version__ = "2.6.1.dev2" +__version__ = "2.6.1.dev3" diff --git a/setup.py b/setup.py index 305b2c0..9d7b91e 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name="adv_optm", - version="2.6.1.dev2", + version="2.6.1.dev3", author="Koratahiu", author_email="hiuhonor@gmail.com", license='Apache 2.0',