Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / aux_loss.py: 98%

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"""Reusable auxiliary-loss functions.""" 

16 

17from __future__ import annotations 

18 

19from contextlib import contextmanager 

20from contextvars import ContextVar 

21from dataclasses import dataclass, field 

22from typing import Any, Iterator 

23 

24import torch # pylint: disable=forbidden-backend-import 

25 

26from hyper_parallel.core.dtensor.dtensor import DTensor, SkipDTensorDispatch 

27 

28 

29@dataclass 

30class _AuxLossScaleState: 

31 """Per-forward scale, shared with its backward and recomputation.""" 

32 

33 scale: torch.Tensor | float = 1.0 

34 active: bool = False 

35 handles: list[Any] = field(default_factory=list) 

36 

37 

38_AUX_LOSS_STATE: ContextVar[_AuxLossScaleState | None] = ContextVar("aux_loss_state", default=None) 

39 

40 

41@contextmanager 

42def aux_loss_scale_context() -> Iterator[None]: 

43 """Isolate auxiliary gradient scaling for one forward/backward microstep. 

44 

45 Call :func:`bind_aux_loss_scale` on the local scalar task loss before 

46 weighting it. Keep this context active through backward so checkpoint 

47 recomputation observes the same scale. Nested contexts are independent. 

48 """ 

49 state = _AuxLossScaleState() 

50 token = _AUX_LOSS_STATE.set(state) 

51 try: 

52 yield 

53 finally: 

54 for handle in state.handles: 

55 handle.remove() 

56 _AUX_LOSS_STATE.reset(token) 

57 

58 

59def bind_aux_loss_scale(loss: torch.Tensor) -> None: 

60 """Make injected losses follow the scalar task loss's backward multiplier. 

61 

62 This is a no-op outside :func:`aux_loss_scale_context` or when no auxiliary 

63 loss was attached. The hook runs before upstream model backward, including 

64 checkpoint recomputation, and never changes the task-loss gradient. 

65 

66 Args: 

67 loss: Unweighted scalar foundation loss connected to the model output. 

68 """ 

69 state = _AUX_LOSS_STATE.get() 

70 if state is None or not state.active or not loss.requires_grad: 

71 return 

72 if loss.numel() != 1: 

73 raise ValueError("Auxiliary gradient scaling requires a scalar task loss") 

74 

75 def save_scale(gradient: torch.Tensor) -> None: 

76 """Capture the actual multiplier after token and parallel weighting. 

77 

78 Args: 

79 gradient: Upstream scalar loss multiplier, possibly a DTensor. 

80 """ 

81 with SkipDTensorDispatch(): 

82 local_gradient = gradient.to_local() if isinstance(gradient, DTensor) else gradient 

83 state.scale = local_gradient.detach() 

84 

85 state.handles.append(loss.register_hook(save_scale)) 

86 

87 

88class _AuxLossAutoScaler(torch.autograd.Function): 

89 """Inject auxiliary-loss gradients without changing the forward loss value.""" 

90 

91 main_loss_backward_scale = torch.tensor(1.0) 

92 

93 @staticmethod 

94 def forward(ctx: Any, output: torch.Tensor, aux_loss: torch.Tensor) -> torch.Tensor: 

95 """Save the auxiliary loss and return the output unchanged. 

96 

97 Args: 

98 ctx: Autograd context for the auxiliary loss and scale state. 

99 output: Forward tensor carrying the injected gradient. 

100 aux_loss: Scalar objective whose gradient is injected in backward. 

101 """ 

102 ctx.save_for_backward(aux_loss) 

103 ctx.scale_state = _AUX_LOSS_STATE.get() 

104 return output 

105 

106 @staticmethod 

107 def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: 

108 """Pass through the output gradient and inject the auxiliary-loss gradient. 

109 

110 Args: 

111 ctx: Context saved by forward. 

112 grad_output: Unchanged gradient returned to the forward tensor. 

113 """ 

114 (aux_loss,) = ctx.saved_tensors 

115 scale = ( 

116 ctx.scale_state.scale if ctx.scale_state is not None 

117 else _AuxLossAutoScaler.main_loss_backward_scale 

118 ) 

119 if isinstance(scale, torch.Tensor): 

120 scale = scale.to(device=aux_loss.device, dtype=aux_loss.dtype) 

121 aux_loss_grad = torch.ones_like(aux_loss) * scale 

122 return grad_output, aux_loss_grad 

123 

124 @staticmethod 

125 def set_loss_scale(scale: torch.Tensor) -> None: 

126 """Set the auxiliary-loss gradient scale.""" 

127 _AuxLossAutoScaler.main_loss_backward_scale = scale 

128 

129 

130def aux_loss_auto_scale(output: torch.Tensor, aux_loss: torch.Tensor) -> torch.Tensor: 

131 """Attach an auxiliary-loss gradient to an unchanged forward tensor. 

132 

133 Args: 

134 output: Main forward output. 

135 aux_loss: Scalar auxiliary loss whose gradient should be injected. 

136 

137 Returns: 

138 ``output`` unchanged in value, with an autograd edge to ``aux_loss``. 

139 """ 

140 state = _AUX_LOSS_STATE.get() 

141 if state is not None: 

142 state.active = True 

143 return _AuxLossAutoScaler.apply(output, aux_loss) 

144 

145 

146def set_aux_loss_scale(scale: torch.Tensor) -> None: 

147 """Set the gradient multiplier used by :func:`aux_loss_auto_scale`. 

148 

149 Args: 

150 scale: Tensor used to scale the injected auxiliary-loss gradient. 

151 """ 

152 _AuxLossAutoScaler.set_loss_scale(scale)