Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_mhc / module.py: 94%

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"""Model-facing MegaMhc module for shifted BF16 mHC.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Callable 

20from typing import Any 

21 

22import torch 

23 

24from .._precision import preserve_fp32_parameters 

25from ..mega_mhc_common import ( 

26 MHC_MAPPINGS, 

27 MHC_SINKHORN_ITERS, 

28 MHC_STREAMS, 

29 MegaMhcOutputs, 

30 validate_mega_mhc_hidden_size, 

31) 

32from .function import hyper_mega_mhc 

33from .graph import DEFAULT_TOKEN_TILE 

34 

35 

36class MegaMhc(torch.nn.Module): 

37 """Own the current layer's mapping predictor and input RMSNorm weight. 

38 

39 Note: 

40 This trainable module requires the native backward hidden-size contract. 

41 Forward-only callers can use ``hyper_mega_mhc(..., need_backward=False)``. 

42 """ 

43 

44 def __init__( 

45 self, 

46 hidden_size: int, 

47 *, 

48 hc_eps: float = 1e-6, 

49 norm_eps: float = 1e-6, 

50 num_iters: int = MHC_SINKHORN_ITERS, 

51 token_tile: int = DEFAULT_TOKEN_TILE, 

52 device: Any | None = None, 

53 ) -> None: 

54 """Initialize the fixed-four-stream shifted mHC parameters. 

55 

56 Args: 

57 hidden_size: Stream width H, a positive multiple of 128, at most 5760. 

58 The module enforces this training limit even in eval/no-grad mode. 

59 hc_eps: Epsilon used by pre-mix and Sinkhorn mapping prediction. 

60 norm_eps: Epsilon used by RMS normalization. 

61 num_iters: Sinkhorn iteration count; must be 20. 

62 token_tile: Positive integer tokens per scheduler task; defaults to 32. 

63 device: Parameter device. None uses Torch's default device; execution 

64 requires moving all parameters and inputs to the same NPU. 

65 

66 Note: 

67 phi [24, 4*H], alpha [3] and bias [24] use FP32; norm_weight [H] 

68 uses BF16. Module dtype conversions preserve the FP32 predictor 

69 parameters and their existing gradients. Runtime tensors require NPU 

70 and the native payload. 

71 

72 Raises: 

73 ValueError: If hidden_size, num_iters or token_tile is unsupported. 

74 """ 

75 super().__init__() 

76 validate_mega_mhc_hidden_size(hidden_size, need_backward=True) 

77 if num_iters != MHC_SINKHORN_ITERS: 

78 raise ValueError(f"MegaMhc requires num_iters={MHC_SINKHORN_ITERS}, got {num_iters}.") 

79 if not isinstance(token_tile, int) or isinstance(token_tile, bool) or token_tile <= 0: 

80 raise ValueError(f"token_tile must be a positive integer, got {token_tile!r}.") 

81 self.hidden_size = hidden_size 

82 self.hc_eps = float(hc_eps) 

83 self.norm_eps = float(norm_eps) 

84 self.num_iters = num_iters 

85 self.token_tile = token_tile 

86 factory_kwargs = {"device": device, "dtype": torch.float32} 

87 self.phi = torch.nn.Parameter( 

88 torch.empty((MHC_MAPPINGS, MHC_STREAMS * hidden_size), **factory_kwargs) 

89 ) 

90 self.alpha = torch.nn.Parameter(torch.ones((3,), **factory_kwargs)) 

91 self.bias = torch.nn.Parameter(torch.zeros((MHC_MAPPINGS,), **factory_kwargs)) 

92 self.norm_weight = torch.nn.Parameter( 

93 torch.ones((hidden_size,), device=device, dtype=torch.bfloat16) 

94 ) 

95 torch.nn.init.normal_( 

96 self.phi, 

97 mean=0.0, 

98 std=(MHC_STREAMS * hidden_size) ** -0.5, 

99 ) 

100 

101 def _apply(self, fn: Callable, recurse: bool = True) -> MegaMhc: 

102 """Keep mapping predictors FP32 when the enclosing model changes dtype or device.""" 

103 return super()._apply(preserve_fp32_parameters(self, fn, ("phi", "alpha", "bias")), recurse=recurse) 

104 

105 @classmethod 

106 def from_module(cls, module: torch.nn.Module) -> MegaMhc: 

107 """Bind an existing four-stream, shifted mHC module without replacing its Parameters. 

