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

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"""Scheme-selected grouped experts using the existing QAT computation strategies.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Mapping 

20from functools import partial 

21from typing import Any 

22 

23from torch import Tensor, nn # pylint: disable=forbidden-backend-import 

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

25from transformers.activations import ACT2FN 

26from transformers.models.deepseek_v4.modeling_deepseek_v4 import DeepseekV4Experts 

27 

28from hyper_parallel.components.modules.grouped_experts import GroupedExperts 

29from hyper_parallel.components.quantization.functional.base_gmm_func import GroupedLinear, _GroupedLinearFunction 

30from hyper_parallel.components.quantization.functional.strategy_factory import _validate_tile_alignment 

31from hyper_parallel.components.quantization.functional.mxfp8_gmm_func import MXFP8GroupedLinear 

32from hyper_parallel.components.quantization.functional.w4a8_gmm_func import W4A8GroupedLinear 

33from hyper_parallel.components.quantization.functional.fake_w4a8_gmm_func import FakeW4A8GroupedLinear 

34from hyper_parallel.components.quantization.ops.npu_mxfp8 import validate_npu_gmm_runtime 

35from hyper_parallel.components.quantization.ops.npu_w4a8 import validate_w4a8_gmm_runtime 

36from hyper_parallel.components.quantization.ops.npu_fake_w4a8 import validate_fake_w4a8_gmm_runtime 

37from hyper_parallel.models.replacement import module_replacement 

38 

39 

40def _swiglu(projected: Tensor, *, limit: float | None = None) -> Tensor: 

41 """Preserve projection dtype and optionally apply the native DeepSeek V4/V4.1 clamp.""" 

42 gate, up = projected.chunk(2, dim=-1) 

43 if limit is not None: 

44 gate = gate.clamp(max=limit) 

45 up = up.clamp(min=-limit, max=limit) 

46 return F.silu(gate) * up 

47 

48 

49class LowPrecisionGroupedExperts(GroupedExperts): 

50 """Preserve high-performance expert layout/routing and bind one existing QAT strategy.""" 

51 

52 def __init__(self, *, module: nn.Module, strategy: GroupedLinear, module_fqn: str = "", 

53 context: Mapping[str, Any]) -> None: 

54 """Build local expert computation from the already resolved dtype scheme. 

55 

56 Args: 

57 module: Source packed experts. 

58 strategy: Computation strategy selected by the Family rule. 

59 module_fqn: Module path used in diagnostics. 

60 context: Existing replacement context containing the resolved low_precision policy. 

61 """ 

62 prefix = f"{module_fqn or '<root>'}: LowPrecisionGroupedExperts" 

63 if any(context.get(axis) for axis in ("tp", "cp", "ep", "pp")) or context.get("ep_size", 1) != 1: 

64 raise NotImplementedError(f"{prefix} requires TP=CP=EP=PP=1") 

65 super().__init__(module=module, module_fqn=module_fqn, context=context) 

66 if (not self.has_gate or self.add_bias or self.use_2d_experts 

67 or not self.router_gating_in_fp32): 

68 raise ValueError(f"{prefix} requires 3D bias-free SwiGLU with routing scores after GMM2") 

69 self.grouped_linear = strategy 

70 _validate_tile_alignment( 

71 "fake_w4a8" if isinstance(strategy, FakeW4A8GroupedLinear) else strategy.FORMAT_NAME, 

72 context["low_precision"].block_size, 

73 (tuple(self.gate_up_proj.shape), tuple(self.down_proj.shape)), 

74 ) 

75 

76 @staticmethod 

77 def _activation(module: nn.Module, config: Any, gated_linear_unit: bool) -> Any: 

78 """Adapt V4.1's native expert activation without retaining its source parameters.""" 

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

80 if hidden_act is None: 

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

82 if hidden_act is None: 

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

84 source_activation = getattr(module, "act_fn", None) 

85 if (hidden_act not in (None, "silu", "swiglu") 

86 or (source_activation is not None and source_activation is not F.silu 

87 and not isinstance(source_activation, (nn.SiLU, type(ACT2FN["silu"]))))): 

88 raise ValueError("LowPrecisionGroupedExperts requires SwiGLU source activation") 

89 if isinstance(module, DeepseekV4Experts): 

90 return partial(_swiglu, limit=module.limit) 

91 # Keep the parent's source validation and the existing non-fused SiLU/multiply numerics. 

92 GroupedExperts._activation(module, config, gated_linear_unit) 

93 return _swiglu 

94 

95 def _grouped_gemm_expert_forward(self, gate_up_proj, down_proj, permuted, 

96 tokens_per_expert, permuted_probs): 

97 """Adapt layout at the shared strategy boundary; outer routing applies scores.""" 

98 del permuted_probs 

99 projected = _GroupedLinearFunction.apply( 

100 permuted, gate_up_proj.transpose(-2, -1), tokens_per_expert, self.grouped_linear, 1) 

101 return _GroupedLinearFunction.apply( 

102 self.activation_func(projected), down_proj.transpose(-2, -1), 

103 tokens_per_expert, self.grouped_linear, 1) 

104 

105 

106@module_replacement 

107def replace_mxfp8_grouped_experts( 

108 *, module: nn.Module, module_fqn: str, context: Mapping[str, Any], 

109) -> LowPrecisionGroupedExperts: 

110 """Bind the MXFP8GroupedLinear strategy selected by the Family rule. 

111 

112 Args: 

113 module: Source packed experts. 

114 module_fqn: Module path for diagnostics. 

115 context: Replacement context containing the resolved dtype scheme. 

116 """ 

117 validate_npu_gmm_runtime() 

118 return LowPrecisionGroupedExperts( 

119 module=module, module_fqn=module_fqn, context=context, 

120 strategy=MXFP8GroupedLinear(), 

121 ) 

122 

123 

124@module_replacement 

125def replace_w4a8_grouped_experts( 

126 *, module: nn.Module, module_fqn: str, context: Mapping[str, Any], 

127) -> LowPrecisionGroupedExperts: 

128 """Bind the W4A8GroupedLinear strategy selected by the Family rule. 

129 

130 Args: 

131 module: Source packed experts. 

132 module_fqn: Module path for diagnostics. 

133 context: Replacement context containing the resolved dtype scheme. 

134 """ 

135 validate_w4a8_gmm_runtime() 

136 return LowPrecisionGroupedExperts( 

137 module=module, module_fqn=module_fqn, context=context, 

138 strategy=W4A8GroupedLinear(context["low_precision"].block_size), 

139 ) 

140 

141 

142@module_replacement 

143def replace_fake_w4a8_grouped_experts( 

144 *, module: nn.Module, module_fqn: str, context: Mapping[str, Any], 

145) -> LowPrecisionGroupedExperts: 

146 """Bind the FakeW4A8GroupedLinear strategy selected by the Family rule. 

147 

148 Args: 

149 module: Source packed experts. 

150 module_fqn: Module path for diagnostics. 

151 context: Replacement context containing the resolved dtype scheme. 

152 """ 

153 validate_fake_w4a8_gmm_runtime(context["low_precision"].block_size) 

154 return LowPrecisionGroupedExperts( 

155 module=module, module_fqn=module_fqn, context=context, 

156 strategy=FakeW4A8GroupedLinear(context["low_precision"].block_size), 

157 )