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

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"""Adamw optimizer.""" 

16 

17from typing import List 

18 

19import torch 

20 

21 

22def adamw( 

23 params: List[torch.Tensor], 

24 grads: List[torch.Tensor], 

25 exp_avgs: List[torch.Tensor], 

26 exp_avg_sqs: List[torch.Tensor], 

27 max_exp_avg_sqs: List[torch.Tensor], 

28 step: int, 

29 *, 

30 amsgrad: bool, 

31 beta1: float, 

32 beta2: float, 

33 lr: float, 

34 weight_decay: float, 

35 eps: float, 

36 maximize: bool 

37) -> None: 

38 r"""Functional API that performs AdamW algorithm computation. 

39 See :class:`~torch.optim.AdamW` for details. 

40 """ 

41 step_tensor = torch.tensor(step, dtype=torch.int64, device=params[0].device) 

42 state_steps = [step_tensor] * len(params) 

43 

44 torch._fused_adamw_( # pylint: disable=protected-access 

45 params, 

46 grads, 

47 exp_avgs, 

48 exp_avg_sqs, 

49 max_exp_avg_sqs if amsgrad else [], 

50 state_steps, 

51 amsgrad=amsgrad, 

52 lr=lr, 

53 beta1=beta1, 

54 beta2=beta2, 

55 weight_decay=weight_decay, 

56 eps=eps, 

57 maximize=maximize 

58 ) 

59 

60 

61class AdamW(torch.optim.Optimizer): 

62 """AdamW optimizer implementation.""" 

63 

64 def __init__( 

65 self, 

66 params, 

67 lr=1e-3, 

68 betas=(0.9, 0.999), 

69 eps=1e-8, 

70 weight_decay=0.01, 

71 amsgrad=False, 

72 maximize=False 

73 ): 

74 defaults = { 

75 "lr": lr, 

76 "betas": betas, 

77 "eps": eps, 

78 "weight_decay": weight_decay, 

79 "amsgrad": amsgrad, 

80 "maximize": maximize 

81 } 

82 super().__init__(params, defaults) 

83 

84 def __setstate__(self, state): 

85 """Set optimizer state.""" 

86 super().__setstate__(state) 

87 for group in self.param_groups: 

88 group.setdefault('amsgrad', False) 

89 group.setdefault('maximize', False) 

90 

91 def __str__(self): 

92 return super().__repr__() 

93 

94 __repr__ = __str__ 

95 

96 def step(self, closure=None): 

97 """Performs a single optimization step.""" 

98 loss = None 

99 if closure is not None: 

100 with torch.enable_grad(): 

101 loss = closure() 

102 

103 for group in self.param_groups: 

104 params_with_grad = [] 

105 grads = [] 

106 exp_avgs = [] 

107 exp_avg_sqs = [] 

108 max_exp_avg_sqs = [] 

109 

110 amsgrad = group['amsgrad'] 

111 beta1, beta2 = group['betas'] 

112 group['step'] = (group.get('step') or 0) + 1 

113 

114 current_rank_params = group['params'] 

115 for p in current_rank_params: 

116 if p.grad is None: 

117 continue 

118 

119 if p.grad.data.is_sparse: 

120 raise RuntimeError('AdamW does not support sparse gradients') 

121 

122 state = self.state[p] 

123 

124 if len(state) == 0: 

125 state['exp_avg'] = torch.zeros_like(p.grad, memory_format=torch.preserve_format) 

126 state['exp_avg_sq'] = torch.zeros_like(p.grad, memory_format=torch.preserve_format) 

127 if amsgrad: 

128 state['max_exp_avg_sq'] = torch.zeros_like(p.grad, memory_format=torch.preserve_format) 

129 

130 params_with_grad.append(p) 

131 grads.append(p.grad) 

132 exp_avgs.append(state['exp_avg']) 

133 exp_avg_sqs.append(state['exp_avg_sq']) 

134 

135 if amsgrad: 

136 max_exp_avg_sqs.append(state['max_exp_avg_sq']) 

137 

138 if params_with_grad: 

139 adamw( 

140 params_with_grad, 

141 grads, 

142 exp_avgs, 

143 exp_avg_sqs, 

144 max_exp_avg_sqs, 

145 group['step'], 

146 amsgrad=amsgrad, 

147 beta1=beta1, 

148 beta2=beta2, 

149 lr=group['lr'], 

150 weight_decay=group['weight_decay'], 

151 eps=group['eps'], 

152 maximize=group['maximize'] 

153 ) 

154 

155 return loss