adv-optm 2.6.dev1__tar.gz → 2.6.1.dev1__tar.gz

{adv_optm-2.6.dev1 → adv_optm-2.6.1.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: adv_optm
-Version: 2.6.dev1
+Version: 2.6.1.dev1
 Summary: A family of highly efficient, lightweight yet powerful optimizers.
 Home-page: https://github.com/Koratahiu/Advanced_Optimizers
 Author: Koratahiu

{adv_optm-2.6.dev1 → adv_optm-2.6.1.dev1}/adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "SinkSGD_adv",
 ]
 
-__version__ = "2.6.dev1"
+__version__ = "2.6.1.dev1"

{adv_optm-2.6.dev1 → adv_optm-2.6.1.dev1}/adv_optm/optim/AdaMuon_adv.py

@@ -137,6 +137,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
         # 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,
@@ -177,8 +178,6 @@ class AdaMuon_adv(torch.optim.Optimizer):
         mars_gamma: float = 0.025,
         # Spectral Normalization
         spectral_normalization: bool = False,
-        # Orthogonalize the weights (Matrix Sign - MSign) every x steps
-        MSign_interval: int | None = None,
         # Centered WD
         centered_wd: float = 0.0,
         centered_wd_mode: str = 'float8',
@@ -229,7 +228,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
         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,
@@ -249,7 +248,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
             # MARS-M
             "approx_mars": approx_mars, "mars_gamma": mars_gamma,
             # Spectral Normalization
-            "spectral_normalization": spectral_normalization, "MSign_interval": MSign_interval,
+            "spectral_normalization": spectral_normalization,
             # Centered WD
             "centered_wd": centered_wd,
             "centered_wd_mode": centered_wd_mode,
@@ -284,11 +283,9 @@ class AdaMuon_adv(torch.optim.Optimizer):
             for device in devices:
                 param_update.set_seed(device)
 
