chore: import upstream snapshot with attribution

swift/optimizers/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+from .base import OptimizerCallback
+from .mapping import optimizers_map

swift/optimizers/base.py

@@ -0,0 +1,56 @@
+from torch.optim import Optimizer
+from transformers.trainer import Trainer as HfTrainer
+from typing import TYPE_CHECKING
+
+try:
+    from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
+except ImportError:
+    from torch.optim.lr_scheduler import LRScheduler
+
+if TYPE_CHECKING:
+    from swift.trainers import Trainer, TrainingArguments
+
+
+class OptimizerCallback:
+    """
+    Callback for creating and managing optimizer and learning rate scheduler.
+
+    This callback provides hooks for customizing the creation of optimizers and
+    learning rate schedulers during the training process. It delegates to the
+    trainer's methods by default but can be subclassed to implement custom
+    optimization strategies.
+
+    Args:
+        args (TrainingArguments): The training arguments containing hyperparameters
+            and configuration settings.
+        trainer (Trainer): The trainer instance that will use this callback.
+    """
+
+    def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
+        self.args = args
+        self.trainer = trainer
+
+    def create_optimizer_and_scheduler(self, num_training_steps: int) -> None:
+        """
+        Create both optimizer and learning rate scheduler for training.
+
+        This method initializes the optimizer and scheduler by calling their
+        respective creation methods and assigns them to the trainer instance.
+
+        Args:
+            num_training_steps (int): The total number of training steps, used
+                for scheduler configuration (e.g., warmup steps, decay schedule).
+
+        Returns:
+            None: The optimizer and scheduler are set directly on the trainer.
+        """
+        trainer = self.trainer
+        trainer.optimizer = self.create_optimizer()
+        trainer.scheduler = self.create_scheduler(num_training_steps, trainer.optimizer)
+
+    def create_optimizer(self, model=None) -> Optimizer:
+        kwargs = {} if model is None else {'model': model}
+        return HfTrainer.create_optimizer(self.trainer, **kwargs)
+
+    def create_scheduler(self, num_training_steps: int, optimizer: Optimizer) -> LRScheduler:
+        return HfTrainer.create_scheduler(self.trainer, num_training_steps, optimizer)

swift/optimizers/galore/__init__.py

@@ -0,0 +1,8 @@
+from transformers.utils import is_bitsandbytes_available
+
+from .adafactor import GaLoreAdafactor
+from .adamw import GaLoreAdamW
+from .utils import GaLoreConfig, GaloreOptimizerCallback
+
+if is_bitsandbytes_available():
+    from .adamw8bit import GaLoreAdamW8bit

swift/optimizers/galore/adafactor.py

