Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / mhc.py: 23%

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 manifold-constrained hyper-connection replacement modules.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Mapping 

20from typing import Any 

21 

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

23import torch.nn.functional as F # pylint: disable=forbidden-backend-import 

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

25 

26from hyper_parallel.components.functional.mhc_post import mhc_post 

27from hyper_parallel.components.functional.mhc_pre import mhc_pre 

28from hyper_parallel.components.functional.sinkhorn import sinkhorn, sinkhorn_knopps 

29from hyper_parallel.models.replacement import module_replacement 

30 

31 

32def _required_attribute(module: nn.Module, name: str) -> Any: 

33 """Return a required source-module attribute.""" 

34 if not hasattr(module, name): 

35 raise TypeError(f"{module.__class__.__name__} is missing required attribute '{name}'") 

36 return getattr(module, name) 

37 

38 

39def _required_config_value(module: nn.Module, name: str) -> Any: 

40 """Return a required value from the source module or its config.""" 

41 if hasattr(module, name): 

42 return getattr(module, name) 

43 config = getattr(module, "config", None) 

44 if config is None or not hasattr(config, name): 

45 raise TypeError(f"{module.__class__.__name__} is missing required MHC setting '{name}'") 

46 return getattr(config, name) 

47 

48 

49def _num_stream(module: nn.Module) -> int: 

50 """Return and validate the source module's residual-stream count.""" 

51 value = getattr(module, "num_stream", None) 

52 if value is None: 

53 value = _required_config_value(module, "mhc_num_stream") 

54 if not isinstance(value, int) or value <= 0: 

55 raise ValueError(f"MHC num_stream must be a positive integer, but got {value}") 

56 return value 

57 

58 

59def _parameter(module: nn.Module, name: str) -> nn.Parameter: 

60 """Return a required source parameter.""" 

61 value = _required_attribute(module, name) 

62 if not isinstance(value, nn.Parameter): 

63 raise TypeError(f"MHC attribute '{name}' must be an nn.Parameter") 

64 return value 

65 

66 

67def _validate_parameter_layout(phi: nn.Linear, parameters: tuple[nn.Parameter, ...]) -> None: 

68 """Validate that MHC parameters share the projection's device and dtype.""" 

69 for parameter in parameters: 

70 if parameter.device != phi.weight.device or parameter.dtype != phi.weight.dtype: 

71 raise ValueError("MHC projection and branch parameters must share device and dtype") 

72 

73 

74@module_replacement 

75class MhcPreModule(nn.Module): 

76 """Prepare hidden states and mixing coefficients for an MHC-wrapped block.""" 

77 

78 def __init__( 

79 self, 

80 *, 

81 module: nn.Module, 

82 module_fqn: str = "", 

83 context: Mapping[str, Any] | None = None, 

84 ) -> None: 

85 """Build the high-performance MHC pre module from an existing module. 

86 

87 Args: 

88 module: Source MHC pre module with the existing parameter layout. 

89 module_fqn: Fully qualified source-module name supplied by replacement. 

90 context: Replacement context supplied by Trainer. 

91 

92 Raises: 

93 TypeError: If required source attributes or settings are missing. 

94 ValueError: If parameter shapes or runtime settings are incompatible. 

95 """ 

96 super().__init__() 

97 del module_fqn, context 

98 self.config = getattr(module, "config", None) 

99 self.num_stream = _num_stream(module) 

100 self.layer_number = getattr(module, "layer_number", 1) 

101 

102 phi = _required_attribute(module, "phi") 

103 if not isinstance(phi, nn.Linear) or phi.bias is not None: 

104 raise TypeError("MhcPreModule requires a bias-free nn.Linear phi projection") 

105 if phi.in_features % self.num_stream != 0: 

106 raise ValueError("MHC pre phi input size must be divisible by num_stream") 

107 if phi.out_features != (self.num_stream + 2) * self.num_stream: 

108 raise ValueError("MHC pre phi output size must equal (num_stream + 2) * num_stream") 

109 self.phi = phi 

110 self.branch_alpha = _parameter(module, "branch_alpha") 

111 self.branch_beta = _parameter(module, "branch_beta") 

112 if self.branch_alpha.numel() != 3: 

113 raise ValueError("MHC pre branch_alpha must contain three values") 

114 expected_beta_size = 2 * self.num_stream + self.num_stream * self.num_stream 

115 if self.branch_beta.numel() != expected_beta_size: 

116 raise ValueError("MHC pre branch_beta has an incompatible size") 

117 

118 self.mhc_use_gamma = bool(_required_config_value(module, "mhc_use_gamma")) 

119 parameters = [self.branch_alpha, self.branch_beta] 

120 if self.mhc_use_gamma: 

121 self.norm_gamma = _parameter(module, "norm_gamma") 

122 if self.norm_gamma.numel() != phi.in_features: 

123 raise ValueError("MHC pre norm_gamma size must equal the phi input size") 

124 parameters.append(self.norm_gamma) 

125 elif hasattr(module, "norm_gamma"): 

126 raise ValueError("MHC pre source has norm_gamma while mhc_use_gamma is disabled") 

127 _validate_parameter_layout(phi, tuple(parameters)) 

128 

129 self.hc_eps = float(_required_attribute(module, "hc_eps")) 

130 self.norm_eps = float(_required_attribute(module, "norm_eps")) 

131 self.mhc_recur_norm = int(_required_config_value(module, "mhc_recur_norm")) 

132 if self.hc_eps <= 0 or self.norm_eps <= 0: 

