[Feature] Implement DistributedMuon optimizer (#418)

* build(deps): docs requirements

* feature: implement DistributedMuon optimizer

* docs: v3.8.0 changelog

* feature: implement StochasticAccumulator

* update: DistributedMuon optimizer

* docs: DistributedMuon optimizer

* docs: v3.8.0 changelog

* update: test recipe

* refactor: remove unused var

* docs: v3.8.0 changelog

* update: test_stochastic_accumulation_hook

* update: recipes

* refactor: exception message format

* update: mse to bce

Author: kozistr

docs/changelogs/v3.8.0.md

@@ -7,6 +7,9 @@
     * You can use this variant by setting `decoupling_c` parameter in the `ScheduleFreeAdamW` optimizer.
 * Add more built-in optimizers, `NAdam`, `RMSProp`, and `LBFGS` optimizers. (#415)
 * Support `cautious` variant for `Muon` optimizer. (#417)
+* Separate distributed functionality from `Muon` to `DistribtuedMuon` optimizer. (#418)
+* Implement `StochasticAccumulator`, which is a gradient hook. (#418)
+    * [stochastic optimizer](https://github.com/lodestone-rock/torchastic/)
 
 ### Update
 

docs/optimizer.md

@@ -304,6 +304,10 @@
     :docstring:
     :members:
 
+::: pytorch_optimizer.DistributedMuon
+    :docstring:
+    :members:
+
 ::: pytorch_optimizer.AdaMuon
     :docstring:
     :members:

pytorch_optimizer/__init__.py

@@ -116,6 +116,7 @@
     DAdaptSGD,
     DeMo,
     DiffGrad,
+    DistributedMuon,
     DynamicLossScaler,
     EmoFact,
     EmoLynx,

pytorch_optimizer/base/optimizer.py

@@ -64,7 +64,7 @@ def set_hessian(param_groups: PARAMETERS, state: STATE, hessian: List[torch.Tens
             for p in group['params']:
                 if p.size() != hessian[i].size():
                     raise ValueError(
-                        f'[-] the shape of parameter and hessian does not match. {p.size()} vs {hessian[i].size()}'
+                        f'the shape of parameter and hessian does not match. {p.size()} vs {hessian[i].size()}'
                     )
 
                 state[p]['hessian'] = hessian[i]
@@ -312,35 +312,35 @@ def get_stable_adamw_rms(grad: torch.Tensor, exp_avg_sq: torch.Tensor, eps: floa
     @staticmethod
     def validate_range(x: float, name: str, low: float, high: float, range_type: str = '[)') -> None:
         if range_type == '[)' and not low <= x < high:
-            raise ValueError(f'[-] {name} must be in the range [{low}, {high})')
+            raise ValueError(f'{name} must be in the range [{low}, {high})')
         if range_type == '[]' and not low <= x <= high:
-            raise ValueError(f'[-] {name} must be in the range [{low}, {high}]')
+            raise ValueError(f'{name} must be in the range [{low}, {high}]')
         if range_type == '(]' and not low < x <= high:
-            raise ValueError(f'[-] {name} must be in the range ({low}, {high}]')
+            raise ValueError(f'{name} must be in the range ({low}, {high}]')
         if range_type == '()' and not low < x < high:
-            raise ValueError(f'[-] {name} must be in the range ({low}, {high})')
+            raise ValueError(f'{name} must be in the range ({low}, {high})')
 
     @staticmethod
     def validate_non_negative(x: Optional[float], name: str) -> None:
         if x is not None and x < 0.0:
-            raise ValueError(f'[-] {name} must be non-negative')
+            raise ValueError(f'{name} must be non-negative')
 
     @staticmethod
     def validate_non_positive(x: Optional[float], name: str) -> None:
         if x is not None and x > 0.0:
-            raise ValueError(f'[-] {name} must be non-positive')
+            raise ValueError(f'{name} must be non-positive')
 
     @staticmethod
     def validate_positive(x: Union[float, int], name: str) -> None:
         if x <= 0:
-            raise ValueError(f'[-] {name} must be positive')
+            raise ValueError(f'{name} must be positive')
 
     @staticmethod
     def validate_boundary(constant: float, boundary: float, bound_type: str = 'upper') -> None:
         if bound_type == 'upper' and constant > boundary:
-            raise ValueError(f'[-] constant {constant} must be in a range of (-inf, {boundary}]')
+            raise ValueError(f'constant {constant} must be in a range of (-inf, {boundary}]')
         if bound_type == 'lower' and constant < boundary:
-            raise ValueError(f'[-] constant {constant} must be in a range of [{boundary}, inf)')
+            raise ValueError(f'constant {constant} must be in a range of [{boundary}, inf)')
 
     @staticmethod
     def validate_step(step: int, step_type: str) -> None:
@@ -351,7 +351,7 @@ def validate_step(step: int, step_type: str) -> None:
     def validate_options(x: str, name: str, options: List[str]) -> None:
         if x not in options:
             opts: str = ' or '.join([f"'{option}'" for option in options]).strip()
-            raise ValueError(f'[-] {name} {x} must be one of ({opts})')
+            raise ValueError(f'{name} {x} must be one of ({opts})')
 
     @staticmethod
     def validate_learning_rate(learning_rate: Optional[float]) -> None:
@@ -361,7 +361,7 @@ def validate_learning_rate(learning_rate: Optional[float]) -> None:
     @staticmethod
     def validate_mod(x: int, y: int) -> None:
         if x % y != 0:
-            raise ValueError(f'[-] {x} must be divisible by {y}')
+            raise ValueError(f'{x} must be divisible by {y}')
 
     def validate_betas(self, betas: BETAS, beta_range_type: str = '[)', beta3_range_type: str = '[]') -> None:
         if betas[0] is not None:

pytorch_optimizer/optimizer/__init__.py

@@ -67,7 +67,7 @@
 from pytorch_optimizer.optimizer.madgrad import MADGRAD
 from pytorch_optimizer.optimizer.mars import MARS
 from pytorch_optimizer.optimizer.msvag import MSVAG
-from pytorch_optimizer.optimizer.muon import AdaMuon, Muon, prepare_muon_parameters
+from pytorch_optimizer.optimizer.muon import AdaMuon, DistributedMuon, Muon, prepare_muon_parameters
 from pytorch_optimizer.optimizer.nero import Nero
 from pytorch_optimizer.optimizer.novograd import NovoGrad
 from pytorch_optimizer.optimizer.orthograd import OrthoGrad
@@ -164,6 +164,7 @@
     DAdaptSGD,
     DeMo,
     DiffGrad,
+    DistributedMuon,
     EXAdam,
     EmoFact,
     EmoLynx,

pytorch_optimizer/optimizer/muon.py

@@ -3,6 +3,7 @@
 
 import torch
 from torch import nn
+from torch.distributed import all_gather, get_rank, get_world_size
 from torch.optim import Optimizer
 
 from pytorch_optimizer.base.exception import NoComplexParameterError, NoSparseGradientError
@@ -216,6 +217,209 @@ def step(self, closure: CLOSURE = None) -> LOSS:
         return loss
 
 
+class DistributedMuon(BaseOptimizer):  # pragma: no cover
+    r"""Distributed Momentum Orthogonalized by Newton-schulz.
+
+    Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-processing step, in which
+    each 2D parameter's update is replaced with the nearest orthogonal matrix. To efficiently orthogonalize each
+    update, we use a Newton-Schulz iteration, which has the advantage that it can be stably run in bfloat16 on the GPU.
+
+    Muon is intended to optimize only the internal ≥2D parameters of a network. Embeddings, classifier heads, and
+    scalar or vector parameters should be optimized using AdamW.
+
+    Some warnings:
+    - We believe this optimizer is unlikely to work well for training with small batch size.
+    - We believe it may not work well for fine-tuning pretrained models, but we haven't tested this.
+
+    Example:
+    -------
+        from pytorch_optimizer import DistributedMuon
+
+        hidden_weights = [p for p in model.body.parameters() if p.ndim >= 2]
+        hidden_gains_biases = [p for p in model.body.parameters() if p.ndim < 2]
+        non_hidden_params = [*model.head.parameters(), *model.embed.parameters()]
+
+        param_groups = [
+            dict(params=hidden_weights, lr=0.02, weight_decay=0.01, use_muon=True),
+            dict(
+                params=hidden_gains_biases + non_hidden_params,
+                lr=3e-4,
+                betas=(0.9, 0.95),
+                weight_decay=0.01,
+                use_muon=False,
+            ),
+        ]
+
+        optimizer = Muon(param_groups)
+
+    :param params: PARAMETERS. the parameters to be optimized by Muon.
+    :param lr: float. learning rate.
+    :param momentum: float. the momentum used by the internal SGD.
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
+    :param nesterov: bool. whether to use nesterov momentum.
+    :param ns_steps: int. the number of Newton-Schulz iterations to run. (5 is probably always enough)
+    :param use_adjusted_lr: bool. whether to use adjusted learning rate, which is from the Moonlight.
+        reference: https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py
+    :param adamw_lr: float. The learning rate for the internal AdamW.
+    :param adamw_betas: The betas for the internal AdamW.
+    :param adamw_wd: float. The weight decay for the internal AdamW.
+    :param adamw_eps: float. The epsilon for the internal AdamW.
+    :param maximize: bool. maximize the objective with respect to the params, instead of minimizing.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 2e-2,
+        momentum: float = 0.95,
+        weight_decay: float = 0.0,
+        weight_decouple: bool = True,
+        nesterov: bool = True,
+        ns_steps: int = 5,
+        use_adjusted_lr: bool = False,
+        adamw_lr: float = 3e-4,
+        adamw_betas: BETAS = (0.9, 0.95),
+        adamw_wd: float = 0.0,
+        adamw_eps: float = 1e-10,
+        maximize: bool = False,
+        **kwargs,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_learning_rate(adamw_lr)
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_range(momentum, 'momentum', 0.0, 1.0, range_type='[)')
+        self.validate_positive(ns_steps, 'ns_steps')
+        self.validate_betas(adamw_betas)
+        self.validate_non_negative(adamw_wd, 'adamw_wd')
+        self.validate_non_negative(adamw_eps, 'adamw_eps')
+
+        self.maximize = maximize
+
+        self.world_size: int = get_world_size()
+        self.rank: int = get_rank()
+
+        for group in params:
+            if 'use_muon' not in group:
+                raise ValueError('`use_muon` must be set.')
+
+            if group['use_muon']:
+                group['lr'] = group.get('lr', lr)
+                group['momentum'] = group.get('momentum', momentum)
+                group['nesterov'] = group.get('nesterov', nesterov)
+                group['weight_decay'] = group.get('weight_decay', weight_decay)
+                group['ns_steps'] = group.get('ns_steps', ns_steps)
+                group['use_adjusted_lr'] = group.get('use_adjusted_lr', use_adjusted_lr)
+            else:
+                group['lr'] = group.get('lr', adamw_lr)
+                group['betas'] = group.get('betas', adamw_betas)
+                group['eps'] = group.get('eps', adamw_eps)
+                group['weight_decay'] = group.get('weight_decay', adamw_wd)
+
+            group['weight_decouple'] = group.get('weight_decouple', weight_decouple)
+
+        super().__init__(params, kwargs)
+
+    def __str__(self) -> str:
+        return 'DistributedMuon'
+
+    def init_group(self, group: GROUP, **kwargs) -> None:
+        for p in group['params']:
+            if p.grad is None:
+                p.grad = torch.zeros_like(p)
+
+            grad = p.grad
+            if grad.is_sparse:
+                raise NoSparseGradientError(str(self))
+
+            if torch.is_complex(p):
+                raise NoComplexParameterError(str(self))
+
+            state = self.state[p]
+
+            if len(state) == 0 and not group['use_muon']:
+                state['exp_avg'] = torch.zeros_like(p)
+                state['exp_avg_sq'] = torch.zeros_like(p)
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None) -> LOSS:
+        loss: LOSS = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            if 'step' not in group:
+                self.init_group(group)
+                group['step'] = 1
+            else:
+                group['step'] += 1
+
+            if group['use_muon']:
+                params = group['params']
+                padded_params = params + [torch.empty_like(params[-1])] * (
+                    self.world_size - len(params) % self.world_size
+                )
+
+                for i in range(len(params))[:: self.world_size]:
+                    if i + self.rank < len(params):
+                        p = params[i + self.rank]
+
+                        grad = p.grad
+
+                        self.maximize_gradient(grad, maximize=self.maximize)
+
+                        state = self.state[p]
+                        if len(state) == 0:
+                            state['momentum_buffer'] = torch.zeros_like(p)
+
+                        self.apply_weight_decay(
+                            p,
+                            grad=grad,
+                            lr=group['lr'],
+                            weight_decay=group['weight_decay'],
+                            weight_decouple=group['weight_decouple'],
+                            fixed_decay=False,
+                        )
+
+                        buf = state['momentum_buffer']
+                        buf.lerp_(grad, weight=1.0 - group['momentum'])
+
+                        update = grad.lerp_(buf, weight=group['momentum']) if group['nesterov'] else buf
+                        if update.ndim > 2:
+                            update = update.view(len(update), -1)
+
+                        update = zero_power_via_newton_schulz_5(update, num_steps=group['ns_steps'])
+
+                        if group.get('cautious'):
+                            self.apply_cautious(update, grad)
+
+                        lr: float = get_adjusted_lr(group['lr'], p.size(), use_adjusted_lr=group['use_adjusted_lr'])
+
+                        p.add_(update.reshape(p.shape), alpha=-lr)
+
+                    all_gather(padded_params[i:i + self.world_size], padded_params[i:i + self.rank])  # fmt: skip
+            else:
+                for p in group['params']:
+                    grad = p.grad
+
+                    exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+
+                    beta1, beta2 = group['betas']
+
+                    bias_correction1: float = self.debias(beta1, group['step'])
+                    bias_correction2_sq: float = math.sqrt(self.debias(beta2, group['step']))
+
+                    exp_avg.lerp_(grad, weight=1.0 - beta1)
+                    exp_avg_sq.lerp_(grad.square(), weight=1.0 - beta2)
+
+                    de_nom = exp_avg_sq.sqrt().add_(group['eps']).div_(bias_correction2_sq)
+
+                    p.addcdiv_(exp_avg / bias_correction1, de_nom, value=-group['lr'])
+
+        return loss
+
+
 class AdaMuon(BaseOptimizer):
     r"""Adaptive Muon optimizer.
 

pytorch_optimizer/optimizer/utils.py

@@ -157,6 +157,57 @@ def load_state_dict(self, state_dict):
             optim.load_state_dict(optim_state_dict)
 
 
+class StochasticAccumulator:
+    r"""Stochastic accumulator.
+
+    Example:
+    -------
+        model = YourModel()
+
+        # apply stochastic grad accumulator hooks
+        StochasticAccumulator.assign_hooks(model)
+
+        while True:
+            loss = model.loss(*your_model_input)
+            for _ in range(grad_accum_length):
+                loss.backward()
+
+            StochasticAccumulator.reassign_grad_buffer(model)
+
+            optimizer.step()
+            optimizer.zero_grad()
+    """
+
+    @staticmethod
+    def stochastic_grad_accum(p: torch.Tensor) -> None:
+        if hasattr(p, 'acc_grad'):
+            acc_grad_fp32 = p.acc_grad.clone().to(torch.float32)
+            acc_grad_fp32.add_(p.grad.to(torch.float32))
+
+            copy_stochastic(p.acc_grad, acc_grad_fp32)
+
+            del acc_grad_fp32
+        else:
+            p.acc_grad = p.grad.clone().to(torch.bfloat16)
+
+        del p.grad
+
+    @staticmethod
+    def reassign_grad_buffer(model: nn.Module) -> None:
+        for _, p in model.named_parameters():
+            if p.requires_grad and hasattr(p, 'acc_grad'):
+                p.grad = p.acc_grad
+                del p.acc_grad
+
+    @staticmethod
+    def assign_hooks(model: nn.Module) -> List:
+        return [
+            p.register_post_accumulate_grad_hook(StochasticAccumulator.stochastic_grad_accum)
+            for _, p in model.named_parameters()
+            if p.requires_grad
+        ]
+
+
 def is_valid_parameters(parameters: PARAMETERS) -> bool:
     r"""Check where the parameters are valid."""
     return isinstance(parameters, (list, tuple)) and len(parameters) > 0 and isinstance(parameters[0], dict)