@@ -0,0 +1,271 @@
+# copy dependencies from transformers/optimization.py
+# code borrowed from https://github.com/jiaweizzhao/GaLore
+import math
+import torch
+from torch.optim import Optimizer
+from transformers.utils.versions import require_version
+
+from .galore_projector import GaLoreProjector
+
+
+class Adafactor(Optimizer):
+    """
+    AdaFactor pytorch implementation can be used as a drop in replacement for Adam original fairseq code:
+    https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py
+
+    Paper: *Adafactor: Adaptive Learning Rates with Sublinear Memory Cost* https://arxiv.org/abs/1804.04235 Note that
+    this optimizer internally adjusts the learning rate depending on the `scale_parameter`, `relative_step` and
+    `warmup_init` options. To use a manual (external) learning rate schedule you should set `scale_parameter=False` and
+    `relative_step=False`.
+
+    Arguments:
+        params (`Iterable[nn.parameter.Parameter]`):
+            Iterable of parameters to optimize or dictionaries defining parameter groups.
+        lr (`float`, *optional*):
+            The external learning rate.
+        eps (`Tuple[float, float]`, *optional*, defaults to `(1e-30, 0.001)`):
+            Regularization constants for square gradient and parameter scale respectively
+        clip_threshold (`float`, *optional*, defaults to 1.0):
+            Threshold of root mean square of final gradient update
+        decay_rate (`float`, *optional*, defaults to -0.8):
+            Coefficient used to compute running averages of square
+        beta1 (`float`, *optional*):
+            Coefficient used for computing running averages of gradient
+        weight_decay (`float`, *optional*, defaults to 0.0):
+            Weight decay (L2 penalty)
+        scale_parameter (`bool`, *optional*, defaults to `True`):
+            If True, learning rate is scaled by root mean square
+        relative_step (`bool`, *optional*, defaults to `True`):
+            If True, time-dependent learning rate is computed instead of external learning rate
+        warmup_init (`bool`, *optional*, defaults to `False`):
+            Time-dependent learning rate computation depends on whether warm-up initialization is being used
+
+    This implementation handles low-precision (FP16, bfloat) values, but we have not thoroughly tested.
+
+    Recommended T5 finetuning settings (https://discuss.huggingface.co/t/t5-finetuning-tips/684/3):
+
+        - Training without LR warmup or clip_threshold is not recommended.
+
+           - use scheduled LR warm-up to fixed LR
+           - use clip_threshold=1.0 (https://arxiv.org/abs/1804.04235)
+        - Disable relative updates
+        - Use scale_parameter=False
+        - Additional optimizer operations like gradient clipping should not be used alongside Adafactor
+
+    Example:
+
+    ```python
+    Adafactor(model.parameters(), scale_parameter=False, relative_step=False, warmup_init=False, lr=1e-3)
+    ```
+
+    Others reported the following combination to work well:
+
+    ```python
+    Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)
+    ```
+
+    When using `lr=None` with [`Trainer`] you will most likely need to use [`~optimization.AdafactorSchedule`]
+    scheduler as following:
+
+    ```python
+    from transformers.optimization import Adafactor, AdafactorSchedule
+
+    optimizer = Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)
+    lr_scheduler = AdafactorSchedule(optimizer)
+    trainer = Trainer(..., optimizers=(optimizer, lr_scheduler))
+    ```
+
+    Usage:
+
+    ```python
+    # replace AdamW with Adafactor
+    optimizer = Adafactor(
+        model.parameters(),
+        lr=1e-3,
+        eps=(1e-30, 1e-3),
+        clip_threshold=1.0,
+        decay_rate=-0.8,
+        beta1=None,
+        weight_decay=0.0,
+        relative_step=False,
+        scale_parameter=False,
+        warmup_init=False,
+    )
+    ```"""
+
+    def __init__(
+        self,
+        params,
+        lr=None,
+        eps=(1e-30, 1e-3),
+        clip_threshold=1.0,
+        decay_rate=-0.8,
+        beta1=None,
+        weight_decay=0.0,
+        scale_parameter=True,
+        relative_step=True,
+        warmup_init=False,
+    ):
+        require_version('torch>=1.5.0')  # add_ with alpha
+        if lr is not None and relative_step:
+            raise ValueError('Cannot combine manual `lr` and `relative_step=True` options')
+        if warmup_init and not relative_step:
+            raise ValueError('`warmup_init=True` requires `relative_step=True`')
+
+        defaults = {
+            'lr': lr,
+            'eps': eps,
+            'clip_threshold': clip_threshold,
+            'decay_rate': decay_rate,
+            'beta1': beta1,
+            'weight_decay': weight_decay,
+            'scale_parameter': scale_parameter,
+            'relative_step': relative_step,
+            'warmup_init': warmup_init,
+        }
+        super().__init__(params, defaults)
+
+    @staticmethod
+    def _get_lr(param_group, param_state):
+        rel_step_sz = param_group['lr']
+        if param_group['relative_step']:
+            min_step = 1e-6 * param_state['step'] if param_group['warmup_init'] else 1e-2
+            rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state['step']))
+        param_scale = 1.0
+        if param_group['scale_parameter']:
+            param_scale = max(param_group['eps'][1], param_state['RMS'])
+        return param_scale * rel_step_sz
+
+    @staticmethod
+    def _get_options(param_group, param_shape):
+        factored = len(param_shape) >= 2
+        use_first_moment = param_group['beta1'] is not None
+        return factored, use_first_moment
+
+    @staticmethod
+    def _rms(tensor):
+        return tensor.norm(2) / (tensor.numel()**0.5)
+
+    @staticmethod
+    def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col):
+        # copy from fairseq's adafactor implementation:
+        # https://github.com/huggingface/transformers/blob/8395f14de6068012787d83989c3627c3df6a252b/src/transformers/optimization.py#L505
+        r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
+        c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
+        return torch.mul(r_factor, c_factor)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """
+        Performs a single optimization step
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad
+                if grad.dtype in {torch.float16, torch.bfloat16}:
+                    grad = grad.float()
+                if grad.is_sparse:
+                    raise RuntimeError('Adafactor does not support sparse gradients.')
+
+                state = self.state[p]
+
+                if 'step' not in state:
+                    state['step'] = 0
+
+                # GaLore Projection
+                if 'rank' in group:
+                    if 'projector' not in state:
+                        state['projector'] = GaLoreProjector(
+                            group['rank'],
+                            update_proj_gap=group['update_proj_gap'],
+                            scale=group['scale'],
+                            proj_type=group['proj_type'])
+
+                    grad = state['projector'].project(grad, state['step'])
+
+                grad_shape = grad.shape
+
+                factored, use_first_moment = self._get_options(group, grad_shape)
+                # State Initialization
+                if 'RMS' not in state:
+                    state['step'] = 0
+
+                    if use_first_moment:
+                        # Exponential moving average of gradient values
+                        state['exp_avg'] = torch.zeros_like(grad)
+                    if factored:
+                        state['exp_avg_sq_row'] = torch.zeros(grad_shape[:-1]).to(grad)
+                        state['exp_avg_sq_col'] = torch.zeros(grad_shape[:-2] + grad_shape[-1:]).to(grad)
+                    else:
+                        state['exp_avg_sq'] = torch.zeros_like(grad)
+
+                    state['RMS'] = 0
+                else:
+                    if use_first_moment:
+                        state['exp_avg'] = state['exp_avg'].to(grad)
+                    if factored:
+                        state['exp_avg_sq_row'] = state['exp_avg_sq_row'].to(grad)
+                        state['exp_avg_sq_col'] = state['exp_avg_sq_col'].to(grad)
+                    else:
+                        state['exp_avg_sq'] = state['exp_avg_sq'].to(grad)
+
+                p_data_fp32 = p
+                if p.dtype in {torch.float16, torch.bfloat16}:
+                    p_data_fp32 = p_data_fp32.float()
+
+                state['step'] += 1
+                state['RMS'] = self._rms(p_data_fp32)
+                lr = self._get_lr(group, state)
+
+                beta2t = 1.0 - math.pow(state['step'], group['decay_rate'])
+                update = (grad**2) + group['eps'][0]
+                if factored:
+                    exp_avg_sq_row = state['exp_avg_sq_row']
+                    exp_avg_sq_col = state['exp_avg_sq_col']
+
+                    exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=(1.0 - beta2t))
+                    exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=(1.0 - beta2t))
+
+                    # Approximation of exponential moving average of square of gradient
+                    update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
+                    update.mul_(grad)
+                else:
+                    exp_avg_sq = state['exp_avg_sq']
+
+                    exp_avg_sq.mul_(beta2t).add_(update, alpha=(1.0 - beta2t))
+                    update = exp_avg_sq.rsqrt().mul_(grad)
+
+                update.div_((self._rms(update) / group['clip_threshold']).clamp_(min=1.0))
+                update.mul_(lr)
+
+                if use_first_moment:
+                    exp_avg = state['exp_avg']
+                    exp_avg.mul_(group['beta1']).add_(update, alpha=(1 - group['beta1']))
+                    update = exp_avg
+
+                # GaLore Projection Back
+                if 'rank' in group:
+                    update = state['projector'].project_back(update)
+
+                if group['weight_decay'] != 0:
+                    p_data_fp32.add_(p_data_fp32, alpha=(-group['weight_decay'] * lr))
+
+                p_data_fp32.add_(-update)
+
+                if p.dtype in {torch.float16, torch.bfloat16}:
+                    p.copy_(p_data_fp32)
+
+        return loss
+
+
+GaLoreAdafactor = Adafactor

swift/optimizers/galore/adamw.py

@@ -0,0 +1,140 @@
+# copy dependencies from transformers/optimization.py
+# code borrowed from https://github.com/jiaweizzhao/GaLore
+import math
+import torch
+from torch import nn
+from torch.optim import Optimizer
+from transformers.utils.versions import require_version
+from typing import Callable, Iterable, Tuple
+
+from .galore_projector import GaLoreProjector
+
+
+class AdamW(Optimizer):
+    """
+    Implements Adam algorithm with weight decay fix as introduced in [Decoupled Weight Decay
+    Regularization](https://arxiv.org/abs/1711.05101).
+
+    Parameters:
+        params (`Iterable[nn.parameter.Parameter]`):
+            Iterable of parameters to optimize or dictionaries defining parameter groups.
+        lr (`float`, *optional*, defaults to 0.001):
+            The learning rate to use.
+        betas (`Tuple[float,float]`, *optional*, defaults to `(0.9, 0.999)`):
+            Adam's betas parameters (b1, b2).
+        eps (`float`, *optional*, defaults to 1e-06):
+            Adam's epsilon for numerical stability.
+        weight_decay (`float`, *optional*, defaults to 0.0):
+            Decoupled weight decay to apply.
+        correct_bias (`bool`, *optional*, defaults to `True`):
+            Whether or not to correct bias in Adam (for instance, in Bert TF repository they use `False`).
+        no_deprecation_warning (`bool`, *optional*, defaults to `False`):
+            A flag used to disable the deprecation warning (set to `True` to disable the warning).
+    """
+
+    def __init__(
+        self,
+        params: Iterable[nn.parameter.Parameter],
+        lr: float = 1e-3,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        eps: float = 1e-6,
+        weight_decay: float = 0.0,
+        correct_bias: bool = True,
+        no_deprecation_warning: bool = False,
+    ):
+        require_version('torch>=1.5.0')  # add_ with alpha
+        if lr < 0.0:
+            raise ValueError(f'Invalid learning rate: {lr} - should be >= 0.0')
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f'Invalid beta parameter: {betas[0]} - should be in [0.0, 1.0)')
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f'Invalid beta parameter: {betas[1]} - should be in [0.0, 1.0)')
+        if not 0.0 <= eps:
+            raise ValueError(f'Invalid epsilon value: {eps} - should be >= 0.0')
+        defaults = {'lr': lr, 'betas': betas, 'eps': eps, 'weight_decay': weight_decay, 'correct_bias': correct_bias}
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure: Callable = None):
+        """
+        Performs a single optimization step.
+
+        Arguments:
+            closure (`Callable`, *optional*): A closure that reevaluates the model and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad
+                if grad.is_sparse:
+                    raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
+
+                state = self.state[p]
+
+                if 'step' not in state:
+                    state['step'] = 0
+
+                # GaLore Projection
+                if 'rank' in group:
+                    if 'projector' not in state:
+                        state['projector'] = GaLoreProjector(
+                            group['rank'],
+                            update_proj_gap=group['update_proj_gap'],
+                            scale=group['scale'],
+                            proj_type=group['proj_type'])
+
+                    grad = state['projector'].project(grad, state['step'])
+
+                # State initialization
+                if 'exp_avg' not in state:
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(grad)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(grad)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+
+                # Decay the first and second moment running average coefficient
+                # In-place operations to update the averages at the same time
+                exp_avg.mul_(beta1).add_(grad, alpha=(1.0 - beta1))
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+                denom = exp_avg_sq.sqrt().add_(group['eps'])
+
+                step_size = group['lr']
+                if group['correct_bias']:  # No bias correction for Bert
+                    bias_correction1 = 1.0 - beta1**state['step']
+                    bias_correction2 = 1.0 - beta2**state['step']
+                    step_size = step_size * math.sqrt(bias_correction2) / bias_correction1
+
+                # compute norm gradient
+                norm_grad = exp_avg / denom
+
+                # GaLore Projection Back
+                if 'rank' in group:
+                    norm_grad = state['projector'].project_back(norm_grad)
+
+                p.add_(norm_grad, alpha=-step_size)
+
+                # Just adding the square of the weights to the loss function is *not*
+                # the correct way of using L2 regularization/weight decay with Adam,
+                # since that will interact with the m and v parameters in strange ways.
+                #
+                # Instead we want to decay the weights in a manner that doesn't interact
+                # with the m/v parameters. This is equivalent to adding the square
+                # of the weights to the loss with plain (non-momentum) SGD.
+                # Add weight decay at the end (fixed version)
+                if group['weight_decay'] > 0.0:
+                    p.add_(p, alpha=(-group['lr'] * group['weight_decay']))
+
+        return loss
+
+
+GaLoreAdamW = AdamW

swift/optimizers/galore/adamw8bit.py

@@ -0,0 +1,113 @@
+# code borrowed from https://github.com/jiaweizzhao/GaLore
+import torch
+from bitsandbytes.optim.optimizer import Optimizer2State
+
+from swift.utils import synchronize
+from .galore_projector import GaLoreProjector
+
+
+class AdamW8bit(Optimizer2State):
+
+    def __init__(self,
+                 params,
+                 lr=1e-3,
+                 betas=(0.9, 0.999),
+                 eps=1e-8,
+                 weight_decay=1e-2,
+                 amsgrad=False,
+                 optim_bits=32,
+                 args=None,
+                 min_8bit_size=4096,
+                 percentile_clipping=100,
+                 block_wise=True,
+                 is_paged=False):
+        super().__init__(
+            'adam',
+            params,
+            lr,
+            betas,
+            eps,
+            weight_decay,
+            8,
+            args,
+            min_8bit_size,
+            percentile_clipping,
+            block_wise,
+            is_paged=is_paged)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        if not self.initialized:
+            self.check_overrides()
+            self.to_gpu()  # needed for fairseq pure fp16 training
+            self.initialized = True
+
+        # if self.is_paged: self.page_mng.prefetch_all()
+        for gindex, group in enumerate(self.param_groups):
+            for pindex, p in enumerate(group['params']):
+                if p.grad is None:
+                    continue
+                state = self.state[p]
+
+                if 'step' not in state:
+                    state['step'] = 0
+
+                # GaLore Projection
+                if 'rank' in group:
+                    if 'projector' not in state:
+                        state['projector'] = GaLoreProjector(
+                            group['rank'],
+                            update_proj_gap=group['update_proj_gap'],
+                            scale=group['scale'],
+                            proj_type=group['proj_type'])
+
+                    if 'weight_decay' in group and group['weight_decay'] > 0:
+                        # ensure that the weight decay is not applied to the norm grad
+                        group['weight_decay_saved'] = group['weight_decay']
+                        group['weight_decay'] = 0
+
+                    grad = state['projector'].project(p.grad, state['step'])
+
+                    # suboptimal implementation
+                    p.saved_data = p.data.clone()
+                    p.data = grad.clone().to(p.data.dtype).to(p.data.device)
+                    p.data.zero_()
+                    p.grad = grad
+
+                if 'state1' not in state:
+                    self.init_state(group, p, gindex, pindex)
+
+                self.prefetch_state(p)
+                self.update_step(group, p, gindex, pindex)
+                synchronize()
+
+                # GaLore Projection Back
+                if 'rank' in group:
+                    p.data = p.saved_data.add_(state['projector'].project_back(p.data))
+
+                    # apply weight decay
+                    if 'weight_decay_saved' in group:
+                        p.data.add_(p.data, alpha=-group['lr'] * group['weight_decay_saved'])
+                        group['weight_decay'] = group['weight_decay_saved']
+                        del group['weight_decay_saved']
+
+        if self.is_paged:
+            # all paged operation are asynchronous, we need
+            # to sync to make sure all tensors are in the right state
+            synchronize()
+
+        return loss
+
+
+GaLoreAdamW8bit = AdamW8bit

swift/optimizers/galore/galore_projector.py

@@ -0,0 +1,109 @@
+# code borrowed from https://github.com/jiaweizzhao/GaLore
+
+import torch
+
+
+class GaLoreProjector:
+
+    def __init__(self, rank, verbose=False, update_proj_gap=200, scale=1.0, proj_type='std'):
+        self.rank = rank
+        self.verbose = verbose
+        self.update_proj_gap = update_proj_gap
+        self.scale = scale
+        self.ortho_matrix = None
+        self.proj_type = proj_type
+
+    def project(self, full_rank_grad, iter):
+
+        if self.proj_type == 'std':
+            if full_rank_grad.shape[0] >= full_rank_grad.shape[1]:
+                if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
+                    self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='right')
+                low_rank_grad = torch.matmul(full_rank_grad, self.ortho_matrix.t())
+            else:
+                if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
+                    self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='left')
+                low_rank_grad = torch.matmul(self.ortho_matrix.t(), full_rank_grad)
+        elif self.proj_type == 'reverse_std':
+            if full_rank_grad.shape[0] >= full_rank_grad.shape[1]:
+                if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
+                    self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='left')
+                low_rank_grad = torch.matmul(self.ortho_matrix.t(), full_rank_grad)
+            else:
+                if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
+                    self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='right')
+                low_rank_grad = torch.matmul(full_rank_grad, self.ortho_matrix.t())
+        elif self.proj_type == 'right':
+            if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
+                self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='right')
+            low_rank_grad = torch.matmul(full_rank_grad, self.ortho_matrix.t())
+        elif self.proj_type == 'left':
+            if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
+                self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='left')
+            low_rank_grad = torch.matmul(self.ortho_matrix.t(), full_rank_grad)
+        elif self.proj_type == 'full':
+            if self.ortho_matrix is None or iter % self.update_proj_gap == 0:
+                self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='full')
+            low_rank_grad = torch.matmul(self.ortho_matrix[0].t(), full_rank_grad) @ self.ortho_matrix[1].t()
+
+        return low_rank_grad
+
+    def project_back(self, low_rank_grad):
+
+        if self.proj_type == 'std':
+            if low_rank_grad.shape[0] >= low_rank_grad.shape[1]:
+                full_rank_grad = torch.matmul(low_rank_grad, self.ortho_matrix)
+            else:
+                full_rank_grad = torch.matmul(self.ortho_matrix, low_rank_grad)
+        elif self.proj_type == 'reverse_std':
+            if low_rank_grad.shape[0] <= low_rank_grad.shape[1]:  # note this is different from std
+                full_rank_grad = torch.matmul(self.ortho_matrix, low_rank_grad)
+            else:
+                full_rank_grad = torch.matmul(low_rank_grad, self.ortho_matrix)
+        elif self.proj_type == 'right':
+            full_rank_grad = torch.matmul(low_rank_grad, self.ortho_matrix)
+        elif self.proj_type == 'left':
+            full_rank_grad = torch.matmul(self.ortho_matrix, low_rank_grad)
+        elif self.proj_type == 'full':
+            full_rank_grad = torch.matmul(self.ortho_matrix[0], low_rank_grad) @ self.ortho_matrix[1]
+
+        return full_rank_grad * self.scale
+
+    # svd decomposition
+    def get_orthogonal_matrix(self, weights, rank, type):
+        module_params = weights
+
+        if module_params.data.dtype != torch.float:
+            float_data = False
+            original_type = module_params.data.dtype
+            original_device = module_params.data.device
+            matrix = module_params.data.float()
+        else:
+            float_data = True
+            matrix = module_params.data
+
+        U, s, Vh = torch.linalg.svd(matrix, full_matrices=False)
+
+        # make the smaller matrix always to be orthogonal matrix
+        if type == 'right':
+            A = U[:, :rank] @ torch.diag(s[:rank])
+            B = Vh[:rank, :]
+
+            if not float_data:
+                B = B.to(original_device).type(original_type)
+            return B
+        elif type == 'left':
+            A = U[:, :rank]
+            B = torch.diag(s[:rank]) @ Vh[:rank, :]
+            if not float_data:
+                A = A.to(original_device).type(original_type)
+            return A
+        elif type == 'full':
+            A = U[:, :rank]
+            B = Vh[:rank, :]
+            if not float_data:
+                A = A.to(original_device).type(original_type)
+                B = B.to(original_device).type(original_type)
+            return [A, B]
+        else:
+            raise ValueError('type should be left, right or full')

