Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / __init__.py: 97%
66 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"""HyperParallel optimizer module."""
16from importlib import import_module as _import_module # pylint: disable=invalid-name
18import inspect
19import logging
20from typing import Any, Dict, List, Optional
22from hyper_parallel.core.optimizer.swap_optimizer import (
23 SwapOptimizer,
24 SwapOptimizerConfig,
25 swap_optimizer,
26)
28logger = logging.getLogger(__name__)
29logger.setLevel(logging.INFO)
31# Optimizer implementations import torch at module load. Keep them off the
32# eager path so importing the package does not pull in torch.
33_LAZY_EXPORTS = {
34 "AdamW": ".adamw",
35 "Muon": ".muon",
36 "Sinkhorn": ".sinkhorn",
37 "ChainedOptimizer": ".optimizer",
38 "detect_dtensor_backend": ".dtensor_compat",
39}
42def __getattr__(name): # pylint: disable=invalid-name
43 """Lazily import torch-only optimizer symbols."""
44 if name not in _LAZY_EXPORTS:
45 raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
46 module = _import_module(_LAZY_EXPORTS[name], __name__)
47 value = getattr(module, name)
48 globals()[name] = value
49 return value
52def __dir__(): # pylint: disable=invalid-name
53 """Include lazy torch-only optimizer exports in ``dir()``."""
54 return sorted(set(globals()) | set(_LAZY_EXPORTS))
57def _load_torch_optimizer_runtime():
58 """Import torch-only optimizer helpers used by the factory APIs."""
59 # pylint: disable=import-outside-toplevel,unused-import
60 import hyper_parallel.core.optimizer.utils # noqa: F401 - install rank0 logging helpers
62 from hyper_parallel.core.optimizer.adamw import AdamW
63 from hyper_parallel.core.optimizer.dtensor_compat import detect_dtensor_backend
64 from hyper_parallel.core.optimizer.muon import Muon
65 from hyper_parallel.core.optimizer.optimizer import ChainedOptimizer
67 return AdamW, Muon, ChainedOptimizer, detect_dtensor_backend
70def _effective_optimizer_config(
71 optimizer_class: Any,
72 configured_values: Dict[str, Any],
73 runtime_optimizer: Any,
74) -> Dict[str, Any]:
75 """Merge constructor defaults, user values, and resolved runtime defaults."""
76 signature = inspect.signature(optimizer_class.__init__)
77 effective_config = {}
78 for name, parameter in signature.parameters.items():
79 if name not in {"self", "params"} and parameter.default is not inspect.Parameter.empty:
80 effective_config[name] = parameter.default
81 effective_config.update(configured_values)
82 effective_config.update(runtime_optimizer.defaults)
83 if hasattr(runtime_optimizer, "hsdp_replica_count"):
84 effective_config["hsdp_replica_count"] = runtime_optimizer.hsdp_replica_count
85 return effective_config
88def _filter_optimizer_config(
89 optimizer_name: str,
90 optimizer_class: Any,
91 configured_values: Optional[Dict[str, Any]],
92) -> Dict[str, Any]:
93 """Normalize prefixed YAML keys and remove unsupported constructor args."""
94 prefix = f"{optimizer_name}_"
95 normalized_config = {
96 key[len(prefix):] if key.startswith(prefix) else key: value
97 for key, value in (configured_values or {}).items()
98 }
99 allowed_keys = (
100 inspect.signature(optimizer_class.__init__).parameters.keys()
101 - {"self", "params"}
102 )
103 allowed_keys = set(allowed_keys) | set(getattr(optimizer_class, "ADDITIONAL_CONFIG_KEYS", ()))
104 filtered_config = {
105 key: value
106 for key, value in normalized_config.items()
107 if key in allowed_keys
108 }
109 excluded_keys = normalized_config.keys() - allowed_keys
110 if excluded_keys:
111 logger.info_rank0(
112 "Excluded %s config: %s",
113 optimizer_name,
114 list(excluded_keys),
115 )
116 return filtered_config
119def _build_configured_optimizer(
120 optimizer_name: str,
121 optimizer_class: Any,
122 param_groups: Any,
123 configured_values: Dict[str, Any],
124) -> Any:
125 """Construct one leaf optimizer and log its effective configuration."""
126 optimizer = optimizer_class(param_groups, **configured_values)
127 logger.info_rank0(
128 f"Effective {optimizer_name} config: %s",
129 _effective_optimizer_config(
130 optimizer_class,
131 configured_values,
132 optimizer,
133 ),
134 )
135 return optimizer
138def get_hyper_optimizer(
139 model: Any,
140 muon_params: List[Dict[str, Any]],
141 adamw_params: List[Dict[str, Any]],
142 muon_kwargs: Optional[Dict[str, Any]] = None,
143 adamw_kwargs: Optional[Dict[str, Any]] = None,
144 *,
145 sinkhorn_params: Optional[List[Dict[str, Any]]] = None,
146 sinkhorn_kwargs: Optional[Dict[str, Any]] = None,
147) -> Any:
148 """Create a chained Muon + AdamW optimizer with optional Sinkhorn groups.
150 Args:
151 model: The neural network model.
152 muon_params: Param groups for Muon. Empty list disables Muon.
153 adamw_params: Param groups for AdamW. Empty list disables AdamW.
154 muon_kwargs: Dedicated configurations dict for Muon.
155 adamw_kwargs: Dedicated configurations dict for AdamW.
156 sinkhorn_params: Optional Sinkhorn groups. Empty or None disables Sinkhorn.
157 sinkhorn_kwargs: Dedicated configurations dict for Sinkhorn.
159 Example:
160 from hyper_parallel.core.optimizer import get_hyper_optimizer
162 _adamw_legacy = {
163 'adamw_lr': 1e-3,
164 'adamw_weight_decay': 1e-2,
165 'adamw_betas': (0.9, 0.95),
166 'adamw_eps': 1e-8,
167 'fused': True
168 }
169 _muon_legacy = {
170 'muon_lr': 2e-2,
171 'muon_weight_decay': 0.1,
172 'muon_momentum': 0.95,
173 'muon_ns_steps': 5,
174 'muon_ns_variant': 'asym5',
175 'muon_nesterov': True,
176 'muon_hsdp_replica_count': 2
177 }
179 optimizer = get_hyper_optimizer(
180 model=model,
181 muon_params=muon_groups,
182 adamw_params=adamw_groups,
183 adamw_kwargs=_adamw_legacy,
184 muon_kwargs=_muon_legacy,
185 )
187 optimizer.step()
188 """
189 AdamW, Muon, ChainedOptimizer, detect_dtensor_backend = _load_torch_optimizer_runtime()
191 filtered_adamw_config = _filter_optimizer_config(
192 "adamw",
193 AdamW,
194 adamw_kwargs,
195 )
196 filtered_muon_config = _filter_optimizer_config(
197 "muon",
198 Muon,
199 muon_kwargs,
200 )
201 optimizers = {}
202 detect_dtensor_backend(adamw_params or sinkhorn_params or [], muon_params)
204 if adamw_params:
205 optimizers["adamw"] = _build_configured_optimizer(
206 "adamw",
207 AdamW,
208 adamw_params,
209 filtered_adamw_config,
210 )
212 if muon_params:
213 optimizers["muon"] = _build_configured_optimizer(
214 "muon",
215 Muon,
216 muon_params,
217 filtered_muon_config,
218 )
220 if sinkhorn_params:
221 sinkhorn_class = __getattr__("Sinkhorn")
222 filtered_sinkhorn_config = _filter_optimizer_config("sinkhorn", sinkhorn_class, sinkhorn_kwargs)
223 optimizers["sinkhorn"] = _build_configured_optimizer(
224 "sinkhorn", sinkhorn_class, sinkhorn_params, filtered_sinkhorn_config,
225 )
227 flatten = len(optimizers) > 1
229 return ChainedOptimizer(
230 model,
231 optimizers=optimizers,
232 flatten=flatten,
233 )
236__all__ = [
237 'SwapOptimizer',
238 'SwapOptimizerConfig',
239 'get_hyper_optimizer',
240 'swap_optimizer',
241]