108 

109 Args: 

110 module: Source with hidden_size, hc_eps, norm_eps, num_iters and Parameters 

111 phi [24, 4H], alpha [3], bias [24], norm_weight [H]. Predictor Parameters must be FP32. 

112 

113 Returns: 

114 Boundary sharing all original Parameters and mathematical configuration. 

115 """ 

116 configuration = {} 

117 for name in ("hidden_size", "hc_eps", "norm_eps", "num_iters"): 

118 if not hasattr(module, name): 

119 raise ValueError(f"source mHC module must expose {name}") 

120 configuration[name] = getattr(module, name) 

121 phi = getattr(module, "phi", None) 

122 if not isinstance(phi, torch.nn.Parameter): 

123 raise TypeError("source phi must be an FP32 Parameter [24, 4H]") 

124 result = cls(**configuration, device=phi.device) 

125 for name in ("phi", "alpha", "bias", "norm_weight"): 

126 parameter = getattr(module, name, None) 

127 expected = getattr(result, name) 

128 if (not isinstance(parameter, torch.nn.Parameter) or parameter.shape != expected.shape 

129 or parameter.dtype != expected.dtype or parameter.device != expected.device): 

130 raise ValueError(f"source {name} must match shape={tuple(expected.shape)}, dtype={expected.dtype}") 

131 setattr(result, name, parameter) 

132 return result 

133 

134 def forward( 

135 self, 

136 previous_output: torch.Tensor, 

137 residual: torch.Tensor, 

138 previous_pre_mix: torch.Tensor, 

139 previous_post_mix: torch.Tensor, 

140 previous_residual_mix: torch.Tensor, 

141 ) -> MegaMhcOutputs: 

142 """Advance one shifted mHC block boundary and predict the next mappings. 

143 

144 Args: 

145 previous_output: BF16 previous block output with shape [*S, H]. 

146 residual: BF16 residual streams with shape [*S, 4, H]. S contains 

147 one or more token dimensions, all inputs share the same NPU. 

148 previous_pre_mix: FP32 coefficients [*S, 4] used to form block_input. 

149 previous_post_mix: FP32 coefficients [*S, 4] that distribute 

150 previous_output into the residual streams. 

151 previous_residual_mix: FP32 residual mixing matrices [*S, 4, 4]. 

152 

153 Returns: 

154 A five-tensor tuple, in order: new_residual (BF16, [*S, 4, H]), 

155 next_pre_mix (FP32, [*S, 4]), next_post_mix (FP32, [*S, 4]), 

156 next_residual_mix (FP32, [*S, 4, 4]), block_input (BF16, [*S, H]). 

157 The new mappings are for the next boundary; block_input uses 

158 previous_pre_mix to combine new_residual, then applies RMSNorm. 

159 

160 Note: 

161 Training uses native autograd. The flattened token count must be 

162 at least max(40, configured AIV core count) when backward is needed; 

163 forward-only execution requires at least one token. Cross-stream 

164 producers must establish dependencies before this call. 

165 

166 Raises: 

167 TypeError: If tensors, dtypes or the NPU device type are invalid. 

168 ValueError: If shapes, devices, hidden size or training token count 

169 violate the native contract. 

170 """ 

171 return hyper_mega_mhc( 

172 previous_output, 

173 residual, 

174 previous_pre_mix, 

175 previous_post_mix, 

176 previous_residual_mix, 

177 self.phi, 

178 self.alpha, 

179 self.bias, 

180 self.norm_weight, 

181 hc_eps=self.hc_eps, 

182 norm_eps=self.norm_eps, 

183 num_iters=self.num_iters, 

184 token_tile=self.token_tile, 

185 )