swift/optimizers/galore/utils.py

@@ -0,0 +1,246 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import importlib
+import torch
+from dataclasses import dataclass
+from torch import nn
+from torch.optim import Optimizer
+from transformers import Trainer as HfTrainer
+from transformers import get_scheduler
+from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union
+
+from swift.trainers import calculate_max_steps
+from swift.utils import get_logger
+from ..base import OptimizerCallback
+
+try:
+    from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
+except ImportError:
+    from torch.optim.lr_scheduler import LRScheduler
+
+if TYPE_CHECKING:
+    from swift.trainers import TrainingArguments
+
+logger = get_logger()
+
+
+@dataclass
+class GaLoreConfig:
+    """
+    The configuration class for the Galore module.
+
+
+    See https://arxiv.org/abs/2403.03507
+
+    Args:
+        rank (`int`): The galore rank
+        target_modules (`Union[str, List[str]]`): The target modules to use, if `None`,
+            will use all attn and mlp linears
+        update_proj_gap(`int`): The projection update interval for galore
+        proj_type(`str`) The project type of Galore, valid values are `std`,
+            `reverse_std`, `right`, `left`, `full`
+        galore_scale(float): the scale of gradient
+        optim_per_parameter(bool): Gives one optimizer per parameter
+    """
+    rank: int = 128
+    target_modules: Union[str, List[str]] = None
+    update_proj_gap: int = 50
+    galore_scale: float = 1.0
+    proj_type: str = 'std'
+    optim_per_parameter: bool = False
+    quantize: bool = False
+    proj_quant: bool = False
+    proj_bits: int = 4
+    proj_group_size: int = 256
+    cos_threshold: float = 0.4
+    gamma_proj: int = 2
+    queue_size: int = 5
+
+
+class GaloreOptimizerWrapper(Optimizer):
+
+    def __init__(self, optimizers: Dict[Any, Optimizer]):
+        self.optimizers = optimizers
+        super().__init__([torch.tensor([1., 2., 3.])], {'lr': 1.})
+
+    def zero_grad(self, *args, **kwargs) -> None:
+        for optim in self.optimizers.values():
+            optim.zero_grad(*args, **kwargs)
+
+    def step(self, *args, **kwargs) -> None:
+        for optim in self.optimizers.values():
+            optim.step(*args, **kwargs)
+
+
+class GaloreSchedulerWrapper(LRScheduler):
+
+    def __init__(self, lr_schedulers: Dict[Any, LRScheduler]):
+        self.lr_schedulers = lr_schedulers
+
+    def step(self, *args, **kwargs) -> None:
+        for lr_scheduler in self.lr_schedulers.values():
+            lr_scheduler.step(*args, **kwargs)
+        self._last_lr = lr_scheduler.get_last_lr()
+
+
+def _create_optimizer_and_scheduler(model: nn.Module, args: 'TrainingArguments', config: GaLoreConfig, max_steps,
+                                    **defaults):
+    galore_params = []
+    for module_name, module in model.named_modules():
+        if not isinstance(module, (nn.Linear, nn.Embedding)) or \
+                not any(target_key in module_name for target_key in config.target_modules):
+            continue
+
+        if not module.weight.requires_grad:
+            continue
+
+        logger.info(f'Enable GaLore for weights in module: {module_name}')
+        galore_params.append(module.weight)
+
+    id_galore_params = [id(p) for p in galore_params]
+    galore_defaults = {
+        'rank': config.rank,
+        'update_proj_gap': config.update_proj_gap,
+        'scale': config.galore_scale,
+        'proj_type': config.proj_type,
+        **defaults
+    }
+    if config.quantize:
+        galore_defaults['quant'] = config.proj_quant
+        galore_defaults['quant_n_bit'] = config.proj_bits
+        galore_defaults['quant_group_size'] = config.proj_group_size
+        galore_defaults['cos_threshold'] = config.cos_threshold
+        galore_defaults['gamma_proj'] = config.gamma_proj
+        galore_defaults['queue_size'] = config.queue_size
+    optim_cls, optim_kwargs = get_optimizer(args, config)
+
+    if config.optim_per_parameter and not config.quantize:
+        # q-galore does not support optim_per_parameter
+        optimizer_dict = {}
+        galore_defaults['update_proj_gap'] = galore_defaults['update_proj_gap'] * 2
+        for p in model.parameters():
+            if p.requires_grad:
+                if id(p) in id_galore_params:
+                    optimizer_dict[p] = optim_cls([{'params': [p], **galore_defaults}], **optim_kwargs)
+                else:
+                    optimizer_dict[p] = optim_cls([{'params': [p], **defaults}], **optim_kwargs)
+
+        # get scheduler dict
+        scheduler_dict = {}
+        for p in model.parameters():
+            if p.requires_grad:
+                scheduler_dict[p] = get_scheduler(
+                    optimizer=optimizer_dict[p],
+                    name=args.lr_scheduler_type,
+                    num_training_steps=max_steps * 2,
+                    num_warmup_steps=args.warmup_steps * 2,
+                    scheduler_specific_kwargs=args.lr_scheduler_kwargs,
+                )
+
+        return GaloreOptimizerWrapper(optimizer_dict), GaloreSchedulerWrapper(scheduler_dict)
+    else:
+        decay_parameters = HfTrainer.get_decay_parameter_names(None, model)
+        param_groups = [{
+            'params': galore_params,
+            **galore_defaults,
+        }]
+        param_groups.extend([
+            {
+                'params': [
+                    p for n, p in model.named_parameters()
+                    if (n in decay_parameters and id(p) not in id_galore_params and p.requires_grad)
+                ],
+                'weight_decay':
+                defaults['weight_decay'],
+            },
+            {
+                'params': [
+                    p for n, p in model.named_parameters()
+                    if (n not in decay_parameters and id(p) not in id_galore_params and p.requires_grad)
+                ],
+                'weight_decay':
+                0.0,
+            },
+        ])
+        optim = optim_cls(param_groups, **optim_kwargs)
+        scheduler = get_scheduler(
+            optimizer=optim,
+            name=args.lr_scheduler_type,
+            num_training_steps=max_steps,
+            num_warmup_steps=args.warmup_steps,
+            scheduler_specific_kwargs=args.lr_scheduler_kwargs,
+        )
+        return optim, scheduler
+
+
+def get_optimizer(args: 'TrainingArguments', config: GaLoreConfig) -> Tuple[Any, Any]:
+    # parse args.optim_args
+    optim_args = {}
+    if args.optim_args:
+        for mapping in args.optim_args.replace(' ', '').split(','):
+            key, value = mapping.split('=')
+            optim_args[key] = value
+
+    optimizer_kwargs = {'lr': args.learning_rate}
+
+    adam_kwargs = {
+        'betas': (args.adam_beta1, args.adam_beta2),
+        'eps': args.adam_epsilon,
+    }
+    if args.optim == 'adafactor':
+        from .adafactor import GaLoreAdafactor
+        optimizer_cls = GaLoreAdafactor
+        optimizer_kwargs.update({'scale_parameter': False, 'relative_step': False})
+    elif args.optim in ('adamw_hf', 'adamw_torch', 'adamw_torch_fused'):
+        if config.quantize:
+            assert importlib.util.find_spec('q_galore_torch') is not None, \
+                'Please install q-galore by `pip install q_galore_torch`'
+            logger.info('If you encounter `absmax2` error, please downgrade your bitsandbytes to 0.40.0')
+            from swift.utils import get_dist_setting
+            _, _, world_size, _ = get_dist_setting()
+            if world_size > 1:
+                # from q_galore_torch import QGaLoreAdamW8bit_simulate as GaLoreAdamW
+                from q_galore_torch import QGaLoreAdamW8bit as GaLoreAdamW
+            else:
+                from q_galore_torch import QGaLoreAdamW8bit as GaLoreAdamW
+        else:
+            from .adamw import GaLoreAdamW
+        optimizer_cls = GaLoreAdamW
+        optimizer_kwargs.update(adam_kwargs)
+    elif 'adamw' in args.optim and '8bit' in args.optim:
+        try:
+            from .adamw8bit import GaLoreAdamW8bit
+            optimizer_cls = GaLoreAdamW8bit
+            optimizer_kwargs.update(adam_kwargs)
+            optimizer_kwargs.update({'optim_bits': 8, 'is_paged': 'paged' in args.optim})
+        except ImportError:
+            raise ValueError('Trainer tried to instantiate bnb optimizer but bnb is not installed!')
+    else:
+        raise ValueError(f'Galore not supported for optimizer type: {args.optim}')
+    return optimizer_cls, optimizer_kwargs
+
+
+class GaloreOptimizerCallback(OptimizerCallback):
+
+    def create_optimizer_and_scheduler(self, num_training_steps: int):
+        trainer = self.trainer
+        args = self.args
+        training_steps = calculate_max_steps(args, trainer.train_dataset)
+        galore_config = GaLoreConfig(
+            target_modules=args.galore_target_modules,
+            rank=args.galore_rank,
+            update_proj_gap=args.galore_update_proj_gap,
+            galore_scale=args.galore_scale,
+            proj_type=args.galore_proj_type,
+            optim_per_parameter=args.galore_optim_per_parameter,
+            quantize=args.galore_quantization,
+            proj_quant=args.galore_proj_quant,
+            proj_bits=args.galore_proj_bits,
+            proj_group_size=args.galore_proj_group_size,
+            cos_threshold=args.galore_cos_threshold,
+            gamma_proj=args.galore_gamma_proj,
+            queue_size=args.galore_queue_size,
+        )
+        optimizer, lr_scheduler = _create_optimizer_and_scheduler(
+            trainer.model, args, galore_config, training_steps, lr=args.learning_rate, weight_decay=args.weight_decay)
+        trainer.optimizer = optimizer
+        trainer.lr_scheduler = lr_scheduler

