Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / lr_scheduler.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"""Learning rate schedule utilities for HyperParallel optimizers.""" 

17 

18import copy 

19import logging 

20import math 

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

22 

23from hyper_parallel.core.optimizer.optimizer import ChainedOptimizer 

24 

25logger = logging.getLogger(__name__) 

26 

27SUPPORTED_DECAY_STYLES = {"constant", "linear", "cosine", "WSD"} 

28SUPPORTED_WSD_STYLES = {"linear", "cosine", "exponential", "minus_sqrt"} 

29 

30 

31class OptimizerParamScheduler: 

32 """Anneals learning rate and weight decay for a SINGLE optimizer. 

33 

34 Aligned with LambdaLR-based scheduler: 

35 - LR values are bit-identical for constant / linear / cosine decay styles. 

36 - Different param groups (e.g. muon vs adamw) naturally get different peak LRs 

37 via ``param_group['initial_lr']`` (set by optimizer constructor), without 

38 injecting extra keys into param_groups. 

39 """ 

40 

41 def __init__( 

42 self, 

43 optimizer: Any, 

44 init_lr: float, 

45 lr_start: float, 

46 min_lr: float, 

47 lr_warmup_steps: int, 

48 lr_decay_steps: int, 

49 lr_decay_style: str, 

50 lr_decay_ratio: float = 1.0, 

51 wsd_decay_steps: Optional[int] = None, 

52 lr_wsd_decay_style: str = "exponential", 

53 override_opt_param_scheduler: bool = False 

54 ): 

55 self.optimizer = optimizer 

56 self.init_lr = init_lr 

57 self.lr_start = lr_start 

58 self.min_lr = min_lr 

59 self.lr_warmup_steps = lr_warmup_steps 

60 self.lr_decay_steps = lr_decay_steps 

61 self.lr_decay_style = lr_decay_style 

62 self.lr_decay_ratio = lr_decay_ratio 

63 

64 self.wsd_decay_steps = wsd_decay_steps 

65 self.lr_wsd_decay_style = lr_wsd_decay_style 

66 

67 self.override_opt_param_scheduler = override_opt_param_scheduler 

68 self.num_steps = 0 

69 

70 # Ensure every param_group has 'initial_lr' (same as PyTorch LambdaLR does). 

71 # This is the peak LR for each group — set by the optimizer constructor 

72 # (e.g. muon lr=0.001, adamw lr=0.0001), so different sub-optimizers 

73 # naturally get different peak LRs without injecting extra keys. 

74 for pg in self.optimizer.param_groups: 

75 pg.setdefault('initial_lr', pg['lr']) 

76 

77 self._validate_params() 

78 

79 # Set the learning rate 

80 self.step(0) 

81 

82 def _validate_params(self) -> None: 

83 """Validate initialization parameters to ensure logical correctness.""" 

84 if self.min_lr < 0.0: 

85 raise ValueError("min_lr must be >= 0.0") 

86 

87 if self.init_lr < self.min_lr: 

88 raise ValueError("init_lr must be >= min_lr") 

89 

90 if self.lr_decay_steps <= 0: 

91 raise ValueError("lr_decay_steps must be > 0") 

92 

93 if self.lr_warmup_steps >= self.lr_decay_steps: 

94 raise ValueError("warmup_steps must be < decay_steps") 

95 

96 if self.lr_decay_style == "WSD": 

97 if self.wsd_decay_steps is None: 

98 raise ValueError("wsd_decay_steps must be not None") 

99 

100 if self.wsd_decay_steps <= 0: 

101 raise ValueError("wsd_decay_steps must be > 0") 

102 

103 if self.wsd_decay_steps > self.lr_decay_steps: 

104 raise ValueError("wsd_decay_steps must be <= lr_decay_steps") 

105 

106 def get_lr(self, param_group: Dict[str, Any]) -> float: 

107 """Calculate and return the learning rate based on current step and decay style.""" 

108 max_lr = param_group['initial_lr'] 

109 init_lr = self.init_lr 

110 min_lr_ratio = self.min_lr / init_lr if init_lr > 0 else 0.0 

111 

112 # 1. Linear warmup (exclusive boundary: step < warmup_steps, same as origin). 

113 if self.lr_warmup_steps > 0 and self.num_steps < self.lr_warmup_steps: 

114 progress = float(self.num_steps) / float(self.lr_warmup_steps) 

115 factor = (self.lr_start + (init_lr - self.lr_start) * progress) / init_lr 

116 return factor * max_lr 

117 

118 # If the learning rate is constant, just return the peak value. 

119 if self.lr_decay_style == 'constant': 

120 return max_lr 

121 

122 # 2. Decay period 

123 # 2.1 WSD decay 

124 if self.lr_decay_style == 'WSD': 

125 # WSD: Warmup -> Stable -> Decay 

126 # For WSD, lr_decay_steps is the total schedule length (no lr_decay_ratio applied). 

127 if self.num_steps > self.lr_decay_steps: 

128 return min_lr_ratio * max_lr 

129 

130 wsd_anneal_start = self.lr_decay_steps - (self.wsd_decay_steps or 0) 

