Spaces:
Running on Zero
Running on Zero
File size: 2,468 Bytes
dcc4594 | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | 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
|