swift/optimizers/lorap.py

@@ -0,0 +1,34 @@
+from torch.optim import Optimizer
+from transformers.trainer import Trainer as HfTrainer
+
+from .base import OptimizerCallback
+
+
+class LorapOptimizerCallback(OptimizerCallback):
+
+    def create_optimizer(self, model=None) -> Optimizer:
+        args = self.args
+        if model is None:
+            model = self.trainer.model
+        optimizer_grouped_parameters = None
+        if hasattr(model, 'create_optimizer_param_groups'):
+            # Lora+ parameter groups
+            optimizer_grouped_parameters = model.create_optimizer_param_groups(
+                lr=args.learning_rate, weight_decay=args.weight_decay)
+
+        if optimizer_grouped_parameters is None:
+            # Default parameter groups
+            decay_parameters = HfTrainer.get_decay_parameter_names(None, model)
+            optimizer_grouped_parameters = [
+                {
+                    'params': [p for n, p in model.named_parameters() if (n in decay_parameters and p.requires_grad)],
+                    'weight_decay': args.weight_decay,
+                },
+                {
+                    'params':
+                    [p for n, p in model.named_parameters() if (n not in decay_parameters and p.requires_grad)],
+                    'weight_decay': 0.0,
+                },
+            ]
+        optimizer_cls, optimizer_kwargs = HfTrainer.get_optimizer_cls_and_kwargs(args)
+        return optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)

