Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_swiglu.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""" 

16SwiGLU distributed operator implementation. 

17 

18SwiGLU splits the input along an axis into gate and up halves, applies silu to the 

19gate, and multiplies element-wise with up. The split is performed on the local shard. 

20When the split axis is sharded, the producer layout strategy is expected to make each 

21local shard contain paired gate/up halves. 

22""" 

23 

24import copy 

25from typing import Tuple 

26 

27from .parallel_ops import DistributedOp 

28 

29 

30def _normalize_swiglu_args(x, axis=-1, dim=None): 

31 """Normalize SwiGLU arguments across platform call sites. 

32 

33 Args: 

34 x: Input tensor. 

35 axis: The split axis (default -1). 

36 dim: Alias for axis; when provided, overrides axis. 

37 

38 Returns: 

39 tuple: ((x, axis), {}) 

40 """ 

41 if dim is not None: 

42 axis = dim 

43 return (x, axis), {} 

44 

45 

46class SwiGLUDistributedOp(DistributedOp): 

47 """ 

48 Distributed implementation for the SwiGLU operator. 

49 

50 Rules: 

51 1. Input must not have Partial status. 

52 2. axis must be an int within the valid range [-ndim, ndim-1]. 

53 3. The split axis may be sharded. The split is performed on the local shard; 

54 when the split axis is sharded, the producer layout strategy is expected 

55 to make each local shard contain paired gate/up halves. 

56 4. The split axis must have an even global size so it can be halved. 

57 5. Output keeps the same sharding mapping as input. 

58 """ 

59 

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

61 """ 

62 Preprocess arguments for SwiGLU. 

63 

64 Args: 

65 args (tuple): Positional arguments (x,), with optional axis or dim in kwargs. 

66 kwargs (dict): Keyword arguments, optionally containing 'axis' or 'dim'. 

67 

68 Returns: 

69 tuple: (local_args, local_kwargs, cache_values) 

70 """ 

71 args, _ = _normalize_swiglu_args(*args, **kwargs) 

72 input_tensor, axis = args 

73 

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

75 local_kwargs = {} 

76 # Cache the global shape to validate the split axis is even in infer_layout. 

77 cache_values = [input_tensor.layout, axis, input_tensor.shape] 

78 return local_args, local_kwargs, cache_values 

79 

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

81 """ 

82 Infer output layout for SwiGLU. 

83 

84 Args: 

85 cache_values (list): [input_layout, axis, input_shape] 

86 

87 Returns: 

88 tuple: ((output_layout,), None) 

89 

90 Raises: 

91 ValueError: If input has Partial status, axis is invalid, or the split 

92 axis length is not divisible by 2. 

93 """ 

94 layout = cache_values[0] 

95 axis = cache_values[1] 

96 input_shape = cache_values[2] 

97 

98 if not self._allow_partial_inputs: 

99 self._check_partial_inputs([layout]) 

100 

101 if not isinstance(axis, int): 

102 raise ValueError( 

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

104 ) 

105 

106 # Use alias_tensor_map to support StridedShard multi-axis mappings. 

107 alias_map = layout.alias_tensor_map 

108 ndim = len(alias_map) 

109 

110 if axis < 0: 

111 axis += ndim 

112 if axis < 0 or axis >= ndim: 

113 raise ValueError( 

114 f"For {self.op_name}, axis out of range " 

115 f"(expected to be in range of [{-ndim}, {ndim - 1}], but got {cache_values[1]})" 

116 ) 

117 

118 # The split axis may be sharded — the split is performed on the local 

119 # shard. When the split axis is sharded, the producer layout strategy is 

120 # expected to make each local shard contain paired gate/up halves. 

121 # 

122 # The split axis must have an even global length so it can be halved into 

123 # gate and up within each shard. 

124 axis_size = input_shape[axis] 

125 if axis_size % 2 != 0: 

126 raise ValueError( 

127 f"For {self.op_name}, the split axis (dim {axis}) must have an even size, " 

128 f"but got {axis_size}. SwiGLU splits the input along this axis " 

129 f"into two equal halves for gate and up." 

130 ) 

131 

132 return (copy.deepcopy(layout),), None