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

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"""Canonical shell for low-precision Dense Linear.""" 

16 

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

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

19 

20from hyper_parallel.components.quantization.functional.base_linear_func import ( 

21 LinearStrategy, 

22 _LinearFunction, 

23) 

24 

25 

26class LowPrecisionLinear(nn.Linear): 

27 """Preserve Linear state while delegating compute to one bound strategy.""" 

28 

29 _hp_linear_compute_kind = "npu_quant" 

30 

31 def __init__( 

32 self, 

33 in_features: int, 

34 out_features: int, 

35 bias: bool = True, 

36 *, 

37 strategy: LinearStrategy, 

38 ) -> None: 

39 """Create a low-precision Linear with newly allocated Parameters.""" 

40 

41 self._validate_strategy(strategy) 

42 super().__init__(in_features, out_features, bias=bias) 

43 self.strategy = strategy 

44 

45 @staticmethod 

46 def _validate_strategy(strategy: LinearStrategy) -> None: 

47 """Require a plain strategy so module registration stays unchanged.""" 

48 

49 if isinstance(strategy, nn.Module) or not isinstance(strategy, LinearStrategy): 

50 raise TypeError("strategy must be a non-Module LinearStrategy instance") 

51 

52 @classmethod 

53 def from_linear( 

54 cls, 

55 linear: nn.Linear, 

56 *, 

57 strategy: LinearStrategy, 

58 ) -> "LowPrecisionLinear": 

59 """Create a no-allocation shell retaining the source Parameters.""" 

60 

61 cls._validate_strategy(strategy) 

62 # object.__new__ is a static method; passing cls explicitly avoids the 

63 # allocation and initialization performed by nn.Linear.__init__. 

64 converted = cls.__new__(cls) # pylint: disable=no-value-for-parameter 

65 nn.Module.__init__(converted) # pylint: disable=unnecessary-dunder-call 

66 converted.in_features = linear.in_features 

67 converted.out_features = linear.out_features 

68 converted.register_parameter("weight", linear.weight) 

69 converted.register_parameter("bias", linear.bias) 

70 converted.strategy = strategy 

71 converted.training = linear.training 

72 return converted 

73 

74 def forward(self, input: torch.Tensor) -> torch.Tensor: # pylint: disable=redefined-builtin 

75 """Apply the selected low-precision compute and high-precision bias.""" 

76 

77 output = _LinearFunction.apply(input, self.weight, self.strategy) 

78 if self.bias is not None: 

79 output = output + self.bias 

80 return output 

81 

82 

83__all__ = ["LowPrecisionLinear"]