adv-optm 2.6.1.dev1__tar.gz → 2.6.1.dev3__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/PKG-INFO +1 -1
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/__init__.py +1 -1
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/util/param_update.py +5 -4
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/util/scaled_optm.py +11 -6
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm.egg-info/PKG-INFO +1 -1
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/setup.py +1 -1
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/LICENSE +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/README.md +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/optim/AdaMuon_adv.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/optim/AdamW_adv.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/optim/Adopt_adv.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/optim/Lion_adv.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/optim/Muon_adv.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/optim/Prodigy_adv.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/optim/SignSGD_adv.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/optim/SinkSGD_adv.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/optim/__init__.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/util/Kourkoutas.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/util/Muon_AuxAdam.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/util/Muon_util.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/util/OrthoGrad.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/util/__init__.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/util/centered_decay.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/util/factorization_util.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/util/lion_k.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/util/signed_util.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/util/sinkhorn.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/util/state_util.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm/util/update_util.py +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm.egg-info/SOURCES.txt +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm.egg-info/dependency_links.txt +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm.egg-info/requires.txt +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/adv_optm.egg-info/top_level.txt +0 -0
- {adv_optm-2.6.1.dev1 → adv_optm-2.6.1.dev3}/setup.cfg +0 -0
|
@@ -106,6 +106,7 @@ def apply_parameter_update(
|
|
|
106
106
|
cwd = group.get("centered_wd", 0.0)
|
|
107
107
|
wd, cwd = adjust_wds(wd, cwd, p)
|
|
108
108
|
scaled_wd = group.get("scaled_wd", False)
|
|
109
|
+
decoupled = scaled_wd or decoupled
|
|
109
110
|
|
|
110
111
|
# Calculate global decay factor for decoupled vs standard
|
|
111
112
|
decay_factor = (lr / self._init_lr) if decoupled else lr
|
|
@@ -113,14 +114,14 @@ def apply_parameter_update(
|
|
|
113
114
|
eff_wd = (wd * decay_factor) if wd != 0 else None
|
|
114
115
|
eff_cwd = (cwd * decay_factor) if cwd != 0 else None
|
|
115
116
|
|
|
117
|
+
if scaled_wd:
|
|
118
|
+
eff_wd = scale_wd(eff_wd, p, skip_vectors=True)
|
|
119
|
+
eff_cwd = scale_wd(eff_cwd, p, skip_vectors=False)
|
|
120
|
+
|
|
116
121
|
if wd_scaler is not None:
|
|
117
122
|
if eff_wd is not None:
|
|
118
|
-
if scaled_wd:
|
|
119
|
-
eff_wd = scale_wd(eff_wd, p)
|
|
120
123
|
eff_wd = eff_wd * wd_scaler
|
|
121
124
|
if eff_cwd is not None:
|
|
122
|
-
if scaled_wd:
|
|
123
|
-
eff_cwd = scale_wd(eff_cwd, p)
|
|
124
125
|
eff_cwd = eff_cwd * wd_scaler
|
|
125
126
|
|
|
126
127
|
state = self.state[p]
|
|
@@ -97,21 +97,26 @@ def adjust_wds(wd: float, cwd: float, p: torch.Tensor) -> tuple[float, float]:
|
|
|
97
97
|
# Centered WD safely regularizes the delta without collapsing base feature variance.
|
|
98
98
|
return wd, cwd
|
|
99
99
|
|
|
100
|
-
def scale_wd(wd: float, p: torch.Tensor) -> float:
|
|
100
|
+
def scale_wd(wd: float | None, p: torch.Tensor, skip_vectors: bool = False) -> float:
|
|
101
101
|
"""
|
|
102
102
|
Scale-invariant, dimension-scaled weight decay.
|
|
103
103
|
"""
|
|
104
|
+
if wd is None:
|
|
105
|
+
return wd
|
|
104
106
|
if getattr(p, '_is_oft', False):
|
|
105
107
|
n_el = p.shape[-1]
|
|
106
108
|
b = (1.0 + math.sqrt(1.0 + 8.0 * n_el)) / 2.0
|
|
107
109
|
wd = (2 * wd) / (b - 1)
|
|
108
110
|
return wd
|
|
109
111
|
|
|
112
|
+
is_vector = p.ndim < 2 or getattr(p, '_is_dora_scale', False) or getattr(p, 'is_vector', False)
|
|
113
|
+
if is_vector:
|
|
114
|
+
return 0.0 if skip_vectors else wd
|
|
115
|
+
|
|
110
116
|
if p.ndim >= 2:
|
|
111
117
|
width = p.numel() // p.shape[0]
|
|
112
118
|
return wd / width
|
|
113
119
|
|
|
114
|
-
|
|
115
120
|
def is_spectral(p: torch.Tensor) -> bool:
|
|
116
121
|
"""Determines if a parameter should undergo spectral normalization updates."""
|
|
117
122
|
if getattr(p, '_is_oft', False):
|
|
@@ -227,14 +232,14 @@ def apply_spectral_riemannian_oft(
|
|
|
227
232
|
# Estimate sigma (The spectral norm) for each block
|
|
228
233
|
sigma = torch.sum(next_u * u_raw, dim=1, keepdim=True)
|
|
229
234
|
|
|
230
|
-
#
|
|
231
|
-
|
|
232
|
-
max_sigma = sigma.max()
|
|
235
|
+
# Squeeze out the last dimension so shape becomes (batch_size, 1)
|
|
236
|
+
sigma = sigma.squeeze(-1)
|
|
233
237
|
|
|
234
238
|
target_scale = 0.5 * scale_factor
|
|
235
239
|
spectral_eps = 1.0 / (2.0 * math.sqrt(block_size))
|
|
236
240
|
|
|
237
|
-
|
|
241
|
+
# Apply the clamp and scaling block-wise
|
|
242
|
+
scale = lr * (target_scale / sigma.clamp_min(spectral_eps))
|
|
238
243
|
|
|
239
244
|
return update_flat.mul_(scale).view(orig_shape)
|
|
240
245
|
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|