Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / custom_shard.py: 100%

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"""shard""" 

16 

17import queue 

18from typing import Callable, Tuple, Optional 

19from hyper_parallel.core.dtensor.layout import DeviceMesh 

20from hyper_parallel.core.dtensor.dtensor import DTensor 

21from hyper_parallel.core.dtensor.placement_types import Placement 

22from hyper_parallel.platform import get_platform 

23platform = get_platform() 

24Tensor = platform.Tensor 

25 

26 

27def _process_custom_shard_inputs(args, in_placements, redistribute_inputs, device_mesh): 

28 """Process input arguments for custom_shard, redistributing DTensors as needed. 

29 

30 Args: 

31 args: Positional arguments to the wrapped function. 

32 in_placements: Expected placements for each input (None entries = non-tensor). 

33 redistribute_inputs: Whether to redistribute inputs to required placements. 

34 device_mesh: Device mesh for redistribution. 

35 

36 Returns: 

37 tuple: (local_args, contain_distributed_arg, args_layout_queue) 

38 

39 Raises: 

40 RuntimeError: If a DTensor arg is found but in_placements is None. 

41 TypeError: If a DTensor arg has None in_placements entry, or a non-DTensor 

42 arg has a non-None in_placements entry. 

43 """ 

44 local_args = [] 

45 contain_distributed_arg = False 

46 args_layout = queue.Queue(len(args)) 

47 

48 for i, arg in enumerate(args): 

49 if isinstance(arg, DTensor): 

50 if in_placements is None: 

51 raise RuntimeError("Found Tensor input but in_placements is None") 

52 

53 required_in_placement = in_placements[i] 

54 if required_in_placement is None: 

55 raise TypeError( 

56 f"Tensor input at position {i} requires Placement, " 

57 "but corresponding in_placements entry is None!" 

58 ) 

59 

60 if redistribute_inputs: 

61 arg = arg.redistribute(device_mesh, required_in_placement) 

62 

63 args_layout.put(arg.layout) 

64 local_tensor = arg.to_local() 

65 local_args.append(local_tensor) 

66 contain_distributed_arg = True 

67 else: 

68 if in_placements is not None and in_placements[i] is not None: 

69 raise TypeError( 

70 f"Non-DTensor input at position {i} requires None in_placements, " 

71 f"but received {in_placements[i]}!" 

72 ) 

73 local_args.append(arg) 

74 

75 return local_args, contain_distributed_arg, args_layout 

76 

77 

78def _wrap_custom_shard_outputs(out, out_placements, contain_distributed_arg, device_mesh): 

79 """Wrap output tensors as DTensors using the specified placements. 

80 

81 Args: 

82 out: Raw output(s) from the wrapped function. 

83 out_placements: Placements for each output tensor. 

84 contain_distributed_arg: Whether any input was a DTensor. 

85 device_mesh: Device mesh for DTensor construction. 

86 

87 Returns: 

88 DTensor or tuple of DTensors, or the raw output if no distributed args. 

89 

90 Raises: 

91 TypeError: If a tensor output has None out_placements, or a non-tensor 

92 output has non-None out_placements. 

93 ValueError: If output count doesn't match out_placements count. 

94 """ 

95 if not contain_distributed_arg: 

96 return out 

97 

98 out_is_tuple = isinstance(out, tuple) 

99 out_tuple = (out,) if not out_is_tuple else out 

100 

101 if len(out_tuple) != len(out_placements): 

102 raise ValueError( 

103 f"Output count {len(out_tuple)} does not match " 

104 f"out_placements count {len(out_placements)}!" 

105 ) 

106 

107 dist_output = [] 

108 for item, out_placement in zip(out_tuple, out_placements): 

109 if isinstance(item, Tensor): 

110 if out_placement is None: 

111 raise TypeError( 

112 "Tensor output requires non-None out_placements!" 

113 ) 

114 dist_output.append( 

115 DTensor.from_local(item, device_mesh=device_mesh, placements=out_placement) 

116 ) 

117 else: 

118 if out_placement is not None: 

119 raise TypeError( 

120 f"Non-tensor output requires None out_placements, got {out_placement}!" 

121 ) 

122 dist_output.append(item) 

123 

124 return dist_output[0] if not out_is_tuple else tuple(dist_output) 

125 

126 

127def custom_shard( 

128 func: Callable, 

129 device_mesh: DeviceMesh, 

130 out_placements: Tuple[Tuple[Placement, ...], ...], 

131 in_placements: Optional[Tuple[Optional[Tuple[Placement, ...]], ...]] = None, 

132 redistribute_inputs: bool = True, 

133) -> Callable: 

134 """ 

135 Wraps a function to handle distributed tensor conversions. 

136 

137 Args: 

138 func (Callable): The function to be wrapped. 

139 device_mesh (DeviceMesh): The device mesh for sharding. 

140 out_placements (Tuple[Tuple[Placement, ...], ...]): Placements for each output tensor. 

141 in_placements (Optional[Tuple[Optional[Tuple[Placement, ...]], ...]], optional): 

142 Placements for each input argument. None entries indicate non-tensor inputs. 

143 redistribute_inputs (bool): Whether to redistribute inputs to required placements. 

144 

145 Returns: 

146 Callable: Wrapped function that handles distributed tensors. 

147 

148 Examples: 

149 >>> mesh = DeviceMesh("npu", (2, 2), mesh_dim_names=("dp", "tp")) 

150 >>> @custom_shard( 

151 ... device_mesh=mesh, 

152 ... out_placements=((Shard(0), Replicate()),), 

153 ... in_placements=((Shard(0), Replicate()), (Replicate(), Shard(1))) 

154 ... ) 

155 ... def my_func(x, y): 

156 ... return x + y 

157 """ 

158 def wrapped(*args, **kwargs): 

159 if in_placements is not None: 

160 if len(in_placements) != len(args): 

161 raise ValueError( 

162 f"in_placements length {len(in_placements)} does not match " 

163 f"the number of input args {len(args)}!" 

164 ) 

165 

166 local_args, contain_distributed_arg, _ = _process_custom_shard_inputs( 

167 args, in_placements, redistribute_inputs, device_mesh 

168 ) 

169 out = func(*local_args, **kwargs) 

170 return _wrap_custom_shard_outputs(out, out_placements, contain_distributed_arg, device_mesh) 

171 

172 return wrapped