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

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

19from typing import Tuple 

20 

21from .parallel_ops import DistributedOp 

22 

23 

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

25 if dim is not None: 

26 axis = dim 

27 return (x, axis), {} 

28 

29 

30class ActivationWithAxisDistributedOp(DistributedOp): 

31 """ 

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

33 

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

35 """ 

36 

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

38 """ 

39 Preprocess arguments for activation-with-axis operators. 

40 

41 Args: 

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

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

44 

45 Returns: 

46 tuple: (local_args, local_kwargs, cache_values) 

47 """ 

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

49 input_tensor = args[0] 

50 axis = args[1] 

51 

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

53 local_kwargs = {} 

54 cache_values = [input_tensor.layout, axis] 

55 return local_args, local_kwargs, cache_values 

56 

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

58 """ 

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

60 

61 Rules: 

62 1. Input must not have Partial status. 

63 2. axis must be an int or tuple. 

64 3. Activation axes must not be sharded. 

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

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

67 

68 Args: 

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

70 

71 Returns: 

72 tuple: ((output_layout,), None) 

73 

74 Raises: 

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

76 axis is sharded. 

77 """ 

78 axis = cache_values[-1] 

79 layouts = cache_values[:-1] 

80 if not layouts: 

81 return None 

82 

83 if not self._allow_partial_inputs: 

84 self._check_partial_inputs(layouts) 

85 

86 self.check_layout(layouts, axis) 

87 

88 first_layout = None 

89 for layout in layouts: 

90 if first_layout is None and layout is not None: 

91 first_layout = layout 

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

93 raise ValueError( 

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

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

96 ) 

97 

98 return (first_layout,), None 

99 

100 def check_layout(self, layouts, axis): 

101 """ 

102 check_layout 

103 """ 

104 min_slice_num = 1 

105 x_dict = layouts[0].to_dict() 

106 x_dev = x_dict["tensor_map"] 

107 

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

109 raise ValueError( 

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

111 ) 

112 

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

114 for axis_index in axes: 

115 tensor_map = x_dev[axis_index] 

116 if tensor_map == -1: 

117 continue 

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

119 if axis_strategy != min_slice_num: 

120 raise ValueError( 

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

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

123 )