Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / optim / builders.py: 95%

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. 

16 

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""" 

21 

22__all__ = ["AdamW", "Muon"] 

23 

24import logging 

25from typing import Any, Dict, List, Optional, Sequence, Tuple 

26 

27from torch import nn # pylint: disable=forbidden-backend-import 

28 

29from hyper_parallel.core.optimizer import get_hyper_optimizer 

30from hyper_parallel.components.optim.parameter_groups import ( 

31 _DEFAULT_ADAMW_NAME_KEYWORDS, 

32 get_adamw_param_groups, 

33 split_muon_adamw_params, 

34) 

35 

36logger = logging.getLogger(__name__) 

37 

38 

39class AdamW: 

40 """Build a core AdamW optimizer from YAML configuration.""" 

41 

42 def __init__( 

43 self, 

44 adamw_config: dict, 

45 model: nn.Module, 

46 no_decay_params: Optional[List[str]] = None, 

47 excluded_param_ids: Optional[Sequence[int]] = None, 

48 ) -> None: 

49 """Initialize AdamW optimizer configuration. 

50 

51 Args: 

52 adamw_config: AdamW hyperparameters resolved from YAML. 

53 model: Module whose trainable parameters are optimized. 

54 no_decay_params: Optional names excluded from weight decay. 

55 """ 

56 self.config = adamw_config 

57 self.model = model 

58 

59 adamw_groups, _ = self.get_adamw_param_groups( 

60 self.model, 

61 weight_decay=adamw_config.get("adamw_weight_decay", 1e-2), 

62 no_decay_params=no_decay_params, 

63 excluded_param_ids=excluded_param_ids, 

64 ) 

65 if not adamw_groups: 

66 raise ValueError("AdamW requires at least one trainable parameter") 

67 

68 self.optimizer = get_hyper_optimizer( 

69 model=self.model, 

70 muon_params=[], 

71 adamw_params=adamw_groups, 

72 adamw_kwargs=adamw_config, 

73 ) 

74 

75 @staticmethod 

76 def get_adamw_param_groups( 

77 model: "nn.Module", 

78 weight_decay: float = 1e-2, 

79 no_decay_params: Optional[Sequence[str]] = None, 

80 param_groups: Optional[Sequence[Dict[str, Any]]] = None, 

81 allowed_param_ids: Optional[Sequence[int]] = None, 

82 excluded_param_ids: Optional[Sequence[int]] = None, 

83 ) -> Tuple[List[Dict[str, Any]], List[str]]: 

84 """Split model parameters into decaying and non-decaying groups. 

85 

86 Delegates to 

87 :func:`hyper_parallel.components.optim.parameter_groups.get_adamw_param_groups`; 

88 retained on the class so existing call sites keep working. 

89 

90 Args: 

91 model: Module whose trainable parameters are routed. 

92 weight_decay: Default decay coefficient for trainable parameters. 

93 no_decay_params: Name keywords excluded from weight decay. 

94 param_groups: Optional prebuilt groups to return unchanged. 

95 allowed_param_ids: Optional parameter identities eligible for routing. 

96 

97 Returns: 

98 Parameter groups and names routed to AdamW. 

99 """ 

100 return get_adamw_param_groups( 

101 model, 

102 weight_decay=weight_decay, 

103 no_decay_params=no_decay_params, 

104 param_groups=param_groups, 

105 allowed_param_ids=allowed_param_ids, 

106 excluded_param_ids=excluded_param_ids, 

107 ) 

108 

109 def get_optimizer(self) -> Any: 

110 """Return the core chained optimizer runtime.""" 

111 return self.optimizer 

112 

113 

114class Muon: 

115 """Build a mixed Muon and AdamW optimizer from YAML configuration.""" 

116 

117 _DEFAULT_ADAMW_NAME_KEYWORDS = _DEFAULT_ADAMW_NAME_KEYWORDS 

118 

119 def __init__( 

120 self, 

121 muon_config: dict, 

122 adamw_config: dict, 

123 model: nn.Module, 

124 extra_adamw_name_keywords: Optional[List[str]] = None, 

125 no_decay_params: Optional[List[str]] = None, 

126 excluded_param_ids: Optional[Sequence[int]] = None, 

127 ) -> None: 

128 """Build a mixed Muon and fallback AdamW runtime for ``model``. 

129 

130 Args: 

131 muon_config: Muon hyperparameters resolved from YAML. 

132 adamw_config: Fallback AdamW hyperparameters resolved from YAML. 

133 model: Module whose trainable parameters are optimized. 

134 extra_adamw_name_keywords: Additional names routed to AdamW. 

135 no_decay_params: Optional names excluded from weight decay. 

136 """ 

137 self.muon_config = muon_config 

138 self.adamw_config = adamw_config 

139 self.model = model 

140 

141 muon_params, adamw_params, muon_names, adamw_names = self.split_muon_adamw_params( 

142 model, 

143 extra_adamw_name_keywords=extra_adamw_name_keywords or (), 

144 excluded_param_ids=excluded_param_ids, 

145 ) 

146 if not muon_params: 

147 raise ValueError("Muon requires at least one eligible matrix parameter") 

148 

149 adamw_groups, _ = AdamW.get_adamw_param_groups( 

150 model, 

151 weight_decay=adamw_config.get("adamw_weight_decay", 1e-2), 

152 no_decay_params=no_decay_params, 

153 allowed_param_ids=[id(parameter) for parameter in adamw_params], 

154 excluded_param_ids=excluded_param_ids, 

155 ) 

156 

157 logger.info_rank0( 

158 "Muon optimizer split: %s Muon parameters, %s AdamW parameters", 

159 len(muon_names), 

160 len(adamw_names), 

161 ) 

162 logger.info_rank0("Muon parameters (first 5): %s", muon_names[:5]) 

163 logger.info_rank0("AdamW parameters (first 5): %s", adamw_names[:5]) 

164 

165 self.optimizer = get_hyper_optimizer( 

166 model=self.model, 

167 muon_params=muon_params, 

168 adamw_params=adamw_groups, 

169 muon_kwargs=muon_config, 

170 adamw_kwargs=adamw_config, 

171 ) 

172 

173 @staticmethod 

174 def split_muon_adamw_params( 

175 model: nn.Module, 

176 extra_adamw_name_keywords: Sequence[str] = (), 

177 excluded_param_ids: Optional[Sequence[int]] = None, 

178 ) -> Tuple[List[nn.Parameter], List[nn.Parameter], List[str], List[str]]: 

179 """Route matrix parameters to Muon and remaining parameters to AdamW. 

180 

181 Delegates to 

182 :func:`hyper_parallel.components.optim.parameter_groups.split_muon_adamw_params`; 

183 retained on the class so existing call sites keep working. 

184 

185 Args: 

186 model: Module whose trainable parameters are routed. 

187 extra_adamw_name_keywords: Additional names reserved for AdamW. 

188 

189 Returns: 

190 Muon parameters, AdamW parameters, and their respective names. 

191 """ 

192 return split_muon_adamw_params( 

193 model, 

194 extra_adamw_name_keywords=extra_adamw_name_keywords, 

195 excluded_param_ids=excluded_param_ids, 

196 ) 

197 

198 def get_optimizer(self) -> Any: 

199 """Return the core chained optimizer runtime.""" 

200 return self.optimizer