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

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"""Root mean square normalization module.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Mapping 

20from typing import Any 

21 

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

23from hyper_parallel.components.checkpoint.weight_conversion import WeightConverter 

24 

25from hyper_parallel.components.checkpoint.conversion_ops import AddScalar 

26from hyper_parallel.models.replacement import module_replacement 

27from hyper_parallel.components.functional import rms_norm 

28 

29 

30@module_replacement 

31class RMSNorm(torch.nn.Module): 

32 """NPU-accelerated replacement for direct-scale RMSNorm.""" 

33 

34 def __init__( 

35 self, 

36 hidden_size: int | None = None, 

37 eps: float = 1e-5, 

38 *, 

39 module: torch.nn.Module | None = None, 

40 module_fqn: str = "", 

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

42 ) -> None: 

43 """Build a direct-scale RMSNorm or replace an existing module. 

44 

45 Args: 

46 hidden_size: Size of the normalized dimension for direct construction. 

47 eps: Epsilon added to the variance. 

48 module: Optional source module whose parameter layout is preserved. 

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

50 context: Replacement context supplied by Trainer. 

51 

52 Raises: 

53 TypeError: If the source module does not expose a supported RMSNorm contract. 

54 ValueError: If neither a source module nor a positive hidden size is supplied. 

55 """ 

56 super().__init__() 

57 del module_fqn, context 

58 if module is not None: 

59 if not isinstance(getattr(module, "weight", None), torch.nn.Parameter): 

60 raise TypeError("RMSNorm source module must expose an nn.Parameter weight") 

61 source_eps = getattr(module, "eps", None) 

62 if source_eps is None: 

63 source_eps = getattr(module, "variance_epsilon", None) 

64 if source_eps is None: 

65 raise TypeError( 

66 "RMSNorm source module must expose eps or variance_epsilon" 

67 ) 

68 self.eps = float(source_eps) 

69 self.weight = module.weight 

70 self.train(module.training) 

71 return 

72 if hidden_size is None or hidden_size <= 0: 

73 raise ValueError("RMSNorm hidden_size must be a positive integer") 

74 self.eps = eps 

75 self.weight = torch.nn.Parameter(torch.ones(hidden_size)) 

76 

77 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: 

78 """Apply NPU-accelerated RMS normalization.""" 

79 return rms_norm(hidden_states, self.weight, self.eps) 

80 

81 

82@module_replacement 

83class OffsetRMSNorm(torch.nn.Module): 

84 """High-performance replacement for RMSNorm using ``1 + weight``.""" 

85 

86 def __init__( 

87 self, 

88 *, 

89 module: torch.nn.Module, 

90 module_fqn: str = "", 

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

92 ) -> None: 

93 """Build a direct-scale RMSNorm from a unit-offset source module.""" 

94 super().__init__() 

95 del module_fqn, context 

96 if not hasattr(module, "weight") or not hasattr(module, "eps"): 

97 raise TypeError( 

98 "OffsetRMSNorm source module must expose weight and eps" 

99 ) 

100 self.eps = module.eps 

101 weight = ( 

102 torch.empty_like(module.weight) 

103 if module.weight.is_meta 

104 else module.weight.detach() + 1.0 

105 ) 

106 self.weight = torch.nn.Parameter( 

107 weight, 

108 requires_grad=module.weight.requires_grad, 

109 ) 

110 self.train(module.training) 

111 

112 def reset_parameters(self) -> None: 

113 """Initialize the direct scale used by randomly initialized models.""" 

114 torch.nn.init.ones_(self.weight) 

115 

116 def make_transforms(self) -> list[WeightConverter]: 

117 """Convert the unit-offset weight to the NPU operator's direct scale.""" 

118 return [ 

119 WeightConverter( 

120 source_patterns=["weight"], 

121 target_patterns="weight", 

122 operations=[AddScalar(1.0)], 

123 ) 

124 ] 

125 

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

127 """Apply unit-offset RMSNorm with the NPU operator.""" 

128 return rms_norm(x, self.weight, self.eps)