-        # Initialize compiled function
-        self._compiled_muon_step_parameter = None
-        self._compiled_adam_step_parameter = None
-        if compiled_optimizer:
-            self.compile(fullgraph=True)
+        # Initialize compiled functions (by parameter shape)
+        self._compiled_muon_step_fns = {}
+        self._compiled_adam_step_fns = {}
 
     def load_state_dict(self, state_dict: dict) -> None:
         """
@@ -446,8 +443,16 @@ class AdaMuon_adv(torch.optim.Optimizer):
             random_int_state_tensor = None
             if is_compiled:
                 step_size = torch.as_tensor(step_size)
-                adam_step_param = self._compiled_adam_step_parameter
-                
+                # Cache compiled function per-shape
+                cache_key = (p.shape, state.get('factored', False))
+                if cache_key not in self._compiled_adam_step_fns:
+                    self._compiled_adam_step_fns[cache_key] = torch.compile(
+                        Muon_AuxAdam._adam_step_parameter,
+                        fullgraph=True,
+                        dynamic=False
+                    )
+                adam_step_param = self._compiled_adam_step_fns[cache_key]
+
                 # Generate state SR random tensor when compiled
                 actual_precision = group.get('adam_actual_state_precision', 'auto')
                 random_int_state_tensor = random_int_tensor
@@ -466,7 +471,15 @@ class AdaMuon_adv(torch.optim.Optimizer):
             random_G_sketch = None
             if is_compiled:
                 lr = torch.as_tensor(group['lr'])
-                muon_step_param = self._compiled_muon_step_parameter
+                # Cache compiled function per-shape
+                cache_key = (p.shape, state.get('factored', False))
+                if cache_key not in self._compiled_muon_step_fns:
+                    self._compiled_muon_step_fns[cache_key] = torch.compile(
+                        self._muon_step_parameter,
+                        fullgraph=True,
+                        dynamic=False
+                    )
+                muon_step_param = self._compiled_muon_step_fns[cache_key]
 
                 # Generate state SR random tensor when compiled
                 actual_precision = group['actual_state_precision']
@@ -484,10 +497,6 @@ class AdaMuon_adv(torch.optim.Optimizer):
 
             muon_step_param(p, grad, state, group, lr, random_int_tensor, random_int_state_tensor, random_G_sketch)
 
-    def compile(self, *args, **kwargs):
-        self._compiled_muon_step_parameter = torch.compile(self._muon_step_parameter, *args, **kwargs)
-        self._compiled_adam_step_parameter = torch.compile(Muon_AuxAdam._adam_step_parameter, *args, **kwargs)
-
     @torch.no_grad()
     def _muon_step_parameter(self, p, grad, state, group, lr, random_int_tensor, random_int_state_tensor, random_G_sketch):
         # Upcast grad for low-precision state modes (non-factored path)
@@ -533,8 +542,9 @@ class AdaMuon_adv(torch.optim.Optimizer):
             else:
                 update = mt_buf.clone()
 
-            # Factorize
-            state['mu_mbuf_nmf'], state['mv_mbuf_nmf'], state['sign_buf'] = _factorize_state(mt_buf, signed=True, shifter=state['shifter'])
+            # Compress new momentum and store factors
+            for key, val in zip(('mu_mbuf_nmf', 'mv_mbuf_nmf', 'sign_buf'), _factorize_state(mt_buf, signed=True, shifter=state['shifter'])):
+                state[key].copy_(val)
             del mt_buf
 
             # Apply update projection
@@ -561,7 +571,8 @@ class AdaMuon_adv(torch.optim.Optimizer):
                 vt_buf = _reconstruct_state((state['mu_vbuf_nmf'], state['mv_vbuf_nmf']), signed=False, shifter=state['shifter'])
                 # Update second momentum in full-size
                 vt_buf.mul_(beta2).addcmul_(update, update, value=1 - beta2)
-                state['mu_vbuf_nmf'], state['mv_vbuf_nmf'] = _factorize_state(vt_buf, signed=False, shifter=state['shifter'])
+                for key, val in zip(('mu_vbuf_nmf', 'mv_vbuf_nmf'), _factorize_state(vt_buf, signed=False, shifter=state['shifter'])):
+                    state[key].copy_(val)
                 # Apply second momentum update (adaptive scaling)
                 if group['use_atan2']:
                     denom = vt_buf.sqrt_()
@@ -620,7 +631,8 @@ class AdaMuon_adv(torch.optim.Optimizer):
                 update_f32 = update.float()
                 vt_buf = _reconstruct_state((state['mu_vbuf_nmf'], state['mv_vbuf_nmf']), signed=False, shifter=state['shifter'])
                 vt_buf.mul_(beta2).addcmul_(update_f32.view(d1, d2), update_f32.view(d1, d2), value=1 - beta2)
-                state['mu_vbuf_nmf'], state['mv_vbuf_nmf'] = _factorize_state(vt_buf, signed=False, shifter=state['shifter'])
+                for key, val in zip(('mu_vbuf_nmf', 'mv_vbuf_nmf'), _factorize_state(vt_buf, signed=False, shifter=state['shifter'])):
+                    state[key].copy_(val)
                 # Apply second moment scaling
                 if group['use_atan2']:
                     denom = vt_buf.sqrt_().view(original_shape)

{adv_optm-2.6.dev1 → adv_optm-2.6.1.dev1}/adv_optm/optim/AdamW_adv.py

@@ -98,6 +98,7 @@ class AdamW_adv(torch.optim.Optimizer):
         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
@@ -121,8 +122,6 @@ class AdamW_adv(torch.optim.Optimizer):
         layer_key_fn: Optional[Callable] = None,
         # Spectral Normed Optimizer
         spectral_normalization: bool = False,
-        # Orthogonalize the weights (Matrix Sign - MSign) every x steps
-        MSign_interval: int | None = None,
         # Centered WD
         centered_wd: float = 0.0,
         centered_wd_mode: str = 'float8',
@@ -158,14 +157,14 @@ class AdamW_adv(torch.optim.Optimizer):
 
         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,
             "compiled_optimizer": compiled_optimizer,
             "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,
-            "spectral_normalization": spectral_normalization, "MSign_interval": MSign_interval,
+            "spectral_normalization": spectral_normalization,
             "centered_wd": centered_wd, "centered_wd_mode": centered_wd_mode,
             "state_precision": state_precision,
             "nnmf_factor": nnmf_factor, "vector_reshape": vector_reshape, "factored_2nd": factored_2nd
@@ -188,10 +187,8 @@ class AdamW_adv(torch.optim.Optimizer):
             for device in devices:
                 param_update.set_seed(device)
 
-        # Initialize compiled function
-        self._compiled_step_parameter = None
-        if compiled_optimizer:
-            self.compile(fullgraph=True)
+        # Initialize compiled function (by parameter shape)
+        self._compiled_step_fns = {}
 
     def load_state_dict(self, state_dict: dict) -> None:
         """
