Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / activation_situ.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"""SiTU eager reference and first-order autograd over CANN ops-nn.""" 

16 

17import importlib 

18from collections.abc import Callable 

19from typing import Any, Optional 

20 

21# Torch is the native framework; the legacy platform-agnostic lint rule is obsolete here. 

22# pylint: disable=forbidden-backend-import 

23import torch 

24from torch.autograd.function import once_differentiable 

25 

26 

27def _situ_eager(gate: torch.Tensor, up: torch.Tensor, beta: float, 

28 linear_beta: Optional[float]) -> torch.Tensor: 

29 gate_float, up_float = gate.float(), up.float() 

30 activated_gate = beta * torch.tanh(gate_float / beta) * torch.sigmoid(gate_float) 

31 if linear_beta is not None and linear_beta > 0: 

32 up_float = linear_beta * torch.tanh(up_float / linear_beta) 

33 return (activated_gate * up_float).to(gate.dtype) 

34 

35 

36def _load_situ_ops() -> tuple[Callable, Callable]: 

37 """Resolve the optional ops-nn training APIs without hiding native errors.""" 

38 try: 

39 # CANN is optional for CPU and eager-only callers. 

40 ops = importlib.import_module("cann_ops_nn") 

41 except (ImportError, OSError) as error: 

42 raise ImportError( 

43 "SiTU backend='ascendc' could not load cann_ops_nn. Install the standalone ops-nn Python package " 

44 "with compatible PyTorch, torch-npu and CANN libraries. " 

45 f"See hyper_parallel/components/functional/README.md. Original error: {error}" 

46 ) from error 

47 situ_glu = getattr(ops, "situ_glu", None) 

48 situ_glu_grad = getattr(ops, "situ_glu_grad", None) 

49 if not callable(situ_glu) or not callable(situ_glu_grad): 

50 raise ImportError( 

51 "SiTU backend='ascendc' requires cann_ops_nn.situ_glu and cann_ops_nn.situ_glu_grad. " 

52 "Install a package containing both training APIs and build both native operators. " 

53 "See hyper_parallel/components/functional/README.md." 

54 ) 

55 return situ_glu, situ_glu_grad 

56 

57 

58class _SituAutograd(torch.autograd.Function): 

59 """Connect ops-nn's separate SituGlu and SituGluGrad entry points.""" 

60 

61 @staticmethod 

62 def forward(ctx: Any, packed: torch.Tensor, beta: float, linear_beta: float) -> torch.Tensor: 

63 """Run dense gate-first rows and retain the native backward inputs. 

64 

65 Args: 

66 ctx: Autograd context holding inputs for backward. 

67 packed: Input tensor with gate channels followed by up channels. 

68 beta: Gate saturation scale. 

69 linear_beta: Up saturation scale; nonpositive values disable saturation. 

70 

71 Returns: 

72 Activated tensor with half the input's last dimension. 

73 """ 

74 situ_glu, situ_glu_grad = _load_situ_ops() 

75 rows = packed.reshape(-1, packed.shape[-1]).contiguous() 

76 ctx.save_for_backward(rows) 

77 ctx.input_shape = packed.shape 

78 ctx.beta = beta 

79 ctx.linear_beta = linear_beta 

80 ctx.situ_glu_grad = situ_glu_grad 

81 output = situ_glu(rows, dim=-1, beta=beta, linear_beta=linear_beta, activate_left=True) 

82 return output.reshape(*packed.shape[:-1], packed.shape[-1] // 2) 

83 

84 @staticmethod 

85 @once_differentiable 

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

87 """Restore the packed gradient with the same saturation parameters. 

88 

89 Args: 

90 ctx: Autograd context populated by forward. 

91 grad_output: Gradient with respect to the activated output. 

92 

93 Returns: 

94 Packed input gradient and no gradients for the scalar attributes. 

95 """ 

96 (rows,) = ctx.saved_tensors 

97 grad_rows = grad_output.reshape(-1, rows.shape[-1] // 2).contiguous() 

98 gradient = ctx.situ_glu_grad( 

99 grad_rows, rows, dim=-1, beta=ctx.beta, linear_beta=ctx.linear_beta, activate_left=True, 

100 ) 

101 return gradient.reshape(ctx.input_shape), None, None 

102 

103 

104def _situ_fused(packed: torch.Tensor, beta: float, linear_beta: float) -> torch.Tensor: 

105 return _SituAutograd.apply(packed, beta, linear_beta) 

106 

107 

108def situ( 

109 gate: torch.Tensor, 

110 up: torch.Tensor, 

111 beta: float = 1.0, 

112 linear_beta: Optional[float] = None, 

113 backend: str = "eager", 

114) -> torch.Tensor: 

115 """Apply SiTU to separate gate and up projections. 

116 

117 Args: 

118 gate: Gate projection with at least one dimension. 

119 up: Up projection with the same shape, dtype, and device as gate. 

120 beta: Positive gate saturation scale. 

121 linear_beta: Up saturation scale; nonpositive values or None disable saturation. 

122 backend: ``eager`` or ``ascendc`` (the fused ops-nn implementation). 

123 

124 Returns: 

125 Activated values with gate's shape and dtype. 

126 """ 

127 # Packing would lose the original split if gate and up had unequal widths. 

128 if gate.shape != up.shape: 

129 raise ValueError("SiTU gate and up must have identical shapes") 

130 if backend == "eager": 

131 return _situ_eager(gate, up, beta, linear_beta) 

132 return situ_packed(torch.cat((gate, up), dim=-1), beta, linear_beta, backend) 

133 

134 

135def situ_packed( 

136 packed: torch.Tensor, 

137 beta: float = 1.0, 

138 linear_beta: Optional[float] = None, 

139 backend: str = "eager", 

140) -> torch.Tensor: 

141 """Apply SiTU to a tensor with gate channels followed by up channels. 

142 

143 Args: 

144 packed: Tensor of shape [..., 2H]; noncontiguous inputs are supported. 

145 beta: Positive gate saturation scale. 

146 linear_beta: Up saturation scale; nonpositive values or None disable saturation. 

147 backend: ``eager`` uses FP32 intermediates. ``ascendc`` calls ops-nn 

148 on NPU with first-order autograd; native errors propagate unchanged. 

149 

150 Raises: 

151 ImportError: The optional CANN package cannot load or lacks a training API. 

152 

153 Returns: 

154 Tensor of shape [..., H] in the input dtype. Empty inputs keep connected 

155 gradients through tensor operations without launching a native kernel. 

156 """ 

157 if backend not in ("eager", "ascendc"): 

158 raise ValueError("SiTU backend must be 'eager' or 'ascendc'") 

159 linear_beta = 0.0 if linear_beta is None else linear_beta 

160 if backend == "eager" or packed.numel() == 0: 

161 gate, up = packed.unflatten(-1, (2, packed.shape[-1] // 2)).unbind(-2) 

162 return _situ_eager(gate, up, beta, linear_beta) 

163 return _situ_fused(packed, beta, linear_beta)