adv-optm 2.1.dev1__tar.gz → 2.1.dev3__tar.gz

{adv_optm-2.1.dev1 → adv_optm-2.1.dev3}/PKG-INFO

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

{adv_optm-2.1.dev1 → adv_optm-2.1.dev3}/adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "AdaMuon_adv",
 ]
 
-__version__ = "2.1.dev1"
+__version__ = "2.1.dev3"

{adv_optm-2.1.dev1 → adv_optm-2.1.dev3}/adv_optm/optim/AdaMuon_adv.py

@@ -9,7 +9,7 @@ from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.Kourkoutas import KourkoutasHelper
 from ..util import Muon_AuxAdam
 
-A = torch.as_tensor(4 / math.pi)
+A = 4 / math.pi
 
 class AdaMuon_adv(torch.optim.Optimizer):
     """
@@ -396,7 +396,7 @@ class AdaMuon_adv(torch.optim.Optimizer):
                     del denom, vt_buf
 
                 # RMS-aligned scaling
-                step_scale = lr * A if group['use_atan2'] else lr
+                step_scale = lr * A if group['use_atan2'] and not group['normuon_variant'] else lr
                 rms_adjustment(update, group['rms_rescaling'], step_scale)
 
                 update = update.reshape(p.shape)
@@ -454,14 +454,18 @@ class AdaMuon_adv(torch.optim.Optimizer):
                     del denom
 
                 # RMS-aligned rescaling
-                step_scale = lr * A if group['use_atan2'] else lr
+                step_scale = lr * A if group['use_atan2'] and not group['normuon_variant'] else lr
                 rms_adjustment(update, group['rms_rescaling'], step_scale)
 
                 update = update.reshape(original_shape)
 
             param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor)
 
-        compiled_muon_step_parameter(state, grad, group, group['lr'], random_int_tensor)
+        if group.get('compiled_optimizer', False):
+            lr = torch.as_tensor(group['lr'])
+        else:
+            lr = group['lr']
+        compiled_muon_step_parameter(state, grad, group, lr, random_int_tensor)
 
     @torch.no_grad()
     def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):

{adv_optm-2.1.dev1 → adv_optm-2.1.dev3}/adv_optm/optim/AdamW_adv.py

@@ -10,7 +10,7 @@ from ..util.update_util import _grams_update, _cautious_update
 from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.Kourkoutas import KourkoutasHelper
 
-A = torch.as_tensor(4 / math.pi)
+A = 4 / math.pi
 
 class AdamW_adv(torch.optim.Optimizer):
     """
@@ -233,7 +233,7 @@ class AdamW_adv(torch.optim.Optimizer):
         current_step = state['step']
         if group.get('kourkoutas_beta', False):
             # Call prepare_step() once at the beginning of the step for all params
-            self.kourkoutas_helper.maybe_prepare_step(current_step)
+            self.kourkoutas_helper.maybe_prepare_step(current_step, p.device)
             # Get the dynamic beta2 calculated in prepare_step()
             beta2 = self.kourkoutas_helper.get_beta2(p, group)
 
@@ -249,6 +249,7 @@ class AdamW_adv(torch.optim.Optimizer):
         random_int_tensor = None
 
         if group.get('compiled_optimizer', False):
+            step_size = torch.as_tensor(step_size)
             if p.dtype == torch.bfloat16 and self.stochastic_rounding:
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)

{adv_optm-2.1.dev1 → adv_optm-2.1.dev3}/adv_optm/optim/Adopt_adv.py

@@ -9,7 +9,7 @@ from ..util.OrthoGrad import _orthogonalize_gradient
 from ..util.Kourkoutas import KourkoutasHelper
 from ..util.update_util import _grams_update, _cautious_update
 