@@ -323,7 +320,15 @@ class AdamW_adv(torch.optim.Optimizer):
                 random_int_state_tensor = param_update._get_random_int_for_sr(p)
             elif group['actual_state_precision'] == 'int8_sr':
                 random_int_state_tensor = param_update._get_random_int_for_8bit_sr(p)
-            step_param_fn = self._compiled_step_parameter
+            # Cache compiled function per-shape
+            cache_key = (p.shape, state.get('factored', False))
+            if cache_key not in self._compiled_step_fns:
+                self._compiled_step_fns[cache_key] = torch.compile(
+                    self._step_parameter,
+                    fullgraph=True,
+                    dynamic=False
+                )
+            step_param_fn = self._compiled_step_fns[cache_key]
         else:
             step_param_fn = self._step_parameter
 
@@ -359,7 +364,8 @@ class AdamW_adv(torch.optim.Optimizer):
                 vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped, value=1.0 - beta2)
 
             # Factorize
-            state['mu_v_nmf'], state['mv_v_nmf'] = _factorize_state(vt, signed=False, shifter=state['shifter'])
+            for key, val in zip(('mu_v_nmf', 'mv_v_nmf'), _factorize_state(vt, signed=False, shifter=state['shifter'])):
+                state[key].copy_(val)
 
             if group['use_atan2']:
                 denom = vt.sqrt_()
@@ -380,7 +386,8 @@ class AdamW_adv(torch.optim.Optimizer):
                 mt.lerp_(grad_reshaped, 1.0 - beta1)
 
                 # Factorize
-                state['mu_m_nmf'], state['mv_m_nmf'], state['sign'] = _factorize_state(mt.clone(), signed=True, shifter=state['shifter'])
+                for key, val in zip(('mu_m_nmf', 'mv_m_nmf', 'sign'), _factorize_state(mt.clone(), signed=True, shifter=state['shifter'])):
+                    state[key].copy_(val)
 
                 update_mt = mt
 
@@ -424,7 +431,8 @@ class AdamW_adv(torch.optim.Optimizer):
                 exp_avg_sq.mul_(beta2).addcmul_(grad_vt, grad_vt, value=1.0 - beta2)
 
             if factored_2nd:
-                state['mu_v_nmf'], state['mv_v_nmf'] = _factorize_state(exp_avg_sq.view(d1, d2), signed=False, shifter=state['shifter'])
+                for key, val in zip(('mu_v_nmf', 'mv_v_nmf'), _factorize_state(exp_avg_sq.view(d1, d2), signed=False, shifter=state['shifter'])):
+                    state[key].copy_(val)
             else:
                 set_state(state, 'exp_avg_sq', exp_avg_sq, actual_precision, random_int_state_tensor, non_neg=True)
 
@@ -471,9 +479,6 @@ class AdamW_adv(torch.optim.Optimizer):
 
         param_update.apply_parameter_update(self, p, group, update, group['lr'], random_int_tensor=random_int_tensor, wd_scaler=wd_scaler)
 
-    def compile(self, *args, **kwargs):
-        self._compiled_step_parameter = torch.compile(self._step_parameter, *args, **kwargs)
-
     @torch.no_grad()
     def step(self, closure=None):
         """Performs a single optimization step."""

{adv_optm-2.6.dev1 → adv_optm-2.6.1.dev1}/adv_optm/optim/Adopt_adv.py

@@ -101,6 +101,7 @@ class Adopt_adv(torch.optim.Optimizer):
         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)
@@ -122,8 +123,6 @@ class Adopt_adv(torch.optim.Optimizer):
         layer_key_fn: Optional[Callable] = None,
         # Spectral Normed Optimizer
         spectral_normalization: bool = False,
