feat: optimizer 参数分组,bias/norm 不做 weight decay

This commit is contained in:
ViperEkura 2026-06-27 16:30:34 +08:00
parent 6715461a36
commit bf09a35c95
1 changed files with 14 additions and 1 deletions

View File

@ -269,7 +269,20 @@ def create_model(config):
def create_optimizer(model, **kwargs) -> optim.Optimizer: def create_optimizer(model, **kwargs) -> optim.Optimizer:
return optim.AdamW(model.parameters(), fused=True, **kwargs) decay_params = []
no_decay_params = []
for name, param in model.named_parameters():
if not param.requires_grad:
continue
if param.dim() < 2 or "norm" in name or "bias" in name:
no_decay_params.append(param)
else:
decay_params.append(param)
param_groups = [
{"params": decay_params, "weight_decay": kwargs.pop("weight_decay", 0.01)},
{"params": no_decay_params, "weight_decay": 0.0},
]
return optim.AdamW(param_groups, fused=True, **kwargs)
def create_scheduler( def create_scheduler(