-A = torch.as_tensor(4 / math.pi)
+A = 4 / math.pi
 
 class Adopt_adv(torch.optim.Optimizer):
     """
@@ -258,7 +258,7 @@ class Adopt_adv(torch.optim.Optimizer):
         current_step = state['step']
         if group.get('kourkoutas_beta', False):
             # Call prepare_step() once at the beginning of the step for all params
-            self.kourkoutas_helper.maybe_prepare_step(current_step)
+            self.kourkoutas_helper.maybe_prepare_step(current_step, p.device)
             # Get the dynamic beta2 calculated in prepare_step()
             beta2 = self.kourkoutas_helper.get_beta2(p, group)
 
@@ -270,14 +270,15 @@ class Adopt_adv(torch.optim.Optimizer):
         random_int_tensor = None
 
         if group.get('compiled_optimizer', False):
+            lr = torch.as_tensor(group['lr'])
             if p.dtype == torch.bfloat16 and self.stochastic_rounding:
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
             step_param_fn = self._compiled_step_parameter
         else:
+            lr = group['lr']
             step_param_fn = self._step_parameter
 
-        lr = group['lr']
 
         step_param_fn(p, grad, state, group, lr, beta1, beta2, random_int_tensor)
 

{adv_optm-2.1.dev1 → adv_optm-2.1.dev3}/adv_optm/optim/Lion_Prodigy_adv.py

@@ -226,6 +226,7 @@ class Lion_Prodigy_adv(torch.optim.Optimizer):
                 random_int_tensor = param_update._get_random_int_for_sr(p)
             # TODO, workaround until pytorch#169634 is fixed
             d = torch.as_tensor(group['d'])
+            dlr = torch.as_tensor(group['dlr'])
             step_param_fn = self._compiled_step_parameter
         else:
             d = group['d']

{adv_optm-2.1.dev1 → adv_optm-2.1.dev3}/adv_optm/optim/Muon_adv.py

@@ -399,7 +399,11 @@ class Muon_adv(torch.optim.Optimizer):
 
             param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor)
 
-        compiled_muon_step_parameter(state, grad, group, group['lr'], random_int_tensor)
+        if group.get('compiled_optimizer', False):
+            lr = torch.as_tensor(group['lr'])
+        else:
+            lr = group['lr']
+        compiled_muon_step_parameter(state, grad, group, lr, random_int_tensor)
 
     @torch.no_grad()
     def step_parameter(self, p: torch.Tensor, group: dict, i: int | None = None):

{adv_optm-2.1.dev1 → adv_optm-2.1.dev3}/adv_optm/optim/Prodigy_adv.py

@@ -11,7 +11,7 @@ from ..util.Kourkoutas import KourkoutasHelper
 from ..util.factorization_util import _get_effective_shape, _reconstruct_state, _factorize_state
 from ..util.update_util import _grams_update, _cautious_update
 
-A = torch.as_tensor(4 / math.pi)
+A = 4 / math.pi
 
 class Prodigy_adv(torch.optim.Optimizer):
     """
@@ -327,7 +327,7 @@ class Prodigy_adv(torch.optim.Optimizer):
         current_step = state['step']
         if group.get('kourkoutas_beta', False):
             # Call prepare_step() once at the beginning of the step for all params
-            self.kourkoutas_helper.maybe_prepare_step(current_step)
+            self.kourkoutas_helper.maybe_prepare_step(current_step, p.device)
             # Get the dynamic beta2 calculated in prepare_step()
             beta2 = self.kourkoutas_helper.get_beta2(p, group)
         else:
@@ -343,6 +343,7 @@ class Prodigy_adv(torch.optim.Optimizer):
                 random_int_tensor = param_update._get_random_int_for_sr(p)
             # TODO, workaround until pytorch#169634 is fixed
             d = torch.as_tensor(group['d'])
+            dlr = torch.as_tensor(dlr)
             step_param_fn = self._compiled_step_parameter
         else:
             d = group['d']

{adv_optm-2.1.dev1 → adv_optm-2.1.dev3}/adv_optm/optim/Simplified_AdEMAMix.py

@@ -211,7 +211,7 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
         current_step = state['step']
         if group.get('kourkoutas_beta', False):
             # Call prepare_step() once at the beginning of the step for all params
-            self.kourkoutas_helper.maybe_prepare_step(current_step)
+            self.kourkoutas_helper.maybe_prepare_step(current_step, p.device)
             # Accumulate current grad's norm for the *next* step
             self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
             # Get the dynamic beta2 calculated in prepare_step()
@@ -244,7 +244,7 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
                 # Pre-generate random tensor for stochastic rounding if needed.
                 random_int_tensor = param_update._get_random_int_for_sr(p)
             # TODO, workaround until pytorch#169634 is fixed
-            sqrt_den_num = torch.as_tensor(sqrt_den_num)
+            lr = torch.as_tensor(lr)
             step_param_fn = self._compiled_step_parameter
         else:
             step_param_fn = self._step_parameter
