Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / __init__.py: 97%

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 

17 

18import inspect 

19import logging 

20from typing import Any, Dict, List, Optional 

21 

22from hyper_parallel.core.optimizer.swap_optimizer import ( 

23 SwapOptimizer, 

24 SwapOptimizerConfig, 

25 swap_optimizer, 

26) 

27 

28logger = logging.getLogger(__name__) 

29logger.setLevel(logging.INFO) 

30 

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 "ChainedOptimizer": ".optimizer", 

37 "detect_dtensor_backend": ".dtensor_compat", 

38} 

39 

40 

41def __getattr__(name): # pylint: disable=invalid-name 

42 """Lazily import torch-only optimizer symbols.""" 

43 if name not in _LAZY_EXPORTS: 

44 raise AttributeError(f"module {__name__!r} has no attribute {name!r}") 

45 module = _import_module(_LAZY_EXPORTS[name], __name__) 

46 value = getattr(module, name) 

47 globals()[name] = value 

48 return value 

49 

50 

51def __dir__(): # pylint: disable=invalid-name 

52 """Include lazy torch-only optimizer exports in ``dir()``.""" 

53 return sorted(set(globals()) | set(_LAZY_EXPORTS)) 

54 

55 

56def _load_torch_optimizer_runtime(): 

57 """Import torch-only optimizer helpers used by the factory APIs.""" 

58 # pylint: disable=import-outside-toplevel,unused-import 

59 import hyper_parallel.core.optimizer.utils # noqa: F401 - install rank0 logging helpers 

60 

61 from hyper_parallel.core.optimizer.adamw import AdamW 

62 from hyper_parallel.core.optimizer.dtensor_compat import detect_dtensor_backend 

63 from hyper_parallel.core.optimizer.muon import Muon 

64 from hyper_parallel.core.optimizer.optimizer import ChainedOptimizer 

65 

66 return AdamW, Muon, ChainedOptimizer, detect_dtensor_backend 

67 

68 

69def _effective_optimizer_config( 

70 optimizer_class: Any, 

71 configured_values: Dict[str, Any], 

72 runtime_optimizer: Any, 

73) -> Dict[str, Any]: 

74 """Merge constructor defaults, user values, and resolved runtime defaults.""" 

75 signature = inspect.signature(optimizer_class.__init__) 

76 effective_config = {} 

77 for name, parameter in signature.parameters.items(): 

78 if name not in {"self", "params"} and parameter.default is not inspect.Parameter.empty: 

79 effective_config[name] = parameter.default 

80 effective_config.update(configured_values) 

81 effective_config.update(runtime_optimizer.defaults) 

82 if hasattr(runtime_optimizer, "hsdp_replica_count"): 

83 effective_config["hsdp_replica_count"] = runtime_optimizer.hsdp_replica_count 

84 return effective_config 

85 

86 

87def _filter_optimizer_config( 

88 optimizer_name: str, 

89 optimizer_class: Any, 

90 configured_values: Optional[Dict[str, Any]], 

91) -> Dict[str, Any]: 

92 """Normalize prefixed YAML keys and remove unsupported constructor args.""" 

93 prefix = f"{optimizer_name}_" 

94 normalized_config = { 

95 key[len(prefix):] if key.startswith(prefix) else key: value 

96 for key, value in (configured_values or {}).items() 

97 } 

98 allowed_keys = ( 

99 inspect.signature(optimizer_class.__init__).parameters.keys() 

100 - {"self", "params"} 

101 ) 

102 allowed_keys = set(allowed_keys) | set(getattr(optimizer_class, "ADDITIONAL_CONFIG_KEYS", ())) 

103 filtered_config = { 

104 key: value 

105 for key, value in normalized_config.items() 

106 if key in allowed_keys 

107 } 

108 excluded_keys = normalized_config.keys() - allowed_keys 

109 if excluded_keys: 