133 raise ValueError("MHC eps values must be positive") 

134 if self.mhc_recur_norm <= 0: 

135 raise ValueError("mhc_recur_norm must be a positive integer") 

136 self.train(module.training) 

137 

138 def forward( 

139 self, x: torch.Tensor 

140 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: 

141 """Prepare hidden states and return the MHC mixing coefficients.""" 

142 gamma = self.norm_gamma if self.mhc_use_gamma else None 

143 return mhc_pre( 

144 x, 

145 self.phi.weight, 

146 self.branch_alpha, 

147 self.branch_beta, 

148 self.num_stream, 

149 self.mhc_recur_norm, 

150 self.norm_eps, 

151 self.hc_eps, 

152 gamma, 

153 ) 

154 

155 

156@module_replacement 

157class MhcPostModule(nn.Module): 

158 """Combine transformed and residual streams after an MHC-wrapped block.""" 

159 

160 def __init__( 

161 self, 

162 *, 

163 module: nn.Module, 

164 module_fqn: str = "", 

165 context: Mapping[str, Any] | None = None, 

166 ) -> None: 

167 """Build the high-performance MHC post module from an existing module. 

168 

169 Args: 

170 module: Source MHC post module. 

171 module_fqn: Fully qualified source-module name supplied by replacement. 

172 context: Replacement context supplied by Trainer. 

173 """ 

174 super().__init__() 

175 del module_fqn, context 

176 self.config = getattr(module, "config", None) 

177 self.num_stream = _num_stream(module) 

178 self.train(module.training) 

179 

180 def forward( 

181 self, 

182 x: torch.Tensor, 

183 residual: torch.Tensor, 

184 h_post: torch.Tensor, 

185 h_res: torch.Tensor, 

186 ) -> torch.Tensor: 

187 """Mix the wrapped block output into the residual streams.""" 

188 return mhc_post( 

189 x, 

190 residual, 

191 h_post, 

192 h_res, 

193 self.num_stream, 

194 ) 

195 

196 

197@module_replacement 

198class PipelinedMhcModule(nn.Module): 

199 """Coefficient module for cross-sublayer pipelined mHC.""" 

200 

201 def __init__( 

202 self, 

203 *, 

204 module: nn.Module, 

205 module_fqn: str = "", 

206 context: Mapping[str, Any] | None = None, 

207 ) -> None: 

208 """Reuse the source module's ``fn/base/scale`` parameter layout.""" 

209 super().__init__() 

210 del module_fqn, context 

211 required = ("input_norm", "fn", "base", "scale", "hc_mult") 

212 missing = [name for name in required if not hasattr(module, name)] 

213 if missing: 

214 raise TypeError(f"pipelined MHC source is missing required attributes: {missing}") 

215 self.input_norm = module.input_norm 

216 self.fn = module.fn 

217 self.base = module.base 

218 self.scale = module.scale 

219 self.hc_mult = int(module.hc_mult) 

220 self.hc_sinkhorn_iters = int(module.hc_sinkhorn_iters) 

221 self.hc_eps = float(module.hc_eps) 

222 expected_mix = (self.hc_mult + 2) * self.hc_mult 

223 if tuple(self.fn.shape)[0] != expected_mix or self.base.numel() != expected_mix: 

224 raise ValueError("pipelined MHC fn/base shapes do not match hc_mult") 

225 if self.scale.numel() != 3: 

226 raise ValueError("pipelined MHC scale must contain three values") 

227 self.train(module.training) 

228 

229 def forward( 

230 self, 

231 hidden_streams: torch.Tensor, 

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

233 """Produce the coefficients consumed by the following sublayer.""" 

234 flattened = self.input_norm(hidden_streams.flatten(start_dim=2).float()) 

235 mix = F.linear( # pylint: disable=not-callable 

236 flattened, self.fn.float() 

237 ) 

238 num_stream = self.hc_mult 

239 pre, post, residual = mix.split( 

240 [num_stream, num_stream, num_stream * num_stream], dim=-1 

241 ) 

242 pre_bias, post_bias, residual_bias = self.base.split( 

243 [num_stream, num_stream, num_stream * num_stream] 

244 ) 

245 pre_scale, post_scale, residual_scale = self.scale.unbind(0) 

246 pre = torch.sigmoid(pre * pre_scale + pre_bias) + self.hc_eps 

247 post = 2 * torch.sigmoid(post * post_scale + post_bias) 

248 residual = residual.view(*residual.shape[:-1], num_stream, num_stream) 

249 residual = residual * residual_scale + residual_bias.view(num_stream, num_stream) 

250 if hidden_streams.device.type == "npu" and hasattr(torch.ops.custom, "npu_sinkhorn"): 

251 residual = sinkhorn(residual, self.hc_sinkhorn_iters, self.hc_eps) 

252 else: 

253 residual = sinkhorn_knopps(residual, self.hc_sinkhorn_iters, self.hc_eps) 

254 return pre, post, residual 

255 

256 

257def pipelined_mhc_post( 

258 sublayer_output: torch.Tensor, 

259 residual: torch.Tensor, 

260 post: torch.Tensor, 

261 combine: torch.Tensor, 

262) -> torch.Tensor: 

263 """Apply the shared high-performance mHC post path to 4-D streams.""" 

264 num_stream = residual.shape[-2] 

265 flattened = mhc_post( 

266 sublayer_output, 

267 residual.flatten(start_dim=2), 

268 post, 

269 combine, 

270 num_stream, 

271 ) 

272 return flattened.unflatten(-1, (num_stream, residual.shape[-1]))