Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / shared_expert.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"""MoE shared expert module.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Mapping 

20from functools import partial 

21from typing import Any, Optional 

22 

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

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

25from transformers.activations import ACT2FN 

26from hyper_parallel.components.checkpoint.weight_conversion import ( 

27 WeightConverter, 

28 WeightRenaming, 

29) 

30 

31from hyper_parallel.components.checkpoint import ConcatenateWithSections 

32from hyper_parallel.models.replacement import module_replacement 

33from hyper_parallel.components.functional import swiglu 

34 

35 

36class LinearWithMatmul(nn.Linear): 

37 """Aligned with ColumnParallelLinear (single device, no parallelism). 

38 

39 name: ColumnParallelLinear 

40 forward path: LinearWithGradAccumulationAndAsyncCommunication.forward 

41 → output = torch.matmul(input, weight.t()) 

42 → output = output + bias (bias explicitly added separately, not addmm fused) 

43 

44 Used for: linear_fc1, linear_fc2 

45 

46 skip_bias_add=True: don't add bias, return (output, bias) for external addition. 

47 skip_bias_add=False: bias explicitly added in forward, return (output, None). 

48 

49 NPU bf16 note: matmul runs in @custom_fwd autograd Function context, 

50 NPU automatically selects the same kernel for non-contiguous weight.t() as weight.t().contiguous(). 

51 calling directly in forward does not trigger this behavior, so explicit .contiguous() is needed. 

52 """ 

53 

54 def __init__( 

55 self, 

56 in_features: int, 

57 out_features: int, 

58 bias: bool = True, 

59 skip_bias_add: bool = True, 

60 device: torch.device | str | None = None, 

61 dtype: torch.dtype | None = None, 

62 ) -> None: 

63 """Initialize the linear parameters and bias-return behavior.""" 

64 super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) 

65 self.skip_bias_add = skip_bias_add 

66 

67 def forward( # pylint: disable=redefined-builtin 

68 self, input: torch.Tensor 

69 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: 

70 """Apply matmul and optionally return bias for external addition.""" 

71 output = torch.matmul(input, self.weight.t().contiguous()) 

72 if self.skip_bias_add: 

73 return output, self.bias 

74 if self.bias is not None: 

75 output = output + self.bias 

76 return output, None 

77 

78 

79@module_replacement 

80class SharedExpert(nn.Module): 

81 """Transformers-compatible parameter-owning MoE shared expert. 

82 

83 Args: 

84 module: Source Transformers shared-expert module. 

85 module_fqn: Fully qualified name supplied by the replacement framework. 

86 context: Additional replacement context. 