-        # Orthogonalize the weights (Matrix Sign - MSign) every x steps
-        MSign_interval: int | None = None,
         # Centered WD
         centered_wd: float = 0.0,
         centered_wd_mode: str = 'float8',
@@ -159,12 +158,12 @@ class Adopt_adv(torch.optim.Optimizer):
             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,
             "tiny_spike": tiny_spike, "k_warmup_steps": k_warmup_steps, "k_logging": k_logging,
-            "spectral_normalization": spectral_normalization, "MSign_interval": MSign_interval,
+            "spectral_normalization": spectral_normalization,
             "centered_wd": centered_wd,
             "centered_wd_mode": centered_wd_mode,
             "state_precision": state_precision,
@@ -191,9 +190,8 @@ class Adopt_adv(torch.optim.Optimizer):
             for device in devices:
                 param_update.set_seed(device)
 
-        self._compiled_step_parameter = None
-        if compiled_optimizer:
-            self.compile(fullgraph=True)
+        # Initialize compiled function (by parameter shape)
+        self._compiled_step_fns = {}
 
     def load_state_dict(self, state_dict: dict) -> None:
         """
@@ -333,7 +331,15 @@ class Adopt_adv(torch.optim.Optimizer):
                 random_int_state_tensor = param_update._get_random_int_for_sr(p)
             elif group['actual_state_precision'] == 'int8_sr':
                 random_int_state_tensor = param_update._get_random_int_for_8bit_sr(p)
-            step_param_fn = self._compiled_step_parameter
+            # Cache compiled function per-shape
+            cache_key = (p.shape, state.get('factored', False))
+            if cache_key not in self._compiled_step_fns:
+                self._compiled_step_fns[cache_key] = torch.compile(
+                    self._step_parameter,
+                    fullgraph=True,
+                    dynamic=False
+                )
+            step_param_fn = self._compiled_step_fns[cache_key]
         else:
             lr = group['lr']
             step_param_fn = self._step_parameter
@@ -375,7 +381,8 @@ class Adopt_adv(torch.optim.Optimizer):
             else:
                 vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped, value=1.0 - beta2)
             # Factorize
-            state['mu_v_nmf'], state['mv_v_nmf'] = _factorize_state(vt, signed=False, shifter=state['shifter'])
+            for key, val in zip(('mu_v_nmf', 'mv_v_nmf'), _factorize_state(vt, signed=False, shifter=state['shifter'])):
+                state[key].copy_(val)
             del vt
 
             if self.use_atan2:
@@ -393,7 +400,8 @@ class Adopt_adv(torch.optim.Optimizer):
                 mt.lerp_(normalized_grad, 1.0 - beta1)
 
                 # Factorize
-                state['mu_m_nmf'], state['mv_m_nmf'], state['sign'] = _factorize_state(mt.clone(), signed=True, shifter=state['shifter'])
+                for key, val in zip(('mu_m_nmf', 'mv_m_nmf', 'sign'), _factorize_state(mt.clone(), signed=True, shifter=state['shifter'])):
+                    state[key].copy_(val)
 
                 update_mt = mt
 
@@ -460,7 +468,8 @@ class Adopt_adv(torch.optim.Optimizer):
                 vt.mul_(beta2).addcmul_(grad_vt, grad_vt, value=1 - beta2)
 
             if factored_2nd:
-                state['mu_v_nmf'], state['mv_v_nmf'] = _factorize_state(vt.view(d1, d2), signed=False, shifter=state['shifter'])
+                for key, val in zip(('mu_v_nmf', 'mv_v_nmf'), _factorize_state(vt.view(d1, d2), signed=False, shifter=state['shifter'])):
+                    state[key].copy_(val)
             else:
                 set_state(state, 'exp_avg_sq', vt, actual_precision, random_int_state_tensor, non_neg=True)
             del random_int_state_tensor
@@ -475,9 +484,6 @@ class Adopt_adv(torch.optim.Optimizer):
         # Parameter Update
         param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor, wd_scaler=wd_scaler)
 
-    def compile(self, *args, **kwargs):
-        self._compiled_step_parameter = torch.compile(self._step_parameter, *args, **kwargs)
-
     @torch.no_grad()
     def step(self, closure=None):
         """Performs a single optimization step."""

{adv_optm-2.6.dev1 → adv_optm-2.6.1.dev1}/adv_optm/optim/Lion_adv.py

@@ -64,6 +64,7 @@ class Lion_adv(torch.optim.Optimizer):
         # 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
@@ -78,8 +79,6 @@ class Lion_adv(torch.optim.Optimizer):
         centered_wd_mode: str = 'float8',
         # Spectral Normed Optimizer
         spectral_normalization: bool = False,
-        # Orthogonalize the weights (Matrix Sign - MSign) every x steps
-        MSign_interval: int | None = None,
         # SMMF factorization
         nnmf_factor: bool = False,
         vector_reshape: bool = False,
@@ -98,13 +97,13 @@ class Lion_adv(torch.optim.Optimizer):
             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,
             auto_kappa_p=auto_kappa_p,
             stochastic_sign=stochastic_sign,
             spectral_normalization=spectral_normalization,
-            MSign_interval=MSign_interval,
             nnmf_factor=nnmf_factor,
             centered_wd= centered_wd,
             centered_wd_mode= centered_wd_mode,
@@ -122,10 +121,8 @@ class Lion_adv(torch.optim.Optimizer):
             for device in devices:
                 param_update.set_seed(device)
 
-        # Initialize compiled function
-        self._compiled_step_parameter = None
-        if compiled_optimizer:
-            self.compile(fullgraph=True)
+        # Initialize compiled function (by parameter shape)
+        self._compiled_step_fns = {}
 
     def load_state_dict(self, state_dict: dict) -> None:
         """
