adv-optm 2.3.dev2__tar.gz → 2.4.dev1__tar.gz

{adv_optm-2.3.dev2 → adv_optm-2.4.dev1}/PKG-INFO

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

{adv_optm-2.3.dev2 → adv_optm-2.4.dev1}/adv_optm/__init__.py

@@ -10,11 +10,6 @@ from .optim import (
     SignSGD_adv,
 )
 
-from .stiefel_optm.Stiefel_LoRA import (
-    Stiefel_LoRA,
-)
-
-
 __all__ = [
     "AdamW_adv",
     "Prodigy_adv",
@@ -25,7 +20,6 @@ __all__ = [
     "Muon_adv",
     "AdaMuon_adv",
     "SignSGD_adv",
-    "Stiefel_LoRA",
 ]
 
-__version__ = "2.3.dev2"
+__version__ = "2.4.dev1"

{adv_optm-2.3.dev2 → adv_optm-2.4.dev1}/adv_optm/optim/AdaMuon_adv.py

@@ -8,6 +8,7 @@ from ..util.factorization_util import _get_effective_shape, _factorize_state, _r
 from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.Kourkoutas import KourkoutasHelper
 from ..util import Muon_AuxAdam
+from ..util.centered_decay import _init_anchor
 
 A = 4 / math.pi
 
@@ -87,6 +88,15 @@ class AdaMuon_adv(torch.optim.Optimizer):
             (default: False)
         mars_gamma (float): The scaling coefficient for MARS gradient correction.
             (default: 0.025)
+        centered_wd (float): Centered Weight Decay coefficient. Instead of decaying weights
+            toward zero, they are decayed toward their initial values (anchors). This
+            can be used together with standard weight decay. (default: 0.0)
+        centered_wd_mode (str): The quantization format used to store the anchor
+            weights to save VRAM. Options include:
+            'full': Stores anchors in the original parameter's precision.
+            'float8': Uses torch.float8_e4m3fn for a balance of precision and memory.
+            'int8': Uses 8-bit block-wise quantization (block size 128).
+            'int4': Uses 4-bit block-wise quantization (block size 32).
         nnmf_factor (bool): whether to use the factorization or disable it to use
             the uncompressed optimizer. (default: False)
         use_muon (bool | None): whether to use Muon or AuxAdamW. MUST be provided
@@ -157,6 +167,9 @@ class AdaMuon_adv(torch.optim.Optimizer):
         # Spectral Normalization
         n_layers: int = 1,
         spectral_normalization: bool = False,
+        # Centered WD
+        centered_wd: float = 0.0,
+        centered_wd_mode: str = 'float8',
         # torch.compile
         compiled_optimizer: bool = False,
         # --- AdamW_adv specific parameters ---
@@ -214,6 +227,9 @@ class AdaMuon_adv(torch.optim.Optimizer):
             "approx_mars": approx_mars, "mars_gamma": mars_gamma,
             # Spectral Normalization
             "n_layers": n_layers, "spectral_normalization": spectral_normalization,
+            # Centered WD
+            "centered_wd": centered_wd,
+            "centered_wd_mode": centered_wd_mode,
             # AdamW_adv defaults
             "adam_betas": adam_betas, "adam_eps": adam_eps, "adam_weight_decay": adam_weight_decay,
             "adam_use_bias_correction": adam_use_bias_correction, "adam_use_atan2": adam_use_atan2,
@@ -261,6 +277,16 @@ class AdaMuon_adv(torch.optim.Optimizer):
         if compiled_optimizer:
             self.compile(fullgraph=True)
 
+    def load_state_dict(self, state_dict: dict) -> None:
+        """
+        Overrides default load_state_dict to implement a workaround for PyTorch's
+        automatic dtype casting. It ensures factorized states remain float32 for 
+        stability, preserves integer/float8 quantized anchor states, and forces 
+        standard states onto the parameter's current dtype/device.
+        """
+        super().load_state_dict(state_dict)
+        param_update.post_process_loaded_state(self)
+
     @property
     def supports_fused_back_pass(self):
         return True
@@ -344,6 +370,8 @@ class AdaMuon_adv(torch.optim.Optimizer):
                 # Note: This requires full-rank memory even if factored
                 state['last_grad'] = torch.zeros_like(p, device=device, dtype=p.dtype)
 
+            _init_anchor(p, state, group)
+
             group['adam_kourkoutas_beta'] = False
             state['is_muon'] = True # Workaround as group was acting weirdly; passing muon params in adam path
 

{adv_optm-2.3.dev2 → adv_optm-2.4.dev1}/adv_optm/optim/AdamW_adv.py

@@ -9,6 +9,8 @@ from ..util.factorization_util import _get_effective_shape, _reconstruct_state,
 from ..util.update_util import _grams_update, _cautious_update
 from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.Kourkoutas import KourkoutasHelper
+from ..util.scaled_optm import scale_update, is_spectral, init_spectral_norm
+from ..util.centered_decay import _init_anchor
 
 A = 4 / math.pi
 
@@ -78,8 +80,19 @@ class AdamW_adv(torch.optim.Optimizer):
             and returns a unique, hashable key representing its "layer" or "bucket".
             If `None`, parameters are bucketed by their memory ID (tensor-wise).
             (default: None)
+        centered_wd (float): Centered Weight Decay coefficient. Instead of decaying weights
+            toward zero, they are decayed toward their initial values (anchors). This
+            can be used together with standard weight decay. (default: 0.0)
+        centered_wd_mode (str): The quantization format used to store the anchor
+            weights to save VRAM. Options include:
+            'full': Stores anchors in the original parameter's precision.
+            'float8': Uses torch.float8_e4m3fn for a balance of precision and memory.
+            'int8': Uses 8-bit block-wise quantization (block size 128).
+            'int4': Uses 4-bit block-wise quantization (block size 32).
         nnmf_factor (bool): whether to use the factorization or disable it to use
             the uncompressed optimizer. (default: False)
+        factored_2nd (bool): whether to keep the first moment uncompressed (dense) 
+            while only factorizing the second moment. (default: True)
     """
 
     def __init__(
@@ -114,9 +127,15 @@ class AdamW_adv(torch.optim.Optimizer):
         k_warmup_steps: int = 0,
         k_logging: int = 0,
         layer_key_fn: Optional[Callable] = None,
+        # Scaled Optimizer
+        scaled_optm: bool = False,
+        # Centered WD
+        centered_wd: float = 0.0,
+        centered_wd_mode: str = 'float8',
         # SMMF factorization
         nnmf_factor: bool = False,
         vector_reshape: bool = False,
+        factored_2nd: bool = False,
         # torch.compile
         compiled_optimizer: bool = False,
     ):
@@ -137,12 +156,14 @@ class AdamW_adv(torch.optim.Optimizer):
 
         defaults = {
             "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
-            "vector_reshape": vector_reshape, "use_atan2": use_atan2,
+            "use_atan2": use_atan2,
             "orthogonal_gradient": orthogonal_gradient, "use_bias_correction": use_bias_correction,
             "beta3_ema": beta3_ema, "alpha": alpha, "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,
-            "nnmf_factor": nnmf_factor
+            "scaled_optm": scaled_optm,
+            "centered_wd": centered_wd, "centered_wd_mode": centered_wd_mode,
+            "nnmf_factor": nnmf_factor, "vector_reshape": vector_reshape, "factored_2nd": factored_2nd
         }
         self.stochastic_rounding = stochastic_rounding
         self.cautious_mask = cautious_mask
@@ -150,6 +171,7 @@ class AdamW_adv(torch.optim.Optimizer):
         self.use_AdEMAMix = use_AdEMAMix
         self.kourkoutas_beta = kourkoutas_beta
         self.layer_key_fn = layer_key_fn
+        self._init_lr = lr
         super().__init__(params, defaults)
 
         if self.kourkoutas_beta:
@@ -167,6 +189,16 @@ class AdamW_adv(torch.optim.Optimizer):
         if compiled_optimizer:
             self.compile(fullgraph=True)
 
+    def load_state_dict(self, state_dict: dict) -> None:
+        """
+        Overrides default load_state_dict to implement a workaround for PyTorch's
+        automatic dtype casting. It ensures factorized states remain float32 for 
+        stability, preserves integer/float8 quantized anchor states, and forces 
+        standard states onto the parameter's current dtype/device.
+        """
+        super().load_state_dict(state_dict)
+        param_update.post_process_loaded_state(self)
+
     @property
     def supports_fused_back_pass(self):
         return True
@@ -194,6 +226,7 @@ class AdamW_adv(torch.optim.Optimizer):
             state['factored'] = (
                 group['nnmf_factor'] and
                 not (len(p.shape) == 1 and not group['vector_reshape'])
+                or group["factored_2nd"]
             )
 
             dtype = torch.float32 if state['factored'] else p.dtype
@@ -203,18 +236,25 @@ class AdamW_adv(torch.optim.Optimizer):
                 state['effective_shape'] = _get_effective_shape(p.numel())
                 d1, d2 = state['effective_shape']
 
-                # First moment (m)
-                if group['betas'][0] > 0:
-                    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)
-                # AdEMAMix slow moment (m_slow)
-                if self.use_AdEMAMix:
-                    state['mu_m_slow_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype)
-                    state['mv_m_slow_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
-                    packed_d2 = (d2 + 7) // 8
-                    state['sign_slow'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=p.device)
+                if not group.get('factored_2nd', False):
+                    # First moment (m)
+                    if group['betas'][0] > 0:
+                        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)
+                    # AdEMAMix slow moment (m_slow)
+                    if self.use_AdEMAMix:
+                        state['mu_m_slow_nmf'] = torch.zeros(d1, device=device, dtype=dtype)
+                        state['mv_m_slow_nmf'] = torch.zeros(d2, device=device, dtype=dtype)
+                        packed_d2 = (d2 + 7) // 8
+                        state['sign_slow'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=device)
+                else:
+                    if group['betas'][0] > 0:
+                        state['exp_avg'] = torch.zeros_like(p, device=device, dtype=dtype)
+                    if self.use_AdEMAMix:
+                        state['exp_avg_slow'] = torch.zeros_like(p, device=device, dtype=dtype)
+
                 # 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)
@@ -228,6 +268,11 @@ class AdamW_adv(torch.optim.Optimizer):
                 # Second moment (v)
                 state['exp_avg_sq'] = torch.zeros_like(p, device=device, dtype=dtype)
 
+            if group.get('scaled_optm', False) and is_spectral(p):
+                init_spectral_norm(group, state, p)
+
+            _init_anchor(p, state, group)
+
         beta1, beta2 = group['betas']
 
         current_step = state['step']
@@ -275,32 +320,42 @@ class AdamW_adv(torch.optim.Optimizer):
             # Accumulate current grad's norm for the *next* step
             self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
 
+        # Determine if we are using dense first-moments alongside a factored second-order second-moment
+        factored_2nd = group.get('factored_2nd', False)
+
         if state['factored']:
             d1, d2 = state['effective_shape']
             grad_reshaped = grad.view(d1, d2)
 
             # Reconstruct momentum from previous step's factors
             if beta1 > 0:
-                mt = _reconstruct_state((state['mu_m_nmf'], state['mv_m_nmf'], state['sign'], d2), signed=True)
+                if factored_2nd:
+                    mt = state['exp_avg'].view(d1, d2)
+                else:
+                    mt = _reconstruct_state((state['mu_m_nmf'], state['mv_m_nmf'], state['sign'], d2), signed=True)
 
                 # Update momentum in full-size
                 mt.lerp_(grad_reshaped, 1.0 - beta1)
 
-                # Factorize
-                state['mu_m_nmf'], state['mv_m_nmf'], state['sign'] = _factorize_state(mt.clone(), signed=True)
+                if not factored_2nd:
+                    # Factorize
+                    state['mu_m_nmf'], state['mv_m_nmf'], state['sign'] = _factorize_state(mt.clone(), signed=True)
 
                 if self.grams_moment:
-                    update_mt = _grams_update(mt, grad_reshaped, inplace=True)
+                    update_mt = _grams_update(mt, grad_reshaped, inplace=not factored_2nd)
                 elif self.cautious_mask:
-                    update_mt = _cautious_update(mt, grad_reshaped, inplace=True)
+                    update_mt = _cautious_update(mt, grad_reshaped, inplace=not factored_2nd)
                 else:
-                    update_mt = mt
+                    update_mt = mt if not factored_2nd else mt.clone()
 
             vt = _reconstruct_state((state['mu_v_nmf'], state['mv_v_nmf']), signed=False)
             vt.mul_(beta2).addcmul_(grad_reshaped, grad_reshaped, value=1.0 - beta2)
 
             if self.use_AdEMAMix:
-                mt_slow = _reconstruct_state((state['mu_m_slow_nmf'], state['mv_m_slow_nmf'], state['sign_slow'], d2), signed=True)
+                if factored_2nd:
+                    mt_slow = state['exp_avg_slow'].view(d1, d2)
+                else:
+                    mt_slow = _reconstruct_state((state['mu_m_slow_nmf'], state['mv_m_slow_nmf'], state['sign_slow'], d2), signed=True)
 
                 mt_slow.lerp_(grad_reshaped, 1.0 - beta3_ema)
 
@@ -308,9 +363,11 @@ class AdamW_adv(torch.optim.Optimizer):
                     update = update_mt.add_(mt_slow, alpha=alpha)
                 else:
                     update = grad_reshaped.add(mt_slow, alpha=alpha)
-                # Factorize
-                state['mu_m_slow_nmf'], state['mv_m_slow_nmf'], state['sign_slow'] = _factorize_state(mt_slow, signed=True)
-                del mt_slow
+                
+                if not factored_2nd:
+                    # Factorize
+                    state['mu_m_slow_nmf'], state['mv_m_slow_nmf'], state['sign_slow'] = _factorize_state(mt_slow, signed=True)
+                    del mt_slow
             else:
                 if beta1 > 0:
                     update = update_mt
@@ -330,8 +387,7 @@ class AdamW_adv(torch.optim.Optimizer):
                 update.div_(denom)
             del vt
 
-            update_scaling = step_size * A if group['use_atan2'] else step_size
-            update = update.view(p.shape).mul_(update_scaling)
+            update = update.view(p.shape)
 
         else:  # Standard AdamW logic for non-factored tensors
             if beta1 > 0:
@@ -369,7 +425,10 @@ class AdamW_adv(torch.optim.Optimizer):
                 update.div_(denom)
             del denom
 
-            update_scaling = step_size * A if group['use_atan2'] else step_size
+        update_scaling = step_size * A if group['use_atan2'] else step_size
+        if group.get('scaled_optm', False):
+            update = scale_update(p, update, update_scaling, vector_state=state.get('spectral_v'))
+        else:
             update.mul_(update_scaling)
 
         param_update.apply_parameter_update(self, p, group, update, step_size, random_int_tensor=random_int_tensor)

{adv_optm-2.3.dev2 → adv_optm-2.4.dev1}/adv_optm/optim/Adopt_adv.py

@@ -8,6 +8,8 @@ from ..util.factorization_util import _get_effective_shape, _reconstruct_state,
 from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.Kourkoutas import KourkoutasHelper
 from ..util.update_util import _grams_update, _cautious_update, _scale_sim_AdEMAMix_update
+from ..util.scaled_optm import scale_update, is_spectral, init_spectral_norm
+from ..util.centered_decay import _init_anchor
 
 A = 4 / math.pi
 
@@ -94,8 +96,19 @@ class Adopt_adv(torch.optim.Optimizer):
             and returns a unique, hashable key representing its "layer" or "bucket".
             If `None`, parameters are bucketed by their memory ID (tensor-wise).
             (default: None)
+        centered_wd (float): Centered Weight Decay coefficient. Instead of decaying weights
+            toward zero, they are decayed toward their initial values (anchors). This
+            can be used together with standard weight decay. (default: 0.0)
+        centered_wd_mode (str): The quantization format used to store the anchor
+            weights to save VRAM. Options include:
+            'full': Stores anchors in the original parameter's precision.
+            'float8': Uses torch.float8_e4m3fn for a balance of precision and memory.
+            'int8': Uses 8-bit block-wise quantization (block size 128).
+            'int4': Uses 4-bit block-wise quantization (block size 32).
         nnmf_factor (bool): whether to use the factorization or disable it to use
             the uncompressed optimizer. (default: False)
+        factored_2nd (bool): whether to keep the first moment uncompressed (dense) 
+            while only factorizing the second moment. (default: True)
     """
 
     def __init__(
@@ -133,9 +146,15 @@ class Adopt_adv(torch.optim.Optimizer):
         k_warmup_steps: int = 0,
         k_logging: int = 0,
         layer_key_fn: Optional[Callable] = None,
+        # Scaled Optimizer
+        scaled_optm: bool = False,
+        # Centered WD
+        centered_wd: float = 0.0,
+        centered_wd_mode: str = 'float8',
         # SMMF factorization
         nnmf_factor: bool = False,
-        vector_reshape: bool = False,
+        vector_reshape: bool = True,
+        factored_2nd: bool = False,
         # torch.compile
         compiled_optimizer: bool = False,
     ):
@@ -163,11 +182,14 @@ class Adopt_adv(torch.optim.Optimizer):
 
         defaults = {
             "lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "cautious_wd": cautious_wd,
-            "vector_reshape": vector_reshape, "beta3_ema": beta3_ema, "alpha": alpha,
+            "beta3_ema": beta3_ema, "alpha": alpha,
             "alpha_grad": alpha_grad,
             "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,
-            "nnmf_factor": nnmf_factor,
+            "scaled_optm": scaled_optm,
+            "centered_wd": centered_wd,
+            "centered_wd_mode": centered_wd_mode,
+            "nnmf_factor": nnmf_factor, "vector_reshape": vector_reshape, "factored_2nd": factored_2nd, 
             "compiled_optimizer": compiled_optimizer,
         }
         self.clip_lambda = clip_lambda
@@ -180,6 +202,7 @@ class Adopt_adv(torch.optim.Optimizer):
         self.Simplified_AdEMAMix = Simplified_AdEMAMix
         self.kourkoutas_beta = kourkoutas_beta
         self.layer_key_fn = layer_key_fn
+        self._init_lr = lr
         super().__init__(params, defaults)
 
         if self.kourkoutas_beta:
@@ -196,6 +219,16 @@ class Adopt_adv(torch.optim.Optimizer):
         if compiled_optimizer:
             self.compile(fullgraph=True)
 
+    def load_state_dict(self, state_dict: dict) -> None:
+        """
+        Overrides default load_state_dict to implement a workaround for PyTorch's
+        automatic dtype casting. It ensures factorized states remain float32 for 
+        stability, preserves integer/float8 quantized anchor states, and forces 
+        standard states onto the parameter's current dtype/device.
+        """
+        super().load_state_dict(state_dict)
+        param_update.post_process_loaded_state(self)
+
     @property
     def supports_fused_back_pass(self): return True
     @property
@@ -218,6 +251,7 @@ class Adopt_adv(torch.optim.Optimizer):
             state['factored'] = (
                 group['nnmf_factor'] and
                 not (len(p.shape) == 1 and not group['vector_reshape'])
+                or group["factored_2nd"]
             )
 
             dtype = torch.float32 if state['factored'] else p.dtype
@@ -226,18 +260,24 @@ class Adopt_adv(torch.optim.Optimizer):
                 state['effective_shape'] = _get_effective_shape(p.numel())
                 d1, d2 = state['effective_shape']
 
-                # First moment (m)
-                if group['betas'][0] > 0:
-                    state['mu_m_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype)
-                    state['mv_m_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
-                    packed_d2 = (d2 + 7) // 8
-                    state['sign'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=p.device)
-                # AdEMAMix slow moment (m_slow)
-                if self.use_AdEMAMix:
-                    state['mu_m_slow_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype)
-                    state['mv_m_slow_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
-                    packed_d2 = (d2 + 7) // 8
-                    state['sign_slow'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=p.device)
+                if not group.get('factored_2nd', False):
+                    # First moment (m)
+                    if group['betas'][0] > 0:
+                        state['mu_m_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype)
+                        state['mv_m_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
+                        packed_d2 = (d2 + 7) // 8
+                        state['sign'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=p.device)
+                    # AdEMAMix slow moment (m_slow)
+                    if self.use_AdEMAMix:
+                        state['mu_m_slow_nmf'] = torch.zeros(d1, device=p.device, dtype=dtype)
+                        state['mv_m_slow_nmf'] = torch.zeros(d2, device=p.device, dtype=dtype)
+                        packed_d2 = (d2 + 7) // 8
+                        state['sign_slow'] = torch.zeros((d1, packed_d2), dtype=torch.uint8, device=p.device)
+                else:
+                    if group['betas'][0] > 0:
+                        state['exp_avg'] = torch.zeros_like(p, device=p.device, dtype=dtype)
+                    if self.use_AdEMAMix:
+                        state['exp_avg_slow'] = torch.zeros_like(p, device=p.device, dtype=dtype)
                 # Second moment (v)
                 vt_init = grad.to(dtype).view(d1, d2).square()
                 # Allocate NMF factors for vt
@@ -253,6 +293,11 @@ class Adopt_adv(torch.optim.Optimizer):
                     state['exp_avg_slow'] = torch.zeros_like(p, device=p.device, dtype=dtype)
                 state['exp_avg_sq'] = grad.to(dtype).square()
 
+            if group.get('scaled_optm', False) and is_spectral(p):
+                init_spectral_norm(group, state, p)
+
+            _init_anchor(p, state, group)
+
         beta1, beta2 = group['betas']
 
         current_step = state['step']
@@ -280,7 +325,7 @@ class Adopt_adv(torch.optim.Optimizer):
             step_param_fn = self._step_parameter
 
         if self.Simplified_AdEMAMix:
-            lr = _scale_sim_AdEMAMix_update(beta1, state['step'] + 1, group["alpha_grad"], lr)
+            lr = _scale_sim_AdEMAMix_update(beta1, state['step'] + 1, group["alpha_grad"], lr, group.get('scaled_optm', False))
 
         step_param_fn(p, grad, state, group, lr, beta1, beta2, random_int_tensor)
 
@@ -302,6 +347,9 @@ class Adopt_adv(torch.optim.Optimizer):
             # Accumulate current grad's norm for the *next* step
             self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
 
+        # Determine if we are using dense first-moments alongside a factored second-order second-moment
+        factored_2nd = group.get('factored_2nd', False)
+
         if state['factored']:
             d1, d2 = state['effective_shape']
             grad_reshaped = grad.view(d1, d2)
@@ -328,35 +376,47 @@ class Adopt_adv(torch.optim.Optimizer):
 
             # ADOPT Step B: Update momentum m_t using normalized gradient
             if beta1 > 0:
-                # Reconstruct m_{t-1}
-                mt = _reconstruct_state((state['mu_m_nmf'], state['mv_m_nmf'], state['sign'], d2), signed=True)
+                if factored_2nd:
+                    mt = state['exp_avg'].view(d1, d2)
+                else:
+                    # Reconstruct m_{t-1}
+                    mt = _reconstruct_state((state['mu_m_nmf'], state['mv_m_nmf'], state['sign'], d2), signed=True)
+
                 if self.Simplified_AdEMAMix:
                     mt.mul_(beta1).add_(normalized_grad, alpha=1.0)
                 else:
                     mt.lerp_(normalized_grad, 1.0 - beta1)
 
-                # Factorize
-                state['mu_m_nmf'], state['mv_m_nmf'], state['sign'] = _factorize_state(mt.clone(), signed=True)
+                if not factored_2nd:
+                    # Factorize
+                    state['mu_m_nmf'], state['mv_m_nmf'], state['sign'] = _factorize_state(mt.clone(), signed=True)
 
                 if self.grams_moment:
-                    update_mt = _grams_update(mt, grad_reshaped, inplace=True)
+                    update_mt = _grams_update(mt, grad_reshaped, inplace=not factored_2nd)
                 elif self.cautious_mask:
-                    update_mt = _cautious_update(mt, grad_reshaped, inplace=True)
+                    update_mt = _cautious_update(mt, grad_reshaped, inplace=not factored_2nd)
                 else:
-                    update_mt = mt
+                    update_mt = mt if not factored_2nd else mt.clone()
 
             if self.use_AdEMAMix:
-                # Reconstruct AdEMAMix EMA
-                mt_slow = _reconstruct_state((state['mu_m_slow_nmf'], state['mv_m_slow_nmf'], state['sign_slow'], d2), signed=True)
+                if factored_2nd:
+                    mt_slow = state['exp_avg_slow'].view(d1, d2)
+                else:
+                    # Reconstruct AdEMAMix EMA
+                    mt_slow = _reconstruct_state((state['mu_m_slow_nmf'], state['mv_m_slow_nmf'], state['sign_slow'], d2), signed=True)
+
                 mt_slow.lerp_(normalized_grad, 1.0 - beta3_ema)
+
                 if beta1 > 0:
                     update = update_mt.add_(mt_slow, alpha=alpha)
                     del normalized_grad
                 else:
                     update = normalized_grad.add_(mt_slow, alpha=alpha)
-                # Factorize
-                state['mu_m_slow_nmf'], state['mv_m_slow_nmf'], state['sign_slow'] = _factorize_state(mt_slow, signed=True)
-                del mt_slow
+                if not factored_2nd:
+                    # Factorize
+                    state['mu_m_slow_nmf'], state['mv_m_slow_nmf'], state['sign_slow'] = _factorize_state(mt_slow, signed=True)
+                    del mt_slow
+
             elif self.Simplified_AdEMAMix:
                 update = update_mt.add_(normalized_grad, alpha=alpha_grad)
                 del normalized_grad
@@ -369,9 +429,6 @@ class Adopt_adv(torch.optim.Optimizer):
 
             update = update.view(p.shape)
 
-            update_scaling = lr * A if self.use_atan2 else lr
-            update.mul_(update_scaling)
-
         else: # Standard ADOPT logic for non-factored tensors
             vt = state['exp_avg_sq'] # v_{t-1}
 
@@ -418,12 +475,17 @@ class Adopt_adv(torch.optim.Optimizer):
                 else:
                     update = normalized_grad
 
-            update_scaling = lr * A if self.use_atan2 else lr
-            update.mul_(update_scaling)
 
             # Update second moment v_t for the next step using raw g_t
             vt.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
 
+        update_scaling = lr * A if self.use_atan2 else lr
+
+        if group.get('scaled_optm', False):
+            update = scale_update(p, update, update_scaling, vector_state=state.get('spectral_v'))
+        else:
+            update.mul_(update_scaling)
+
         # Parameter Update
         param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor)