Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / low_precision_linear.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"""Low-precision Linear factories selected by Module Family rules.""" 

16 

17from collections.abc import Mapping 

18from typing import Any 

19 

20from torch import nn 

21 

22from hyper_parallel.components.quantization.functional.mxfp8_linear_func import MXFP8LinearStrategy 

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

24from hyper_parallel.components.quantization.ops.npu_mxfp8 import validate_npu_runtime 

25from hyper_parallel.models.replacement import module_replacement 

26 

27 

28@module_replacement 

29def replace_mxfp8_linear( 

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

31) -> LowPrecisionLinear: 

32 """Bind the MXFP8 Dense strategy selected by the Family rule. 

33 

34 Args: 

35 module: Source Linear whose loaded parameters are retained. 

36 module_fqn: Module path for diagnostics. 

37 context: Replacement context containing parallel topology. 

38 

39 Returns: 

40 The existing low-precision Linear with an MXFP8 strategy. 

41 """ 

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

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

44 if module.in_features % 32 or module.out_features % 32: 

45 raise ValueError(f"{module_fqn!r}: MXFP8 Linear requires feature dimensions divisible by 32") 

46 if context.get("pp"): 

47 raise NotImplementedError(f"{module_fqn!r}: MXFP8 Linear does not support pipeline parallelism") 

48 validate_npu_runtime() 

49 return LowPrecisionLinear.from_linear(module, strategy=MXFP8LinearStrategy())