Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / functional / base_linear_func.py: 92%

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"""Shared Dense low-precision autograd flow and format strategy interface.""" 

16 

17from abc import ABC, abstractmethod 

18from typing import Optional 

19 

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

21from torch.autograd.function import once_differentiable # pylint: disable=forbidden-backend-import 

22 

23from hyper_parallel.components.quantization.config import MATRIX_NDIM 

24from hyper_parallel.components.quantization.tensor import QuantizedTensorStorage 

25from hyper_parallel.components.quantization.tensor.quantized_tensor import QuantizedTensor 

26from hyper_parallel.components.quantization.functional._saved_quantized import ( 

27 restore_quantized_operands, 

28 save_quantized_operands, 

29) 

30 

31 

32def _as_matrix(tensor: torch.Tensor) -> torch.Tensor: 

33 """Flatten leading dimensions while preserving the contracting axis.""" 

34 

35 if tensor.ndim == MATRIX_NDIM: 

36 return tensor 

37 return tensor.reshape(-1, tensor.shape[-1]) 

38 

39 

40class LinearStrategy(ABC): 

41 """Own the common Dense forward/backward flow and expose format hooks.""" 

42 

43 @abstractmethod 

44 def quantize_input( 

45 self, 

46 inputs: torch.Tensor, 

47 *, 

48 rowwise: bool, 

49 colwise: bool, 

50 ) -> QuantizedTensorStorage: 

51 """Create the format-specific input representation.""" 

52 

53 @abstractmethod 

54 def quantize_weight( 

55 self, 

56 weight: torch.Tensor, 

57 *, 

58 rowwise: bool, 

59 colwise: bool, 

60 ) -> QuantizedTensorStorage: 

61 """Create the format-specific weight representation.""" 

62 

63 @abstractmethod 

64 def quantize_grad_output( 

65 self, 

66 grad_output: torch.Tensor, 

67 *, 

68 rowwise: bool, 

69 colwise: bool, 

70 ) -> QuantizedTensorStorage: 

71 """Create the format-specific gradient-output representation.""" 

72 

73 @abstractmethod 

74 def matmul( 

75 self, 

76 left: QuantizedTensorStorage, 

77 right: QuantizedTensorStorage, 

78 *, 

79 layout: str, 

80 output_dtype: torch.dtype, 

81 ) -> torch.Tensor: 

82 """Execute one format-specific Dense matrix multiplication.""" 

83 

84 def forward( 

85 self, 

86 ctx: torch.autograd.function.FunctionCtx, 

87 inputs: torch.Tensor, 

88 weight: torch.Tensor, 

89 ) -> torch.Tensor: 

90 """Run the complete bias-free Dense forward lifecycle.""" 

91 

92 input_matrix = _as_matrix(inputs) 

93 needs_grad_input = inputs.requires_grad 

94 needs_grad_weight = weight.requires_grad 

95 input_quant = self.quantize_input( 

96 input_matrix, 

97 rowwise=True, 

98 colwise=needs_grad_weight, 

99 ) 

100 weight_quant = self.quantize_weight( 

101 weight, 

102 rowwise=True, 

103 colwise=needs_grad_input, 

104 ) 

105 output = self.matmul( 

106 input_quant, 

107 weight_quant, 

108 layout="NT", 

109 output_dtype=inputs.dtype, 

110 ) 

111 

112 ctx.input_shape = inputs.shape 

113 ctx.weight_dtype = weight.dtype 

114 ctx.strategy = self 

115 input_quant.update_usage(rowwise=False, colwise=needs_grad_weight) 

116 weight_quant.update_usage(rowwise=False, colwise=needs_grad_input) 

117 operands = ( 

118 input_quant if needs_grad_weight else None, 

119 weight_quant if needs_grad_input else None, 

120 ) 

121 # Keep physical payloads under PyTorch's saved-tensor lifecycle when 

122 # the production quantizer returned a wrapper Tensor. CPU contract 

123 # strategies may use a light-weight storage double; retain those 

124 # objects directly so the shared seam remains testable without NPU. 

125 ctx.quantized_operands_saved = all( 

126 operand is None or isinstance(operand, QuantizedTensor) 

127 for operand in operands 

128 ) 

129 if ctx.quantized_operands_saved: 

130 save_quantized_operands(ctx, *operands) 

131 ctx.input_quant = None 

132 ctx.weight_quant = None 

133 else: 

134 ctx.input_quant, ctx.weight_quant = operands 

135 if inputs.ndim != MATRIX_NDIM: 

136 output = output.reshape(*inputs.shape[:-1], output.shape[-1]) 

137 return output 

138 

139 def backward( 

140 self, 

141 ctx: torch.autograd.function.FunctionCtx, 

142 grad_output: torch.Tensor, 

143 ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: 

144 """Run the complete Dense dgrad/wgrad lifecycle.""" 

145 

146 grad_matrix = _as_matrix(grad_output) 

147 needs_grad_input = ctx.needs_input_grad[0] 

148 needs_grad_weight = ctx.needs_input_grad[1] 

149 if ctx.quantized_operands_saved: 

150 _, input_quant, weight_quant = restore_quantized_operands(ctx) 

151 else: 

152 input_quant = ctx.input_quant 

153 weight_quant = ctx.weight_quant 

154 grad_quant = self.quantize_grad_output( 

155 grad_matrix, 

156 rowwise=needs_grad_input, 

157 colwise=needs_grad_weight, 

158 ) 

159 

160 grad_input = None 

161 grad_weight = None 

162 if needs_grad_input: 

163 grad_input = self.matmul( 

164 grad_quant, 

165 weight_quant, 

166 layout="NN", 

167 output_dtype=grad_output.dtype, 

168 ) 

169 if len(ctx.input_shape) != MATRIX_NDIM: 

170 grad_input = grad_input.reshape(ctx.input_shape) 

171 if needs_grad_weight: 

172 grad_weight = self.matmul( 

173 grad_quant, 

174 input_quant, 

175 layout="TN", 

176 output_dtype=ctx.weight_dtype, 

177 ) 

178 

179 grad_quant.update_usage(rowwise=False, colwise=False) 

180 ctx.input_quant = None 

181 ctx.weight_quant = None 

182 return grad_input, grad_weight 

183 

184 

185# PyTorch's Function base exposes variadic callbacks; this bridge intentionally 

186# specializes them for the dense-linear ``apply`` contract. 

187class _LinearFunction(torch.autograd.Function): # pylint: disable=abstract-method,arguments-differ 

188 """Bridge PyTorch autograd callbacks to one bound Dense strategy.""" 

189 

190 @staticmethod 

191 def forward( # pylint: disable=arguments-differ 

192 ctx: torch.autograd.function.FunctionCtx, 

193 inputs: torch.Tensor, 

194 weight: torch.Tensor, 

195 strategy: LinearStrategy, 

196 ) -> torch.Tensor: 

197 """Delegate the complete forward lifecycle to the strategy.""" 

198 

199 return strategy.forward(ctx, inputs, weight) 

200 

201 @staticmethod 

202 @once_differentiable 

203 def backward( # pylint: disable=arguments-differ 

204 ctx: torch.autograd.function.FunctionCtx, 

205 grad_output: torch.Tensor, 

206 ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], None]: 

207 """Delegate the complete backward lifecycle to the strategy.""" 

208 

209 grad_input, grad_weight = ctx.strategy.backward(ctx, grad_output) 

210 return grad_input, grad_weight, None 

211 

212 

213__all__ = ["LinearStrategy", "_LinearFunction"]