@@ -208,7 +205,15 @@ class Lion_adv(torch.optim.Optimizer):
             if group.get('stochastic_sign', False):
                 random_noise_tensor = param_update._get_random_noise_for_sso(p)
             lr = torch.as_tensor(lr)
-            step_param_fn = self._compiled_step_parameter
+            # Cache compiled function per-shape
+            cache_key = (p.shape, state.get('factored', False))
+            if cache_key not in self._compiled_step_fns:
+                self._compiled_step_fns[cache_key] = torch.compile(
+                    self._step_parameter,
+                    fullgraph=True,
+                    dynamic=False
+                )
+            step_param_fn = self._compiled_step_fns[cache_key]
         else:
             step_param_fn = self._step_parameter
 
@@ -281,9 +286,6 @@ class Lion_adv(torch.optim.Optimizer):
 
         param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor)
 
-    def compile(self, *args, **kwargs):
-        self._compiled_step_parameter = torch.compile(self._step_parameter, *args, **kwargs)
-
     @torch.no_grad()
     def step(self, closure: Optional[callable] = None):
         """Performs a single optimization step."""

{adv_optm-2.6.dev1 → adv_optm-2.6.1.dev1}/adv_optm/optim/Muon_adv.py

@@ -111,6 +111,7 @@ class Muon_adv(torch.optim.Optimizer):
         # 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,
@@ -146,8 +147,6 @@ class Muon_adv(torch.optim.Optimizer):
         mars_gamma: float = 0.025,
         # Spectral Normalization
         spectral_normalization: bool = False,
-        # Orthogonalize the weights (Matrix Sign - MSign) every x steps
-        MSign_interval: int | None = None,
         # Centered WD
         centered_wd: float = 0.0,
         centered_wd_mode: str = 'float8',
@@ -203,7 +202,7 @@ class Muon_adv(torch.optim.Optimizer):
         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,
@@ -220,7 +219,7 @@ class Muon_adv(torch.optim.Optimizer):
             # MARS-M
             "approx_mars": approx_mars, "mars_gamma": mars_gamma,
             # Spectral Normalization
-            "spectral_normalization": spectral_normalization, "MSign_interval": MSign_interval,
+            "spectral_normalization": spectral_normalization,
             # Centered WD
             "centered_wd": centered_wd,
             "centered_wd_mode": centered_wd_mode,
@@ -256,11 +255,9 @@ class Muon_adv(torch.optim.Optimizer):
             for device in devices:
                 param_update.set_seed(device)
 
