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

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 

16"""HyperParallel optimizer module.""" 

17 

18import inspect 

19import logging 

20from typing import Any, Dict, List, Optional 

21 

22from torch import nn 

23 

24from hyper_parallel.core.optimizer.adamw import AdamW 

25from hyper_parallel.core.optimizer.lr_scheduler import get_hyper_lr_scheduler 

26from hyper_parallel.core.optimizer.muon import Muon 

27from hyper_parallel.core.optimizer.optimizer import ChainedOptimizer 

28 

29logging.basicConfig(level=logging.INFO) 

30logger = logging.getLogger(__name__) 

31 

32__all__ = ['get_hyper_optimizer', 'get_hyper_lr_scheduler'] 

33 

34 

35def get_hyper_optimizer( 

36 model: nn.Module, 

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

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

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

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

41) -> ChainedOptimizer: 

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

43 

44 Args: 

45 model: The neural network model. 

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

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

48 muon_kwargs: Dedicated configurations dict for Muon. 

49 adamw_kwargs: Dedicated configurations dict for AdamW. 

50 

51 Example: 

52 _adamw_legacy = { 

53 'adamw_lr': 1e-3,  

54 'adamw_weight_decay': 1e-2,  

55 'adamw_betas': (0.9, 0.95),  

56 'adamw_eps': 1e-8, 

57 'fused': True 

58 } 

59 _muon_legacy = { 

60 'muon_lr': 2e-2,  

61 'muon_weight_decay': 0.1,  

62 'muon_momentum': 0.95,  

63 'muon_ns_steps': 5, 

64 'muon_nesterov': True, 

65 'muon_hsdp_replica_count': 2 

66 } 

67 

68 optimizer = get_hyper_optimizer( 

69 model=model, 

70 muon_params=muon_groups, 

71 adamw_params=adamw_groups, 

72 adamw_kwargs=_adamw_legacy, 

73 muon_kwargs=_muon_legacy, 

74 )  

75 """ 

76 # 1. Arguments Preparation 

77 # 1.1 adamw 

78 adamw_raw = adamw_kwargs or {} 

79 adamw_config = { 

80 k[6:] if k.startswith("adamw_") else k: v 

81 for k, v in adamw_raw.items() 

82 } 

83 allowed_keys_adamw = inspect.signature(AdamW.__init__).parameters.keys() - {'self', 'params'} 

84 filtered_adamw_config = {k: v for k, v in adamw_config.items() if k in allowed_keys_adamw} 

85 if excluded_adamw_keys := adamw_config.keys() - allowed_keys_adamw: 

86 logger.info_rank0("Excluded adamw config: %s", list(excluded_adamw_keys)) 

87 

88 # 1.2 muon 

89 muon_raw = muon_kwargs or {} 

90 muon_config = { 

91 k[5:] if k.startswith("muon_") else k: v 

92 for k, v in muon_raw.items() 

93 } 

94 allowed_keys_muon = inspect.signature(Muon.__init__).parameters.keys() - {'self', 'params'} 

95 filtered_muon_config = {k: v for k, v in muon_config.items() if k in allowed_keys_muon} 

96 if excluded_muon_keys := muon_config.keys() - allowed_keys_muon: 

97 logger.info_rank0("Excluded muon config: %s", list(excluded_muon_keys)) 

98 

99 # 2. Optimizer Creation 

100 optimizers = {} 

101 

102 if adamw_params: 

103 optimizers["adamw"] = AdamW(adamw_params, **filtered_adamw_config) 

104 logger.info_rank0("Using adamw config: %s", filtered_adamw_config) 

105 

106 if muon_params: 

107 optimizers["muon"] = Muon(muon_params, **filtered_muon_config) 

108 logger.info_rank0("Using muon config: %s", filtered_muon_config) 

109 

110 flatten = bool(adamw_params and muon_params) 

111 

112 return ChainedOptimizer(model, optimizers=optimizers, flatten=flatten)