@@ -289,10 +289,7 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
             state['mu_v_nmf'], state['mv_v_nmf'] = _factorize_state(vt, signed=False)
             del vt
 
-            if group['use_bias_correction']:
-                update.mul_(sqrt_den_num)
-
-            update = update.view(p.shape).mul_(lr)
+            update = update.view(p.shape).mul_(lr * sqrt_den_num)
 
         else:  # Standard optimizer logic for non-factored tensors
             exp_avg_sq = state['exp_avg_sq']
@@ -308,8 +305,7 @@ class Simplified_AdEMAMix(torch.optim.Optimizer):
             update.div_(denom)
             del denom
 
-            update_scaling = lr * sqrt_den_num if group['use_bias_correction'] else lr
-            update.mul_(update_scaling)
+            update.mul_(lr * sqrt_den_num)
 
         param_update.apply_parameter_update(self, p, group, update, lr, random_int_tensor=random_int_tensor)
 

adv_optm-2.1.dev3/adv_optm/util/Kourkoutas.py

@@ -0,0 +1,196 @@
+import torch
+from torch.optim import Optimizer
+
+class KourkoutasHelper:
+    """
+    A helper class to add layer-wise Kourkoutas-β functionality to a PyTorch optimizer.
+    """
+    def __init__(self, optimizer: Optimizer):
+        # We need a reference to the optimizer to access its param_groups and state
+        if not hasattr(optimizer, 'param_groups'):
+            raise TypeError("optimizer must be a valid torch.optim.Optimizer instance.")
+        self.optimizer = optimizer
+        self.layer_state = {}
+
+        self.layer_info = {}
+        self._layer_info_built = False
+        self._current_step_prepared = -1
+
+        # Store stats for external logging (e.g., TensorBoard)
+        self.last_beta2_stats = {}
+
+        # This ensures the map is complete before the first backward pass,
+        # making it compatible with fused back pass mechanisms.
+        self._build_layer_info_if_needed()
+
+    def _build_layer_info_if_needed(self):
+        """Builds a map of layers and the parameters they contain."""
+        if self._layer_info_built:
+            return
+
+        if hasattr(self.optimizer, 'layer_key_fn') and self.optimizer.layer_key_fn is not None:
+            # A custom key function was provided by the user. We will use it.
+            pass
+        else:
+            # No key function was provided. Default to coarse, shape-based bucketing.
+            self.optimizer.layer_key_fn = lambda p: \
+                (id(p),) if p.dim() == 2 and 1 <= p.shape[0] <= 10 and p.shape[1] in {768, 1280, 4096} \
+                else tuple(p.shape)
+            # This ensures that we won't mix embeddings with tokens (1 to 10)
+            # TODO find a better way to safeguard the embeddings
+
+        for group in self.optimizer.param_groups:
+            if not group.get('kourkoutas_beta', False) and not group.get('adam_kourkoutas_beta', False):
+                continue
+
+            for p in group['params']:
+                # The mapping is static and should not depend on the presence of a gradient.
+                layer_key = self.optimizer.layer_key_fn(p)
+                if layer_key not in self.layer_info:
+                    self.layer_info[layer_key] = {'params': [], 'group_ref': group}
+                self.layer_info[layer_key]['params'].append(p)
+
+        self._layer_info_built = True
+
+    def _get_or_init_layer_ema_tensor(self, layer_key, layer_params, device):
+        """
+        Retrieves the EMA tensor for this layer.
+        It handles synchronization between the internal layer_state and 
+        the external optimizer.state (which is required for state_dict saving/loading).
+        """
+        # Initialize container in layer_state if missing
+        if layer_key not in self.layer_state:
+            self.layer_state[layer_key] = {
+                'sum_sq_accumulator': torch.tensor(0.0, device=device, dtype=torch.float32)
+            }
+
+        internal_ema = self.layer_state[layer_key].get('kourkoutas_r_ema')
+
+        # Check optimizer.state for any existing state (e.g. from a loaded checkpoint)
+        # We check the first parameter in the list to see if it has state.
+        # If a checkpoint was loaded, optimizer.state[p] will contain the tensor.
+        representative_p = layer_params[0]
+        external_ema = self.optimizer.state[representative_p].get('kourkoutas_r_ema')
+
+        # Case A: Desync detected (Optimizer has state, but Internal doesn't, or they differ).
+        # This usually happens after load_state_dict(). We trust the optimizer.state.
+        if external_ema is not None and (internal_ema is None or internal_ema is not external_ema):
+            # Adopt the external tensor as our working tensor
+            self.layer_state[layer_key]['kourkoutas_r_ema'] = external_ema
+
+            # Ensure ALL params in this layer point to this exact tensor object
+            # (Fixes any fragmentation if only some params had state)
+            for p in layer_params:
+                self.optimizer.state[p]['kourkoutas_r_ema'] = external_ema
+
+            return external_ema
+
+        # Case B: No state anywhere. Create new.
+        if internal_ema is None:
+            new_ema = torch.tensor(0.0, device=device, dtype=torch.float32)
+            self.layer_state[layer_key]['kourkoutas_r_ema'] = new_ema
+
+            # Register this tensor in optimizer.state for ALL params so it gets saved
+            for p in layer_params:
+                self.optimizer.state[p]['kourkoutas_r_ema'] = new_ema
+
+            return new_ema
+
+        # Case C: Internal state exists and looks valid.
+        # We just need to ensure the link to optimizer.state is maintained (just in case).
+        # This is a cheap reference assignment.
+        for p in layer_params:
+            if 'kourkoutas_r_ema' not in self.optimizer.state[p]:
+                 self.optimizer.state[p]['kourkoutas_r_ema'] = internal_ema
+
+        return internal_ema
+
+    def prepare_step(self, current_step: int, device):
+        """
+        Calculates dynamic beta2 for all layers using the completed scalar accumulators
+        from the PREVIOUS step. Should be called once at the start of an optimizer step.
+        """
+        beta2_log = []
+        master_defaults = self.optimizer.defaults
+
+        for layer_key, info in self.layer_info.items():
+            group = info['group_ref']
+
+            if not group.get('kourkoutas_beta', False) and not group.get('adam_kourkoutas_beta', False):
+                continue
+
+            # Retrieve the EMA tensor. This function ensures the tensor is present
+            # in self.optimizer.state[p] for all parameters, ensuring state_dict support.
+            r_ema_tensor = self._get_or_init_layer_ema_tensor(layer_key, info['params'], device)
+
+            # Get accumulator
+            accumulator = self.layer_state[layer_key]['sum_sq_accumulator']
+            pooled_grad_norm = torch.sqrt(accumulator)
+
+            # Use group-specific K-b settings, falling back to the optimizer's master defaults.
+            # This makes the helper robust against param groups that enable kourkoutas_beta
+            # but are missing the other required hyperparameters.
+            # In hybrid optimizers like Muon_adv, the Kourkoutas-related keys in the
+            # defaults and param_groups are prefixed with 'adam_' to avoid conflicts.
+            # We must detect this case and use the correct key names.
+            prefix = 'adam_' if group.get('adam_kourkoutas_beta', False) else ''
+
+            ema_alpha = group.get(f'{prefix}ema_alpha', master_defaults[f'{prefix}ema_alpha'])
+            betas_tuple = group.get(f'{prefix}betas', master_defaults[f'{prefix}betas'])
+            beta2_max = betas_tuple[1]
+            beta2_min = group.get(f'{prefix}beta2_min', master_defaults[f'{prefix}beta2_min'])
+            tiny_spike = group.get(f'{prefix}tiny_spike', master_defaults[f'{prefix}tiny_spike'])
+            k_warmup_steps = group.get(f'{prefix}k_warmup_steps', master_defaults[f'{prefix}k_warmup_steps'])
+
+            # Update the persistent EMA tensor in-place.
+            r_ema_tensor.mul_(ema_alpha).add_(pooled_grad_norm, alpha=1.0 - ema_alpha)
+
+            # Calculate Beta2
+            if current_step < k_warmup_steps:
+                beta2 = beta2_max
+            else:
+                raw = pooled_grad_norm / (r_ema_tensor + tiny_spike)
+                sun = raw / (1.0 + raw)
+                beta2 = beta2_max - (beta2_max - beta2_min) * sun
+
+            # Store the final calculated beta2 in the helper's transient state for this step.
+            self.layer_state[layer_key]['dynamic_beta2'] = beta2.item() if isinstance(beta2, torch.Tensor) and not group.get('compiled_optimizer', False) else beta2
+
+            # Reset the accumulator for the next optimizer step.
+            accumulator.zero_()
+
+            beta2_log.append(self.layer_state[layer_key]['dynamic_beta2'])
+
+        # Compute stats for TensorBoard
+        if beta2_log:
+            beta2_tensor = torch.as_tensor(beta2_log, device='cpu')
+            self.last_beta2_stats = {
+                'mean': beta2_tensor.mean().item()
+                }
+
+    def maybe_prepare_step(self, current_step: int, device):
+        """
+        A universal guard that calls prepare_step() exactly once per training step.
+        """
+        if self._current_step_prepared < current_step:
+            self.prepare_step(current_step, device)
+            self._current_step_prepared = current_step
+
+    def accumulate_gradient_sq_norm(self, p: torch.Tensor, grad: torch.Tensor):
+        """
+        Accumulates the squared L2 norm of a single gradient for the next step's calculation.
+        """
+        layer_key = self.optimizer.layer_key_fn(p)
+
+        if layer_key in self.layer_info and layer_key in self.layer_state:
+            # Accumulate for the *next* step's prepare_step call
+            self.layer_state[layer_key]['sum_sq_accumulator'] += torch.sum(grad.detach().pow(2)).float()
+
+    def get_beta2(self, p: torch.Tensor, group: dict) -> float:
+        """
+        Gets the appropriate beta2 for the current parameter, handling warmup and dynamic value fetching.
+        """
+        layer_key = self.optimizer.layer_key_fn(p)
+        # The default is the max value, which is correct for unmapped params or edge cases
+        beta2_default = group.get('betas', group.get('adam_betas'))[1] if group.get('betas', group.get('adam_betas')) else 0.999
+        return self.layer_state.get(layer_key, {}).get('dynamic_beta2', beta2_default)