-        # Initialize compiled function
-        self._compiled_muon_step_parameter = None
-        self._compiled_adam_step_parameter = None
-        if compiled_optimizer:
-            self.compile(fullgraph=True)
+        # Initialize compiled functions (by parameter shape)
+        self._compiled_muon_step_fns = {}
+        self._compiled_adam_step_fns = {}
 
     def load_state_dict(self, state_dict: dict) -> None:
         """
@@ -398,7 +395,15 @@ class Muon_adv(torch.optim.Optimizer):
             random_int_state_tensor = None
             if is_compiled:
                 step_size = torch.as_tensor(step_size)
-                adam_step_param = self._compiled_adam_step_parameter
+                # Cache compiled function per-shape
+                cache_key = (p.shape, state.get('factored', False))
+                if cache_key not in self._compiled_adam_step_fns:
+                    self._compiled_adam_step_fns[cache_key] = torch.compile(
+                        Muon_AuxAdam._adam_step_parameter,
+                        fullgraph=True,
+                        dynamic=False
+                    )
+                adam_step_param = self._compiled_adam_step_fns[cache_key]
                 
                 actual_precision = group.get('adam_actual_state_precision', 'auto')
                 random_int_state_tensor = random_int_tensor
@@ -417,7 +422,15 @@ class Muon_adv(torch.optim.Optimizer):
             random_G_sketch = None
             if is_compiled:
                 lr = torch.as_tensor(group['lr'])
-                muon_step_param = self._compiled_muon_step_parameter
+                # Cache compiled function per-shape
+                cache_key = (p.shape, state.get('factored', False))
+                if cache_key not in self._compiled_muon_step_fns:
+                    self._compiled_muon_step_fns[cache_key] = torch.compile(
+                        self._muon_step_parameter,
+                        fullgraph=True,
+                        dynamic=False
+                    )
+                muon_step_param = self._compiled_muon_step_fns[cache_key]
 
                 # Generate state SR random tensor when compiled
                 actual_precision = group['actual_state_precision']
@@ -435,10 +448,6 @@ class Muon_adv(torch.optim.Optimizer):
 
             muon_step_param(p, grad, state, group, lr, random_int_tensor, random_int_state_tensor, random_G_sketch)
 
-    def compile(self, *args, **kwargs):
-        self._compiled_muon_step_parameter = torch.compile(self._muon_step_parameter, *args, **kwargs)
-        self._compiled_adam_step_parameter = torch.compile(Muon_AuxAdam._adam_step_parameter, *args, **kwargs)
-
     @torch.no_grad()
     def _muon_step_parameter(self, p, grad, state, group, lr, random_int_tensor, random_int_state_tensor, random_G_sketch):
         # Upcast grad for low-precision state modes (non-factored path)
@@ -477,8 +486,9 @@ class Muon_adv(torch.optim.Optimizer):
                 # Standard momentum
                 update = mt_buf.clone()
 
-            # Factorize
-            state['mu_mbuf_nmf'], state['mv_mbuf_nmf'], state['sign_buf'] = _factorize_state(mt_buf, signed=True, shifter=state['shifter'])
+            # Compress new momentum and store factors
+            for key, val in zip(('mu_mbuf_nmf', 'mv_mbuf_nmf', 'sign_buf'), _factorize_state(mt_buf, signed=True, shifter=state['shifter'])):
+                state[key].copy_(val)
             del mt_buf
 
             # Orthogonalization step

{adv_optm-2.6.dev1 → adv_optm-2.6.1.dev1}/adv_optm/optim/Prodigy_adv.py

@@ -115,6 +115,7 @@ class Prodigy_adv(torch.optim.Optimizer):
         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)
@@ -156,8 +157,6 @@ class Prodigy_adv(torch.optim.Optimizer):
         centered_wd_mode: str = 'float8',
         # Spectral Normalization
         spectral_normalization: bool = False,
-        # Orthogonalize the weights (Matrix Sign - MSign) every x steps
-        MSign_interval: int | None = None,
     ):
         if not (lr >= 0.0):
             raise ValueError(f"Learning-rate should be >= 0.0. Got {lr}")
@@ -183,7 +182,7 @@ class Prodigy_adv(torch.optim.Optimizer):
 
         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,
@@ -195,7 +194,7 @@ class Prodigy_adv(torch.optim.Optimizer):
             "tiny_spike": tiny_spike, "k_warmup_steps": k_warmup_steps, "k_logging": k_logging,
             "centered_wd": centered_wd, "centered_wd_mode": centered_wd_mode,
             "nnmf_factor": nnmf_factor, "vector_reshape": vector_reshape, "factored_2nd": factored_2nd,
-            "spectral_normalization": spectral_normalization, "MSign_interval": MSign_interval,
+            "spectral_normalization": spectral_normalization,
         }
         self.stochastic_rounding = stochastic_rounding
         self.fsdp_in_use = fsdp_in_use
@@ -220,11 +219,8 @@ class Prodigy_adv(torch.optim.Optimizer):
             for device in devices:
                 param_update.set_seed(device)
 
-        # Initialize compiled function
-        self._compiled_step_parameter = None
-
-        if compiled_optimizer:
-            self.compile(fullgraph=True)
+        # Initialize compiled function (by parameter shape)
+        self._compiled_step_fns = {}
 
     def load_state_dict(self, state_dict: dict) -> None:
         """
