Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_activation_with_axis.py: 98%

1# Copyright 2025-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""" 

16Activation with axis distributed operator implementation. 

17""" 

18 

19import copy 

20from typing import Tuple 

21 

22from .parallel_ops import DistributedOp 

23 

24 

25def _normalize_activation_with_axis_args(x, axis=-1, dim=None): 

26 if dim is not None: 

27 axis = dim 

28 return (x, axis), {} 

29 

30 

31class ActivationWithAxisDistributedOp(DistributedOp): 

32 """ 

33 Distributed implementation for activation-with-axis operators (e.g., softmax). 

34 

35 Inherits from DistributedOp and provides activation-with-axis specific implementations. 

36 """ 

37 

38 def preprocess(self, args: tuple, kwargs: dict) -> tuple: 

39 """ 

40 Preprocess arguments for activation-with-axis operators. 

41 

42 Args: 

43 args (tuple): Input arguments, first element is the input tensor. 

44 kwargs (dict): Keyword arguments, optionally containing axis/dim. 

45 

46 Returns: 

47 tuple: (local_args, local_kwargs, cache_values) 

48 """ 

49 args, _ = _normalize_activation_with_axis_args(*args, **kwargs) 

50 input_tensor = args[0] 

51 axis = args[1] 

52 

53 local_args = (input_tensor.to_local(), axis) 

54 local_kwargs = {} 

55 cache_values = [input_tensor.layout, axis] 

56 return local_args, local_kwargs, cache_values 

57 

58 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221 

59 """ 

60 Infer output layouts for activation-with-axis operations. 

61 

62 Rules: 

63 1. Input must not have Partial status. 

64 2. axis must be an int or tuple. 

65 3. Activation axes must not be sharded. 

66 4. If multiple input layouts are provided, all tensor inputs must share the same layout. 

67 5. Output layout is identical to the input layout. 

68 

69 Args: 

70 cache_values (list): [input_layout, axis], or [input_layouts..., axis] 

71 

72 Returns: 

73 tuple: ((output_layout,), None) 

74 

75 Raises: 

76 ValueError: If input has Partial status, axis is invalid, or an activation 

77 axis is sharded. 

78 """ 

79 axis = cache_values[-1] 

80 layouts = cache_values[:-1] 

81 if not layouts: 

82 return None 

83 

84 if not self._allow_partial_inputs: 

85 self._check_partial_inputs(layouts) 

86 

87 self.check_layout(layouts, axis) 

88 

89 first_layout = None 

90 for layout in layouts: 

91 if first_layout is None and layout is not None: 

92 first_layout = layout 

93 if layout is not None and first_layout is not None and layout != first_layout: 

94 raise ValueError( 

95 f"For {self.op_name}, requires all tensor inputs to have the same layout. " 

96 f"Input a: {first_layout}, Input b: {layout}" 

97 ) 

98 

99 return (copy.deepcopy(first_layout),), None 

100 

101 def check_layout(self, layouts, axis): 

102 """ 

103 check_layout 

104 """ 

105 min_slice_num = 1 

106 x_dict = layouts[0].to_dict() 

107 x_dev = x_dict["tensor_map"] 

108 

109 if not isinstance(axis, (int, tuple)): 

110 raise ValueError( 

111 f"For {self.op_name}, axis should be int or tuple, but got {type(axis)}" 

112 ) 

113 

114 axes = (axis,) if isinstance(axis, int) else axis 

115 for axis_index in axes: 

116 tensor_map = x_dev[axis_index] 

117 if tensor_map == -1: 

118 continue 

119 axis_strategy = x_dict["mesh_shape"][len(x_dict["mesh_shape"]) - tensor_map - 1] 

120 if axis_strategy != min_slice_num: 

121 raise ValueError( 

122 f"For {self.op_name}, the axis dimension (in dim {axis_index}) is sharded " 

123 f"(strategy is {axis_strategy}). This operation requires the reduction axis to be un-sharded." 

124 )