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

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 .parallel_ops import DistributedOp 

20 

21 

22class ActivationWithAxisDistributedOp(DistributedOp): 

23 """ 

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

25 

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

27 """ 

28 

29 def infer_layout(self, layouts, extra_args=None): 

30 """ 

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

32 

33 For activation-with-axis operations, all inputs should have the same layout, 

34 and the output will have the same layout. 

35 

36 Args: 

37 primitive: Primitive instance 

38 layouts: Layouts of input tensors 

39 

40 Returns: 

41 tuple: Layout for output tensor. 

42 

43 Raises: 

44 ValueError: If input layouts are not compatible or have partial status. 

45 """ 

46 if not layouts: 

47 return None 

48 

49 # Check partial inputs 

50 if not self._allow_partial_inputs: 

51 self._check_partial_inputs(layouts) 

52 

53 self.check_layout(layouts, extra_args) 

54 # Verify all layouts are the same 

55 first_layout = None 

56 for layout in layouts: 

57 if first_layout is None and layout is not None: 

58 first_layout = layout 

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

60 raise ValueError( 

61 f"Operation {self.op_name} requires all tensor inputs to have the same layout. " 

62 f"Input a: {first_layout}, Input b: {layout}") 

63 return first_layout 

64 

65 def check_layout(self, layouts, extra_args): 

66 """ 

67 check_layout 

68 """ 

69 min_slice_num = 1 

70 x_dict = layouts[0].to_dict() 

71 x_dev = x_dict["tensor_map"] 

72 extra_args = extra_args[0] 

73 if not isinstance(extra_args, (int, tuple)): 

74 raise ValueError( 

75 f"Operation {self.op_name}: The extra args should be int or tuple, but got ({type(extra_args)})") 

76 extra_args = (extra_args,) if isinstance(extra_args, int) else extra_args 

77 for axis_index in extra_args: 

78 tensor_map = x_dev[axis_index] 

79 if tensor_map == -1: 

80 continue 

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

82 if axis_strategy != min_slice_num: 

83 raise ValueError( 

84 f"Operation {self.op_name}: The axis dimension (in dim {axis_index}) is sharded " 

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