@@ -366,7 +362,15 @@ class Prodigy_adv(torch.optim.Optimizer):
                 random_int_state_tensor = param_update._get_random_int_for_sr(p)
             elif group['actual_state_precision'] == 'int8_sr':
                 random_int_state_tensor = param_update._get_random_int_for_8bit_sr(p)
-            step_param_fn = self._compiled_step_parameter
+            # Cache compiled function per-shape
+            cache_key = (p.shape, state.get('factored', False))
+            if cache_key not in self._compiled_step_fns:
+                self._compiled_step_fns[cache_key] = torch.compile(
+                    self._step_parameter,
+                    fullgraph=True,
+                    dynamic=False
+                )
+            step_param_fn = self._compiled_step_fns[cache_key]
         else:
             d = group['d']
             step_param_fn = self._step_parameter
@@ -402,7 +406,8 @@ class Prodigy_adv(torch.optim.Optimizer):
                 mt.mul_(self.beta1).add_(grad_reshaped, alpha=d * (1.0 - self.beta1))
 
                 # Factorize
-                state['mu_m_nmf'], state['mv_m_nmf'], state['sign'] = _factorize_state(mt.clone(), signed=True, shifter=state['shifter'])
+                for key, val in zip(('mu_m_nmf', 'mv_m_nmf', 'sign'), _factorize_state(mt.clone(), signed=True, shifter=state['shifter'])):
+                    state[key].copy_(val)
 
                 update_mt = mt
 
@@ -423,7 +428,8 @@ class Prodigy_adv(torch.optim.Optimizer):
                 update = grad_reshaped.mul(d)
 
             # Factorize
-            state['mu_v_nmf'], state['mv_v_nmf'] = _factorize_state(vt, signed=False, shifter=state['shifter'])
+            for key, val in zip(('mu_v_nmf', 'mv_v_nmf'), _factorize_state(vt, signed=False, shifter=state['shifter'])):
+                state[key].copy_(val)
 
             if group['use_atan2']:
                 denom = vt.sqrt_()
@@ -475,7 +481,8 @@ class Prodigy_adv(torch.optim.Optimizer):
                 exp_avg_sq.mul_(beta2).addcmul_(grad_vt, grad_vt, value=d * d * (1.0 - beta2))
 
             if factored_2nd:
-                state['mu_v_nmf'], state['mv_v_nmf'] = _factorize_state(exp_avg_sq.view(d1, d2), signed=False, shifter=state['shifter'])
+                for key, val in zip(('mu_v_nmf', 'mv_v_nmf'), _factorize_state(exp_avg_sq.view(d1, d2), signed=False, shifter=state['shifter'])):
+                    state[key].copy_(val)
             else:
                 set_state(state, 'exp_avg_sq', exp_avg_sq, actual_precision, random_int_state_tensor, non_neg=True)
             del random_int_state_tensor
@@ -524,9 +531,6 @@ class Prodigy_adv(torch.optim.Optimizer):
 
         param_update.apply_parameter_update(self, p, group, update, dlr, random_int_tensor=random_int_tensor, wd_scaler=wd_scaler)
 
-    def compile(self, *args, **kwargs):
-        self._compiled_step_parameter = torch.compile(self._step_parameter, *args, **kwargs)
-
     @torch.no_grad()
     def step(self, closure=None):
         """Performs a single optimization step."""