87 """ 

88 

89 @staticmethod 

90 def _separate_projections(module: nn.Module) -> tuple[nn.Linear, nn.Linear, nn.Linear]: 

91 """Validate and return separate source projections.""" 

92 source_linears = (module.gate_proj, module.up_proj, module.down_proj) 

93 if any(not isinstance(layer, nn.Linear) for layer in source_linears): 

94 raise TypeError("SharedExpert source projections must be nn.Linear instances") 

95 gate_proj, up_proj, down_proj = source_linears 

96 if gate_proj.in_features != up_proj.in_features: 

97 raise ValueError("SharedExpert gate and up projections must have the same input size") 

98 if gate_proj.out_features != up_proj.out_features: 

99 raise ValueError("SharedExpert gate and up projections must have the same output size") 

100 if down_proj.in_features != gate_proj.out_features: 

101 raise ValueError("SharedExpert down projection input size is incompatible with gate/up") 

102 if down_proj.out_features != gate_proj.in_features: 

103 raise ValueError("SharedExpert down projection output size is incompatible with gate/up") 

104 if gate_proj.weight.requires_grad != up_proj.weight.requires_grad: 

105 raise ValueError("SharedExpert gate and up weights must use the same training policy") 

106 if (gate_proj.bias is None) != (up_proj.bias is None): 

107 raise ValueError("SharedExpert gate and up projections must use the same bias policy") 

108 if gate_proj.bias is not None and gate_proj.bias.requires_grad != up_proj.bias.requires_grad: 

109 raise ValueError("SharedExpert gate and up biases must use the same training policy") 

110 return gate_proj, up_proj, down_proj 

111 

112 def _initialize_fused_source(self, module: nn.Module) -> None: 

113 """Reuse an already fused source module without converting parameters.""" 

114 self.linear_fc1 = module.linear_fc1 

115 self.linear_fc2 = module.linear_fc2 

116 if not all(hasattr(layer, "skip_bias_add") for layer in (self.linear_fc1, self.linear_fc2)): 

117 raise TypeError("fused SharedExpert projections must return output and bias separately") 

118 if not hasattr(module, "activation_func"): 

119 raise TypeError("fused SharedExpert source module is missing activation_func") 

120 self.activation_func = module.activation_func 

121 self._return_bias_tuple = True 

122 self._gate_size = 0 

123 self._up_size = 0 

124 self.train(module.training) 

125 

126 def _initialize_separate_linears( 

127 self, 

128 gate_proj: nn.Linear, 

129 up_proj: nn.Linear, 

130 down_proj: nn.Linear, 

131 ) -> None: 

132 """Construct fused linear layers from separate source projections.""" 

133 self._gate_size = gate_proj.out_features 

134 self._up_size = up_proj.out_features 

135 self.linear_fc1 = LinearWithMatmul( 

136 gate_proj.in_features, 

137 self._gate_size + self._up_size, 

138 bias=gate_proj.bias is not None, 

139 device=gate_proj.weight.device, 

140 dtype=gate_proj.weight.dtype, 

141 ) 

142 self.linear_fc1.weight.requires_grad_(gate_proj.weight.requires_grad) 

143 if gate_proj.bias is not None: 

144 self.linear_fc1.bias.requires_grad_(gate_proj.bias.requires_grad) 

145 self.linear_fc2 = LinearWithMatmul( 

146 down_proj.in_features, 

147 down_proj.out_features, 

148 bias=down_proj.bias is not None, 

149 device=down_proj.weight.device, 

150 dtype=down_proj.weight.dtype, 

151 ) 

152 self.linear_fc2.weight.requires_grad_(down_proj.weight.requires_grad) 

153 if down_proj.bias is not None: 

154 self.linear_fc2.bias.requires_grad_(down_proj.bias.requires_grad) 

155 

156 def _initialize_activation(self, module: nn.Module, config: Any) -> None: 

157 """Select the source-compatible gated activation implementation.""" 

158 hidden_act = getattr(config, "hidden_act", None) 

159 if hidden_act is None: 

160 hidden_act = getattr(config, "hidden_activation", None) 

161 if hidden_act is None: 

162 hidden_act = getattr(config, "mlp_hidden_act", None) 

163 activation_func = getattr(module, "act_fn", None) 

164 if activation_func is None: 

165 if hidden_act is None: 

166 hidden_act = "silu" 

167 activation_func = ACT2FN[hidden_act] 

168 if hidden_act == "silu" and bool(getattr(config, "use_fused_swiglu", True)): 

169 self.activation_func = partial(swiglu, dim=-1) 

170 return 

171 

172 def glu(x: torch.Tensor) -> torch.Tensor: 

173 """Apply the configured gated linear unit.""" 

174 gate, value = torch.chunk(x, 2, dim=-1) 

175 return activation_func(gate) * value 

176 

177 self.activation_func = glu 

178 

179 def __init__( 

180 self, 

181 *, 

182 module: nn.Module, 

183 module_fqn: str = "", 

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

185 ) -> None: 

186 """Build the high-performance module from a Transformers shared expert.""" 

187 super().__init__() 

188 del module_fqn, context 

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

190 self.config = config 

191 self._source_is_fused = all( 

192 hasattr(module, name) for name in ("linear_fc1", "linear_fc2") 

193 ) 

194 self._source_is_separate = all( 

195 hasattr(module, name) for name in ("gate_proj", "up_proj", "down_proj") 

196 ) 

197 if not self._source_is_fused and not self._source_is_separate: 

198 raise TypeError( 

199 "SharedExpert requires linear_fc1/linear_fc2 or gate_proj/up_proj/down_proj" 

200 ) 

201 

202 if self._source_is_fused: 

203 self._initialize_fused_source(module) 

204 return 

205 

206 gate_proj, up_proj, down_proj = self._separate_projections(module) 

207 self._initialize_separate_linears(gate_proj, up_proj, down_proj) 

208 self._initialize_activation(module, config) 

209 self._return_bias_tuple = False 

210 self.train(module.training) 

211 

212 def make_transforms(self) -> list[WeightRenaming | WeightConverter]: 

213 """Describe reversible source-checkpoint to high-performance conversion.""" 

214 if self._source_is_fused: 

215 return [] 

216 transforms: list[WeightRenaming | WeightConverter] = [ 

217 WeightConverter( 

218 source_patterns=["gate_proj.weight", "up_proj.weight"], 

219 target_patterns="linear_fc1.weight", 

220 operations=[ 

221 ConcatenateWithSections( 

222 (self._gate_size, self._up_size), 

223 dim=0, 

224 ) 

225 ], 

226 ), 

227 WeightRenaming( 

228 source_patterns="down_proj.weight", 

229 target_patterns="linear_fc2.weight", 

230 ), 

231 ] 

232 if self.linear_fc1.bias is not None: 

233 transforms.insert( 

234 1, 

235 WeightConverter( 

236 source_patterns=["gate_proj.bias", "up_proj.bias"], 

237 target_patterns="linear_fc1.bias", 

238 operations=[ 

239 ConcatenateWithSections( 

240 (self._gate_size, self._up_size), 

241 dim=0, 

242 ) 

243 ], 

244 ), 

245 ) 

246 if self.linear_fc2.bias is not None: 

247 transforms.append( 

248 WeightRenaming( 

249 source_patterns="down_proj.bias", 

250 target_patterns="linear_fc2.bias", 

251 ) 

252 ) 

253 return transforms 

254 

255 def forward( 

256 self, 

257 hidden_states: torch.Tensor, 

258 ) -> torch.Tensor | tuple[torch.Tensor, Optional[torch.Tensor]]: 

259 """Apply the shared-expert MLP while preserving the source return contract.""" 

260 intermediate, bias = self.linear_fc1(hidden_states) 

261 if bias is not None: 

262 intermediate = intermediate + bias 

263 intermediate = self.activation_func(intermediate) 

264 output, output_bias = self.linear_fc2(intermediate) 

265 if not self._return_bias_tuple: 

266 if output_bias is not None: 

267 output = output + output_bias 

268 return output 

269 return output, output_bias