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.tensor import QuantizedTensorStorage 

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

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

26 restore_quantized_operands, 

27 save_quantized_operands, 

28) 

29 

30 

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

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

33 

34 if tensor.ndim == 2: 

35 return tensor 

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

37 

38 

39class LinearStrategy(ABC): 

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

41 

42 @abstractmethod 

43 def quantize_input( 

44 self, 

45 inputs: torch.Tensor, 

46 *, 

47 rowwise: bool, 

48 colwise: bool, 

49 ) -> QuantizedTensorStorage: 

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

51 

52 @abstractmethod 

53 def quantize_weight( 

54 self, 

55 weight: torch.Tensor, 

56 *, 

57 rowwise: bool, 

58 colwise: bool, 

59 ) -> QuantizedTensorStorage: 

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

61 

62 @abstractmethod 

63 def quantize_grad_output( 

64 self, 

65 grad_output: torch.Tensor, 

66 *, 

67 rowwise: bool, 

68 colwise: bool, 

69 ) -> QuantizedTensorStorage: 

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

71 

72 @abstractmethod 

73 def matmul( 

74 self, 

75 left: QuantizedTensorStorage, 

76 right: QuantizedTensorStorage, 

77 *, 

78 layout: str, 

79 output_dtype: torch.dtype, 

80 ) -> torch.Tensor: 

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

82 

83 def forward( 

84 self, 

85 ctx: torch.autograd.function.FunctionCtx, 

86 inputs: torch.Tensor, 

87 weight: torch.Tensor, 

88 ) -> torch.Tensor: 

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

90 

91 input_matrix = _as_matrix(inputs) 

92 needs_grad_input = inputs.requires_grad 

93 needs_grad_weight = weight.requires_grad 

94 input_quant = self.quantize_input( 

95 input_matrix, 

96 rowwise=True, 

97 colwise=needs_grad_weight, 

98 ) 

99 weight_quant = self.quantize_weight( 

100 weight, 

101 rowwise=True, 

102 colwise=needs_grad_input, 

103 ) 

104 output = self.matmul( 

105 input_quant, 

106 weight_quant, 

107 layout="NT", 

108 output_dtype=inputs.dtype, 

109 ) 

110 

111 ctx.input_shape = inputs.shape 

112 ctx.weight_dtype = weight.dtype 

113 ctx.strategy = self 

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

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

116 operands = ( 

117 input_quant if needs_grad_weight else None, 

118 weight_quant if needs_grad_input else None, 

119 ) 

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

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

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

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

124 ctx.quantized_operands_saved = all( 

125 operand is None or isinstance(operand, QuantizedTensor) 

126 for operand in operands 

127 ) 

128 if ctx.quantized_operands_saved: 

129 save_quantized_operands(ctx, *operands) 

130 ctx.input_quant = None 

131 ctx.weight_quant = None 

132 else: 

133 ctx.input_quant, ctx.weight_quant = operands 

134 if inputs.ndim != 2: 

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

136 return output 

137 

138 def backward( 

139 self, 

140 ctx: torch.autograd.function.FunctionCtx, 

141 grad_output: torch.Tensor, 

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

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

144 

145 grad_matrix = _as_matrix(grad_output) 

146 needs_grad_input = ctx.needs_input_grad[0] 

147 needs_grad_weight = ctx.needs_input_grad[1] 

148 if ctx.quantized_operands_saved: 

149 _, input_quant, weight_quant = restore_quantized_operands(ctx) 

150 else: 

151 input_quant = ctx.input_quant 

152 weight_quant = ctx.weight_quant 

153 grad_quant = self.quantize_grad_output( 

154 grad_matrix, 

155 rowwise=needs_grad_input, 

156 colwise=needs_grad_weight, 

157 ) 

158 

159 grad_input = None 

160 grad_weight = None 

161 if needs_grad_input: 

162 grad_input = self.matmul( 

163 grad_quant, 

164 weight_quant, 

165 layout="NN", 

166 output_dtype=grad_output.dtype, 

167 ) 

168 if len(ctx.input_shape) != 2: 

169 grad_input = grad_input.reshape(ctx.input_shape) 

170 if needs_grad_weight: 

171 grad_weight = self.matmul( 

172 grad_quant, 

173 input_quant, 

174 layout="TN", 

175 output_dtype=ctx.weight_dtype, 

176 ) 

177 

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

179 ctx.input_quant = None 

180 ctx.weight_quant = None 

181 return grad_input, grad_weight 

182 

183 

184class _LinearFunction(torch.autograd.Function): 

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

186 

187 @staticmethod 

188 def forward( 

189 ctx: torch.autograd.function.FunctionCtx, 

190 inputs: torch.Tensor, 

191 weight: torch.Tensor, 

192 strategy: LinearStrategy, 

193 ) -> torch.Tensor: 

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

195 

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

197 

198 @staticmethod 

199 @once_differentiable 

200 def backward( 

201 ctx: torch.autograd.function.FunctionCtx, 

202 grad_output: torch.Tensor, 

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

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

205 

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

207 return grad_input, grad_weight, None 

208 

209 

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