adv_optm-2.1.dev3/adv_optm/util/Muon_AuxAdam.py

@@ -0,0 +1,194 @@
+import torch
+
+import math
+
+from ..util import param_update
+from ..util.OrthoGrad import _orthogonalize_gradient
+from ..util.factorization_util import _get_effective_shape, _reconstruct_state, _factorize_state
+from ..util.update_util import _grams_update, _cautious_update
+
+A = 4 / math.pi
+
+@torch.no_grad()
+def _init_auxadam_state(self, p, group):
+    state = self.state[p]
+
+    state['step'] = 0
+
+    state['factored'] = (
+        group['adam_nnmf_factor'] and
+        not (len(p.shape) == 1 and not group['vector_reshape'])
+    )
+    dtype = torch.float32 if state['factored'] else p.dtype
+    device = p.device
+
+    if state['factored']:
+        state['effective_shape'] = _get_effective_shape(p.numel())
+        d1, d2 = state['effective_shape']
+        # First moment (m)
+        if group['adam_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)
+        if group.get('adam_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)
+        # 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)
+    else:  # Fallback to standard AdamW for non-factored tensors
+        if group['adam_betas'][0] > 0:
+            state['exp_avg'] = torch.zeros_like(p, device=device, dtype=dtype)
+        if group.get('adam_use_AdEMAMix'):
+            state['exp_avg_slow'] = torch.zeros_like(p, device=device, dtype=dtype)
+        state['exp_avg_sq'] = torch.zeros_like(p, device=device, dtype=dtype)
+
+
+@torch.no_grad()
+def _adam_step_parameter(self, p, grad, state, group, is_compiled, random_int_tensor):
+
+    step = state['step']
+
+    beta1_adam, beta2_adam = group['adam_betas']
+
+    if self.kourkoutas_helper:
+        # Prepare Kourkoutas-β once per optimizer step.
+        self.kourkoutas_helper.maybe_prepare_step(step, p.device)
+        # Get the dynamic beta2_adam calculated in prepare_step()
+        beta2_adam = self.kourkoutas_helper.get_beta2(p, group)
+
+    if group['adam_use_bias_correction']:
+        current_step = step + 1
+        beta1_adam, beta2_adam = group['adam_betas']
+        bias_correction1 = 1.0 - beta1_adam ** current_step
+        sqrt_bias_correction2 = (1.0 - beta2_adam ** current_step)**0.5
+    else:
+        bias_correction1 = 1.0
+        sqrt_bias_correction2 = 1.0
+
+    state['step'] += 1
+
+    step_size = group['lr'] / bias_correction1
+
+    if group.get('compiled_optimizer', False):
+        step_size = torch.as_tensor(step_size)
+
+    @torch.compile(fullgraph=True, disable= not is_compiled)
+    def compiled_muon_step_parameter(state, grad, group, step_size, sqrt_bias_correction2, random_int_tensor):
+        if grad.dtype != torch.float32 and state.get('factored', False):
+            grad = grad.float()
+        if group.get("adam_orthogonal_gradient"):
+            grad = _orthogonalize_gradient(p, grad)
+
+        if self.kourkoutas_helper:
+            # Accumulate current grad's norm for the *next* step
+            self.kourkoutas_helper.accumulate_gradient_sq_norm(p, grad)
+
+        if group.get('adam_use_AdEMAMix'):
+            beta3_ema = group['adam_beta3_ema']
+            alpha = group['adam_alpha']
+
+        if state['factored']:
+            d1, d2 = state['effective_shape']
+            grad_reshaped = grad.view(d1, d2)
+
+            # Reconstruct momentum from previous step's factors
+            if beta1_adam > 0:
+                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_adam)
+
+                # Factorize
+                state['mu_m_nmf'], state['mv_m_nmf'], state['sign'] = _factorize_state(mt.clone(), signed=True)
+
+                if group.get('adam_grams_moment'):
+                    update_mt = _grams_update(mt, grad_reshaped, inplace=True)
+                elif group.get('adam_cautious_mask'):
+                    update_mt = _cautious_update(mt, grad_reshaped, inplace=True)
+                else:
+                    update_mt = mt
+
+            vt = _reconstruct_state((state['mu_v_nmf'], state['mv_v_nmf']), signed=False)
+            vt.mul_(beta2_adam).addcmul_(grad_reshaped, grad_reshaped, value=1.0 - beta2_adam)
+
+            if group.get('adam_use_AdEMAMix'):
+                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)
+
+                if beta1_adam > 0:
+                    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
+            else:
+                if beta1_adam > 0:
+                    update = update_mt
+                else:
+                    update = grad_reshaped.clone()
+
+            if group['adam_use_atan2']:
+                denom = vt.sqrt()
+                denom.div_(sqrt_bias_correction2)
+                update.atan2_(denom)
+            else:
+                denom = vt.sqrt()
+                denom.div_(sqrt_bias_correction2).add_(group['adam_eps'])
+                update.div_(denom)
+            del denom
+
+            # Factorize
+            state['mu_v_nmf'], state['mv_v_nmf'] = _factorize_state(vt, signed=False)
+            del vt
+
+            update_scaling = step_size * A if group['use_atan2'] else step_size
+            update = update.view(p.shape).mul_(update_scaling)
+
+        else:  # Standard AdamW logic for non-factored tensors
+            if beta1_adam > 0:
+                exp_avg = state['exp_avg']
+                exp_avg.lerp_(grad, 1.0 - beta1_adam)
+
+                if group.get('adam_grams_moment'):
+                    update_mt = _grams_update(exp_avg, grad)
+                elif group.get('adam_cautious_mask'):
+                    update_mt = _cautious_update(exp_avg, grad)
+                else:
+                    update_mt = exp_avg.clone()
+
+            if group.get('adam_use_AdEMAMix'):
+                exp_avg_slow = state['exp_avg_slow']
+                exp_avg_slow.lerp_(grad, 1.0 - beta3_ema)
+
+                if beta1_adam > 0:
+                    update = update_mt.add_(exp_avg_slow, alpha=alpha)
+                else:
+                    update = torch.add(grad, exp_avg_slow, alpha=alpha)
+            else:
+                update = update_mt if beta1_adam > 0 else grad.clone()
+
+            exp_avg_sq = state['exp_avg_sq']
+            exp_avg_sq.mul_(beta2_adam).addcmul_(grad, grad, value=1 - beta2_adam)
+
+            if group.get('adam_use_atan2'):
+                denom = exp_avg_sq.sqrt()
+                denom.div_(sqrt_bias_correction2)
+                update.atan2_(denom)
+            else:
+                denom = exp_avg_sq.sqrt()
+                denom.div_(sqrt_bias_correction2).add_(group['adam_eps'])
+                update.div_(denom)
+            del denom
+
+            update_scaling = step_size * A if group['adam_use_atan2'] else step_size
+            update.mul_(update_scaling)
+
+        param_update.apply_parameter_update(self, p, group, update, step_size, group["adam_weight_decay"], random_int_tensor=random_int_tensor)
+
+    compiled_muon_step_parameter(state, grad, group, step_size, sqrt_bias_correction2, random_int_tensor)