Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / swiglu_mlp.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"""Transformers-compatible SwiGLU MLP with a fused Gate/Up projection.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Mapping 

20from typing import Any 

21 

22# This package provides PyTorch-specific high-performance modules. 

23# pylint: disable=forbidden-backend-import 

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

25import torch.nn.functional as F 

26from torch import nn 

27from hyper_parallel.components.checkpoint.weight_conversion import ( 

28 WeightConverter, 

29 WeightRenaming, 

30) 

31 

32from hyper_parallel.components.checkpoint import ConcatenateWithSections 

33from hyper_parallel.models.replacement import module_replacement 

34from hyper_parallel.components.functional import swiglu 

35 

36 

37@module_replacement 

38class SwiGLUMLP(nn.Module): 

39 """Transformers-compatible SwiGLU MLP using one fused Gate/Up matmul.""" 

40 

41 @staticmethod 

42 def _source_projections(module: nn.Module) -> tuple[nn.Linear, nn.Linear, nn.Linear]: 

43 """Validate and return the source Gate, Up, and Down projections.""" 

44 required = ("gate_proj", "up_proj", "down_proj") 

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

46 if missing: 

47 raise TypeError(f"SwiGLUMLP source module is missing: {missing}") 

48 projections = (module.gate_proj, module.up_proj, module.down_proj) 

49 if not all(isinstance(projection, nn.Linear) for projection in projections): 

50 raise TypeError("SwiGLUMLP requires Gate, Up, and Down projections to be nn.Linear") 

51 return projections 

52 

53 @staticmethod 

54 def _validate_projection_contract( 

55 gate_proj: nn.Linear, 

56 up_proj: nn.Linear, 

57 down_proj: nn.Linear, 

58 ) -> None: 

59 """Validate source projection shapes, placement, and training policy.""" 

60 if gate_proj.in_features != up_proj.in_features: 

61 raise ValueError("Gate and Up projections must have the same input size") 

62 if gate_proj.out_features != up_proj.out_features: 

63 raise ValueError("Gate and Up projections must have the same output size") 

64 if down_proj.in_features != gate_proj.out_features: 

65 raise ValueError("Down projection input size must equal the SwiGLU intermediate size") 

66 if down_proj.out_features != gate_proj.in_features: 

67 raise ValueError("Down projection output size must equal the MLP hidden size") 

68 if gate_proj.weight.device != up_proj.weight.device or gate_proj.weight.dtype != up_proj.weight.dtype: 

69 raise ValueError("Gate and Up projection weights must share device and dtype") 

70 if gate_proj.weight.device != down_proj.weight.device or gate_proj.weight.dtype != down_proj.weight.dtype: 

71 raise ValueError("Gate, Up, and Down projection weights must share device and dtype") 

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

73 raise ValueError("Gate and Up projection weights must share a training policy") 

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

75 raise ValueError("Gate and Up projections must use the same bias policy") 

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

77 raise ValueError("Gate and Up projection biases must share a training policy") 

78 

79 def _initialize_fc1(self, gate_proj: nn.Linear, up_proj: nn.Linear) -> None: 

80 """Create the packed Gate/Up projection with source parameter values.""" 

81 pack = ConcatenateWithSections((self.intermediate_size, self.intermediate_size)) 

82 packed_weight = pack.convert( 

83 {"gate": gate_proj.weight.detach(), "up": up_proj.weight.detach()}, 

84 ["gate", "up"], 

85 ["linear_fc1.weight"], 

86 )["linear_fc1.weight"] 

87 self.linear_fc1 = nn.Linear( 

88 self.hidden_size, 

89 2 * self.intermediate_size, 

90 bias=gate_proj.bias is not None, 

91 device=gate_proj.weight.device, 

92 dtype=gate_proj.weight.dtype, 

93 ) 

94 self.linear_fc1.weight = nn.Parameter(packed_weight, requires_grad=gate_proj.weight.requires_grad) 

95 if gate_proj.bias is not None: 

96 packed_bias = pack.convert( 

97 {"gate": gate_proj.bias.detach(), "up": up_proj.bias.detach()}, 

98 ["gate", "up"], 

99 ["linear_fc1.bias"], 

100 )["linear_fc1.bias"] 

101 self.linear_fc1.bias = nn.Parameter(packed_bias, requires_grad=gate_proj.bias.requires_grad) 

102 

103 def _initialize_fc2(self, down_proj: nn.Linear) -> None: 

104 """Create the Down projection and reuse its source parameters.""" 

105 self.linear_fc2 = nn.Linear( 

106 down_proj.in_features, 

107 down_proj.out_features, 

108 bias=down_proj.bias is not None, 

109 device=down_proj.weight.device, 

110 dtype=down_proj.weight.dtype, 

111 ) 

112 self.linear_fc2.weight = down_proj.weight 

113 self.linear_fc2.bias = down_proj.bias 

114 

115 def __init__( 

116 self, 

117 *, 

118 module: nn.Module, 

119 module_fqn: str = "", 

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

121 ) -> None: 

122 """Build the high-performance MLP from separate source projections. 

123 

124 Args: 

125 module: Source MLP exposing ``gate_proj``, ``up_proj``, and ``down_proj``. 

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

127 context: Replacement context supplied by Trainer. 

128 

129 Raises: 

130 TypeError: If the source projection modules are missing or unsupported. 

131 ValueError: If projection layouts, training policies, or activation differ. 

132 """ 

133 super().__init__() 

134 del module_fqn, context 

135 gate_proj, up_proj, down_proj = self._source_projections(module) 

136 self._validate_projection_contract(gate_proj, up_proj, down_proj) 

137 

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

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

140 if hidden_act is not None and hidden_act not in ("silu", "swiglu"): 

141 raise ValueError(f"SwiGLUMLP requires a SiLU activation, but got {hidden_act}") 

142 self.config = config 

143 self.hidden_size = gate_proj.in_features 

144 self.intermediate_size = gate_proj.out_features 

145 

146 self._initialize_fc1(gate_proj, up_proj) 

147 self._initialize_fc2(down_proj) 

148 self.train(module.training) 

149 

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

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

152 transforms: list[WeightRenaming | WeightConverter] = [ 

153 WeightConverter( 

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

155 target_patterns="linear_fc1.weight", 

156 operations=[ 

157 ConcatenateWithSections((self.intermediate_size, self.intermediate_size)) 

158 ], 

159 ) 

160 ] 

161 if self.linear_fc1.bias is not None: 

162 transforms.append( 

163 WeightConverter( 

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

165 target_patterns="linear_fc1.bias", 

166 operations=[ 

167 ConcatenateWithSections((self.intermediate_size, self.intermediate_size)) 

168 ], 

169 ) 

170 ) 

171 transforms.append(WeightRenaming("down_proj.weight", "linear_fc2.weight")) 

172 if self.linear_fc2.bias is not None: 

173 transforms.append(WeightRenaming("down_proj.bias", "linear_fc2.bias")) 

174 return transforms 

175 

176 def forward(self, x: torch.Tensor) -> torch.Tensor: 

177 """Apply the fused Gate/Up projection, SwiGLU, and Down projection.""" 

178 intermediate_parallel = self.linear_fc1(x) 

179 if intermediate_parallel.device.type == "npu": 

180 intermediate_parallel = swiglu( # pylint: disable=not-callable 

181 intermediate_parallel 

182 ) 

183 else: 

184 gate, up = intermediate_parallel.chunk(2, dim=-1) 

185 intermediate_parallel = F.silu(gate) * up 

186 return self.linear_fc2(intermediate_parallel)