Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / adamw.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"""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 device = torch.npu.current_device() if torch.npu.is_available() else torch.cuda.current_device() 

42 step_tensor = torch.tensor(step, dtype=torch.int64, device=device) 

43 state_steps = [step_tensor] * len(params) 

44 

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

46 params, 

47 grads, 

48 exp_avgs, 

49 exp_avg_sqs, 

50 max_exp_avg_sqs if amsgrad else [], 

51 state_steps, 

52 amsgrad=amsgrad, 

53 lr=lr, 

54 beta1=beta1, 

55 beta2=beta2, 

56 weight_decay=weight_decay, 

57 eps=eps, 

58 maximize=maximize 

59 ) 

60 

61 

62class AdamW(torch.optim.Optimizer): 

63 """AdamW optimizer implementation.""" 

64 

65 def __init__( 

66 self, 

67 params, 

68 lr=1e-3, 

69 betas=(0.9, 0.999), 

70 eps=1e-8, 

71 weight_decay=0.01, 

72 amsgrad=False, 

73 maximize=False 

74 ): 

75 defaults = { 

76 "lr": lr, 

77 "betas": betas, 

78 "eps": eps, 

79 "weight_decay": weight_decay, 

80 "amsgrad": amsgrad, 

81 "maximize": maximize 

82 } 

83 super().__init__(params, defaults) 

84 

85 def __setstate__(self, state): 

86 """Set optimizer state.""" 

87 super().__setstate__(state) 

88 for group in self.param_groups: 

89 group.setdefault('amsgrad', False) 

90 group.setdefault('maximize', False) 

91 

92 def step(self, closure=None): 

93 """Performs a single optimization step.""" 

94 loss = None 

95 if closure is not None: 

96 with torch.enable_grad(): 

97 loss = closure() 

98 

99 for group in self.param_groups: 

100 params_with_grad = [] 

101 grads = [] 

102 exp_avgs = [] 

103 exp_avg_sqs = [] 

104 max_exp_avg_sqs = [] 

105 

106 amsgrad = group['amsgrad'] 

107 beta1, beta2 = group['betas'] 

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

109 

110 current_rank_params = group['params'] 

111 for p in current_rank_params: 

112 if p.grad is None: 

113 continue 

114 

115 if p.grad.data.is_sparse: 

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

117 

118 state = self.state[p] 

119 

120 if len(state) == 0: 

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

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

123 if amsgrad: 

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

125 

126 params_with_grad.append(p) 

127 grads.append(p.grad) 

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

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

130 

131 if amsgrad: 

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

133 

134 if params_with_grad: 

135 adamw( 

136 params_with_grad, 

137 grads, 

138 exp_avgs, 

139 exp_avg_sqs, 

140 max_exp_avg_sqs, 

141 group['step'], 

142 amsgrad=amsgrad, 

143 beta1=beta1, 

144 beta2=beta2, 

145 lr=group['lr'], 

146 weight_decay=group['weight_decay'], 

147 eps=group['eps'], 

148 maximize=group['maximize'] 

149 ) 

150 

151 return loss