Update Wall-X to 1.1.0 (#104)
This commit is contained in:
Starrick Liu
@@ -0,0 +1,9 @@
|
||||
from .utils import get_optimizer, register_optimizer
|
||||
|
||||
# DMuon is optional (external package). Skip registration if dmuon isn't
|
||||
# installed so wall_x still imports; the registry will simply not have
|
||||
# "dmuon" and get_optimizer("dmuon", ...) will raise a clear KeyError.
|
||||
try:
|
||||
from . import dmuon # noqa: F401 - side-effect registers "dmuon"
|
||||
except ImportError:
|
||||
pass
|
||||
@@ -0,0 +1,16 @@
|
||||
import torch.nn as nn
|
||||
|
||||
from .utils import get_dmuon_optimizer
|
||||
|
||||
|
||||
def is_dmuon_model(model: nn.Module) -> bool:
|
||||
"""True if ``dmuon.dedicate_params()`` has been applied to this model.
|
||||
|
||||
Checked via an attribute the external ``dmuon`` package attaches to
|
||||
the root module, so this predicate works without importing ``dmuon``
|
||||
and returns ``False`` for ordinary (non-DMuon) models.
|
||||
"""
|
||||
return hasattr(model, "_dedicated_comm_ctx")
|
||||
|
||||
|
||||
__all__ = ["get_dmuon_optimizer", "is_dmuon_model"]
|
||||
@@ -0,0 +1,156 @@
|
||||
"""DMuon optimizer builder."""
|
||||
|
||||
import inspect
|
||||
import logging
|
||||
|
||||
from ..utils import register_optimizer
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _emit(log_fn, message, *args, level=logging.INFO):
|
||||
if args:
|
||||
message = message % args
|
||||
if log_fn is not None:
|
||||
log_fn(message, level=level)
|
||||
else:
|
||||
_logger.log(level, message)
|
||||
|
||||
|
||||
def _is_rank0():
|
||||
try:
|
||||
import torch.distributed as dist
|
||||
|
||||
return (
|
||||
not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0
|
||||
)
|
||||
except Exception:
|
||||
return True
|
||||
|
||||
|
||||
def _build_ns_backend(dmuon, opt_cfg):
|
||||
coefficients = getattr(opt_cfg, "ns_coefficients", "default")
|
||||
if coefficients in (None, "default"):
|
||||
return opt_cfg.ns_backend
|
||||
|
||||
if coefficients != "wallx_muon":
|
||||
raise ValueError(
|
||||
"Unsupported DMuon ns_coefficients="
|
||||
f"{coefficients!r}. Supported: 'default', 'wallx_muon'."
|
||||
)
|
||||
if opt_cfg.ns_backend != "direct":
|
||||
raise ValueError(
|
||||
"ns_coefficients='wallx_muon' is intended to match Wall-X's "
|
||||
"direct-space Muon implementation. Set ns_backend='direct'."
|
||||
)
|
||||
|
||||
wallx_coefficients = [[3.4445, -4.7750, 2.0315] for _ in range(opt_cfg.ns_steps)]
|
||||
return dmuon.NewtonSchulz(
|
||||
backend="direct",
|
||||
coefficients=wallx_coefficients,
|
||||
)
|
||||
|
||||
|
||||
def get_dmuon_optimizer(model, *, opt_cfg, param_groups=None, log_fn=None):
|
||||
"""Build dmuon.Muon from a DMuonConfig.
|
||||
|
||||
When ``param_groups`` is provided, Wall-X expects DMuon to preserve the
|
||||
PyTorch optimizer group semantics and then split each user group into
|
||||
dedicated/Muon and non-dedicated/AdamW subgroups internally.
|
||||
"""
|
||||
from wall_x.config.hyperparams_config import DMuonConfig
|
||||
|
||||
if not isinstance(opt_cfg, DMuonConfig):
|
||||
raise TypeError(
|
||||
f"get_dmuon_optimizer expects DMuonConfig, got {type(opt_cfg).__name__}"
|
||||
)
|
||||
import dmuon
|
||||
|
||||
muon_signature = inspect.signature(dmuon.Muon)
|
||||
supports_param_groups = "param_groups" in muon_signature.parameters
|
||||
if param_groups is not None and not supports_param_groups:
|
||||
raise RuntimeError(
|
||||
"Wall-X built optimizer param_groups for DMuon, but the installed "
|
||||
"dmuon.Muon does not accept a param_groups= argument. Please update "
|
||||
"DMuon to the param-group-aware implementation before enabling "
|
||||
"action_expert_learning_rate with optimizer_type='dmuon'."
|
||||
)
|
||||
|
||||
ns_backend = _build_ns_backend(dmuon, opt_cfg)
|
||||
|
||||
_emit(
|
||||
log_fn,
|
||||
"DMuon: Muon lr=%s momentum=%s ns_steps=%s; "
|
||||
"AdamW lr=%s betas=%s wd=%s; "
|
||||
"ns_backend=%s ns_coefficients=%s nesterov=%s",
|
||||
opt_cfg.muon_lr,
|
||||
opt_cfg.momentum,
|
||||
opt_cfg.ns_steps,
|
||||
opt_cfg.adamw_lr,
|
||||
opt_cfg.adamw_betas,
|
||||
opt_cfg.adamw_weight_decay,
|
||||
opt_cfg.ns_backend,
|
||||
opt_cfg.ns_coefficients,
|
||||
opt_cfg.nesterov,
|
||||
)
|
||||
if param_groups is not None:
|
||||
_emit(
|
||||
log_fn,
|
||||
"DMuon param_groups enabled: %s",
|
||||
[
|
||||
{
|
||||
"group_name": group.get("group_name", f"group_{idx}"),
|
||||
"lr": group.get("lr"),
|
||||
"num_params": len(group.get("params", [])),
|
||||
}
|
||||
for idx, group in enumerate(param_groups)
|
||||
],
|
||||
)
|
||||
|
||||
kwargs = {}
|
||||
if param_groups is not None:
|
||||
kwargs["param_groups"] = param_groups
|
||||
|
||||
optimizer = dmuon.Muon(
|
||||
model,
|
||||
lr=opt_cfg.muon_lr,
|
||||
momentum=opt_cfg.momentum,
|
||||
weight_decay=opt_cfg.muon_weight_decay,
|
||||
ns_steps=opt_cfg.ns_steps,
|
||||
adamw_lr=opt_cfg.adamw_lr,
|
||||
adamw_betas=tuple(opt_cfg.adamw_betas),
|
||||
adamw_weight_decay=opt_cfg.adamw_weight_decay,
|
||||
adamw_eps=opt_cfg.adamw_eps,
|
||||
ns_backend=ns_backend,
|
||||
nesterov=opt_cfg.nesterov,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if param_groups is not None and _is_rank0():
|
||||
summarize = getattr(dmuon, "summarize_param_groups", None)
|
||||
format_summary = getattr(dmuon, "format_param_group_summary", None)
|
||||
if summarize is None or format_summary is None:
|
||||
_emit(
|
||||
log_fn,
|
||||
"DMuon param_groups are enabled, but the installed DMuon package "
|
||||
"does not expose param-group diagnostics. Update DMuon if you need "
|
||||
"startup verification of the Muon/AdamW subgroup split.",
|
||||
level=logging.WARNING,
|
||||
)
|
||||
else:
|
||||
try:
|
||||
summary = summarize(model, optimizer, max_rows=80)
|
||||
_emit(log_fn, "%s", format_summary(summary))
|
||||
except Exception as exc:
|
||||
_logger.exception("Failed to summarize DMuon param_groups")
|
||||
_emit(
|
||||
log_fn,
|
||||
"Failed to summarize DMuon param_groups: %s",
|
||||
exc,
|
||||
level=logging.WARNING,
|
||||
)
|
||||
|
||||
return optimizer
|
||||
|
||||
|
||||
register_optimizer("dmuon", get_dmuon_optimizer)
|
||||
@@ -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