Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / quantizers / hifloat8.py: 41%

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"""Pangu-compatible HiFloat8 current-scaling quantizer.""" 

16 

17from typing import Any, Optional 

18 

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

20 

21from hyper_parallel.components.quantization.ops.npu_hifloat8 import ( 

22 HiFloat8NpuOps, 

23 _get_hifloat8_npu_ops, 

24) 

25from hyper_parallel.components.quantization.quantizers.base import Quantizer 

26from hyper_parallel.components.quantization.tensor.hifloat8_tensor import ( 

27 HiFloat8Tensor, 

28) 

29from hyper_parallel.components.quantization.tensor.quantized_tensor import ( 

30 QuantizedTensor, 

31) 

32 

33INPUT_WEIGHT_FORMAT_MAX = 15.0 

34GRADIENT_FORMAT_MAX = 224.0 

35_SUPPORTED_FORMAT_MAX = (INPUT_WEIGHT_FORMAT_MAX, GRADIENT_FORMAT_MAX) 

36 

37 

38class HiFloat8Quantizer(Quantizer): 

39 """Quantize one HiFloat8 role with stateless current tensorwise scaling.""" 

40 

41 def __init__( 

42 self, 

43 *, 

44 fp8_max: float, 

45 npu_ops: Optional[HiFloat8NpuOps] = None, 

46 ) -> None: 

47 """Create one frozen-role quantizer compatible with Pangu's structure.""" 

48 

49 if float(fp8_max) not in _SUPPORTED_FORMAT_MAX: 

50 raise ValueError( 

51 "HiFloat8 current scaling supports format max 15 or 224, " 

52 f"got {fp8_max}." 

53 ) 

54 self.fp8_max = float(fp8_max) 

55 self._npu_ops = npu_ops 

56 

57 @property 

58 def npu_ops(self) -> HiFloat8NpuOps: 

59 """Return the injected or lazily resolved NPU operator adapter.""" 

60 

61 if self._npu_ops is None: 

62 self._npu_ops = _get_hifloat8_npu_ops() 

63 return self._npu_ops 

64 

65 @property 

66 def quant_dtype(self) -> Any: 

67 """Return the runtime-frozen logical HiFloat8 dtype identifier.""" 

68 

69 return self.npu_ops.quant_dtype 

70 

71 @staticmethod 

72 def _validate_input( 

73 tensor: torch.Tensor, 

74 *, 

75 group_list: Optional[torch.Tensor], 

76 group_list_type: int, 

77 rowwise: bool, 

78 colwise: bool, 

79 ) -> None: 

80 if not isinstance(tensor, torch.Tensor) or isinstance( 

81 tensor, 

82 QuantizedTensor, 

83 ): 

84 raise TypeError( 

85 "HiFloat8Quantizer expects a high-precision torch.Tensor." 

86 ) 

87 if tensor.dtype not in (torch.float16, torch.bfloat16): 

88 raise TypeError( 

89 "HiFloat8 current quantization requires float16 or bfloat16 " 

90 f"input, got {tensor.dtype}." 

91 ) 

92 if tensor.ndim not in (2, 3): 

93 raise ValueError( 

94 "HiFloat8 quantization requires a matrix or grouped weight, " 

95 "got shape " 

96 f"{tuple(tensor.shape)}." 

97 ) 

98 if not rowwise and not colwise: 

99 raise ValueError( 

100 "HiFloat8Quantizer requires rowwise=True or colwise=True." 

101 ) 

102 if group_list_type not in (0, 1): 

103 raise ValueError( 

104 "HiFloat8Quantizer group_list_type must be 0 or 1, " 

105 f"got {group_list_type}." 

106 ) 

107 if group_list is not None: 

108 if tensor.ndim != 2: 

109 raise ValueError( 

110 "Grouped HiFloat8 quantization requires a 2D operand." 

111 ) 

112 if not isinstance(group_list, torch.Tensor) or group_list.ndim != 1: 

113 raise ValueError( 

114 "HiFloat8Quantizer group_list must be one-dimensional." 

115 ) 

116 if group_list.dtype != torch.int64: 

117 raise TypeError( 

118 "HiFloat8Quantizer group_list must use torch.int64." 

119 ) 

120 

121 def quantize( 

122 self, 

123 tensor: torch.Tensor, 

124 *, 

125 group_list: Optional[torch.Tensor] = None, 

126 group_list_type: int = 0, 

127 rowwise: bool = False, 

128 colwise: bool = False, 

129 ) -> HiFloat8Tensor: 

130 """Quantize Dense or grouped operands with current scaling. 

131 

132 Two-dimensional activation/gradient operands use one tensorwise scale. 

133 GMM callers expand that scalar to one identical entry per expert. 

134 Three-dimensional expert weights use one scale per leading expert. 

135 """ 

136 

137 self._validate_input( 

138 tensor, 

139 group_list=group_list, 

140 group_list_type=group_list_type, 

141 rowwise=rowwise, 

142 colwise=colwise, 

143 ) 

144 if tensor.ndim == 3: 

145 flattened = tensor.reshape(tensor.shape[0], -1) 

146 data, scale = self.npu_ops.dynamic_quant( 

147 flattened, 

148 dst_type_max=self.fp8_max, 

149 quant_mode="pertoken", 

150 ) 

151 data = data.reshape_as(tensor) 

152 else: 

153 data, scale = self.npu_ops.dynamic_quant( 

154 tensor, 

155 dst_type_max=self.fp8_max, 

156 quant_mode="pertensor", 

157 ) 

158 scale = torch.where( 

159 torch.isfinite(scale) & (scale > 0), 

160 scale, 

161 torch.ones_like(scale), 

162 ) 

163 if group_list is not None and scale.numel() == 1: 

164 scale = scale.reshape(1).expand(group_list.numel()).contiguous() 

165 return HiFloat8Tensor( 

166 shape=tensor.shape, 

167 dtype=tensor.dtype, 

168 quantizer=self, 

169 row_data=data if rowwise else None, 

170 row_scale=scale if rowwise else None, 

171 col_data=data if colwise else None, 

172 col_scale=scale if colwise else None, 

173 device=tensor.device, 

174 ) 

175 

176 

177__all__ = [ 

178 "GRADIENT_FORMAT_MAX", 

179 "HiFloat8Quantizer", 

180 "INPUT_WEIGHT_FORMAT_MAX", 

181]