swift/optimizers/mapping.py

@@ -0,0 +1,16 @@
+from .base import OptimizerCallback
+from .galore import GaloreOptimizerCallback
+from .lorap import LorapOptimizerCallback
+from .multimodal import MultimodalOptimizerCallback
+from .muon import MuonOptimizerCallback
+from .muonclip import MuonClipOptimizerCallback
+
+# Add your own optimizers here, use --optimizer xxx to train
+optimizers_map = {
+    'default': OptimizerCallback,
+    'galore': GaloreOptimizerCallback,
+    'lorap': LorapOptimizerCallback,
+    'muon': MuonOptimizerCallback,
+    'muonclip': MuonClipOptimizerCallback,
+    'multimodal': MultimodalOptimizerCallback,
+}

swift/optimizers/multimodal.py

@@ -0,0 +1,74 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch.nn as nn
+from peft import PeftModel
+from transformers import Trainer as HfTrainer
+from typing import List, Optional, Tuple
+
+from swift.utils import get_logger
+from .base import OptimizerCallback
+
+logger = get_logger()
+
+
+def get_param_startswith(model,
+                         chosen_prefix: List[str],
+                         rejected_prefix: Optional[List[str]] = None) -> List[Tuple[str, nn.Parameter]]:
+    chosen_prefix = chosen_prefix or []
+    rejected_prefix = rejected_prefix or []
+    res = []
+    if not chosen_prefix:
+        return res
+    is_peft_model = isinstance(model, PeftModel)
+    if is_peft_model:
+        model = model.model
+    for n, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        is_rejected = False
+        for prefix in rejected_prefix:
+            if n.startswith(prefix):
+                is_rejected = True
+                break
+        if is_rejected:
+            continue
+        for prefix in chosen_prefix:
+            if n.startswith(prefix):
+                if is_peft_model:
+                    n = f'base_model.model.{n}'
+                res.append((n, p))
+                break
+    return res
+
+
+class MultimodalOptimizerCallback(OptimizerCallback):
+
+    def create_optimizer(self, model=None):
+        """ViT/Aligner/LLM use different learning rates."""
+        args = self.args
+        if model is None:
+            model = self.trainer.model
+        decay_parameters = set(HfTrainer.get_decay_parameter_names(None, model))
+        model_arch = model.model_meta.model_arch
+        vit_parameters = get_param_startswith(model, model_arch.vision_tower, model_arch.aligner)
+        aligner_parameters = get_param_startswith(model, model_arch.aligner)
+        llm_parameters = get_param_startswith(model, model_arch.language_model)
+        optimizer_grouped_parameters = []
+        vit_lr = args.vit_lr if args.vit_lr is not None else args.learning_rate
+        aligner_lr = args.aligner_lr if args.aligner_lr is not None else args.learning_rate
+        logger.info(f'vit_lr: {vit_lr}, aligner_lr: {aligner_lr}, llm_lr: {args.learning_rate}')
+        for lr, parameters in zip([vit_lr, aligner_lr, args.learning_rate],
+                                  [vit_parameters, aligner_parameters, llm_parameters]):
+            for use_wd, wd in zip([False, True], [0., args.weight_decay]):
+                if use_wd:
+                    params = [p for n, p in parameters if n in decay_parameters]
+                else:
+                    params = [p for n, p in parameters if n not in decay_parameters]
+                if not params:
+                    continue
+                optimizer_grouped_parameters.append({
+                    'params': params,
+                    'weight_decay': wd,
+                    'lr': lr,
+                })
+        optimizer_cls, optimizer_kwargs = HfTrainer.get_optimizer_cls_and_kwargs(args, model)
+        return optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)

