Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / modules / replacement.py: 100%

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"""Legacy policy-based replacement factories shared across models. 

16 

17New configurations can select fixed strategies through Module Family rules. 

18These factories retain the original low-precision module layout and policy API. 

19""" 

20 

21from collections.abc import Mapping 

22from typing import Any 

23 

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

25 

26from hyper_parallel.components.quantization.functional import build_low_precision_strategy 

27from hyper_parallel.components.quantization.functional.linear_strategy_factory import ( 

28 build_linear_strategy, 

29) 

30from hyper_parallel.components.quantization.modules.grouped_experts import ( 

31 GroupedExperts, 

32) 

33from hyper_parallel.components.quantization.modules.linear import LowPrecisionLinear 

34from hyper_parallel.models.replacement import module_replacement 

35 

36 

37def _check_ep1_only(context: Mapping[str, Any], factory_name: str) -> None: 

38 """Reject any active model-parallel axis (EP=1 packed experts only).""" 

39 active_model_parallel_axes = [ 

40 axis.upper() 

41 for axis in ("tp", "cp", "ep", "pp") 

42 if context.get(axis) 

43 ] 

44 if active_model_parallel_axes: 

45 raise NotImplementedError( 

46 f"{factory_name} currently require TP=CP=EP=PP=1; " 

47 f"active axes: {active_model_parallel_axes}." 

48 ) 

49 

50 

51@module_replacement 

52def replace_linear( 

53 *, 

54 module: nn.Module, 

55 module_fqn: str, 

56 context: Mapping[str, Any], 

57) -> LowPrecisionLinear: 

58 """Replace one exact Dense Linear using its resolved low-precision policy.""" 

59 

60 if type(module) is not nn.Linear: # pylint: disable=unidiomatic-typecheck 

61 raise TypeError( 

62 f"{module_fqn!r} must be exact nn.Linear, got {type(module).__name__}" 

63 ) 

64 if context.get("pp"): 

65 raise NotImplementedError( 

66 "Low-precision Linear training is not yet supported with pipeline parallelism." 

67 ) 

68 try: 

69 strategy = build_linear_strategy( 

70 context.get("low_precision"), 

71 in_features=module.in_features, 

72 out_features=module.out_features, 

73 ) 

74 except ValueError as error: 

75 raise ValueError( 

76 f"Low-precision Linear target {module_fqn!r} is invalid: {error}" 

77 ) from error 

78 return LowPrecisionLinear.from_linear(module, strategy=strategy) 

79 

80 

81@module_replacement 

82def replace_grouped_experts( 

83 *, 

84 module: nn.Module, 

85 module_fqn: str, 

86 context: Mapping[str, Any], 

87) -> GroupedExperts: 

88 """Replace packed experts with the policy-selected grouped-linear strategy.""" 

89 

90 _check_ep1_only(context, "Low-precision grouped experts") 

91 parameters = tuple(module.parameters(recurse=False)) 

92 strategy = build_low_precision_strategy( 

93 context.get("low_precision"), 

94 tile_shapes=tuple(tuple(parameter.shape) for parameter in parameters), 

95 ) 

96 return GroupedExperts.from_module( 

97 module, 

98 fqn=module_fqn, 

99 grouped_linear=strategy, 

100 ) 

101 

102 

103__all__ = ["replace_linear", "replace_grouped_experts"]