Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / optim / builders.py: 93%
134 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""YAML-targeted optimizer builders: AdamW/Muon selection.
17Parameter-name and parameter-group logic lives in
18:mod:`hyper_parallel.components.optim.parameter_groups`; the optimizer
19algorithm implementations stay in ``hyper_parallel.core.optimizer``.
20"""
22__all__ = ["AdamW", "Muon", "Sinkhorn"]
24import logging
25import re
26from typing import Any, Dict, List, Optional, Sequence, Tuple
28from torch import nn # pylint: disable=forbidden-backend-import
30from hyper_parallel.core.optimizer import get_hyper_optimizer
31from hyper_parallel.components.optim.parameter_groups import (
32 _DEFAULT_ADAMW_NAME_KEYWORDS,
33 get_adamw_param_groups,
34 split_muon_adamw_params,
35)
37logger = logging.getLogger(__name__)
40class AdamW:
41 """Build a core AdamW optimizer from YAML configuration."""
43 def __init__(
44 self,
45 adamw_config: dict,
46 model: nn.Module,
47 no_decay_params: Optional[List[str]] = None,
48 excluded_param_ids: Optional[Sequence[int]] = None,
49 ) -> None:
50 """Initialize AdamW optimizer configuration.
52 Args:
53 adamw_config: AdamW hyperparameters resolved from YAML.
54 model: Module whose trainable parameters are optimized.
55 no_decay_params: Optional names excluded from weight decay.
56 """
57 self.config = adamw_config
58 self.model = model
60 adamw_groups, _ = self.get_adamw_param_groups(
61 self.model,
62 weight_decay=adamw_config.get("adamw_weight_decay", 1e-2),
63 no_decay_params=no_decay_params,
64 excluded_param_ids=excluded_param_ids,
65 )
66 if not adamw_groups:
67 raise ValueError("AdamW requires at least one trainable parameter")
69 self.optimizer = get_hyper_optimizer(
70 model=self.model,
71 muon_params=[],
72 adamw_params=adamw_groups,
73 adamw_kwargs=adamw_config,
74 )
76 @staticmethod
77 def get_adamw_param_groups(
78 model: "nn.Module",
79 weight_decay: float = 1e-2,
80 no_decay_params: Optional[Sequence[str]] = None,
81 param_groups: Optional[Sequence[Dict[str, Any]]] = None,
82 allowed_param_ids: Optional[Sequence[int]] = None,
83 excluded_param_ids: Optional[Sequence[int]] = None,
84 ) -> Tuple[List[Dict[str, Any]], List[str]]:
85 """Split model parameters into decaying and non-decaying groups.
87 Delegates to
88 :func:`hyper_parallel.components.optim.parameter_groups.get_adamw_param_groups`;
89 retained on the class so existing call sites keep working.
91 Args:
92 model: Module whose trainable parameters are routed.
93 weight_decay: Default decay coefficient for trainable parameters.
94 no_decay_params: Name keywords excluded from weight decay.
95 param_groups: Optional prebuilt groups to return unchanged.
96 allowed_param_ids: Optional parameter identities eligible for routing.
98 Returns:
99 Parameter groups and names routed to AdamW.
100 """
101 return get_adamw_param_groups(
102 model,
103 weight_decay=weight_decay,
104 no_decay_params=no_decay_params,
105 param_groups=param_groups,
106 allowed_param_ids=allowed_param_ids,
107 excluded_param_ids=excluded_param_ids,
108 )
110 def get_optimizer(self) -> Any:
111 """Return the core chained optimizer runtime."""
112 return self.optimizer
115class Muon:
116 """Build a mixed Muon and AdamW optimizer from YAML configuration."""
118 _DEFAULT_ADAMW_NAME_KEYWORDS = _DEFAULT_ADAMW_NAME_KEYWORDS
120 def __init__(
121 self,
122 muon_config: dict,
123 adamw_config: dict,
124 model: nn.Module,
125 extra_adamw_name_keywords: Optional[List[str]] = None,
126 no_decay_params: Optional[List[str]] = None,
127 excluded_param_ids: Optional[Sequence[int]] = None,
128 ) -> None:
129 """Build a mixed Muon and fallback AdamW runtime for ``model``.
131 Args:
132 muon_config: Muon hyperparameters resolved from YAML.
133 adamw_config: Fallback AdamW hyperparameters resolved from YAML.
134 model: Module whose trainable parameters are optimized.
135 extra_adamw_name_keywords: Additional names routed to AdamW.
136 no_decay_params: Optional names excluded from weight decay.
137 """
138 self.muon_config = muon_config
139 self.adamw_config = adamw_config
140 self.model = model
142 muon_params, adamw_params, muon_names, adamw_names = self.split_muon_adamw_params(
143 model,
144 extra_adamw_name_keywords=extra_adamw_name_keywords or (),
145 excluded_param_ids=excluded_param_ids,
146 )
147 if not muon_params:
148 raise ValueError("Muon requires at least one eligible matrix parameter")
150 adamw_groups, _ = AdamW.get_adamw_param_groups(
151 model,
152 weight_decay=adamw_config.get("adamw_weight_decay", 1e-2),
153 no_decay_params=no_decay_params,
154 allowed_param_ids=[id(parameter) for parameter in adamw_params],
155 excluded_param_ids=excluded_param_ids,
156 )
158 logger.info_rank0(
159 "Muon optimizer split: %s Muon parameters, %s AdamW parameters",
160 len(muon_names),
161 len(adamw_names),
162 )
163 logger.info_rank0("Muon parameters (first 5): %s", muon_names[:5])
164 logger.info_rank0("AdamW parameters (first 5): %s", adamw_names[:5])
166 self.optimizer = get_hyper_optimizer(
167 model=self.model,
168 muon_params=muon_params,
169 adamw_params=adamw_groups,
170 muon_kwargs=muon_config,
171 adamw_kwargs=adamw_config,
172 )
174 @staticmethod
175 def split_muon_adamw_params(
176 model: nn.Module,
177 extra_adamw_name_keywords: Sequence[str] = (),
178 excluded_param_ids: Optional[Sequence[int]] = None,
179 ) -> Tuple[List[nn.Parameter], List[nn.Parameter], List[str], List[str]]:
180 """Route matrix parameters to Muon and remaining parameters to AdamW.
182 Delegates to
183 :func:`hyper_parallel.components.optim.parameter_groups.split_muon_adamw_params`;
184 retained on the class so existing call sites keep working.
186 Args:
187 model: Module whose trainable parameters are routed.
188 extra_adamw_name_keywords: Additional names reserved for AdamW.
190 Returns:
191 Muon parameters, AdamW parameters, and their respective names.
192 """
193 return split_muon_adamw_params(
194 model,
195 extra_adamw_name_keywords=extra_adamw_name_keywords,
196 excluded_param_ids=excluded_param_ids,
197 )
199 def get_optimizer(self) -> Any:
200 """Return the core chained optimizer runtime."""
201 return self.optimizer
204class Sinkhorn:
205 """Build Sinkhorn for embeddings/heads, Muon for matrices, and AdamW otherwise."""
207 def __init__(
208 self,
209 sinkhorn_config: dict,
210 muon_config: dict,
211 adamw_config: dict,
212 model: nn.Module,
213 no_decay_params: Optional[List[str]] = None,
214 excluded_param_ids: Optional[Sequence[int]] = None,
215 ) -> None:
216 """Partition trainable parameters and build the existing chained runtime.
218 Args:
219 sinkhorn_config: Sinkhorn kwargs, optional ``param_patterns`` and
220 ``param_groups`` containing selectors and group overrides.
221 muon_config: Muon kwargs, including model-owned callbacks.
222 adamw_config: AdamW kwargs and optional ``param_patterns``.
223 model: Model whose parameters are partitioned by identity.
224 no_decay_params: Name keywords exempt from AdamW decay. Defaults
225 to bias and scale parameters; normalization weights decay.
226 excluded_param_ids: Parameters managed by another optimizer.
227 """
228 configs = {"sinkhorn": dict(sinkhorn_config), "muon": dict(muon_config), "adamw": dict(adamw_config)}
229 selectors = {name: self._patterns(config.pop("param_patterns", ())) for name, config in configs.items()}
230 override_patterns, sinkhorn_groups = self._sinkhorn_overrides(configs["sinkhorn"].pop("param_groups", ()))
231 aliases = self._aliases(model, excluded_param_ids)
232 if not aliases:
233 raise ValueError("Sinkhorn builder requires at least one trainable parameter")
234 roles = self._roles(model)
235 selected = {name: [] for name in configs}
236 for param, names in aliases.items():
237 family, group_index = self._route(param, names, roles, selectors, override_patterns)
238 selected[family].append(param)
239 if family == "sinkhorn":
240 sinkhorn_groups[group_index]["params"].append(param)
241 adamw_groups = self._adamw_groups(selected["adamw"], aliases, roles[1], no_decay_params)
242 logger.info_rank0("Optimizer split: %s", {name: len(params) for name, params in selected.items()})
243 self.optimizer = get_hyper_optimizer(
244 model=model,
245 muon_params=selected["muon"],
246 adamw_params=adamw_groups,
247 sinkhorn_params=[group for group in sinkhorn_groups if group["params"]],
248 muon_kwargs=configs["muon"],
249 adamw_kwargs=configs["adamw"],
250 sinkhorn_kwargs=configs["sinkhorn"],
251 )
253 @classmethod
254 def _sinkhorn_overrides(cls, overrides):
255 """Separate pattern selectors from Sinkhorn parameter-group overrides."""
256 patterns = []
257 groups = [{"params": []}]
258 for group in overrides:
259 if "params" in group:
260 raise ValueError("Sinkhorn builder groups use param_patterns, not params")
261 selector = cls._patterns(group.get("param_patterns", ()))
262 if not selector:
263 raise ValueError("Sinkhorn parameter group overrides require param_patterns")
264 patterns.append(selector)
265 groups.append({**{key: value for key, value in group.items() if key != "param_patterns"}, "params": []})
266 return patterns, groups
268 @classmethod
269 def _route(cls, param, names, roles, selectors, override_patterns):
270 """Resolve one parameter's optimizer and reject conflicting selectors."""
271 matches = [index for index, patterns in enumerate(override_patterns) if cls._matches(names, patterns)]
272 if len(matches) > 1:
273 raise ValueError(f"Overlapping Sinkhorn parameter groups for {names}")
274 explicit = {name for name, patterns in selectors.items() if cls._matches(names, patterns)}
275 if matches:
276 explicit.add("sinkhorn")
277 if len(explicit) > 1:
278 raise ValueError(f"Overlapping optimizer selectors for {names}")
279 family = next(iter(explicit)) if explicit else cls._default_family(param, names, *roles)
280 if family == "sinkhorn" and param.ndim != 2:
281 raise ValueError(f"Sinkhorn requires a 2D matrix: {names}")
282 if family == "muon" and param.ndim < 2:
283 raise ValueError(f"Muon requires a matrix: {names}")
284 return family, matches[0] + 1 if matches else 0
286 @staticmethod
287 def _patterns(patterns):
288 if isinstance(patterns, str):
289 patterns = [patterns]
290 return [re.compile(pattern) for pattern in patterns]
292 @staticmethod
293 def _matches(names, patterns):
294 return any(pattern.search(name) for pattern in patterns for name in names)
296 @staticmethod
297 def _aliases(model, excluded_param_ids):
298 excluded = set(excluded_param_ids or ())
299 aliases = {}
300 for name, param in model.named_parameters(remove_duplicate=False):
301 if param.requires_grad and id(param) not in excluded:
302 aliases.setdefault(param, []).append(name)
303 return aliases
305 @staticmethod
306 def _roles(model):
307 embeddings, norms = set(), set()
308 for module in model.modules():
309 if isinstance(module, nn.Embedding):
310 embeddings.add(id(module.weight))
311 if re.search(r"norm(?:[123]d|gated)?$", type(module).__name__, re.IGNORECASE):
312 norms.update(id(param) for param in module.parameters(recurse=False))
313 return embeddings, norms
315 @staticmethod
316 def _default_family(param, names, embedding_ids, norm_ids):
317 if id(param) in norm_ids:
318 return "adamw"
319 if id(param) in embedding_ids or any(
320 re.search(r"(?:^|\.)(?:lm_head|out_head|output_head)\.weight$", name) for name in names):
321 return "sinkhorn"
322 return "muon" if param.ndim >= 2 else "adamw"
324 @staticmethod
325 def _adamw_groups(params, aliases, norm_ids, no_decay_params):
326 """Split AdamW parameters by decay policy, including normalization weights."""
327 decay, no_decay = [], []
328 for param in params:
329 names = aliases[param]
330 if no_decay_params is None:
331 excluded = (param.ndim < 2 and id(param) not in norm_ids) or any(
332 name.rsplit(".", 1)[-1] in ("bias", "scale", "scales", "scaling_factor") for name in names
333 )
334 else:
335 excluded = any(keyword.lower() in name.lower() for keyword in no_decay_params for name in names)
336 (no_decay if excluded else decay).append(param)
337 groups = [{"params": decay}] if decay else []
338 if no_decay:
339 groups.append({"params": no_decay, "weight_decay": 0.0})
340 return groups
342 def get_optimizer(self) -> Any:
343 """Return the core chained optimizer runtime."""
344 return self.optimizer