131 

132 if self.num_steps <= wsd_anneal_start: 

133 return max_lr # Stable Phase: keep max_lr without decaying 

134 

135 # Final decay phase of WSD 

136 wsd_decay_ratio = float(self.num_steps - wsd_anneal_start) / float(self.wsd_decay_steps or 1) 

137 

138 if self.lr_wsd_decay_style == "linear": 

139 coeff = 1.0 - wsd_decay_ratio 

140 elif self.lr_wsd_decay_style == "cosine": 

141 coeff = 0.5 * (math.cos(math.pi * wsd_decay_ratio) + 1.0) 

142 elif self.lr_wsd_decay_style == "exponential": 

143 coeff = (2.0 * (0.5 ** wsd_decay_ratio)) - 1.0 

144 else: # minus_sqrt fallback 

145 coeff = 1.0 - math.sqrt(wsd_decay_ratio) 

146 

147 factor = max(0.0, coeff) * (1 - min_lr_ratio) + min_lr_ratio 

148 return factor * max_lr 

149 

150 # 2.2 Non-WSD decay: use lr_decay_ratio to compute effective decay_steps 

151 lr_decay_steps = int(self.lr_decay_steps * self.lr_decay_ratio) 

152 if self.num_steps > lr_decay_steps: 

153 return min_lr_ratio * max_lr 

154 

155 decay_ratio = float(self.num_steps - self.lr_warmup_steps) / float( 

156 max(1, lr_decay_steps - self.lr_warmup_steps) 

157 ) 

158 

159 decay_ratio = max(0.0, min(1.0, decay_ratio)) 

160 

161 if self.lr_decay_style == 'linear': 

162 factor = max(min_lr_ratio, 1.0 - decay_ratio) 

163 return factor * max_lr 

164 

165 if self.lr_decay_style == 'cosine': 

166 coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) 

167 factor = coeff * (1 - min_lr_ratio) + min_lr_ratio 

168 return max(0.0, factor) * max_lr 

169 

170 raise ValueError(f"Unsupported decay style: {self.lr_decay_style}") 

171 

172 def step(self, increment: int = 1) -> None: 

173 """Advance the scheduler steps and set lr for all parameters groups.""" 

174 self.num_steps += increment 

175 for param_group in self.optimizer.param_groups: 

176 param_group['lr'] = self.get_lr(param_group) 

177 

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

179 """Return the current learning rate for all parameter groups.""" 

180 return [group['lr'] for group in self.optimizer.param_groups] 

181 

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

183 """Return the state of the scheduler as a dict.""" 

184 num_groups = len(self.optimizer.param_groups) 

185 

186 # LambdaLR.load_state_dict compatible layout 

187 state_dict = { 

188 "last_epoch": self.num_steps, 

189 "base_lrs": [pg['initial_lr'] for pg in self.optimizer.param_groups], 

190 "_last_lr": [pg['lr'] for pg in self.optimizer.param_groups], 

191 "_step_count": self.num_steps + 1, 

192 "lr_lambdas": [None] * num_groups, 

193 } 

194 

195 # OptimizerParamScheduler config 

196 state_dict.update({ 

197 "initial_lr": self.init_lr, 

198 "lr_warmup_steps": self.lr_warmup_steps, 

199 "lr_decay_steps": self.lr_decay_steps, 

200 "lr_decay_style": self.lr_decay_style, 

201 "lr_decay_ratio": self.lr_decay_ratio, 

202 "min_lr": self.min_lr, 

203 "wsd_decay_steps": self.wsd_decay_steps, 

204 "lr_wsd_decay_style": self.lr_wsd_decay_style 

205 }) 

206 

207 return state_dict 

208 

209 def _check_and_set(self, cls_value: Any, sd_value: Any, name: str) -> Any: 

210 """Strong validation logic during checkpoint loading.""" 

211 if self.override_opt_param_scheduler: 

212 return cls_value 

213 if cls_value != sd_value: 

214 logger.warning("Scheduler Config Override: %s changed from %s to %s.", name, sd_value, cls_value) 

215 return cls_value 

216 

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

218 """Load state and immediately refresh the underlying optimizer's learning rate. 

219 

220 Compatible with both: 

221 - OptimizerParamScheduler state dicts (last_epoch, lr_warmup_steps, ...) 

222 - PyTorch LambdaLR state dicts (last_epoch, base_lrs, _step_count, ...) 

223 """ 

224 # Restore num_steps from either format 

225 if 'last_epoch' in state_dict: 

226 self.num_steps = state_dict['last_epoch'] 

227 elif 'num_steps' in state_dict: 

228 self.num_steps = state_dict['num_steps'] 

229 else: 

230 self.num_steps = 0 

231 

232 # Restore scheduler config if present (OptimizerParamScheduler format only; 

233 # LambdaLR checkpoints lack these keys and will be silently skipped). 

