| from __future__ import annotations |
| import torch |
|
|
|
|
| def parameter_group_name(name: str, parameter: torch.nn.Parameter) -> str: |
| """Return the frozen AdamW group for one named parameter.""" |
| if getattr(parameter, "_no_weight_decay", False): |
| return "no_decay" |
| if parameter.ndim < 2 or name.endswith("bias"): |
| return "no_decay" |
| return "decay" |
|
|
|
|
| def optimizer_group_manifest(model) -> dict: |
| rows = {"decay": [], "no_decay": []} |
| for name, parameter in model.named_parameters(): |
| if parameter.requires_grad: |
| rows[parameter_group_name(name, parameter)].append(name) |
| return { |
| group: { |
| "names": names, |
| "tensors": len(names), |
| "parameters": sum(dict(model.named_parameters())[name].numel() for name in names), |
| } |
| for group, names in rows.items() |
| } |
|
|
|
|
| def build_adamw(model, lr=1e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1, fused=True): |
| decay, no_decay = [], [] |
| for name, p in model.named_parameters(): |
| if not p.requires_grad: |
| continue |
| if parameter_group_name(name, p) == "no_decay": |
| no_decay.append(p) |
| else: |
| decay.append(p) |
| groups = [ |
| {"params": decay, "weight_decay": weight_decay}, |
| {"params": no_decay, "weight_decay": 0.0}, |
| ] |
| kwargs = dict(lr=lr, betas=betas, eps=eps) |
| if fused and torch.cuda.is_available(): |
| kwargs["fused"] = True |
| return torch.optim.AdamW(groups, **kwargs) |
|
|