swift/optimizers/muon.py

@@ -0,0 +1,46 @@
+import os
+import sys
+
+from swift.utils import git_clone_github
+from .base import OptimizerCallback
+
+
+class MuonOptimizerCallback(OptimizerCallback):
+
+    def create_optimizer(self, model=None):
+        args = self.args
+        if model is None:
+            model = self.trainer.model
+        if not args.local_repo_path:
+            args.local_repo_path = git_clone_github('https://github.com/MoonshotAI/Moonlight.git')
+        sys.path.append(os.path.join(args.local_repo_path, 'examples'))
+        from toy_train import Muon
+
+        # parse args.optim_args
+        optim_args = {}
+        if args.optim_args:
+            for mapping in args.optim_args.replace(' ', '').split(','):
+                key, value = mapping.split('=')
+                optim_args[key] = value
+
+        model_arch = model.model_meta.model_arch
+        embed_key = getattr(model_arch, 'embedding', None) or 'embed_tokens'
+        lm_head_key = getattr(model_arch, 'lm_head', None) or 'lm_head'
+        muon_params = [
+            p for n, p in model.named_parameters()
+            if p.requires_grad and p.ndim >= 2 and embed_key not in n and lm_head_key not in n
+        ]
+        adamw_params = [
+            p for n, p in model.named_parameters()
+            if p.requires_grad and not (p.ndim >= 2 and embed_key not in n and lm_head_key not in n)
+        ]
+
+        return Muon(
+            lr=args.learning_rate,
+            wd=args.weight_decay,
+            muon_params=muon_params,
+            adamw_params=adamw_params,
+            adamw_betas=(args.adam_beta1, args.adam_beta2),
+            adamw_eps=args.adam_epsilon,
+            **optim_args,
+        )

swift/optimizers/muonclip.py