234 config_keys = [ 

235 ('lr_warmup_steps', 'warmup_steps'), 

236 ('lr_decay_steps', 'decay_steps'), 

237 ('lr_decay_style', 'decay_style'), 

238 ('lr_decay_ratio', 'decay_ratio'), 

239 ('min_lr', 'min_lr'), 

240 ('init_lr', 'initial_lr'), 

241 ('wsd_decay_steps', 'wsd_decay_steps'), 

242 ('lr_wsd_decay_style', 'lr_wsd_decay_style'), 

243 ] 

244 for attr, log_name in config_keys: 

245 if attr in state_dict: 

246 setattr(self, attr, self._check_and_set(getattr(self, attr), state_dict[attr], log_name)) 

247 

248 # Recompute LR for current step without advancing the counter 

249 self.step(0) 

250 

251 

252class LRSchedulersContainer: 

253 """Container for multiple learning rate schedulers. 

254 

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

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

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

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

259 """ 

260 

261 def __init__(self, optimizers: ChainedOptimizer, scheduler_kwargs: Dict[str, Any]) -> None: 

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

263 self._schedulers_by_name: Dict[str, OptimizerParamScheduler] = {} 

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

265 self._schedulers_by_name[name] = OptimizerParamScheduler( 

266 optimizer=opt, **scheduler_kwargs, 

267 ) 

268 self.schedulers: List[OptimizerParamScheduler] = [ 

269 self._schedulers_by_name[name] for name in self._names 

270 ] 

271 

272 def __iter__(self) -> Iterator[OptimizerParamScheduler]: 

273 """Iterate over the registered schedulers.""" 

274 return iter(self.schedulers) 

275 

276 def __len__(self) -> int: 

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

278 return len(self.schedulers) 

279 

280 def step(self) -> None: 

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

282 for scheduler in self.schedulers: 

283 scheduler.step() 

284 

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

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

287 param_last_lr: List[float] = [] 

288 for scheduler in self.schedulers: 

289 param_last_lr.extend(scheduler.get_last_lr()) 

290 return param_last_lr 

291 

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

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

294 

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

296 """ 

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

298 

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

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

301 

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

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

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

305 """ 

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

307 raise RuntimeError( 

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

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

310 ) 

311 for name in self._names: 

312 if name not in state_dict: 

313 raise RuntimeError( 

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

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

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

317 ) 

318 self._schedulers_by_name[name].load_state_dict( 

319 copy.deepcopy(state_dict[name]), 

320 ) 

321 

322 

323def get_hyper_lr_scheduler( 

324 optimizer: ChainedOptimizer, 

325 total_steps: int, 

326 warmup_steps: int = 0, 

327 warmup_ratio: float = 0.0, 

328 decay_style: str = "cosine", 

329 lr: float = 1e-4, 

330 lr_min: float = 1e-7, 

331 lr_start: float = 0.0, 

332 lr_decay_ratio: float = 1.0, 

333 wsd_decay_steps: Optional[int] = None, 

334 lr_wsd_decay_style: str = "exponential", 

335 override_opt_param_scheduler: bool = False, 

336) -> LRSchedulersContainer: 

337 """Create a learning rate scheduler compatible with HyperParallel optimizers. 

338 

339 Example: 

340 from hyper_parallel.core.optimizer.lr_scheduler import get_hyper_lr_scheduler 

341 

342 lr_scheduler = get_hyper_lr_scheduler( 

343 optimizer=optimizer, 

344 total_steps=train_steps, 

345 warmup_steps=0, 

346 warmup_ratio=lr_warmup_ratio, 

347 decay_style=lr_decay_style, 

348 lr_decay_ratio=lr_decay_ratio, 

349 lr_min=lr_min, 

350 lr=lr, 

351 lr_start=lr_start, 

352 ) 

353 return lr_scheduler 

354 """ 

355 

356 if decay_style not in SUPPORTED_DECAY_STYLES: 

357 raise ValueError( 

358 f"Unknown decay_style '{decay_style}'. " 

359 f"Supported: {sorted(SUPPORTED_DECAY_STYLES)}" 

360 ) 

361 

362 if decay_style == "WSD" and lr_wsd_decay_style not in SUPPORTED_WSD_STYLES: 

363 raise ValueError( 

364 f"Unknown lr_wsd_decay_style '{lr_wsd_decay_style}'. " 

365 f"Supported: {sorted(SUPPORTED_WSD_STYLES)}" 

366 ) 

367 

368 # pylint: disable=chained-comparison 

369 if warmup_steps <= 0 and warmup_ratio > 0: 

370 warmup_steps = int(total_steps * warmup_ratio) 

371 

372 scheduler_kwargs = { 

373 "init_lr": lr, 

374 "lr_start": lr_start, 

375 "min_lr": lr_min, 

376 "lr_warmup_steps": warmup_steps, 

377 "lr_decay_steps": total_steps, 

378 "lr_decay_style": decay_style, 

379 "lr_decay_ratio": lr_decay_ratio, 

380 "wsd_decay_steps": wsd_decay_steps, 

381 "lr_wsd_decay_style": lr_wsd_decay_style, 

382 "override_opt_param_scheduler": override_opt_param_scheduler, 

383 } 

384 

385 return LRSchedulersContainer(optimizers=optimizer, scheduler_kwargs=scheduler_kwargs)