Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / quantization / functional / _saved_quantized.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"""Save physical operands so autograd hooks do not need quantized-wrapper dispatch.""" 

16 

17from typing import Optional 

18 

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

20 

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

22 

23 

24# Saved layout: [group_list, *input_buffers, *weight_buffers]. 

25# Each operand reserves row_data, row_scale, col_data, col_scale in that order, 

26# including None placeholders, so missing operands do not shift later slots. 

27_GROUP_LIST_SLOTS = 1 

28_BUFFERS_PER_OPERAND = 4 

29 

30 

31def save_quantized_operands( 

32 ctx: torch.autograd.function.FunctionCtx, 

33 inputs: Optional[QuantizedTensor], 

34 weight: Optional[QuantizedTensor], 

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

36) -> None: 

37 """Save only physical tensors; keep reconstruction metadata outside hooks. 

38 

39 Args: 

40 ctx: Context for one Dense or grouped projection. 

41 inputs: Input operand needed for the weight gradient, or None. 

42 weight: Weight operand needed for the input gradient, or None. 

43 group_list: Group boundaries or counts for grouped projections. 

44 """ 

45 metadata = [] 

46 tensors = [group_list] 

47 for operand in (inputs, weight): 

48 if operand is None: 

49 metadata.append(None) 

50 tensors.extend((None,) * _BUFFERS_PER_OPERAND) 

51 else: 

52 metadata.append((type(operand), tuple(operand.shape), operand.dtype, operand.quantizer)) 

53 tensors.extend((operand.row_data, operand.row_scale, operand.col_data, operand.col_scale)) 

54 ctx.quantized_metadata = tuple(metadata) 

55 ctx.save_for_backward(*tensors) 

56 

57 

58def restore_quantized_operands( 

59 ctx: torch.autograd.function.FunctionCtx, 

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

61 """Rebuild wrappers from tensors restored by autograd, without requantization. 

62 

63 Args: 

64 ctx: Context populated by save_quantized_operands. 

65 

66 Returns: 

67 Grouping tensor, input operand and weight operand, in that order. 

68 """ 

69 tensors = ctx.saved_tensors 

70 operands = [] 

71 for index, metadata in enumerate(ctx.quantized_metadata): 

72 if metadata is None: 

73 operands.append(None) 

74 continue 

75 tensor_type, shape, dtype, quantizer = metadata 

76 # Skip the grouping prefix and the buffers of preceding operands. 

77 offset = _GROUP_LIST_SLOTS + _BUFFERS_PER_OPERAND * index 

78 row_data, row_scale, col_data, col_scale = tensors[offset:offset + _BUFFERS_PER_OPERAND] 

79 operands.append(tensor_type( 

80 shape, dtype, quantizer=quantizer, 

81 row_data=row_data, row_scale=row_scale, col_data=col_data, col_scale=col_scale, 

82 )) 

83 return tensors[0], operands[0], operands[1]