Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import torch | |
| import torch.optim as optim | |
| def create_optimizer(model, args): | |
| from torch.nn.parallel import DistributedDataParallel as DDP | |
| raw_model = model.module if isinstance(model, DDP) else model | |
| optimizer_name = getattr(args, "optimizer", "adamw") | |
| if optimizer_name == "muon": | |
| return _create_muon_split(raw_model, args) | |
| else: | |
| return _create_adamw(raw_model, args) | |
| def _create_adamw(model, args): | |
| decay, no_decay = _split_decay_params(model) | |
| optimizer = optim.AdamW([ | |
| {"params": decay, "weight_decay": 0.01}, | |
| {"params": no_decay, "weight_decay": 0}, | |
| ], lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-8, fused=True) | |
| return [optimizer] | |
| def _create_muon_split(model, args): | |
| from models.muon import Muon | |
| muon_params = [] | |
| adamw_decay = [] | |
| adamw_no_decay = [] | |
| for name, param in model.named_parameters(): | |
| if not param.requires_grad: | |
| continue | |
| # Muon only makes sense for 2D weight matrices (not embeddings, not biases) | |
| if param.ndim == 2 and "embed" not in name: | |
| muon_params.append(param) | |
| elif param.ndim == 1 or name.endswith(".bias") or "norm" in name: | |
| adamw_no_decay.append(param) | |
| else: | |
| adamw_decay.append(param) | |
| lr = args.learning_rate | |
| momentum = getattr(args, "muon_momentum", 0.95) | |
| backend_steps = getattr(args, "muon_backend_steps", 5) | |
| optimizers = [] | |
| if muon_params: | |
| muon_opt = Muon( | |
| muon_params, lr=lr, momentum=momentum, | |
| backend_steps=backend_steps, weight_decay=0.01, | |
| ) | |
| optimizers.append(muon_opt) | |
| # AdamW for the rest | |
| adamw_groups = [] | |
| if adamw_decay: | |
| adamw_groups.append({"params": adamw_decay, "weight_decay": 0.01}) | |
| if adamw_no_decay: | |
| adamw_groups.append({"params": adamw_no_decay, "weight_decay": 0}) | |
| if adamw_groups: | |
| adamw_opt = optim.AdamW(adamw_groups, lr=lr, betas=(0.9, 0.999), eps=1e-8, fused=True) | |
| optimizers.append(adamw_opt) | |
| return optimizers | |
| def _split_decay_params(model): | |
| decay = [] | |
| no_decay = [] | |
| for name, param in model.named_parameters(): | |
| if param.requires_grad: | |
| if len(param.shape) == 1 or name.endswith(".bias") or "norm" in name: | |
| no_decay.append(param) | |
| else: | |
| decay.append(param) | |
| return decay, no_decay | |