Spaces:
Running on Zero
Running on Zero
File size: 2,268 Bytes
be88765 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | from __future__ import annotations
from functools import partial
from typing import Iterable, Callable
import torch
def _separate_params_by_dim(parameters: Iterable[torch.nn.Parameter]) -> tuple[list[torch.nn.Parameter], list[torch.nn.Parameter]]:
"""Split parameters into matrix-like (ndim>=2) and vector-like (ndim<2)."""
matrix_like: list[torch.nn.Parameter] = []
vector_like: list[torch.nn.Parameter] = []
for p in parameters:
if not p.requires_grad:
continue
(matrix_like if p.ndim >= 2 else vector_like).append(p)
return matrix_like, vector_like
def build_adamw(lr: float, weight_decay: float = 0.0, betas: tuple[float, float] = (0.9, 0.999), eps: float = 1e-8, lr_muon: float = 0.0) -> Callable[[Iterable[torch.nn.Parameter]], torch.optim.Optimizer]:
"""Return a factory that creates AdamW on provided parameters."""
return partial(torch.optim.AdamW, lr=lr, weight_decay=weight_decay, betas=betas, eps=eps)
def build_muon_with_aux_adam(
parameters: Iterable[torch.nn.Parameter],
lr_muon: float = 0.02,
lr: float = 3e-4,
weight_decay: float = 0.01,
betas: tuple[float, float] = (0.9, 0.95),
) -> torch.optim.Optimizer:
"""
Construct MuonWithAuxAdam with two param groups from an iterable of parameters:
- matrix-like parameters (ndim>=2) use Muon with lr_muon and weight decay
- vector-like parameters (bias/gain etc.) use AdamW with lr_adam/betas/weight_decay
Designed to be used with Hydra as a partial, e.g., self.optimizer(self.parameters()).
"""
try:
from muon import MuonWithAuxAdam # type: ignore
except Exception as e: # pragma: no cover
raise ImportError("Muon is not installed. Install with: pip install git+https://github.com/KellerJordan/Muon") from e
matrix_like, vector_like = _separate_params_by_dim(parameters)
param_groups = [
dict(params=matrix_like, use_muon=True, lr=lr_muon, weight_decay=weight_decay*0.1), # according to Jianlin Su's blog, the rule of thumb is to set lr_muon as 10*lr_adam and set weight_decay_muon as 0.1 * weight_decay_adam
dict(params=vector_like, use_muon=False, lr=lr, betas=betas, weight_decay=weight_decay),
]
return MuonWithAuxAdam(param_groups)
|