110 logger.info_rank0( 

111 "Excluded %s config: %s", 

112 optimizer_name, 

113 list(excluded_keys), 

114 ) 

115 return filtered_config 

116 

117 

118def _build_configured_optimizer( 

119 optimizer_name: str, 

120 optimizer_class: Any, 

121 param_groups: Any, 

122 configured_values: Dict[str, Any], 

123) -> Any: 

124 """Construct one leaf optimizer and log its effective configuration.""" 

125 optimizer = optimizer_class(param_groups, **configured_values) 

126 logger.info_rank0( 

127 f"Effective {optimizer_name} config: %s", 

128 _effective_optimizer_config( 

129 optimizer_class, 

130 configured_values, 

131 optimizer, 

132 ), 

133 ) 

134 return optimizer 

135 

136 

137def get_hyper_optimizer( 

138 model: Any, 

139 muon_params: List[Dict[str, Any]], 

140 adamw_params: List[Dict[str, Any]], 

141 muon_kwargs: Optional[Dict[str, Any]] = None, 

142 adamw_kwargs: Optional[Dict[str, Any]] = None, 

143) -> Any: 

144 """Create a chained Muon + AdamW optimizer. 

145 

146 Args: 

147 model: The neural network model. 

148 muon_params: Param groups for Muon. Empty list disables Muon. 

149 adamw_params: Param groups for AdamW. Empty list disables AdamW. 

150 muon_kwargs: Dedicated configurations dict for Muon. 

151 adamw_kwargs: Dedicated configurations dict for AdamW. 

152 

153 Example: 

154 from hyper_parallel.core.optimizer import get_hyper_optimizer 

155 

156 _adamw_legacy = { 

157 'adamw_lr': 1e-3, 

158 'adamw_weight_decay': 1e-2, 

159 'adamw_betas': (0.9, 0.95), 

160 'adamw_eps': 1e-8, 

161 'fused': True 

162 } 

163 _muon_legacy = { 

164 'muon_lr': 2e-2, 

165 'muon_weight_decay': 0.1, 

166 'muon_momentum': 0.95, 

167 'muon_ns_steps': 5, 

168 'muon_ns_variant': 'asym5', 

169 'muon_nesterov': True, 

170 'muon_hsdp_replica_count': 2 

171 } 

172 

173 optimizer = get_hyper_optimizer( 

174 model=model, 

175 muon_params=muon_groups, 

176 adamw_params=adamw_groups, 

177 adamw_kwargs=_adamw_legacy, 

178 muon_kwargs=_muon_legacy, 

179 ) 

180 

181 optimizer.step() 

182 """ 

183 AdamW, Muon, ChainedOptimizer, detect_dtensor_backend = _load_torch_optimizer_runtime() 

184 

185 filtered_adamw_config = _filter_optimizer_config( 

186 "adamw", 

187 AdamW, 

188 adamw_kwargs, 

189 ) 

190 filtered_muon_config = _filter_optimizer_config( 

191 "muon", 

192 Muon, 

193 muon_kwargs, 

194 ) 

195 optimizers = {} 

196 detect_dtensor_backend(adamw_params, muon_params) 

197 

198 if adamw_params: 

199 optimizers["adamw"] = _build_configured_optimizer( 

200 "adamw", 

201 AdamW, 

202 adamw_params, 

203 filtered_adamw_config, 

204 ) 

205 

206 if muon_params: 

207 optimizers["muon"] = _build_configured_optimizer( 

208 "muon", 

209 Muon, 

210 muon_params, 

211 filtered_muon_config, 

212 ) 

213 

214 flatten = bool(adamw_params and muon_params) 

215 

216 return ChainedOptimizer( 

217 model, 

218 optimizers=optimizers, 

219 flatten=flatten, 

220 ) 

221 

222 

223__all__ = [ 

224 'SwapOptimizer', 

225 'SwapOptimizerConfig', 

226 'get_hyper_optimizer', 

227 'swap_optimizer', 

228]