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

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-callable learning-rate scheduler factories.""" 

16 

17from functools import partial 

18from typing import Any 

19 

20from hyper_parallel.core.optimizer.lr_scheduler import ( 

21 LRSchedulersContainer, 

22 get_constant_schedule_with_warmup, 

23 get_cosine_schedule_with_warmup, 

24 get_linear_schedule_with_warmup, 

25) 

26 

27 

28class MultiLRScheduler: 

29 """Build one learning-rate scheduler for each child optimizer.""" 

30 

31 def __init__( 

32 self, 

33 optimizer: Any, 

34 lr_decay_style: str, 

35 train_iters: int, 

36 lr_config: dict, 

37 ) -> None: 

38 """Initialize schedulers for every child optimizer. 

39 

40 Args: 

41 optimizer: Optimizer or collection of optimizers to schedule. 

42 lr_decay_style: One of ``constant``, ``linear`` or ``cosine``. 

43 train_iters: Total number of training iterations. 

44 lr_config: Mapping with warmup/decay settings such as 

45 ``lr_warmup_ratio``, ``lr_warmup_steps``, ``lr_start``, 

46 ``lr_decay_ratio`` and ``min_lr``. 

47 """ 

48 self.config = lr_config 

49 self.optimizer = optimizer 

50 self.lr_decay_style = lr_decay_style 

51 self.train_iters = train_iters 

52 

53 lr_start = self.config.get("lr_start", 0.0) 

54 init_lr = self.config.get("lr", 1e-3) 

55 min_lr = self.config.get("min_lr", self.config.get("lr_min", 1e-7)) 

56 

57 lr_warmup_ratio = self.config.get("lr_warmup_ratio", 0.0) 

58 lr_warmup_steps = self.config.get("lr_warmup_steps") 

59 if lr_warmup_steps is None: 

60 lr_warmup_steps = int(train_iters * lr_warmup_ratio) 

61 lr_decay_ratio = self.config.get("lr_decay_ratio", 1.0) 

62 

63 if self.lr_decay_style == "constant": 

64 lr_scheduler = partial( 

65 get_constant_schedule_with_warmup, 

66 num_warmup_steps=lr_warmup_steps, 

67 init_lr=init_lr, 

68 lr_start=lr_start, 

69 ) 

70 elif self.lr_decay_style == "linear": 

71 lr_scheduler = partial( 

72 get_linear_schedule_with_warmup, 

73 num_warmup_steps=lr_warmup_steps, 

74 num_training_steps=train_iters, 

75 init_lr=init_lr, 

76 min_lr=min_lr, 

77 lr_start=lr_start, 

78 ) 

79 elif self.lr_decay_style == "cosine": 

80 lr_scheduler = partial( 

81 get_cosine_schedule_with_warmup, 

82 num_warmup_steps=lr_warmup_steps, 

83 num_training_steps=train_iters, 

84 init_lr=init_lr, 

85 lr_decay_ratio=lr_decay_ratio, 

86 min_lr=min_lr, 

87 lr_start=lr_start, 

88 ) 

89 else: 

90 raise ValueError( 

91 f"Unsupported lr_decay_style {self.lr_decay_style!r}; expected 'constant', 'linear' or 'cosine'" 

92 ) 

93 

94 self.lr_scheduler = LRSchedulersContainer( 

95 optimizers=self.optimizer, 

96 scheduler=lr_scheduler, 

97 ) 

98 

99 def get_lr_scheduler(self) -> Any: 

100 """Return the scheduler container used by the trainer.""" 

101 return self.lr_scheduler