chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,132 @@
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.optim import Optimizer
|
||||
from torch.optim.optimizer import ParamsT
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Type, Callable, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class OptimizerSpec:
|
||||
"""Spec for creating an optimizer that is part of a `ChainedOptimizer`."""
|
||||
|
||||
class_type: Type[Optimizer]
|
||||
init_args: Dict[str, Any]
|
||||
param_filter: Optional[Callable[[Tensor], bool]]
|
||||
|
||||
|
||||
class ChainedOptimizer(Optimizer):
|
||||
"""
|
||||
A wrapper around multiple optimizers that allows for chaining them together.
|
||||
The optimizers are applied in the order they are passed in the constructor.
|
||||
Each optimizer is responsible for updating a subset of the parameters, which
|
||||
is determined by the `param_filter` function. If no optimizer is found for a
|
||||
parameter group, an exception is raised.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
params: ParamsT,
|
||||
optimizer_specs: List[OptimizerSpec],
|
||||
lr: float,
|
||||
weight_decay: float = 0.0,
|
||||
optimizer_selection_callback: Optional[Callable[[Tensor, int], None]] = None,
|
||||
**common_kwargs,
|
||||
):
|
||||
self.optimizer_specs = optimizer_specs
|
||||
self.optimizer_selection_callback = optimizer_selection_callback
|
||||
self.optimizers: List[Optimizer] = []
|
||||
defaults = dict(lr=lr, weight_decay=weight_decay)
|
||||
super().__init__(params, defaults)
|
||||
|
||||
# Split the params for each optimizer
|
||||
params_for_optimizers = [[] for _ in optimizer_specs]
|
||||
for param_group in self.param_groups:
|
||||
params = param_group["params"]
|
||||
indices = param_group["optimizer_and_param_group_indices"] = set()
|
||||
for param in params:
|
||||
assert isinstance(param, Tensor), f"Expected a Tensor, got {type(param)}"
|
||||
found_optimizer = False
|
||||
for index, spec in enumerate(optimizer_specs):
|
||||
if spec.param_filter is None or spec.param_filter(param):
|
||||
if self.optimizer_selection_callback is not None:
|
||||
self.optimizer_selection_callback(param, index)
|
||||
params_for_optimizers[index].append(param)
|
||||
indices.add((index, 0))
|
||||
found_optimizer = True
|
||||
break
|
||||
if not found_optimizer:
|
||||
raise ValueError("No valid optimizer found for the given parameter")
|
||||
|
||||
# Initialize the optimizers
|
||||
for spec, selected_params in zip(optimizer_specs, params_for_optimizers):
|
||||
optimizer_args = {
|
||||
'lr': lr,
|
||||
'weight_decay': weight_decay,
|
||||
}
|
||||
optimizer_args.update(common_kwargs)
|
||||
optimizer_args.update(spec.init_args)
|
||||
optimizer = spec.class_type(selected_params, **optimizer_args)
|
||||
self.optimizers.append(optimizer)
|
||||
|
||||
def state_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"optimizers": [opt.state_dict() for opt in self.optimizers],
|
||||
**super().state_dict(),
|
||||
}
|
||||
|
||||
def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
|
||||
optimizers = state_dict.pop("optimizers")
|
||||
super().load_state_dict(state_dict)
|
||||
for i in range(len(self.optimizers)):
|
||||
self.optimizers[i].load_state_dict(optimizers[i])
|
||||
|
||||
def zero_grad(self, set_to_none: bool = True) -> None:
|
||||
for opt in self.optimizers:
|
||||
opt.zero_grad(set_to_none=set_to_none)
|
||||
|
||||
def _copy_lr_to_optimizers(self) -> None:
|
||||
for param_group in self.param_groups:
|
||||
indices = param_group["optimizer_and_param_group_indices"]
|
||||
for optimizer_idx, param_group_idx in indices:
|
||||
self.optimizers[optimizer_idx].param_groups[param_group_idx]["lr"] = param_group["lr"]
|
||||
|
||||
def step(self, closure=None) -> None:
|
||||
loss = None
|
||||
if closure is not None:
|
||||
with torch.enable_grad():
|
||||
loss = closure()
|
||||
self._copy_lr_to_optimizers()
|
||||
for opt in self.optimizers:
|
||||
opt.step(closure=None)
|
||||
return loss
|
||||
|
||||
def add_param_group(self, param_group: Dict[str, Any]) -> None:
|
||||
super().add_param_group(param_group)
|
||||
|
||||
# If optimizer has not been initialized, skip adding the param groups
|
||||
if not self.optimizers:
|
||||
return
|
||||
|
||||
# Split the params for each optimizer
|
||||
params_for_optimizers = [[] for _ in self.optimizer_specs]
|
||||
params = param_group["params"]
|
||||
indices = param_group["optimizer_and_param_group_indices"] = set()
|
||||
for param in params:
|
||||
assert isinstance(param, Tensor), f"Expected a Tensor, got {type(param)}"
|
||||
found_optimizer = False
|
||||
for index, spec in enumerate(self.optimizer_specs):
|
||||
if spec.param_filter is None or spec.param_filter(param):
|
||||
if self.optimizer_selection_callback is not None:
|
||||
self.optimizer_selection_callback(param, index)
|
||||
params_for_optimizers[index].append(param)
|
||||
indices.add((index, len(self.optimizers[index].param_groups)))
|
||||
found_optimizer = True
|
||||
break
|
||||
if not found_optimizer:
|
||||
raise ValueError("No valid optimizer found for the given parameter group")
|
||||
|
||||
# Add the selected param group to the optimizers
|
||||
for optimizer, selected_params in zip(self.optimizers, params_for_optimizers):
|
||||
if selected_params:
|
||||
optimizer.add_param_group({"params": selected_params})
|
||||
@@ -0,0 +1,200 @@
|
||||
import collections
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from torch.nn import Parameter
|
||||
from typing import List
|
||||
from .chained_optimizer import ChainedOptimizer, OptimizerSpec
|
||||
|
||||
from modules.commons.common_layers import AdamWLinear, AdamWConv1d
|
||||
|
||||
|
||||
def zeropower_via_newtonschulz5(G: Tensor, steps: int) -> Tensor:
|
||||
"""
|
||||
Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
|
||||
quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
|
||||
of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
|
||||
zero even beyond the point where the iteration no longer converges all the way to one everywhere
|
||||
on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
|
||||
where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
|
||||
performance at all relative to UV^T, where USV^T = G is the SVD.
|
||||
"""
|
||||
assert G.ndim == 3 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
|
||||
a, b, c = (3.4445, -4.7750, 2.0315)
|
||||
|
||||
X = G.to(torch.float32)
|
||||
|
||||
# Ensure spectral norm is at most 1
|
||||
X = F.normalize(X, p=2.0, dim=(-2, -1), eps=1e-7)
|
||||
|
||||
X = X.to(torch.float16)
|
||||
|
||||
# Perform the NS iterations
|
||||
if X.size(-2) < X.size(-1):
|
||||
for _ in range(steps):
|
||||
A = torch.bmm(X, X.mT)
|
||||
A = torch.baddbmm(A, A, A, beta=b, alpha=c)
|
||||
X = torch.baddbmm(X, A, X, beta=a, alpha=1)
|
||||
else:
|
||||
for _ in range(steps):
|
||||
A = torch.bmm(X.mT, X)
|
||||
A = torch.baddbmm(A, A, A, beta=b, alpha=c)
|
||||
X = torch.baddbmm(X, X, A, beta=a, alpha=1)
|
||||
|
||||
return X
|
||||
|
||||
|
||||
def gram_newton_schulz(G: Tensor, steps: int) -> Tensor:
|
||||
"""
|
||||
Refer to:
|
||||
Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon
|
||||
Authors: Jack Zhang, Noah Amsel, Berlin Chen, Tri Dao
|
||||
Blogpost: https://dao-ailab.github.io/blog/2026/gram-newton-schulz/
|
||||
|
||||
Gram Newton-Schulz iteration to compute the orthogonalization of G.
|
||||
Mathematically identical to standard Newton-Schulz but computes iterating
|
||||
on the smaller NxN Gram matrix to save up to 50% FLOPs.
|
||||
"""
|
||||
assert G.ndim == 3
|
||||
reset_iterations = [2]
|
||||
original_shape = G.shape
|
||||
dtype = G.dtype
|
||||
|
||||
X = G.to(torch.float32)
|
||||
X = F.normalize(X, p=2.0, dim=(-2, -1), eps=1e-7)
|
||||
should_transpose = X.size(-2) > X.size(-1)
|
||||
if should_transpose:
|
||||
X = X.mT
|
||||
X = X.to(torch.float16)
|
||||
|
||||
a, b, c = (3.4445, -4.7750, 2.0315)
|
||||
|
||||
if X.size(-2) != X.size(-1):
|
||||
R = torch.bmm(X, X.mT)
|
||||
Q = None
|
||||
for i in range(steps):
|
||||
if i in reset_iterations and i != 0:
|
||||
X = torch.bmm(Q, X)
|
||||
R = torch.bmm(X, X.mT)
|
||||
Q = None
|
||||
Z = torch.baddbmm(R, R, R, beta=b, alpha=c)
|
||||
if i != 0 and i not in reset_iterations:
|
||||
Q = torch.baddbmm(Q, Q, Z, beta=a, alpha=1.0)
|
||||
else:
|
||||
Q = Z.clone()
|
||||
Q.diagonal(dim1=-2, dim2=-1).add_(a)
|
||||
if i < steps - 1 and (i + 1) not in reset_iterations:
|
||||
RZ = torch.baddbmm(R, R, Z, beta=a, alpha=1.0)
|
||||
R = torch.baddbmm(RZ, Z, RZ, beta=a, alpha=1.0)
|
||||
X = torch.bmm(Q, X) if not should_transpose else torch.bmm(X.mT, Q)
|
||||
else:
|
||||
for _ in range(steps):
|
||||
A = torch.bmm(X, X.mT)
|
||||
B = torch.baddbmm(A, A, A, beta=b, alpha=c)
|
||||
X = torch.baddbmm(X, B, X, beta=a, alpha=1.0)
|
||||
|
||||
return X.to(dtype).view(original_shape)
|
||||
|
||||
|
||||
class Muon(torch.optim.Optimizer):
|
||||
"""
|
||||
Muon - MomentUm Orthogonalized by Newton-schulz
|
||||
|
||||
https://kellerjordan.github.io/posts/muon/
|
||||
|
||||
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 float16 on the GPU.
|
||||
|
||||
Some warnings:
|
||||
- This optimizer should not be used for the embedding layer, the final fully connected layer,
|
||||
or any {0,1}-D parameters; those should all be optimized by a standard method (e.g., AdamW).
|
||||
- To use it with 4D convolutional filters, it works well to just flatten their last 3 dimensions.
|
||||
|
||||
Arguments:
|
||||
lr: The learning rate used by the internal SGD.
|
||||
momentum: The momentum used by the internal SGD.
|
||||
nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended)
|
||||
ns_steps: The number of Newton-Schulz iteration steps to use.
|
||||
"""
|
||||
|
||||
def __init__(self, params, lr=5e-4, weight_decay=0.1, momentum=0.95, nesterov=True, ns_steps=5):
|
||||
defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps)
|
||||
super().__init__(params, defaults)
|
||||
|
||||
@torch.no_grad()
|
||||
def step(self, closure=None):
|
||||
for group in self.param_groups:
|
||||
shape_groups = {}
|
||||
for p in filter(lambda p: p.grad is not None, group["params"]):
|
||||
g = p.grad
|
||||
state = self.state[p]
|
||||
if "momentum_buffer" not in state:
|
||||
state["momentum_buffer"] = torch.zeros_like(g)
|
||||
key = (p.shape, p.device, p.dtype)
|
||||
if key not in shape_groups:
|
||||
shape_groups[key] = {"params": [], "grads": [], "buffers": []}
|
||||
shape_groups[key]["params"].append(p)
|
||||
shape_groups[key]["grads"].append(g)
|
||||
shape_groups[key]["buffers"].append(state["momentum_buffer"])
|
||||
for key in shape_groups:
|
||||
group_data = shape_groups[key]
|
||||
p, g, buf, m = group_data["params"], group_data["grads"], group_data["buffers"], group["momentum"]
|
||||
torch._foreach_lerp_(buf, g, 1-m)
|
||||
if group["nesterov"]:
|
||||
torch._foreach_lerp_(g, buf, m)
|
||||
g = torch.stack(g)
|
||||
else:
|
||||
g = torch.stack(buf)
|
||||
original_shape = g.shape
|
||||
if g.ndim >= 4: # for the case of conv filters
|
||||
g = g.view(g.size(0), g.size(1), -1)
|
||||
g = gram_newton_schulz(g, steps=group["ns_steps"])
|
||||
|
||||
if group["weight_decay"] > 0:
|
||||
torch._foreach_mul_(p, 1 - group["lr"] * group["weight_decay"])
|
||||
torch._foreach_add_(p, g.view(original_shape).unbind(0), alpha=-group["lr"] * max(g[0].size()) ** 0.5)
|
||||
|
||||
|
||||
def get_params_for_muon(model) -> List[Parameter]:
|
||||
"""
|
||||
Filter parameters of a module into two groups: those that can be optimized by Muon,
|
||||
and those that should be optimized by a standard optimizer.
|
||||
Args:
|
||||
module: The module to filter parameters for.
|
||||
Returns:
|
||||
A list of parameters that should be optimized with muon.
|
||||
"""
|
||||
excluded_module_classes = (nn.Embedding, AdamWLinear, AdamWConv1d)
|
||||
muon_params = []
|
||||
# BFS through all submodules and exclude parameters from certain module types
|
||||
queue = collections.deque([model])
|
||||
while queue:
|
||||
module = queue.popleft()
|
||||
if isinstance(module, excluded_module_classes):
|
||||
continue
|
||||
for param in module.parameters(recurse=False):
|
||||
if not param.requires_grad:
|
||||
continue
|
||||
if param.ndim >= 2:
|
||||
muon_params.append(param)
|
||||
queue.extend(list(module.children()))
|
||||
return muon_params
|
||||
|
||||
|
||||
class Muon_AdamW(ChainedOptimizer):
|
||||
def __init__(self, model, lr=0.0005, weight_decay=0.0, muon_args=None, adamw_args=None, verbose=False):
|
||||
muon_args = {} if muon_args is None else muon_args
|
||||
adamw_args = {} if adamw_args is None else adamw_args
|
||||
muon_params_id_set = set(id(p) for p in get_params_for_muon(model))
|
||||
spec_muon = OptimizerSpec(Muon, muon_args, lambda param: id(param) in muon_params_id_set)
|
||||
spec_adamw = OptimizerSpec(torch.optim.AdamW, adamw_args, None)
|
||||
specs = [spec_muon, spec_adamw]
|
||||
callback = None
|
||||
if verbose:
|
||||
callback = lambda p, spec_idx: print(
|
||||
f"Adding param {p.shape} to optimizer{spec_idx} {str(specs[spec_idx].class_type)}"
|
||||
)
|
||||
super().__init__(model.parameters(), specs, lr=lr, weight_decay=weight_decay, optimizer_selection_callback=callback)
|
||||
Reference in New Issue
Block a user