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

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"""Manifold hyper-connection pre-processing functions.""" 

16 

17from typing import Any, Optional, Tuple 

18 

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

20 

21try: 

22 import omni_training_custom_ops # noqa: F401 # pylint: disable=unused-import 

23except ImportError: 

24 omni_training_custom_ops = None 

25 

26from hyper_parallel.components.functional.sinkhorn import sinkhorn 

27 

28 

29def _mhc_pre_backward(ctx, dh_in, dh_post, dh_res, dh_x): 

30 """Invoke the MHC gradient operator with the saved forward state.""" 

31 x, phi, alpha, gamma, h_post, inv_rms, h_mix, h_pre = ctx.saved_tensors 

32 return torch.ops.custom.npu_manifold_constrained_hyper_connection_pre_grad( 

33 x, phi, alpha, dh_in, dh_post, dh_res, inv_rms, h_mix, h_pre, h_post, 

34 gamma=gamma, 

35 hc_eps=ctx.hc_eps, 

36 grad_x_post=dh_x, 

37 ) 

38 

39 

40class _MhcPre(torch.autograd.Function): 

41 """Autograd bridge for the NPU MHC pre custom operator.""" 

42 

43 @staticmethod 

44 def forward( 

45 ctx: Any, 

46 x: torch.Tensor, 

47 phi: torch.Tensor, 

48 alpha: torch.Tensor, 

49 bias: torch.Tensor, 

50 gamma: Optional[torch.Tensor] = None, 

51 out_flag: bool = True, 

52 norm_eps: float = 1e-6, 

53 hc_eps: float = 1e-6, 

54 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: 

55 """Run the NPU MHC pre forward operator.""" 

56 if omni_training_custom_ops is None: 

57 raise ImportError("MHC AscendC requires omni_training_custom_ops") 

58 hin, h_post, h_comb_before, inv_rms, h_mix, h_pre = ( 

59 torch.ops.custom.npu_manifold_constrained_hyper_connection_pre( 

60 x, 

61 phi, 

62 alpha, 

63 bias, 

64 gamma=gamma, 

65 out_flag=out_flag, 

66 norm_eps=norm_eps, 

67 hc_eps=hc_eps, 

68 ) 

69 ) 

70 ctx.save_for_backward(x, phi, alpha, gamma, h_post, inv_rms, h_mix, h_pre) 

71 ctx.hc_eps = hc_eps 

72 ctx.has_gamma = gamma is not None 

73 return hin, h_pre, h_post, h_comb_before, x 

74 

75 @staticmethod 

76 def backward( 

77 ctx: Any, 

78 dh_in: torch.Tensor, 

79 dh_pre: torch.Tensor, 

80 dh_post: torch.Tensor, 

81 dh_res: torch.Tensor, 

82 dh_x: torch.Tensor, 

83 ) -> tuple: 

84 """Run the NPU MHC pre backward operator.""" 

85 del dh_pre 

86 dx, dphi, dalpha, dbias, dgamma = _mhc_pre_backward(ctx, dh_in, dh_post, dh_res, dh_x) 

87 if not ctx.has_gamma: 

88 dgamma = None 

89 grads = [dx, dphi, dalpha, dbias, dgamma, None, None, None] 

90 return tuple(grads) 

91 

92 

93def mhc_pre( 

94 x: torch.Tensor, 

95 phi: torch.Tensor, 

96 branch_alpha: torch.Tensor, 

97 branch_beta: torch.Tensor, 

98 num_stream: int, 

99 sinkhorn_iters: int = 20, 

100 norm_eps: float = 1e-6, 

101 hc_eps: float = 1e-6, 

102 gamma: Optional[torch.Tensor] = None, 

103) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: 

104 """Prepare an input and its mixing coefficients for an MHC-wrapped block. 

105 

106 Args: 

107 x: Flattened residual streams. 

108 phi: MHC projection weight. 

109 branch_alpha: Learned scale for each coefficient group. 

110 branch_beta: Learned bias for each coefficient group. 

111 num_stream: Number of residual streams. 

112 sinkhorn_iters: Number of Sinkhorn iterations. 

113 norm_eps: RMS normalization epsilon. 

114 hc_eps: Hyper-connection numerical stability epsilon. 

115 gamma: Optional RMS normalization scale. 

116 

117 Returns: 

118 Block input, post coefficients, residual coefficients, and residual buffer. 

119 """ 

120 shape = x.size() 

121 x = x.reshape(shape[0], shape[1], num_stream, -1) 

122 if gamma is not None: 

123 gamma = gamma.reshape(num_stream, -1).float() 

124 y, _, h_post, h_comb_before, residual = _MhcPre.apply( 

125 x, 

126 phi.float(), 

127 branch_alpha.float(), 

128 branch_beta.float(), 

129 gamma, 

130 True, 

131 norm_eps, 

132 hc_eps, 

133 ) 

134 h_res = sinkhorn(h_comb_before, sinkhorn_iters, hc_eps) 

135 residual = residual.reshape(shape[0], shape[1], -1) 

136 return y, h_post, h_res, residual