@@ -0,0 +1,458 @@
+import math
+import threading
+import torch
+import torch.nn.functional as F
+from contextlib import suppress
+from torch.optim import Optimizer
+from typing import TYPE_CHECKING, Optional
+
+from .base import OptimizerCallback
+
+if TYPE_CHECKING:
+    from swift.trainers import TrainingArguments
+
+
+class _MaxLogitsTracker:
+    """
+    Collect a per-step scalar max logits value even when training loop can't pass it into optimizer.step().
+
+    - Eager attention: patch torch.softmax / F.softmax to capture exact softmax input max (attention scores).
+    - SDPA / FlashAttention: logits not exposed; record conservative upper bound via norms:
+        max(qk^T * scale) <= max||q|| * max||k|| * scale
+
+    Note: This is a GLOBAL scalar for the whole step (not per-layer, not per-head).
+    """
+
+    _tls = threading.local()
+    _enabled = False
+    _patched_softmax = False
+    _patched_sdpa = False
+    _patched_flash = False
+
+    _orig_torch_softmax = None
+    _orig_F_softmax = None
+    _orig_sdpa = None
+    _orig_flash_attn_func = None
+
+    @classmethod
+    def _get_and_reset(cls) -> Optional[float]:
+        v = getattr(cls._tls, 'max_logits', None)
+        cls._tls.max_logits = None
+        return v
+
+    @classmethod
+    def _update(cls, v: float):
+        if v is None:
+            return
+        cur = getattr(cls._tls, 'max_logits', None)
+        if cur is None or v > cur:
+            cls._tls.max_logits = float(v)
+
+    @classmethod
+    def enable_softmax(cls):
+        if cls._patched_softmax:
+            return
+        cls._patched_softmax = True
+
+        cls._orig_torch_softmax = torch.softmax
+        cls._orig_F_softmax = F.softmax
+
+        def _maybe_capture(x: torch.Tensor, dim):
+            # attention scores softmax: usually [B,H,Lq,Lk], dim=-1
+            if not isinstance(x, torch.Tensor):
+                return
+            if x.dim() != 4:
+                return
+            if dim is None or not (dim == -1 or dim == x.dim() - 1):
+                return
+
+            with suppress(Exception):
+                cls._update(float(x.detach().float().amax().item()))
+
+        def _torch_softmax(x, dim=None, dtype=None):
+            with suppress(Exception):
+                _maybe_capture(x, dim)
+            return cls._orig_torch_softmax(x, dim=dim, dtype=dtype)
+
+        def _F_softmax(x, dim=None, _stacklevel=3, dtype=None):
+            with suppress(Exception):
+                _maybe_capture(x, dim)
+            return cls._orig_F_softmax(x, dim=dim, _stacklevel=_stacklevel, dtype=dtype)
+
+        torch.softmax = _torch_softmax
+        F.softmax = _F_softmax
+
+    @classmethod
+    def enable_sdpa(cls):
+        if cls._patched_sdpa:
+            return
+        cls._patched_sdpa = True
+
+        if not hasattr(F, 'scaled_dot_product_attention'):
+            return
+
+        cls._orig_sdpa = F.scaled_dot_product_attention
+
+        def _sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False):
+            with suppress(Exception):
+                if isinstance(query, torch.Tensor) and isinstance(key, torch.Tensor):
+                    q = query.detach()
+                    k = key.detach()
+
+                    # upper bound using vector norms
+                    qn = q.float().norm(p=2, dim=-1).max().item()
+                    kn = k.float().norm(p=2, dim=-1).max().item()
+                    d = q.size(-1)
+                    s = float(scale) if scale is not None else (1.0 / math.sqrt(float(d)))
+                    cls._update(qn * kn * s)
+
+            return cls._orig_sdpa(
+                query,
+                key,
+                value,
+                attn_mask=attn_mask,
+                dropout_p=dropout_p,
+                is_causal=is_causal,
+                scale=scale,
+                enable_gqa=enable_gqa,
+            )
+
+        F.scaled_dot_product_attention = _sdpa
+
+    @classmethod
+    def enable_flash_attn(cls):
+        if cls._patched_flash:
+            return
+        cls._patched_flash = True
+
+        try:
+            import flash_attn.flash_attn_interface as _fai
+            flash_attn_func = _fai.flash_attn_func
+        except Exception:
+            return
+
+        cls._orig_flash_attn_func = flash_attn_func
+
+        def _flash_attn(q,
+                        k,
+                        v,
+                        dropout_p=0.0,
+                        softmax_scale=None,
+                        causal=False,
+                        window_size=(-1, -1),
+                        alibi_slopes=None,
+                        deterministic=False,
+                        return_attn_probs=False):
+            with suppress(Exception):
+                if isinstance(q, torch.Tensor) and isinstance(k, torch.Tensor):
+                    qn = q.detach().float().norm(p=2, dim=-1).max().item()
+                    kn = k.detach().float().norm(p=2, dim=-1).max().item()
+                    d = q.size(-1)
+                    s = float(softmax_scale) if softmax_scale is not None else (1.0 / math.sqrt(float(d)))
+                    cls._update(qn * kn * s)
+
+            return cls._orig_flash_attn_func(
+                q,
+                k,
+                v,
+                dropout_p=dropout_p,
+                softmax_scale=softmax_scale,
+                causal=causal,
+                window_size=window_size,
+                alibi_slopes=alibi_slopes,
+                deterministic=deterministic,
+                return_attn_probs=return_attn_probs,
+            )
+
+        _fai.flash_attn_func = _flash_attn
+
+    @classmethod
+    def enable_all(cls):
+        if cls._enabled:
+            return
+        cls._enabled = True
+        cls.enable_softmax()
+        cls.enable_sdpa()
+        cls.enable_flash_attn()
+
+    @classmethod
+    def consume(cls) -> Optional[float]:
+        return cls._get_and_reset()
+
+
+class MuonClip(Optimizer):
+    """
+    MuonClip (stable version):
+      - Muon-style update for apply_muon=True (2D weights): momentum buffer + Moonlight polynomial NS orthogonalization.
+      - Other params (apply_muon=False): simple momentum SGD (kept minimal; you can switch to AdamW if needed).
+      - QK-Clip uses a scalar max_logits (exact in eager, upper bound in sdpa/flash) and applies gamma_sqrt scaling
+        to Q/K weights marked with is_qk=True.
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float = 2e-4,
+        momentum: float = 0.95,
+        weight_decay: float = 0.1,
+        nesterov: bool = False,
+        newton_schulz_steps: int = 5,
+        qk_clip_tau: float = 10000.0,
+        qk_clip_enabled: bool = True,
+        rms_scale_factor: float = 0.2,
+    ):
+        defaults = dict(
+            lr=lr,
+            momentum=momentum,
+            weight_decay=weight_decay,
+            nesterov=nesterov,
+            newton_schulz_steps=newton_schulz_steps,
+            qk_clip_tau=qk_clip_tau,
+            qk_clip_enabled=qk_clip_enabled,
+            rms_scale_factor=rms_scale_factor,
+        )
+        super().__init__(params, defaults)
+        _MaxLogitsTracker.enable_all()
+
+    @staticmethod
+    @torch.no_grad()
+    def newton_schulz(G: torch.Tensor, steps: int = 5, eps: float = 1e-7) -> torch.Tensor:
+        """
+        Moonlight/Muon polynomial Newton-Schulz iteration (stable).
+        Works for rectangular matrices by transposing when needed.
+        """
+        # constants from your previous stable implementation
+        a, b, c = (3.4445, -4.7750, 2.0315)
+
+        X = G.bfloat16() / (G.norm() + eps)
+        transposed = False
+        if G.size(0) > G.size(1):
+            X = X.T
+            transposed = True
+
+        for _ in range(steps):
+            A = X @ X.T
+            B = b * A + c * A @ A
+            X = a * X + B @ X
+
+        if transposed:
+            X = X.T
+
+        return X.to(G.dtype)
+
+    def _is_qk_weight(self, group) -> bool:
+        return bool(group.get('is_qk', False))
+
+    @torch.no_grad()
+    def step(self, closure=None, max_logits: Optional[float] = None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        # fallback: collect scalar max_logits from tracker if not provided
+        if max_logits is None:
+            max_logits = _MaxLogitsTracker.consume()
+
+        for group in self.param_groups:
+            lr = float(group['lr'])
+            momentum = float(group['momentum'])
+            weight_decay = float(group['weight_decay'])
+            nesterov = bool(group.get('nesterov', False))
+            ns_steps = int(group.get('newton_schulz_steps', 5))
+            qk_clip_tau = float(group.get('qk_clip_tau', 10000.0))
+            qk_clip_enabled = bool(group.get('qk_clip_enabled', True))
+            apply_muon = bool(group.get('apply_muon', True))
+            is_qk_group = self._is_qk_weight(group)
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state['momentum_buffer'] = torch.zeros_like(p)
+                    state['step'] = 0
+
+                buf = state['momentum_buffer']
+                state['step'] += 1
+
+                buf.mul_(momentum).add_(grad)
+
+                # build update
+                if apply_muon and p.ndim >= 2:
+                    orth = self.newton_schulz(buf, steps=ns_steps)
+                    n, m = p.shape[0], p.shape[1]
+                    rms_scale_factor = float(group.get('rms_scale_factor', 0.2))
+                    rms_scale = math.sqrt(max(n, m)) * rms_scale_factor
+                    update = orth * rms_scale
+                else:
+                    update = buf
+
+                if nesterov:
+                    update = grad.add(update, alpha=momentum)
+
+                # decoupled-ish weight decay
+                if weight_decay != 0:
+                    p.mul_(1 - lr * weight_decay)
+
+                # QK-Clip (scalar)
+                if qk_clip_enabled and is_qk_group and (max_logits is not None):
+                    if max_logits > qk_clip_tau:
+                        gamma = qk_clip_tau / float(max_logits)
+                        gamma_sqrt = math.sqrt(gamma)
+                        # scale weight and update (matches your previous stable version)
+                        p.mul_(gamma_sqrt)
+                        update = update * gamma_sqrt
+
+                # apply update
+                p.add_(update, alpha=-lr)
+
+        return loss
+
+
+class MuonClipOptimizerCallback(OptimizerCallback):
+
+    def create_optimizer(self, model=None):
+        args = self.args
+        if model is None:
+            model = self.trainer.model
+        # parse args.optim_args
+        optim_args = {}
+        raw = getattr(args, 'optim_args', None)
+        if raw:
+            for mapping in raw.replace(' ', '').split(','):
+                if not mapping:
+                    continue
+                if '=' not in mapping:
+                    continue
+                key, value = mapping.split('=', 1)
+                if not key:
+                    continue
+
+                lv = value.lower()
+                if lv in ('true', 'false'):
+                    value = (lv == 'true')
+                else:
+                    try:
+                        f = float(value)
+                        value = int(f) if f.is_integer() else f
+                    except ValueError:
+                        pass
+
+                optim_args[key] = value
+
+        # resolve keys like create_muon_optimizer
+        model_arch = model.model_meta.model_arch
+        embed_key = getattr(model_arch, 'embedding', None) or 'embed_tokens'
+        lm_head_key = getattr(model_arch, 'lm_head', None) or 'lm_head'
+
+        # hyperparams (single-source of truth)
+        lr = args.learning_rate
+        weight_decay = optim_args.get('weight_decay', args.weight_decay)
+        momentum = optim_args.get('momentum', 0.95)
+        nesterov = optim_args.get('nesterov', False)
+        newton_schulz_steps = optim_args.get('newton_schulz_steps', 5)
+        qk_clip_tau = optim_args.get('qk_clip_tau', 100.0)
+        qk_clip_enabled = optim_args.get('qk_clip_enabled', True)
+        rms_scale_factor = optim_args.get('rms_scale_factor', 0.2)
+
+        # collect trainable params and group them
+        muon_named = []
+        rest_named = []
+
+        for name, p in model.named_parameters():
+            if not p.requires_grad:
+                continue
+
+            is_muon_candidate = (p.ndim >= 2 and embed_key not in name and lm_head_key not in name)
+            if is_muon_candidate:
+                muon_named.append((name, p))
+            else:
+                rest_named.append((name, p))
+
+        def _is_qk_name(name: str) -> bool:
+            ln = name.lower()
+            # qwen2.5/qwen3 common patterns
+            return ('q_proj' in ln) or ('k_proj' in ln) or ('.wq' in ln) or ('.wk' in ln) or ('/wq' in ln) or ('/wk'
+                                                                                                               in ln)
+
+        qk_muon_params = []
+        other_muon_params = []
+        for name, p in muon_named:
+            (qk_muon_params if _is_qk_name(name) else other_muon_params).append(p)
+
+        rest_params = [p for _, p in rest_named]
+
+        # build param groups
+        base_group_config = {
+            'lr': lr,
+            'momentum': momentum,
+            'weight_decay': weight_decay,
+            'nesterov': nesterov,
+            'newton_schulz_steps': newton_schulz_steps,
+            'qk_clip_tau': qk_clip_tau,
+            'qk_clip_enabled': qk_clip_enabled,
+            'rms_scale_factor': rms_scale_factor,
+        }
+        param_groups = []
+
+        if qk_muon_params:
+            group = base_group_config.copy()
+            group.update({
+                'params': qk_muon_params,
+                'apply_muon': True,
+                'is_qk': True,
+            })
+            param_groups.append(group)
+
+        if other_muon_params:
+            group = base_group_config.copy()
+            group.update({
+                'params': other_muon_params,
+                'apply_muon': True,
+                'is_qk': False,
+            })
+            param_groups.append(group)
+
+        if rest_params:
+            group = base_group_config.copy()
+            group.update({
+                'params': rest_params,
+                'apply_muon': False,
+                'is_qk': False,
+            })
+            param_groups.append(group)
+
+        # safety fallback
+        if not param_groups:
+            all_params = [p for _, p in model.named_parameters() if p.requires_grad]
+            param_groups = [{
+                'params': all_params,
+                'lr': lr,
+                'momentum': momentum,
+                'weight_decay': weight_decay,
+                'nesterov': nesterov,
+                'newton_schulz_steps': newton_schulz_steps,
+                'qk_clip_tau': qk_clip_tau,
+                'qk_clip_enabled': qk_clip_enabled,
+                'apply_muon': True,
+                'is_qk': False,
+            }]
+
+        # Only pass supported init kwargs; real behavior comes from param_groups
+        optimizer = MuonClip(
+            param_groups,
+            lr=lr,
+            momentum=momentum,
+            weight_decay=weight_decay,
+            nesterov=nesterov,
+            newton_schulz_steps=newton_schulz_steps,
+            qk_clip_tau=qk_clip_tau,
+            qk_clip_enabled=qk_clip_enabled,
+            rms_scale_factor=rms_scale_factor,
+        )
+
+        return optimizer