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

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"""Learning rate schedule utilities for HyperParallel optimizers.""" 

17 

18import copy 

19import logging 

20import math 

21from typing import TYPE_CHECKING, Any, Dict, Iterator, List 

22 

23from torch.optim.lr_scheduler import LambdaLR 

24 

25from hyper_parallel.core.optimizer.optimizer import ChainedOptimizer 

26 

27if TYPE_CHECKING: 

28 from torch.optim import Optimizer 

29 

30logger = logging.getLogger(__name__) 

31 

32 

33 

34def get_constant_schedule_with_warmup( 

35 optimizer: "Optimizer", 

36 num_warmup_steps: int, 

37 init_lr: float, 

38 last_epoch: int = -1, 

39 lr_start: float = 0.0, 

40): 

41 """ 

42 Creates a schedule with a constant learning rate preceded by a warmup period during which the learning rate 

43 increases linearly between 0 and the initial lr set in the optimizer. 

44 """ 

45 

46 def _lr_lambda(current_step: int): 

47 if current_step < num_warmup_steps: 

48 warmup_progress = current_step / max(1, num_warmup_steps) 

49 return (lr_start + (init_lr - lr_start) * warmup_progress) / init_lr 

50 

51 return 1.0 

52 

53 return LambdaLR(optimizer, _lr_lambda, last_epoch=last_epoch) 

54 

55 

56def get_linear_schedule_with_warmup( 

57 optimizer: "Optimizer", 

58 num_warmup_steps: int, 

59 num_training_steps: int, 

60 init_lr: float, 

61 last_epoch: int = -1, 

62 min_lr: float = 1e-7, 

63 lr_start: float = 0.0, 

64): 

65 """ 

66 Creates a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, 

67 after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer. 

68 """ 

69 

70 def _lr_lambda(current_step: int): 

71 if current_step < num_warmup_steps: 

72 warmup_progress = current_step / max(1, num_warmup_steps) 

73 return (lr_start + (init_lr - lr_start) * warmup_progress) / init_lr 

74 

75 min_lr_ratio = min_lr / init_lr if init_lr != 0.0 else 0.0 

76 return max( 

77 min_lr_ratio, 

78 float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps)), 

79 ) 

80 

81 return LambdaLR(optimizer, _lr_lambda, last_epoch) 

82 

83 

84def get_cosine_schedule_with_warmup( 

85 optimizer: "Optimizer", 

86 num_warmup_steps: int, 

87 num_training_steps: int, 

88 init_lr: float, 

89 num_cycles: float = 0.5, 

90 last_epoch: int = -1, 

91 lr_decay_ratio: float = 1.0, 

92 min_lr: float = 1e-7, 

93 lr_start: float = 0.0, 

94): 

95 """ 

96 Creates a schedule with a learning rate that decreases following the values of the cosine function between 

97 the initial lr set in the optimizer to min_lr, after a warmup period during which it increases linearly between 0 

98 and the initial lr set in the optimizer. 

99 """ 

100 

101 def lr_lambda(current_step: int): 

102 lr_decay_steps = int(num_training_steps * lr_decay_ratio) 

103 if current_step < num_warmup_steps: 

104 warmup_progress = current_step / max(1, num_warmup_steps) 

105 return (lr_start + (init_lr - lr_start) * warmup_progress) / init_lr 

106 

107 min_lr_ratio = min_lr / init_lr if init_lr != 0.0 else 0.0 

108 if current_step > lr_decay_steps: 

109 return min_lr_ratio 

110 

111 progress = float(current_step - num_warmup_steps) / float(max(1, lr_decay_steps - num_warmup_steps)) 

112 if not 0 <= progress <= 1: 

113 raise ValueError(f"cosine schedule progress must be in [0, 1], but got {progress}") 

114 factor = 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)) 

115 factor = factor * (1 - min_lr_ratio) + min_lr_ratio 

116 return max(0, factor) 

117 

118 return LambdaLR(optimizer, lr_lambda, last_epoch) 

119 

120 

121class LRSchedulersContainer: 

122 """Container for multiple learning rate schedulers. 

123 

124 Each scheduler is keyed by the same name as its corresponding sub-optimizer 

125 in ``ChainedOptimizer.optimizers_dict`` (e.g. ``"muon"``, ``"adamw"``). 

126 This ensures that ``state_dict`` / ``load_state_dict`` are robust to 

127 insertion-order differences between the save and load environments. 

128 """ 

129 # , **scheduler_kwargs,, scheduler_kwargs: Dict[str, Any] 

130 def __init__(self, optimizers: ChainedOptimizer, scheduler) -> None: 

131 self._names: List[str] = list(optimizers.optimizers_dict.keys()) 

132 self._schedulers_by_name = {} 

133 for name, opt in optimizers.optimizers_dict.items(): 

134 self._schedulers_by_name[name] = scheduler( 

135 optimizer=opt 

136 ) 

137 self.schedulers = [self._schedulers_by_name[name] for name in self._names] 

138 

139 def __iter__(self) -> Iterator: 

140 """Iterate over the registered schedulers.""" 

141 return iter(self.schedulers) 

142 

143 def __len__(self) -> int: 

144 """Return the total number of schedulers in the container.""" 

145 return len(self.schedulers) 

146 

147 def step(self) -> None: 

148 """Advance the step for all schedulers.""" 

149 for scheduler in self.schedulers: 

150 scheduler.step() 

151 

152 def get_last_lr(self) -> List[float]: 

153 """Return a flattened list of the last learning rates across all sub-schedulers.""" 

154 param_last_lr: List[float] = [] 

155 for scheduler in self.schedulers: 

156 param_last_lr.extend(scheduler.get_last_lr()) 

157 return param_last_lr 

158 

159 def state_dict(self) -> Dict[str, Any]: 

160 """Return scheduler states keyed by sub-optimizer name. 

161 

162 Compatible with veomni's ``{'muon': ..., 'adamw': ...}`` format. 

163 """ 

164 return {name: self._schedulers_by_name[name].state_dict() for name in self._names} 

165 

166 def load_state_dict(self, state_dict: Dict[str, Any]) -> None: 

167 """Load scheduler states keyed by sub-optimizer name. 

168 

169 Matches by name rather than position, so the checkpoint key order 

170 (e.g. ``muon, adamw``) need not match the local creation order 

171 (e.g. ``adamw, muon``). 

172 """ 

173 if len(self._names) != len(state_dict): 

174 raise RuntimeError( 

175 f"Scheduler count mismatch! Current has {len(self._names)}, " 

176 f"but checkpoint contains {len(state_dict)} states." 

177 ) 

178 for name in self._names: 

179 if name not in state_dict: 

180 raise RuntimeError( 

181 f"Missing state for scheduler '{name}' in state_dict. " 

182 f"Available keys: {sorted(state_dict.keys())}, " 

183 f"expected keys: {sorted(self._names)}." 

184 ) 

185 self._schedulers_by_name[name].load_state_dict( 

186 copy.deepcopy(state_dict[name]), 

187 ) 

188 scheduler = self._schedulers_by_name[name] 

189 last_lrs = scheduler.get_last_lr() 

190 if len(last_lrs) != len(scheduler.optimizer.param_groups): 

191 raise RuntimeError( 

192 f"Scheduler '{name}' restored {len(last_lrs)} learning rates " 

193 f"for {len(scheduler.optimizer.param_groups)} parameter groups." 

194 ) 

195 for parameter_group, learning_rate in zip( 

196 scheduler.optimizer.param_groups, 

197 last_lrs, 

198 ): 

199 parameter_group["lr"] = learning_rate