Update Wall-X to 1.1.0 (#104)
This commit is contained in:
Starrick Liu
@@ -0,0 +1,230 @@
|
||||
import inspect
|
||||
|
||||
from torch.optim import AdamW
|
||||
|
||||
from wall_x.config.hyperparams_config import LRGroupConfig
|
||||
|
||||
_OPTIMIZERS = {}
|
||||
|
||||
|
||||
def register_optimizer(name, optimizer_fn):
|
||||
_OPTIMIZERS[name] = optimizer_fn
|
||||
|
||||
|
||||
def get_optimizer(name, *args, **kwargs):
|
||||
if name not in _OPTIMIZERS:
|
||||
raise KeyError(f"Unknown optimizer '{name}'. Registered: {sorted(_OPTIMIZERS)}")
|
||||
return _OPTIMIZERS[name](*args, **kwargs)
|
||||
|
||||
|
||||
def _group_weight_decay(opt_cfg):
|
||||
return getattr(opt_cfg, "weight_decay", None)
|
||||
|
||||
|
||||
def resolve_lr_group_configs(opt_cfg, default_action_lr_keywords):
|
||||
"""Return structured LR groups, with legacy action config fallback.
|
||||
|
||||
``optimizer.lr_groups`` is the preferred path. The legacy
|
||||
``action_expert_learning_rate`` fields are still converted into a single
|
||||
action group so older configs keep working.
|
||||
"""
|
||||
if opt_cfg.lr_groups:
|
||||
if (
|
||||
opt_cfg.action_expert_learning_rate is not None
|
||||
or opt_cfg.action_lr_keywords is not None
|
||||
):
|
||||
raise ValueError(
|
||||
"Use either optimizer.lr_groups or legacy "
|
||||
"action_expert_learning_rate/action_lr_keywords, not both."
|
||||
)
|
||||
return opt_cfg.lr_groups
|
||||
|
||||
if opt_cfg.action_expert_learning_rate is None:
|
||||
return []
|
||||
|
||||
action_lr_keywords = (
|
||||
opt_cfg.action_lr_keywords
|
||||
if opt_cfg.action_lr_keywords is not None
|
||||
else default_action_lr_keywords
|
||||
)
|
||||
return [
|
||||
LRGroupConfig(
|
||||
name="action_lr_group",
|
||||
lr=opt_cfg.action_expert_learning_rate,
|
||||
include=action_lr_keywords,
|
||||
fail_on_empty=True,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def uses_legacy_action_lr_groups(opt_cfg) -> bool:
|
||||
return not opt_cfg.lr_groups and opt_cfg.action_expert_learning_rate is not None
|
||||
|
||||
|
||||
def _make_param_group(name, params, lr, opt_cfg):
|
||||
group = {
|
||||
"params": params,
|
||||
"lr": lr,
|
||||
"group_name": name,
|
||||
}
|
||||
weight_decay = _group_weight_decay(opt_cfg)
|
||||
if weight_decay is not None:
|
||||
group["weight_decay"] = weight_decay
|
||||
return group
|
||||
|
||||
|
||||
def _validate_lr_group(group: LRGroupConfig, *, base_group_name: str):
|
||||
if not group.name:
|
||||
raise ValueError("optimizer.lr_groups entries must have a non-empty name")
|
||||
if group.name == base_group_name:
|
||||
raise ValueError(
|
||||
f"optimizer.lr_groups name {group.name!r} is reserved for the base group"
|
||||
)
|
||||
if "/" in group.name:
|
||||
raise ValueError(
|
||||
f"optimizer.lr_groups name {group.name!r} must not contain '/'. "
|
||||
"DMuon appends '/muon' and '/adamw' to group names."
|
||||
)
|
||||
if not group.include:
|
||||
raise ValueError(
|
||||
f"optimizer.lr_groups.{group.name} must define at least one include keyword"
|
||||
)
|
||||
|
||||
|
||||
def build_lr_param_groups(model, opt_cfg, lr_groups, *, base_group_name="base"):
|
||||
"""Split trainable params into named LR groups plus a base group.
|
||||
|
||||
``lr_groups`` is a list of :class:`LRGroupConfig`. Each group matches
|
||||
parameter names by substring. A parameter may match at most one explicit
|
||||
group; unmatched trainable parameters remain in the ``base`` group using
|
||||
``opt_cfg.learning_rate``.
|
||||
|
||||
Returns a list of torch.optim-compatible param_group dicts. The
|
||||
``group_name`` key is non-standard but preserved by torch.optim via
|
||||
``setdefault`` in ``add_param_group`` and is consumed downstream for
|
||||
per-group lr logging.
|
||||
|
||||
For AdamW / native Muon, each returned group includes ``weight_decay``.
|
||||
For DMuon, weight decay is split between Muon and AdamW defaults, so the
|
||||
groups only carry lr and metadata; DMuon applies its own per-route defaults.
|
||||
The caller is expected to gate on ``opt_cfg.optimizer_type`` before calling
|
||||
this.
|
||||
"""
|
||||
if not lr_groups:
|
||||
return None
|
||||
|
||||
names = [group.name for group in lr_groups]
|
||||
duplicate_names = sorted({name for name in names if names.count(name) > 1})
|
||||
if duplicate_names:
|
||||
raise ValueError(
|
||||
f"optimizer.lr_groups contains duplicate names: {duplicate_names}"
|
||||
)
|
||||
|
||||
for group in lr_groups:
|
||||
_validate_lr_group(group, base_group_name=base_group_name)
|
||||
|
||||
base_params = []
|
||||
grouped_params = {group.name: [] for group in lr_groups}
|
||||
ambiguous = []
|
||||
|
||||
for name, param in model.named_parameters():
|
||||
if not param.requires_grad:
|
||||
continue
|
||||
matches = [
|
||||
group.name
|
||||
for group in lr_groups
|
||||
if any(keyword in name for keyword in group.include)
|
||||
]
|
||||
if len(matches) > 1:
|
||||
ambiguous.append((name, matches))
|
||||
continue
|
||||
if matches:
|
||||
grouped_params[matches[0]].append(param)
|
||||
else:
|
||||
base_params.append(param)
|
||||
|
||||
if ambiguous:
|
||||
examples = ", ".join(f"{name} -> {matches}" for name, matches in ambiguous[:10])
|
||||
raise ValueError(
|
||||
"Some parameters match multiple optimizer.lr_groups. Make group "
|
||||
f"include patterns disjoint. Examples: {examples}"
|
||||
)
|
||||
|
||||
if opt_cfg.train_action_expert_only:
|
||||
assert len(base_params) == 0, (
|
||||
f"Expected 0 base_params after pre-wrap freeze, got {len(base_params)}. "
|
||||
"Ensure base params are frozen before building the optimizer."
|
||||
)
|
||||
|
||||
param_groups = []
|
||||
if len(base_params) > 0:
|
||||
param_groups.append(
|
||||
_make_param_group(
|
||||
base_group_name, base_params, opt_cfg.learning_rate, opt_cfg
|
||||
)
|
||||
)
|
||||
|
||||
for group in lr_groups:
|
||||
params = grouped_params[group.name]
|
||||
if len(params) == 0:
|
||||
if group.fail_on_empty:
|
||||
raise ValueError(
|
||||
f"No params found for optimizer.lr_groups.{group.name}. "
|
||||
f"Please check include={group.include!r}."
|
||||
)
|
||||
continue
|
||||
param_groups.append(_make_param_group(group.name, params, group.lr, opt_cfg))
|
||||
|
||||
return param_groups
|
||||
|
||||
|
||||
def build_action_expert_param_groups(model, opt_cfg, action_lr_keywords):
|
||||
"""Compatibility wrapper for the legacy action-expert LR config."""
|
||||
return build_lr_param_groups(
|
||||
model,
|
||||
opt_cfg,
|
||||
[
|
||||
LRGroupConfig(
|
||||
name="action_lr_group",
|
||||
lr=opt_cfg.action_expert_learning_rate,
|
||||
include=action_lr_keywords,
|
||||
fail_on_empty=True,
|
||||
)
|
||||
],
|
||||
base_group_name="base_lr_group",
|
||||
)
|
||||
|
||||
|
||||
def get_adamw_optimizer(model, *, opt_cfg, param_groups=None):
|
||||
"""Build AdamW from AdamWConfig."""
|
||||
from wall_x.config.hyperparams_config import AdamWConfig
|
||||
|
||||
if not isinstance(opt_cfg, AdamWConfig):
|
||||
raise TypeError(
|
||||
f"get_adamw_optimizer expects AdamWConfig, got {type(opt_cfg).__name__}"
|
||||
)
|
||||
|
||||
if param_groups is None:
|
||||
params = [p for p in model.parameters() if p.requires_grad]
|
||||
else:
|
||||
# Per-group lr / weight_decay are preserved by torch.optim (via
|
||||
# setdefault in add_param_group), so top-level values act only as
|
||||
# defaults. Extra keys like ``group_name`` are kept in-place and used
|
||||
# downstream for per-group lr logging.
|
||||
params = param_groups
|
||||
|
||||
kw = {
|
||||
"lr": opt_cfg.learning_rate,
|
||||
"weight_decay": opt_cfg.weight_decay,
|
||||
"betas": tuple(opt_cfg.betas),
|
||||
"eps": opt_cfg.eps,
|
||||
}
|
||||
sig_params = inspect.signature(AdamW.__init__).parameters
|
||||
if "foreach" in sig_params and opt_cfg.foreach is not None:
|
||||
kw["foreach"] = opt_cfg.foreach
|
||||
if "fused" in sig_params:
|
||||
kw["fused"] = opt_cfg.fused
|
||||
return AdamW(params, **kw)
|
||||
|
||||
|
||||
register_optimizer("adamw", get_adamw_optimizer)
|
||||
Reference in New Issue
Block a user