Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / optim / parameter_groups.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"""Parameter-name classification and optimizer parameter-group construction.""" 

16 

17__all__ = [ 

18 "get_adamw_param_groups", 

19 "get_parameter_names", 

20 "split_muon_adamw_params", 

21] 

22 

23import logging 

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

25 

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

27 

28logger = logging.getLogger(__name__) 

29 

30_DEFAULT_ADAMW_NAME_KEYWORDS: Tuple[str, ...] = ( 

31 "embed", 

32 "lm_head", 

33 "out_head", 

34 "out_proj.bias", 

35) 

36 

37 

38def get_parameter_names( 

39 model: nn.Module, 

40 forbidden_param_names: Optional[Sequence[str]], 

41) -> List[str]: 

42 """Return canonical parameter names that do not match no-decay keywords. 

43 

44 Transparent module wrappers may expose different internal names through 

45 ``named_children()`` while normalizing their public ``named_parameters()`` 

46 names. Deriving and consuming names from the same public traversal keeps 

47 optimizer grouping stable when those wrappers are installed. 

48 

49 Args: 

50 model: Model whose parameters should be classified. 

51 forbidden_param_names: Case-insensitive keywords excluded from weight decay. 

52 

53 Returns: 

54 Parameter names eligible for weight decay. 

55 """ 

56 forbidden_names = tuple(name.lower() for name in forbidden_param_names or ()) 

57 return [ 

58 name 

59 for name, _ in model.named_parameters() 

60 if not any(forbidden in name.lower() for forbidden in forbidden_names) 

61 ] 

62 

63 

64def get_adamw_param_groups( 

65 model: "nn.Module", 

66 weight_decay: float = 1e-2, 

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

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

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

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

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

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

73 

74 The function also adds ``model_name`` optimizer metadata to each selected 

75 parameter. 

76 

77 Args: 

78 model: Module whose trainable parameters are routed. 

79 weight_decay: Default decay coefficient for trainable parameters. 

80 no_decay_params: Name keywords excluded from weight decay. 

81 param_groups: Optional prebuilt groups to return unchanged. 

82 allowed_param_ids: Optional parameter identities eligible for routing. 

83 

84 Returns: 

85 Parameter groups and names routed to AdamW. 

86 """ 

87 adamw_names = [] 

88 no_decay_keywords = ("bias", "norm", "ln_") if no_decay_params is None else no_decay_params 

89 no_decay_params = tuple(keyword.lower() for keyword in no_decay_keywords) 

90 

91 if param_groups is None: 

92 decay_param_names = set(get_parameter_names(model, no_decay_params)) 

93 allowed_ids = set(allowed_param_ids) if allowed_param_ids is not None else None 

94 excluded_ids = set(excluded_param_ids or ()) 

95 decay_parameters = [] 

96 no_decay_parameters = [] 

97 no_decay_parameter_names = [] 

98 for n, p in model.named_parameters(): 

99 if not p.requires_grad or id(p) in excluded_ids or (allowed_ids is not None and id(p) not in allowed_ids): 

100 continue 

101 setattr(p, "model_name", n) 

102 adamw_names.append(n) 

103 if n in decay_param_names: 

104 decay_parameters.append(p) 

105 else: 

106 no_decay_parameter_names.append(n) 

107 no_decay_parameters.append(p) 

108 

109 param_groups = [] 

110 if decay_parameters: 

111 param_groups.append({"params": decay_parameters, "weight_decay": weight_decay}) 

112 if no_decay_parameters: 

113 logger.info_rank0(f"Parameters without weight decay: {no_decay_parameter_names}") 

114 param_groups.append({"params": no_decay_parameters, "weight_decay": 0.0}) 

115 

116 return param_groups, adamw_names 

117 

118 

119def split_muon_adamw_params( 

120 model: nn.Module, 

121 extra_adamw_name_keywords: Sequence[str] = (), 

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

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

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

125 

126 Args: 

127 model: Module whose trainable parameters are routed. 

128 extra_adamw_name_keywords: Additional names reserved for AdamW. 

129 

130 Returns: 

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

132 """ 

133 adamw_keywords = [] 

134 for keyword in (*_DEFAULT_ADAMW_NAME_KEYWORDS, *extra_adamw_name_keywords): 

135 adamw_keywords.append(keyword.lower()) 

136 adamw_keywords = tuple(adamw_keywords) 

137 muon_params, adamw_params = [], [] 

138 muon_names, adamw_names = [], [] 

139 excluded_ids = set(excluded_param_ids or ()) 

140 

141 for name, parameter in model.named_parameters(): 

142 setattr(parameter, "model_name", name) 

143 if not parameter.requires_grad or id(parameter) in excluded_ids: 

144 continue 

145 

146 is_matrix = parameter.ndim >= 2 

147 is_adamw = any(keyword in name.lower() for keyword in adamw_keywords) 

148 if is_matrix and not is_adamw: 

149 muon_params.append(parameter) 

150 muon_names.append(name) 

151 else: 

152 adamw_params.append(parameter) 

153 adamw_names.append(name) 

154 

155 return muon_params, adamw_params, muon_names, adamw_names